Replace UnitFinder/TreeMap with KDtree (d2) for a large performance improvement

Author: Bytekeeper

build.gradle.kts

@@ -2,7 +2,7 @@ plugins {
     java
     id("maven")
     id("me.champeau.gradle.jmh") version "0.4.8"
-    id("com.github.ben-manes.versions") version "0.21.0"
+    id("com.github.ben-manes.versions") version "0.22.0"
 }
 
 group = "org.bk"

src/jmh/java/org/bk/ass/query/PositionQueriesBenchmark.java

@@ -0,0 +1,64 @@
+package org.bk.ass.query;
+
+import org.bk.ass.path.Position;
+import org.openbw.bwapi4j.test.BWDataProvider;
+import org.openjdk.jmh.annotations.*;
+
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.List;
+import java.util.SplittableRandom;
+import java.util.function.Function;
+
+@Measurement(iterations = 2, time = 2)
+@Warmup(iterations = 2, time = 2)
+@Fork(2)
+public class PositionQueriesBenchmark {
+
+    @State(Scope.Thread)
+    public static class MyState {
+
+        PositionQueries<Position> positionQueries;
+        List<PositionAndId> entities;
+        List<Position> items;
+
+        @Setup
+        public void setup() {
+            SplittableRandom rnd = new SplittableRandom(815);
+            entities = new ArrayList<>();
+            items = new ArrayList<>();
+            for (int i = 0; i < 1000; i++) {
+                int x = rnd.nextInt(0, 10000);
+                int y = rnd.nextInt(0, 10000);
+                entities.add(new PositionAndId(i, x, y));
+                items.add(new Position(x, y));
+            }
+            positionQueries = new PositionQueries<>(items, Function.identity());
+        }
+    }
+
+    static {
+        try {
+            BWDataProvider.injectValues();
+        } catch (Exception e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+
+    @Benchmark
+    public PositionQueries<Position> queryCreation(MyState state) {
+        return new PositionQueries<>(state.items, Function.identity());
+    }
+
+
+    @Benchmark
+    public Collection<Position> inRadius(MyState state) {
+        return state.positionQueries.inRadius(5000, 5000, 1000);
+    }
+
+    @Benchmark
+    public Position nearest(MyState state) {
+        return state.positionQueries.nearest(2500, 2500);
+    }
+}

src/jmh/java/org/bk/ass/query/UnitFinderBenchmark.java

@@ -1,7 +1,7 @@
 package org.bk.ass.cluster;
 
 import org.bk.ass.collection.UnorderedCollection;
-import org.bk.ass.query.UnitFinder;
+import org.bk.ass.query.PositionQueries;
 
 import java.util.*;
 import java.util.Map.Entry;
@@ -14,7 +14,7 @@
  * is split into multiple clusters. In that case, only a part of the elements are retained in the
  * previous cluster. <br>
  * Make sure to call {@link #scan(int)} to actually perform clustering. Calling {@link
- * #updateDB(UnitFinder, int)} or {@link #updateDB(Collection, Function)} will reset the clustering
+ * #updateDB(PositionQueries, int)} or {@link #updateDB(Collection, Function)} will reset the clustering
  * process. Make sure to always cluster all elements or to check if the clustering is done with
  * {@link #isComplete()}.<br>
  * The clustering is ongoing, each call to {@link #scan(int)} will continue or restart the
@@ -35,8 +35,8 @@ public class StableDBScanner<U> {
   private Map<U, Cluster<U>> elementToCluster = Collections.emptyMap();
   private ClusterSurrogate<U> currentCluster;
 
-  public StableDBScanner(UnitFinder<U> unitFinder, int minPoints, int radius) {
-    this(unitFinder, minPoints, u -> unitFinder.inRadius(u, radius));
+  public StableDBScanner(PositionQueries<U> positionQueries, int minPoints, int radius) {
+    this(positionQueries, minPoints, u -> positionQueries.inRadius(u, radius));
   }
 
   public StableDBScanner(int minPoints) {
@@ -65,8 +65,8 @@ public StableDBScanner<U> updateDB(Collection<U> db, Function<U, Collection<U>>
     return this;
   }
 
-  public StableDBScanner<U> updateDB(UnitFinder<U> unitFinder, int radius) {
-    return updateDB(unitFinder, u -> unitFinder.inRadius(u, radius));
+  public StableDBScanner<U> updateDB(PositionQueries<U> positionQueries, int radius) {
+    return updateDB(positionQueries, u -> positionQueries.inRadius(u, radius));
   }
 
   /**

src/main/java/org/bk/ass/cluster/StableDBScanner.java

@@ -0,0 +1,38 @@
+package org.bk.ass.query;
+
+public class Distances {
+    public static final DistanceProvider EUCLIDEAN_DISTANCE =
+            (ax, ay, bx, by) -> (int) Math.sqrt((float) (bx - ax) * (bx - ax) + (by - ay) * (by - ay));
+
+    /**
+     * When using this, all radius queries need to be made with the squared radius
+     */
+    public static final DistanceProvider EUCLIDEAN_DISTANCE_SQUARED =
+            (ax, ay, bx, by) -> (bx - ax) * (bx - ax) + (by - ay) * (by - ay);
+
+    /**
+     * Use this to query using the distance approximation used in OpenBW
+     */
+    public static final DistanceProvider BW_DISTANCE_APPROXIMATION =
+            (ax, ay, bx, by) -> {
+                int min = Math.abs(ax - bx);
+                int max = Math.abs(ay - by);
+                int minCalc;
+                if (max < min) {
+                    minCalc = max;
+                    max = min;
+                    min = minCalc;
+                }
+
+                if (min < max >> 2) {
+                    return max;
+                } else {
+                    minCalc = 3 * min >> 3;
+                    return (minCalc >> 5) + minCalc + max - (max >> 4) - (max >> 6);
+                }
+            };
+
+    private Distances() {
+        // Utility class
+    }
+}

src/main/java/org/bk/ass/query/Distances.java

@@ -1,5 +1,7 @@
 package org.bk.ass.query;
 
+import org.bk.ass.path.Position;
+
 import java.util.Objects;
 
 /**
@@ -16,6 +18,10 @@ public PositionAndId(int id, int x, int y) {
     this.id = id;
   }
 
+  public Position toPosition() {
+    return new Position(x, y);
+  }
+
   @Override
   public String toString() {
     return "id: " + id + " (" + x + ", " + y + ")";