UnitFinder now allows additional criterias for searching.

Author: Bytekeeper

build.gradle.kts

@@ -2,7 +2,7 @@ plugins {
     java
     id("maven")
     id("me.champeau.gradle.jmh") version "0.4.8"
-    id("org.sonarqube") version "2.7"
+    id("org.sonarqube") version "2.7.1"
     id("com.github.ben-manes.versions") version "0.21.0"
 }
 
@@ -21,11 +21,11 @@ repositories {
 
 dependencies {
     implementation(fileTree("lib").include("*.jar"))
-    implementation("com.github.JasperGeurtz:JBWAPI:develop-SNAPSHOT")
+    implementation("com.github.JasperGeurtz:JBWAPI:0.7")
 
-    testImplementation("org.junit.jupiter:junit-jupiter:5.4.2")
-    testImplementation("org.assertj:assertj-core:3.12.2")
-    testImplementation("io.jenetics:jenetics:4.4.0")
+    testImplementation("org.junit.jupiter:junit-jupiter:5.5.1")
+    testImplementation("org.assertj:assertj-core:3.13.2")
+    testImplementation("io.jenetics:jenetics:5.0.1")
 }
 
 configure<JavaPluginConvention> {

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

@@ -3,6 +3,7 @@
 import java.util.*;
 import java.util.Map.Entry;
 import java.util.function.Function;
+import java.util.function.Predicate;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
@@ -109,12 +110,29 @@ public Collection<U> inArea(int ax, int ay, int bx, int by) {
     return subMapOfArea(ax, ay, bx, by).map(Entry::getValue).collect(Collectors.toList());
   }
 
+  /**
+   * Returns a collection of all units in the area with top left (ax, ay) and bottom right (bx, by)
+   * including those on the boundary. Only units passing the given criteria are returned.
+   */
+  public Collection<U> inArea(int ax, int ay, int bx, int by, Predicate<U> criteria) {
+    return subMapOfArea(ax, ay, bx, by, criteria).map(Entry::getValue).collect(Collectors.toList());
+  }
+
   private Stream<Entry<PositionAndId, U>> subMapOfArea(int ax, int ay, int bx, int by) {
     return map
             .subMap(
                     new PositionAndId(-1, ax, ay), true, new PositionAndId(Integer.MAX_VALUE, bx, by), true)
             .entrySet().stream()
-        .filter(u -> u.getKey().y <= by && u.getKey().y >= ay);
+            .filter(u -> u.getKey().y <= by && u.getKey().y >= ay);
+  }
+
+  private Stream<Entry<PositionAndId, U>> subMapOfArea(
+          int ax, int ay, int bx, int by, Predicate<U> criteria) {
+    return map
+            .subMap(
+                    new PositionAndId(-1, ax, ay), true, new PositionAndId(Integer.MAX_VALUE, bx, by), true)
+            .entrySet().stream()
+            .filter(u -> u.getKey().y <= by && u.getKey().y >= ay && criteria.test(u.getValue()));
   }
 
   /**
@@ -168,8 +186,42 @@ public Optional<U> closestTo(int x, int y) {
       return Optional.empty();
     }
     return subMapOfArea(x - squareHSize, y - squareHSize, x + squareHSize, y + squareHSize)
-        .min(Comparator.comparing(e -> distanceProvider.distance(e.getKey().x, e.getKey().y, x, y)))
-        .map(Entry::getValue);
+            .min(Comparator.comparing(e -> distanceProvider.distance(e.getKey().x, e.getKey().y, x, y)))
+            .map(Entry::getValue);
+  }
+
+  /**
+   * Returns the closest unit to the given position matching the given criteria or {@link
+   * Optional#empty()} if no units are present.
+   */
+  public Optional<U> closestTo(int x, int y, Predicate<U> criteria) {
+    PositionAndId query = new PositionAndId(-1, x, y);
+    int squareHSize = Integer.MAX_VALUE;
+    Entry<PositionAndId, U> lowerBound = map.lowerEntry(query);
+    while (lowerBound != null && !criteria.test(lowerBound.getValue())) {
+      lowerBound = map.lowerEntry(lowerBound.getKey());
+    }
+    if (lowerBound != null) {
+      int lx = lowerBound.getKey().x;
+      int ly = lowerBound.getKey().y;
+      squareHSize = max(abs(lx - x), abs(ly - y));
+    }
+    Entry<PositionAndId, U> higherBound = map.higherEntry(query);
+    while (higherBound != null && !criteria.test(higherBound.getValue())) {
+      higherBound = map.higherEntry(higherBound.getKey());
+    }
+    if (higherBound != null) {
+      int hx = higherBound.getKey().x;
+      int hy = higherBound.getKey().y;
+      squareHSize = min(squareHSize, max(abs(hx - x), abs(hy - y)));
+    }
+    if (squareHSize == Integer.MAX_VALUE) {
+      return Optional.empty();
+    }
+    return subMapOfArea(
+            x - squareHSize, y - squareHSize, x + squareHSize, y + squareHSize, criteria)
+            .min(Comparator.comparing(e -> distanceProvider.distance(e.getKey().x, e.getKey().y, x, y)))
+            .map(Entry::getValue);
   }
 
   @Override

src/test/java/org/bk/ass/query/UnitFinderTest.java

@@ -1,17 +1,19 @@
 package org.bk.ass.query;
 
-import static org.assertj.core.api.Assertions.assertThat;
+import org.junit.jupiter.api.Test;
 
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.Collections;
+import java.util.Objects;
 import java.util.function.Function;
-import org.junit.jupiter.api.Test;
+
+import static org.assertj.core.api.Assertions.assertThat;
 
 class UnitFinderTest {
 
   @Test
-  public void shouldNotDieIfEmpty() {
+  void shouldNotDieIfEmpty() {
     // GIVEN
     UnitFinder<Object> finder =
         new UnitFinder<>(
@@ -30,7 +32,7 @@ public void shouldNotDieIfEmpty() {
   }
 
   @Test
-  public void shouldFindClosest() {
+  void shouldFindClosest() {
     // GIVEN
     UnitFinder<PositionAndId> finder =
         new UnitFinder<>(
@@ -45,7 +47,43 @@ public void shouldFindClosest() {
   }
 
   @Test
-  public void shouldFindAllInArea() {
+  void shouldFindClosestWithCriteria() {
+    // GIVEN
+    UnitFinder<Unit> finder =
+            new UnitFinder<>(
+                    Arrays.asList(
+                            new Unit(-1, new PositionAndId(-1, -30, 0)),
+                            new Unit(0, new PositionAndId(0, 0, -30)),
+                            new Unit(1, new PositionAndId(1, 30, 0))),
+                    u -> u.positionAndId);
+
+    // WHEN
+    Unit closest = finder.closestTo(0, 0, u -> u.id > 0).get();
+
+    // THEN
+    assertThat(closest).isEqualTo(new Unit(1, new PositionAndId(1, 30, 0)));
+  }
+
+  @Test
+  void shouldFindDifferentItemIfCriteriaFails() {
+    // GIVEN
+    UnitFinder<Unit> finder =
+            new UnitFinder<>(
+                    Arrays.asList(
+                            new Unit(-1, new PositionAndId(-1, -30, 0)),
+                            new Unit(0, new PositionAndId(0, 0, -30)),
+                            new Unit(1, new PositionAndId(1, 30, 0))),
+                    u -> u.positionAndId);
+
+    // WHEN
+    Unit closest = finder.closestTo(30, 0, u -> u.id < 0).get();
+
+    // THEN
+    assertThat(closest).isEqualTo(new Unit(-1, new PositionAndId(-1, -30, 0)));
+  }
+
+  @Test
+  void shouldFindAllInArea() {
     // GIVEN
     UnitFinder<PositionAndId> finder =
         new UnitFinder<>(
@@ -67,7 +105,7 @@ public void shouldFindAllInArea() {
   }
 
   @Test
-  public void shouldFindAllInRadius() {
+  void shouldFindAllInRadius() {
     // GIVEN
     UnitFinder<PositionAndId> finder =
         new UnitFinder<>(
@@ -88,7 +126,7 @@ public void shouldFindAllInRadius() {
   }
 
   @Test
-  public void shouldFindAllInRadiusOfUnit() {
+  void shouldFindAllInRadiusOfUnit() {
     // GIVEN
     PositionAndId unit = new PositionAndId(-1, 0, 0);
     UnitFinder<PositionAndId> finder =
@@ -108,4 +146,35 @@ public void shouldFindAllInRadiusOfUnit() {
     // THEN
     assertThat(result).hasSize(5);
   }
+
+  private static class Unit {
+    public final int id;
+    public final PositionAndId positionAndId;
+
+    private Unit(int id, PositionAndId positionAndId) {
+      this.id = id;
+      this.positionAndId = positionAndId;
+    }
+
+    @Override
+    public boolean equals(Object o) {
+      if (this == o) return true;
+      if (o == null || getClass() != o.getClass()) return false;
+      Unit unit = (Unit) o;
+      return id == unit.id && Objects.equals(positionAndId, unit.positionAndId);
+    }
+
+    @Override
+    public int hashCode() {
+      return Objects.hash(id, positionAndId);
+    }
+
+    @Override
+    public String toString() {
+      return "Unit{" +
+              "id=" + id +
+              ", positionAndId=" + positionAndId +
+              '}';
+    }
+  }
 }