cache calls to connected&loaded units per frame

Author: JasperGeurtz

src/main/java/bwapi/Game.java

@@ -43,6 +43,11 @@ public final class Game {
     private static final int REGION_DATA_SIZE = 5000;
 
     private final Set<Integer> visibleUnits = new HashSet<>();
+    private final Map<Unit, List<Unit>> connectedUnits = new HashMap<>();
+    private int lastConnectedUnitsUpdate = -1;
+    private final Map<Unit, List<Unit>> loadedUnits = new HashMap<>();
+    private int lastLoadedUnitsUpdate = -1;
+
     private List<Unit> allUnits;
     private final ClientData clientData;
 
@@ -146,6 +151,10 @@ private static boolean hasPower(final int x, final int y, final UnitType unitTyp
     */
     void init() {
         visibleUnits.clear();
+        connectedUnits.clear();
+        lastConnectedUnitsUpdate = -1;
+        loadedUnits.clear();
+        lastLoadedUnitsUpdate = -1;
 
         final int forceCount = gameData().getForceCount();
         forces = new Force[forceCount];
@@ -310,6 +319,60 @@ void onFrame(final int frame) {
         getAllUnits().forEach(u -> u.updatePosition(frame));
     }
 
+    /**
+     * Lazily update connectedUnits. Only users of the calls pay for it, and only
+     * pay once per frame.
+     * Avoids previous O(n^2) implementation which would be costly for
+     * lategame carrier fights
+     */
+    public List<Unit> getConnected(final Unit unit) {
+        final int frame = getFrameCount();
+        if (lastConnectedUnitsUpdate < frame) {
+            connectedUnits.clear();
+            for (final Unit u : getAllUnits()) {
+                Unit owner = u.getCarrier();
+                if (owner == null) {
+                    owner = u.getHatchery();
+                }
+                if (owner != null) {
+                    if (!connectedUnits.containsKey(owner)) {
+                        connectedUnits.put(owner, new ArrayList<>());
+                    }
+                    connectedUnits.get(owner).add(u);
+                }
+            }
+            lastConnectedUnitsUpdate = frame;
+        }
+        if (!connectedUnits.containsKey(unit)) {
+            return Collections.emptyList();
+        }
+        return new ArrayList<>(connectedUnits.get(unit));
+    }
+
+    /**
+     * @see #getConnected
+     */
+    public List<Unit> getLoadedUnits(final Unit unit) {
+        final int frame = getFrameCount();
+        if (lastLoadedUnitsUpdate < frame) {
+            loadedUnits.clear();
+            for (final Unit u : getAllUnits()) {
+                final Unit owner = u.getTransport();
+                if (owner != null) {
+                    if (!loadedUnits.containsKey(owner)) {
+                        loadedUnits.put(owner, new ArrayList<>());
+                    }
+                    loadedUnits.get(owner).add(u);
+                }
+            }
+            lastLoadedUnitsUpdate = frame;
+        }
+        if (!loadedUnits.containsKey(unit)) {
+            return Collections.emptyList();
+        }
+        return new ArrayList<>(loadedUnits.get(unit));
+    }
+
     /**
      * Retrieves the set of all teams/forces. Forces are commonly seen in @UMS
      * game types and some others such as @TvB and the team versions of game types.

src/main/java/bwapi/Unit.java

@@ -1157,9 +1157,7 @@ public List<Unit> getLoadedUnits() {
         if (getType().spaceProvided() < 1) {
             return Collections.emptyList();
         }
-        return game.getAllUnits().stream()
-                .filter(u -> equals(u.getTransport()))
-                .collect(Collectors.toList());
+        return game.getLoadedUnits(this);
     }
 
     /**
@@ -1199,9 +1197,7 @@ public List<Unit> getInterceptors() {
         if (getType() != Protoss_Carrier && getType() != Hero_Gantrithor) {
             return Collections.emptyList();
         }
-        return game.getAllUnits().stream()
-                .filter(u -> equals(u.getCarrier()))
-                .collect(Collectors.toList());
+        return game.getConnected(this);
     }
 
     /**
@@ -1229,9 +1225,7 @@ public List<Unit> getLarva() {
         if (!getType().producesLarva()) {
             return Collections.emptyList();
         }
-        return game.getAllUnits().stream()
-                .filter(u -> equals(u.getHatchery()))
-                .collect(Collectors.toList());
+        return game.getConnected(this);
     }
 
     public List<Unit> getUnitsInRadius(final int radius) {