Use Multipolygon for all r-tree like structures

README.md

@@ -31,7 +31,7 @@ Country information and boundary data comes from [GeoNames](http://download.geon
 
 ## Performance
 
-~4 μs average to look up city from the GeoNames data set on a 13" MacBook Pro (i5-5257U) using a single thread. The retained heap usage for `ReverseGeocoder` is ~5.54 MB.
+~3 μs average to look up city from the GeoNames data set on a 13" MacBook Pro (i5-5257U) using a single thread. The retained heap size for `ReverseGeocoder` is ~7.5 MB.
 
 ## Algorithms
 

src/main/java/uk/recurse/geocoding/reverse/Feature.java

@@ -6,6 +6,7 @@
 import com.fasterxml.jackson.annotation.JsonProperty;
 
 import java.util.Map;
+import java.util.stream.Stream;
 
 @JsonIgnoreProperties(ignoreUnknown = true)
 class Feature {
@@ -24,15 +25,11 @@ class Feature {
         this.geometry = geometry;
     }
 
-    boolean contains(float lat, float lon) {
-        return geometry.contains(lat, lon);
-    }
-
     Country country() {
         return country;
     }
 
-    Geometry geometry() {
-        return geometry;
+    Stream<Geometry> geometries() {
+        return geometry.flatten(country);
     }
 }

src/main/java/uk/recurse/geocoding/reverse/FeatureCollection.java

@@ -4,58 +4,27 @@
 import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
 import com.fasterxml.jackson.annotation.JsonProperty;
 
-import java.util.Comparator;
-import java.util.Map;
-import java.util.function.Function;
 import java.util.stream.Stream;
 
-import static java.util.stream.Collectors.groupingBy;
-import static java.util.stream.Collectors.toMap;
-
 @JsonIgnoreProperties(ignoreUnknown = true)
 class FeatureCollection {
 
-    private static final Comparator<Feature> POPULATION =
-            Comparator.comparingInt(feature -> feature.country().population());
-
-    private final Map<BoundingBox, Feature[]> featureMap;
+    private final Geometry world;
+    private final Country[] countries;
 
     @JsonCreator
     FeatureCollection(@JsonProperty("features") Feature[] features) {
-        featureMap = partitionByContinent(features);
-    }
-
-    private Map<BoundingBox, Feature[]> partitionByContinent(Feature[] features) {
-        return Stream.of(features)
-                .sorted(POPULATION.reversed()) // performance bias towards populous countries
-                .collect(groupingBy(feature -> feature.country().continent()))
-                .values()
-                .stream()
-                .map(continent -> continent.toArray(new Feature[0]))
-                .collect(toMap(this::continentBoundingBox, Function.identity()));
-    }
-
-    private BoundingBox continentBoundingBox(Feature[] continent) {
-        Geometry[] countries = Stream.of(continent)
-                .map(Feature::geometry)
-                .toArray(Geometry[]::new);
-        return new BoundingBox(countries);
+        world = new MultiPolygon(Stream.of(features).flatMap(Feature::geometries));
+        countries = Stream.of(features)
+                .map(Feature::country)
+                .toArray(Country[]::new);
     }
 
-    Feature getFeature(float lat, float lon) {
-        for (Map.Entry<BoundingBox, Feature[]> entry: featureMap.entrySet()) {
-            if (entry.getKey().contains(lat, lon)) {
-                for (Feature feature : entry.getValue()) {
-                    if (feature.contains(lat, lon)) {
-                        return feature;
-                    }
-                }
-            }
-        }
-        return null;
+    Country getCountry(float lat, float lon) {
+        return world.getCountry(lat, lon);
     }
 
-    Stream<Feature> stream() {
-         return featureMap.values().stream().flatMap(Stream::of);
+    Stream<Country> countries() {
+        return Stream.of(countries);
     }
 }

src/main/java/uk/recurse/geocoding/reverse/Geometry.java

@@ -5,11 +5,17 @@
 import com.fasterxml.jackson.annotation.JsonTypeInfo;
 import com.fasterxml.jackson.annotation.JsonTypeInfo.Id;
 
+import java.util.stream.Stream;
+
 @JsonTypeInfo(use = Id.NAME, property = "type")
 @JsonSubTypes({@Type(Polygon.class), @Type(MultiPolygon.class)})
 interface Geometry {
 
     boolean contains(float lat, float lon);
 
+    Country getCountry(float lat, float lon);
+
     BoundingBox boundingBox();
+
+    Stream<Geometry> flatten(Country country);
 }

src/main/java/uk/recurse/geocoding/reverse/MultiPolygon.java

@@ -3,22 +3,29 @@
 import com.fasterxml.jackson.annotation.JsonCreator;
 import com.fasterxml.jackson.annotation.JsonProperty;
 
+import java.util.ArrayList;
+import java.util.Comparator;
+import java.util.List;
 import java.util.stream.Stream;
 
+import static java.util.stream.Collectors.toList;
+
 class MultiPolygon implements Geometry {
 
     private final Geometry[] geometries;
     private final BoundingBox boundingBox;
 
     @JsonCreator
     MultiPolygon(@JsonProperty("coordinates") Ring[][] rings) {
-        this(RTree.sortTileRecursive(Stream.of(rings)
-                .map(Polygon::new)
-                .toArray(Polygon[]::new)));
+        this(Stream.of(rings).map(Polygon::new));
+    }
+
+    MultiPolygon(Stream<? extends Geometry> geometries) {
+        this(SortTileRecursive.pack(geometries.collect(toList())));
     }
 
-    MultiPolygon(Geometry[] geometries) {
-        this.geometries = geometries;
+    private MultiPolygon(List<Geometry> geometries) {
+        this.geometries = geometries.toArray(new Geometry[geometries.size()]);
         boundingBox = new BoundingBox(this.geometries);
     }
 
@@ -34,8 +41,73 @@ public boolean contains(float lat, float lon) {
         return false;
     }
 
+    @Override
+    public Country getCountry(float lat, float lon) {
+        if (boundingBox.contains(lat, lon)) {
+            for (Geometry geometry : geometries) {
+                Country country = geometry.getCountry(lat, lon);
+                if (country != null) {
+                    return country;
+                }
+            }
+        }
+        return null;
+    }
+
     @Override
     public BoundingBox boundingBox() {
         return boundingBox;
     }
+
+    @Override
+    public Stream<Geometry> flatten(Country country) {
+        return Stream.of(geometries).flatMap(geometry -> geometry.flatten(country));
+    }
+
+    // algorithm paper: http://www.dtic.mil/dtic/tr/fulltext/u2/a324493.pdf
+    static class SortTileRecursive {
+
+        private static final int PAGE_SIZE = 16;
+        private static final Comparator<Geometry> X_COORDINATE =
+                Comparator.comparingDouble(geometry -> geometry.boundingBox().longitude());
+        private static final Comparator<Geometry> Y_COORDINATE =
+                Comparator.comparingDouble(geometry -> geometry.boundingBox().latitude());
+
+        static List<Geometry> pack(List<Geometry> rectangles) {
+            int leafPages = ceilDiv(rectangles.size(), PAGE_SIZE);
+            if (leafPages <= 1) {
+                return rectangles;
+            }
+            int verticalSlices = ceilSqrt(leafPages);
+            List<Geometry> nodes = new ArrayList<>(leafPages);
+            rectangles.sort(X_COORDINATE);
+            for (List<Geometry> verticalSlice : partition(rectangles, verticalSlices)) {
+                verticalSlice.sort(Y_COORDINATE);
+                int runs = ceilDiv(verticalSlice.size(), PAGE_SIZE);
+                for (List<Geometry> run : partition(verticalSlice, runs)) {
+                    nodes.add(new MultiPolygon(run));
+                }
+            }
+            return pack(nodes);
+        }
+
+        private static <T> List<List<T>> partition(List<T> list, int n) {
+            int size = ceilDiv(list.size(), n);
+            List<List<T>> partitions = new ArrayList<>(n);
+            for (int i = 0; i < n; i++) {
+                int start = i * size;
+                int end = Math.min(start + size, list.size());
+                partitions.add(list.subList(start, end));
+            }
+            return partitions;
+        }
+
+        private static int ceilDiv(double x, double y) {
+            return (int) Math.ceil(x / y);
+        }
+
+        private static int ceilSqrt(double a) {
+            return (int) Math.ceil(Math.sqrt(a));
+        }
+    }
 }

src/main/java/uk/recurse/geocoding/reverse/Polygon.java

@@ -3,35 +3,44 @@
 import com.fasterxml.jackson.annotation.JsonCreator;
 import com.fasterxml.jackson.annotation.JsonProperty;
 
-import java.util.Arrays;
+import java.util.stream.Stream;
 
 class Polygon implements Geometry {
 
     private final Ring ring;
-    private final Geometry[] holes;
+    private final Geometry holes;
+    private final Country country;
 
     @JsonCreator
     Polygon(@JsonProperty("coordinates") Ring[] rings) {
         ring = rings[0];
-        holes = RTree.sortTileRecursive(Arrays.copyOfRange(rings, 1, rings.length));
+        holes = new MultiPolygon(Stream.of(rings).skip(1));
+        country = null;
+    }
+
+    private Polygon(Polygon polygon, Country country) {
+        ring = polygon.ring;
+        holes = polygon.holes;
+        this.country = country;
     }
 
     @Override
     public boolean contains(float lat, float lon) {
-        return ring.contains(lat, lon) && !inHoles(lat, lon);
+        return ring.contains(lat, lon) && !holes.contains(lat, lon);
+    }
+
+    @Override
+    public Country getCountry(float lat, float lon) {
+        return contains(lat, lon) ? country : null;
     }
 
     @Override
     public BoundingBox boundingBox() {
         return ring.boundingBox();
     }
 
-    private boolean inHoles(float lat, float lon) {
-        for (Geometry hole : holes) {
-            if (hole.contains(lat, lon)) {
-                return true;
-            }
-        }
-        return false;
+    @Override
+    public Stream<Geometry> flatten(Country country) {
+        return Stream.of(new Polygon(this, country));
     }
 }

src/main/java/uk/recurse/geocoding/reverse/ReverseGeocoder.java

@@ -64,7 +64,7 @@ public Optional<Country> getCountry(double lat, double lon) {
      * @return the country at the given coordinate
      */
     public Optional<Country> getCountry(float lat, float lon) {
-        return Optional.ofNullable(featureCollection.getFeature(lat, lon)).map(Feature::country);
+        return Optional.ofNullable(featureCollection.getCountry(lat, lon));
     }
 
     /**
@@ -73,7 +73,7 @@ public Optional<Country> getCountry(float lat, float lon) {
      * @return stream of countries
      */
     public Stream<Country> countries() {
-        return featureCollection.stream().map(Feature::country);
+        return featureCollection.countries();
     }
 
 }

src/main/java/uk/recurse/geocoding/reverse/Ring.java

@@ -28,13 +28,23 @@ public boolean contains(float lat, float lon) {
         return boundingBox.contains(lat, lon) && pnpoly(lat, lon);
     }
 
+    @Override
+    public Country getCountry(float lat, float lon) {
+        return null;
+    }
+
     @Override
     public BoundingBox boundingBox() {
         return boundingBox;
     }
 
+    @Override
+    public Stream<Geometry> flatten(Country country) {
+        return Stream.of(this);
+    }
+
     // algorithm notes: https://www.ecse.rpi.edu/Homepages/wrf/Research/Short_Notes/pnpoly.html
-    private boolean pnpoly(double lat, double lon) {
+    private boolean pnpoly(float lat, float lon) {
         boolean contains = false;
         for (int i = 0, j = latitude.length - 1; i < latitude.length; j = i++) {
             if (((latitude[i] > lat) != (latitude[j] > lat))