Add lazy per-frame UnitListCache for filtered unit lists

Author: Jasper Geurtz

src/main/java/bwapi/Cache.java

@@ -1,5 +1,14 @@
 package bwapi;
 
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.function.Function;
+import java.util.function.Predicate;
+import java.util.function.Supplier;
+import java.util.stream.Collectors;
 
 class Cache<T> {
     private int frame = -1;
@@ -35,4 +44,83 @@ class BooleanCache extends Cache<Boolean>{}
 class OrderCache extends Cache<Order>{}
 class UnitTypeCache extends Cache<UnitType>{}
 class UpgradeTypeCache extends Cache<UpgradeType>{}
-class TechTypeCache extends Cache<TechType>{}
+class TechTypeCache extends Cache<TechType>{}
+
+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) {
+                    connectedUnits.computeIfAbsent(owner, k -> new ArrayList<>()).add(u);
+                }
+            }
+            lastUpdate = frame;
+        }
+        List<Unit> connected = connectedUnits.get(unit);
+        return connected == null ?
+                Collections.emptyList() :
+                Collections.unmodifiableList(connected);
+    }
+
+    void reset() {
+        lastUpdate = -1;
+        connectedUnits.clear();
+    }
+}
+
+/**
+ * Lazily filters and caches an unmodifiable List<Unit> at most once per frame.
+ * Only callers of the corresponding getter pay for the filtering, and at most once per frame.
+ */
+class UnitListCache {
+    private final Game game;
+    private final Supplier<List<Unit>> source;
+    private final Predicate<Unit> filter;
+    private int lastFrame = -1;
+    private List<Unit> cached = Collections.emptyList();
+
+    UnitListCache(final Game game, final Predicate<Unit> filter) {
+        this(game, game::getAllUnits, filter);
+    }
+
+    UnitListCache(final Game game, final Supplier<List<Unit>> source, final Predicate<Unit> filter) {
+        this.game = game;
+        this.source = source;
+        this.filter = filter;
+    }
+
+    List<Unit> get() {
+        final int frame = game.getFrameCount();
+        if (lastFrame != frame) {
+            cached = source.get().stream()
+                    .filter(filter)
+                    .collect(Collectors.collectingAndThen(Collectors.toList(), Collections::unmodifiableList));
+            lastFrame = frame;
+        }
+        return cached;
+    }
+
+    void reset() {
+        lastFrame = -1;
+        cached = Collections.emptyList();
+    }
+}

src/main/java/bwapi/ConnectedUnitCache.java

@@ -102,6 +102,10 @@ public class Game {
     final ConnectedUnitCache loadedUnitsCache = new ConnectedUnitCache(this, Unit::getTransport);
     final ConnectedUnitCache interceptorsCache = new ConnectedUnitCache(this, Unit::getCarrier);
     final ConnectedUnitCache larvaCache = new ConnectedUnitCache(this, Unit::getHatchery);
+    private final UnitListCache minerals = new UnitListCache(this, u -> u.getType().isMineralField());
+    private final UnitListCache geysers = new UnitListCache(this, u -> u.getType() == Resource_Vespene_Geyser);
+    private final UnitListCache neutralUnits = new UnitListCache(this, u -> u.getPlayer().equals(neutral()));
+    private final UnitListCache selfPylons = new UnitListCache(this, () -> self().getUnits(), u -> u.getType() == Protoss_Pylon);
 
     final SideEffectQueue sideEffects = new SideEffectQueue();
 
@@ -155,6 +159,10 @@ void init() {
         loadedUnitsCache.reset();
         interceptorsCache.reset();
         larvaCache.reset();
+        minerals.reset();
+        geysers.reset();
+        neutralUnits.reset();
+        selfPylons.reset();
 
         final int forceCount = gameData().getForceCount();
         forces = new Force[forceCount];
@@ -368,9 +376,7 @@ public List<Unit> getAllUnits() {
      * @return List<Unit> containing @minerals
      */
     public List<Unit> getMinerals() {
-        return getAllUnits().stream()
-                .filter(u -> u.getType().isMineralField())
-                .collect(Collectors.toList());
+        return minerals.get();
     }
 
     /**
@@ -379,9 +385,7 @@ public List<Unit> getMinerals() {
      * @return List<Unit> containing @geysers
      */
     public List<Unit> getGeysers() {
-        return getAllUnits().stream()
-                .filter(u -> u.getType() == Resource_Vespene_Geyser)
-                .collect(Collectors.toList());
+        return geysers.get();
     }
 
     /**
@@ -391,9 +395,7 @@ public List<Unit> getGeysers() {
      * @return List<Unit> containing all neutral units.
      */
     public List<Unit> getNeutralUnits() {
-        return getAllUnits().stream()
-                .filter(u -> u.getPlayer().equals(neutral()))
-                .collect(Collectors.toList());
+        return neutralUnits.get();
     }
 
     /**
@@ -1018,7 +1020,7 @@ public boolean hasPowerPrecise(final int x, final int y) {
      * @return true if the type at the given position will have power, false if the type at the given position will be unpowered.
      */
     public boolean hasPowerPrecise(final int x, final int y, final UnitType unitType) {
-        return isValidPosition(x, y) && hasPower(x, y, unitType, self().getUnits().stream().filter(u -> u.getType() == Protoss_Pylon).collect(Collectors.toList()));
+        return isValidPosition(x, y) && hasPower(x, y, unitType, selfPylons.get());
     }
 
     public boolean hasPowerPrecise(final Position position) {

src/main/java/bwapi/Game.java

@@ -4,7 +4,6 @@
 
 import java.util.List;
 import java.util.Objects;
-import java.util.stream.Collectors;
 
 import static bwapi.UnitType.*;
 import static bwapi.UpgradeType.*;
@@ -28,6 +27,7 @@ public class Player implements Comparable<Player> {
     private final PlayerType playerType;
     private final TilePosition startLocation;
     private final Color color;
+    private final UnitListCache units;
     private PlayerSelf self = null;
     PlayerSelf self() {
         if (self == null) {
@@ -44,6 +44,7 @@ PlayerSelf self() {
         this.playerType = PlayerType.idToEnum[playerData.getType()];
         this.startLocation = new TilePosition(playerData.getStartLocationX(), playerData.getStartLocationY());
         this.color = new Color(playerData.getColor());
+        this.units = new UnitListCache(game, u -> equals(u.getPlayer()));
     }
 
     /**
@@ -77,9 +78,7 @@ public String getName() {
      * This does not include units that are loaded into transports, @Bunkers, @Refineries, @Assimilators, or @Extractors.
      */
     public List<Unit> getUnits() {
-        return game.getAllUnits().stream()
-                .filter(u -> equals(u.getPlayer()))
-                .collect(Collectors.toList());
+        return units.get();
     }
 
     /**

src/main/java/bwapi/Player.java

@@ -17,6 +17,7 @@
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertSame;
 import static org.junit.Assert.assertTrue;
 import static org.mockito.BDDMockito.given;
 import static org.mockito.Mockito.doReturn;
@@ -210,4 +211,35 @@ public void onFrameShouldRefreshPlayersAndExcludeDefeatedOrLeftPlayers() {
 
         assertTrue(game.enemies().isEmpty());
     }
+
+    @Test
+    public void unitListCachesShouldBeLazyAndCachedPerFrame() throws IOException {
+        Game game = GameBuilder.createGame();
+        Player self = game.self();
+        Player neutral = game.neutral();
+
+        Unit mineral = mock(Unit.class);
+        when(mineral.getType()).thenReturn(UnitType.Resource_Mineral_Field);
+        when(mineral.getPlayer()).thenReturn(neutral);
+
+        Unit geyser = mock(Unit.class);
+        when(geyser.getType()).thenReturn(UnitType.Resource_Vespene_Geyser);
+        when(geyser.getPlayer()).thenReturn(neutral);
+
+        Unit pylon = mock(Unit.class);
+        when(pylon.getType()).thenReturn(UnitType.Protoss_Pylon);
+        when(pylon.getPlayer()).thenReturn(self);
+
+        game.setAllUnits(Arrays.asList(mineral, geyser, pylon));
+
+        assertEquals(Collections.singletonList(mineral), game.getMinerals());
+        assertSame(game.getMinerals(), game.getMinerals());
+        assertEquals(Collections.singletonList(geyser), game.getGeysers());
+        assertSame(game.getGeysers(), game.getGeysers());
+        assertEquals(Arrays.asList(mineral, geyser), game.getNeutralUnits());
+        assertSame(game.getNeutralUnits(), game.getNeutralUnits());
+        assertEquals(Collections.singletonList(pylon), self.getUnits());
+        assertSame(self.getUnits(), self.getUnits());
+        assertEquals(Arrays.asList(mineral, geyser), neutral.getUnits());
+    }
 }