Accelerate repeated spatial unit queries

src/main/java/bwapi/Game.java

@@ -5,6 +5,7 @@
 import java.util.*;
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
+import java.util.stream.Stream;
 
 import static bwapi.CommandType.*;
 import static bwapi.Point.TILE_WALK_FACTOR;
@@ -51,6 +52,8 @@ public final class Game {
     private List<Unit> staticGeysers;
     private List<Unit> staticNeutralUnits;
 
+    final UnitFinder unitFinder = new UnitFinder(this);
+
     // CONSTANT
     private Player[] players;
     private Region[] regions;
@@ -680,7 +683,7 @@ public List<Unit> getUnitsOnTile(final int tileX, final int tileY) {
      * given build tile.
      */
     public List<Unit> getUnitsOnTile(final int tileX, final int tileY, final UnitFilter pred) {
-        return getAllUnits().stream().filter(u -> {
+        return unitFinder.onTile(tileX, tileY).filter(u -> {
             final TilePosition tp = u.getTilePosition();
             return tp.x == tileX && tp.y == tileY && pred.test(u);
         }).collect(Collectors.toList());
@@ -702,9 +705,13 @@ public List<Unit> getUnitsInRectangle(final int left, final int top, final int r
      * given rectangle bounds.
      */
     public List<Unit> getUnitsInRectangle(final int left, final int top, final int right, final int bottom, final UnitFilter pred) {
-        return getAllUnits().stream()
-                .filter(u -> left <= u.getRight() && top <= u.getBottom() && right >= u.getLeft() && bottom >= u.getTop() && pred.test(u))
-                .collect(Collectors.toList());
+        return unitsInRectangle(left, top, right, bottom, pred).collect(Collectors.toList());
+    }
+
+    private Stream<Unit> unitsInRectangle(final int left, final int top, final int right, final int bottom, final UnitFilter pred) {
+        return unitFinder.inRectangle((long) left - UnitType.maxUnitWidth(), (long) top - UnitType.maxUnitHeight(),
+                (long) right + UnitType.maxUnitWidth(), (long) bottom + UnitType.maxUnitHeight())
+                .filter(u -> left <= u.getRight() && top <= u.getBottom() && right >= u.getLeft() && bottom >= u.getTop() && pred.test(u));
     }
 
     public List<Unit> getUnitsInRectangle(final Position leftTop, final Position rightBottom) {
@@ -739,9 +746,13 @@ public List<Unit> getUnitsInRadius(final Position center, final int radius) {
     }
 
     public List<Unit> getUnitsInRadius(final Position center, final int radius, final UnitFilter pred) {
-        return getAllUnits().stream()
-                .filter(u -> center.getApproxDistance(u.getPosition()) <= radius && pred.test(u))
-                .collect(Collectors.toList());
+        return unitsInRadius(center, radius, pred).collect(Collectors.toList());
+    }
+
+    private Stream<Unit> unitsInRadius(final Position center, final int radius, final UnitFilter pred) {
+        return unitFinder.inRectangle((long) center.x - radius, (long) center.y - radius,
+                (long) center.x + radius, (long) center.y + radius)
+                .filter(u -> center.getApproxDistance(u.getPosition()) <= radius && pred.test(u));
     }
 
     public Unit getClosestUnitInRectangle(final Position center, final int left, final int top, final int right, final int bottom) {
@@ -761,7 +772,7 @@ public Unit getClosestUnitInRectangle(final Position center, final int left, fin
      * @see UnitFilter
      */
     public Unit getClosestUnitInRectangle(final Position center, final int left, final int top, final int right, final int bottom, final UnitFilter pred) {
-        return getUnitsInRectangle(left, top, right, bottom, pred).stream()
+        return unitsInRectangle(left, top, right, bottom, pred)
                 .min(Comparator.comparingInt(u -> u.getDistance(center)))
                 .orElse(null);
     }
@@ -789,7 +800,7 @@ public Unit getClosestUnit(final Position center, final int radius) {
      * @see UnitFilter
      */
     public Unit getClosestUnit(final Position center, final int radius, final UnitFilter pred) {
-        return getUnitsInRadius(center, radius, pred).stream()
+        return unitsInRadius(center, radius, pred)
                 .min(Comparator.comparingInt(u -> u.getDistance(center)))
                 .orElse(null);
     }

src/main/java/bwapi/Unit.java

@@ -6390,6 +6390,7 @@ void updatePosition(final int frame) {
         if (frame > lastPositionUpdate) {
             lastPositionUpdate = frame;
             position = new Position(unitData.getPositionX(), unitData.getPositionY());
+            game.unitFinder.reset();
         }
     }
 }

src/main/java/bwapi/UnitFinder.java

@@ -0,0 +1,104 @@
+package bwapi;
+
+import java.util.*;
+import java.util.stream.Stream;
+
+/**
+ * Lazily indexes unit centers once per frame. Bounds and filters are evaluated
+ * at query time, including any latency-compensated type changes.
+ */
+class UnitFinder {
+    private static final int CELL_SHIFT = 7;
+    private static final int MAX_TILE_X = Arrays.stream(UnitType.values()).mapToInt(UnitType::tileWidth).max().getAsInt() * 16;
+    private static final int MAX_TILE_Y = Arrays.stream(UnitType.values()).mapToInt(UnitType::tileHeight).max().getAsInt() * 16;
+
+    private final Game game;
+    private final Map<Long, List<Integer>> cells = new HashMap<>();
+    private List<Unit> indexedUnits;
+    private List<Unit> queriedUnits;
+
+    UnitFinder(final Game game) {
+        this.game = game;
+    }
+
+    void reset() {
+        indexedUnits = null;
+        queriedUnits = null;
+    }
+
+    Stream<Unit> inRectangle(final long left, final long top, final long right, final long bottom) {
+        final List<Unit> units = game.getAllUnits();
+        if (left > right || top > bottom) {
+            return Stream.empty();
+        }
+        // A broad query is cheaper as a single pass, and needs no index.
+        if (cellCount(left, top, right, bottom) >= units.size() || !prepare(units)) {
+            return units.stream();
+        }
+        final BitSet candidates = new BitSet();
+        addCandidates(candidates, left, top, right, bottom);
+        // Preserve getAllUnits order, including closest-unit ties and filter order.
+        return candidates.stream().mapToObj(units::get);
+    }
+
+    Stream<Unit> onTile(final int tileX, final int tileY) {
+        if (tileX < 0 || tileY < 0) {
+            return Stream.empty();
+        }
+        final List<Unit> units = game.getAllUnits();
+        final long left = tileX * 32L - MAX_TILE_X;
+        final long top = tileY * 32L - MAX_TILE_Y;
+        final long right = tileX * 32L + 31 + MAX_TILE_X;
+        final long bottom = tileY * 32L + 31 + MAX_TILE_Y;
+        if (4 * cellCount(left, top, right, bottom) >= units.size() || !prepare(units)) {
+            return units.stream();
+        }
+        final BitSet candidates = new BitSet();
+        // getTilePosition uses absolute coordinates; include reflected centers.
+        addCandidates(candidates, left, top, right, bottom);
+        addCandidates(candidates, -right, top, -left, bottom);
+        addCandidates(candidates, left, -bottom, right, -top);
+        addCandidates(candidates, -right, -bottom, -left, -top);
+        return candidates.stream().mapToObj(units::get);
+    }
+
+    private boolean prepare(final List<Unit> units) {
+        if (indexedUnits == units) {
+            return true;
+        }
+        // Do not build an index for a single local query in a frame.
+        if (queriedUnits != units) {
+            queriedUnits = units;
+            return false;
+        }
+        cells.clear();
+        for (int i = 0; i < units.size(); i++) {
+            final Position position = units.get(i).getPosition();
+            cells.computeIfAbsent(key(position.x >> CELL_SHIFT, position.y >> CELL_SHIFT), unused -> new ArrayList<>()).add(i);
+        }
+        indexedUnits = units;
+        return true;
+    }
+
+    private void addCandidates(final BitSet candidates, final long left, final long top, final long right, final long bottom) {
+        for (long x = left >> CELL_SHIFT; x <= right >> CELL_SHIFT; x++) {
+            for (long y = top >> CELL_SHIFT; y <= bottom >> CELL_SHIFT; y++) {
+                final List<Integer> cell = cells.get(key(x, y));
+                if (cell != null) {
+                    for (final int i : cell) {
+                        candidates.set(i);
+                    }
+                }
+            }
+        }
+    }
+
+    private static long cellCount(final long left, final long top, final long right, final long bottom) {
+        return ((right >> CELL_SHIFT) - (left >> CELL_SHIFT) + 1) * ((bottom >> CELL_SHIFT) - (top >> CELL_SHIFT) + 1);
+    }
+
+    private static long key(final long x, final long y) {
+        // Mix the axes to avoid diagonal cells sharing the same Long.hashCode.
+        return (x << 32) ^ ((y ^ (x * 0x9e3779b9L)) & 0xffffffffL);
+    }
+}

src/main/java/bwapi/UnitType.java

@@ -258,14 +258,19 @@ public enum UnitType {
         this.id = id;
     }
 
+    private static class MaxDimensions {
+        private static final int width = Arrays.stream(UnitType.values()).mapToInt(UnitType::width).max().getAsInt();
+        private static final int height = Arrays.stream(UnitType.values()).mapToInt(UnitType::height).max().getAsInt();
+    }
+
     /**
      * Retrieves the maximum unit width from the set of all units. Used
      * internally to search through unit positions efficiently.
      *
      * @return The maximum width of all unit types, in pixels.
      */
     public static int maxUnitWidth() {
-        return Arrays.stream(UnitType.values()).max(Comparator.comparingInt(UnitType::width)).get().width();
+        return MaxDimensions.width;
     }
 
     /**
@@ -275,7 +280,7 @@ public static int maxUnitWidth() {
      * @return The maximum height of all unit types, in pixels.
      */
     public static int maxUnitHeight() {
-        return Arrays.stream(UnitType.values()).max(Comparator.comparingInt(UnitType::height)).get().height();
+        return MaxDimensions.height;
     }
 
     /**