Give each unittype their own cache

src/main/java/bwapi/ConnectedUnitCache.java

@@ -0,0 +1,47 @@
+package bwapi;
+
+import java.util.*;
+import java.util.function.Function;
+
+class ConnectedUnitCache {
+    private int lastUpdate = -1;
+    private final Map<Unit, List<Unit>> connectedUnits = new HashMap<>();
+    private final Function<Unit, Unit> condition;
+    private final Game game;
+
+    ConnectedUnitCache(final Game game, final Function<Unit, Unit> condition) {
+        this.game = game;
+        this.condition = condition;
+    }
+
+    /**
+     * 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
+     */
+    List<Unit> getConnected(final Unit unit) {
+        final int frame = game.getFrameCount();
+        if (lastUpdate < frame) {
+            connectedUnits.values().forEach(List::clear);
+            for (final Unit u : game.getAllUnits()) {
+                final Unit owner = condition.apply(u);
+                if (owner != null) {
+                    if (!connectedUnits.containsKey(owner)) {
+                        connectedUnits.put(owner, new ArrayList<>());
+                    }
+                    connectedUnits.get(owner).add(u);
+                }
+            }
+            lastUpdate = frame;
+        }
+        if (!connectedUnits.containsKey(unit)) {
+            return Collections.emptyList();
+        }
+        return Collections.unmodifiableList(connectedUnits.get(unit));
+    }
+    void reset() {
+        lastUpdate = -1;
+        connectedUnits.clear();
+    }
+}

src/main/java/bwapi/Game.java

@@ -43,10 +43,6 @@ 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;
@@ -103,6 +99,10 @@ public final class Game {
     private BWClientConfiguration configuration = new BWClientConfiguration();
     private boolean latcom = true;
 
+    final ConnectedUnitCache loadedUnitsCache = new ConnectedUnitCache(this, Unit::getTransport);
+    final ConnectedUnitCache interceptorsCache = new ConnectedUnitCache(this, Unit::getCarrier);
+    final ConnectedUnitCache larvaCache = new ConnectedUnitCache(this, Unit::getHatchery);
+
     final SideEffectQueue sideEffects = new SideEffectQueue();
 
     Game() {
@@ -151,10 +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;
+
+        loadedUnitsCache.reset();
+        interceptorsCache.reset();
+        larvaCache.reset();
 
         final int forceCount = gameData().getForceCount();
         forces = new Force[forceCount];
@@ -319,60 +319,6 @@ 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
-     */
-    List<Unit> getConnected(final Unit unit) {
-        final int frame = getFrameCount();
-        if (lastConnectedUnitsUpdate < frame) {
-            connectedUnits.values().forEach(List::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 Collections.unmodifiableList(connectedUnits.get(unit));
-    }
-
-    /**
-     * @see #getConnected
-     */
-    List<Unit> getLoadedUnits(final Unit unit) {
-        final int frame = getFrameCount();
-        if (lastLoadedUnitsUpdate < frame) {
-            loadedUnits.values().forEach(List::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 Collections.unmodifiableList(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

@@ -1172,7 +1172,7 @@ public List<Unit> getLoadedUnits() {
         if (getType().spaceProvided() < 1) {
             return Collections.emptyList();
         }
-        return game.getLoadedUnits(this);
+        return game.loadedUnitsCache.getConnected(this);
     }
 
     /**
@@ -1212,7 +1212,7 @@ public List<Unit> getInterceptors() {
         if (getType() != Protoss_Carrier && getType() != Hero_Gantrithor) {
             return Collections.emptyList();
         }
-        return game.getConnected(this);
+        return game.interceptorsCache.getConnected(this);
     }
 
     /**
@@ -1240,7 +1240,7 @@ public List<Unit> getLarva() {
         if (!getType().producesLarva()) {
             return Collections.emptyList();
         }
-        return game.getConnected(this);
+        return game.larvaCache.getConnected(this);
     }
 
     public List<Unit> getUnitsInRadius(final int radius) {