introduce EdgeTree writer and reader

server/src/main/java/org/elasticsearch/common/geo/EdgeTreeReader.java

@@ -0,0 +1,193 @@
+/*
+ * Licensed to Elasticsearch under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch licenses this file to you under
+ * the Apache License, Version 2.0 (the "License"); you may
+ * not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.elasticsearch.common.geo;
+
+import org.apache.lucene.util.BytesRef;
+import org.elasticsearch.common.io.stream.ByteBufferStreamInput;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+
+import static org.apache.lucene.geo.GeoUtils.lineCrossesLine;
+
+public class EdgeTreeReader {
+    final BytesRef bytesRef;
+
+    public EdgeTreeReader(BytesRef bytesRef) {
+        this.bytesRef = bytesRef;
+    }
+
+    /**
+     * Returns true if the rectangle query and the edge tree's shape overlap
+     */
+    public boolean containedInOrCrosses(int minX, int minY, int maxX, int maxY) throws IOException {
+        return this.containsBottomLeft(minX, minY, maxX, maxY) || this.crosses(minX, minY, maxX, maxY);
+    }
+
+    boolean containsBottomLeft(int minX, int minY, int maxX, int maxY) throws IOException {
+        ByteBufferStreamInput input = new ByteBufferStreamInput(ByteBuffer.wrap(bytesRef.bytes, bytesRef.offset, bytesRef.length));
+        int thisMinX = input.readInt();
+        int thisMinY = input.readInt();
+        int thisMaxX = input.readInt();
+        int thisMaxY = input.readInt();
+
+        if (thisMinY > maxY || thisMaxX < minX || thisMaxY < minY || thisMinX > maxX) {
+            return false; // tree and bbox-query are disjoint
+        }
+
+        if (minX <= thisMinX && minY <= thisMinY && maxX >= thisMaxX && maxY >= thisMaxY) {
+            return true; // bbox-query fully contains tree's extent.
+        }
+
+        return containsBottomLeft(input, readRoot(input, input.position()), minX, minY, maxX, maxY);
+    }
+
+    public boolean crosses(int minX, int minY, int maxX, int maxY) throws IOException {
+        ByteBufferStreamInput input = new ByteBufferStreamInput(ByteBuffer.wrap(bytesRef.bytes, bytesRef.offset, bytesRef.length));
+        int thisMinX = input.readInt();
+        int thisMinY = input.readInt();
+        int thisMaxX = input.readInt();
+        int thisMaxY = input.readInt();
+
+        if (thisMinY > maxY || thisMaxX < minX || thisMaxY < minY || thisMinX > maxX) {
+            return false; // tree and bbox-query are disjoint
+        }
+
+        if (minX <= thisMinX && minY <= thisMinY && maxX >= thisMaxX && maxY >= thisMaxY) {
+            return true; // bbox-query fully contains tree's extent.
+        }
+
+        return crosses(input, readRoot(input, input.position()), minX, minY, maxX, maxY);
+    }
+
+    public Edge readRoot(ByteBufferStreamInput input, int position) throws IOException {
+        return readEdge(input, position);
+    }
+
+    private static Edge readEdge(ByteBufferStreamInput input, int position) throws IOException {
+        input.position(position);
+        int minY = input.readInt();
+        int maxY = input.readInt();
+        int x1 = input.readInt();
+        int y1 = input.readInt();
+        int x2 = input.readInt();
+        int y2 = input.readInt();
+        int rightOffset = input.readInt();
+        return new Edge(input.position(), x1, y1, x2, y2, minY, maxY, rightOffset);
+    }
+
+
+    Edge readLeft(ByteBufferStreamInput input, Edge root) throws IOException {
+        return readEdge(input, root.streamOffset);
+    }
+
+    Edge readRight(ByteBufferStreamInput input, Edge root) throws IOException {
+        return readEdge(input, root.streamOffset + root.rightOffset);
+    }
+
+    /**
+     * Returns true if the bottom-left point of the rectangle query is contained within the
+     * tree's edges.
+     */
+    private boolean containsBottomLeft(ByteBufferStreamInput input, Edge root, int minX, int minY, int maxX, int maxY) throws IOException {
+        boolean res = false;
+        if (root.maxY >= minY) {
+            // is bbox-query contained within linearRing
+            // cast infinite ray to the right from bottom-left of bbox-query to see if it intersects edge
+            if (lineCrossesLine(root.x1, root.y1, root.x2, root.y2,minX, minY, Integer.MAX_VALUE, minY)) {
+                res = true;
+            }
+
+            if (root.rightOffset > 0) { /* has left node */
+                res ^= containsBottomLeft(input, readLeft(input, root), minX, minY, maxX, maxY);
+            }
+
+            if (root.rightOffset > 0 && maxY >= root.minY) { /* no right node if rightOffset == -1 */
+                res ^= containsBottomLeft(input, readRight(input, root), minX, minY, maxX, maxY);
+            }
+        }
+        return res;
+    }
+
+    /**
+     * Returns true if the box crosses any edge in this edge subtree
+     * */
+    private boolean crosses(ByteBufferStreamInput input, Edge root, int minX, int minY, int maxX, int maxY) throws IOException {
+        boolean res = false;
+        // we just have to cross one edge to answer the question, so we descend the tree and return when we do.
+        if (root.maxY >= minY) {
+
+            // does rectangle's edges intersect or reside inside polygon's edge
+            if (lineCrossesLine(root.x1, root.y1, root.x2, root.y2, minX, minY, maxX, minY) ||
+                lineCrossesLine(root.x1, root.y1, root.x2, root.y2, maxX, minY, maxX, maxY) ||
+                lineCrossesLine(root.x1, root.y1, root.x2, root.y2, maxX, maxY, minX, maxY) ||
+                lineCrossesLine(root.x1, root.y1, root.x2, root.y2, minX, maxY, minX, minY)) {
+                return true;
+            }
+
+            if (root.rightOffset > 0) { /* has left node */
+                if (crosses(input, readLeft(input, root), minX, minY, maxX, maxY)) {
+                    return true;
+                }
+            }
+
+            if (root.rightOffset > 0 && maxY >= root.minY) { /* no right node if rightOffset == -1 */
+                if (crosses(input, readRight(input, root), minX, minY, maxX, maxY)) {
+                    return true;
+                }
+            }
+        }
+        return false;
+    }
+
+
+    private static class Edge {
+        int streamOffset;
+        int x1;
+        int y1;
+        int x2;
+        int y2;
+        int minY;
+        int maxY;
+        int rightOffset;
+
+        /**
+         * Object representing an edge node read from bytes
+         *
+         * @param streamOffset offset in byte-reference where edge terminates
+         * @param x1 x-coordinate of first point in segment
+         * @param y1 y-coordinate of first point in segment
+         * @param x2 x-coordinate of second point in segment
+         * @param y2 y-coordinate of second point in segment
+         * @param minY minimum y-coordinate in this edge-node's tree
+         * @param maxY maximum y-coordinate in this edge-node's tree
+         * @param rightOffset the start offset in the byte-reference of the right edge-node
+         */
+        Edge(int streamOffset, int x1, int y1, int x2, int y2, int minY, int maxY, int rightOffset) {
+            this.streamOffset = streamOffset;
+            this.x1 = x1;
+            this.y1 = y1;
+            this.x2 = x2;
+            this.y2 = y2;
+            this.minY = minY;
+            this.maxY = maxY;
+            this.rightOffset = rightOffset;
+        }
+    }
+}

server/src/main/java/org/elasticsearch/common/geo/EdgeTreeWriter.java

@@ -0,0 +1,162 @@
+/*
+ * Licensed to Elasticsearch under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch licenses this file to you under
+ * the Apache License, Version 2.0 (the "License"); you may
+ * not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.elasticsearch.common.geo;
+
+import org.apache.lucene.util.BytesRef;
+import org.elasticsearch.common.io.stream.BytesStreamOutput;
+import org.elasticsearch.common.io.stream.StreamOutput;
+
+import java.io.IOException;
+import java.util.Arrays;
+
+/**
+ * Shape edge-tree writer for use in doc-values
+ */
+public class EdgeTreeWriter {
+
+    /**
+     * | minY | maxY | x1 | y1 | x2 | y2 | right_offset |
+     */
+    static final int EDGE_SIZE_IN_BYTES = 28;
+
+    int minX;
+    int minY;
+    int maxX;
+    int maxY;
+    final Edge tree;
+
+    public EdgeTreeWriter(int[] x, int[] y) {
+        minX = minY = Integer.MAX_VALUE;
+        maxX = maxY = Integer.MIN_VALUE;
+        Edge edges[] = new Edge[y.length - 1];
+        for (int i = 1; i < y.length; i++) {
+            int y1 = y[i-1];
+            int x1 = x[i-1];
+            int y2 = y[i];
+            int x2 = x[i];
+            int minY, maxY;
+            if (y1 < y2) {
+                minY = y1;
+                maxY = y2;
+            } else {
+                minY = y2;
+                maxY = y1;
+            }
+            edges[i - 1] = new Edge(x1, y1, x2, y2, minY, maxY);
+            this.minX = Math.min(this.minX, Math.min(x1, x2));
+            this.minY = Math.min(this.minY, Math.min(y1, y2));
+            this.maxX = Math.max(this.maxX, Math.max(x1, x2));
+            this.maxY = Math.max(this.maxY, Math.max(y1, y2));
+        }
+        Arrays.sort(edges);
+        this.tree = createTree(edges, 0, edges.length - 1);
+    }
+
+    public BytesRef toBytesRef() throws IOException {
+        BytesStreamOutput output = new BytesStreamOutput(4 * 4 + EDGE_SIZE_IN_BYTES * tree.size);
+        // write extent of edges
+        output.writeInt(minX);
+        output.writeInt(minY);
+        output.writeInt(maxX);
+        output.writeInt(maxY);
+        // write edge-tree itself
+        writeTree(tree, output);
+        output.close();
+        return output.bytes().toBytesRef();
+    }
+
+    private void writeTree(Edge edge, StreamOutput output) throws IOException {
+        if (edge == null) {
+            return;
+        }
+        output.writeInt(edge.minY);
+        output.writeInt(edge.maxY);
+        output.writeInt(edge.x1);
+        output.writeInt(edge.y1);
+        output.writeInt(edge.x2);
+        output.writeInt(edge.y2);
+        // left node is next node, write offset of right node
+        if (edge.left != null) {
+            output.writeInt(edge.left.size * EDGE_SIZE_IN_BYTES);
+        } else if (edge.right == null){
+            output.writeInt(-1);
+        } else {
+            output.writeInt(0);
+        }
+        writeTree(edge.left, output);
+        writeTree(edge.right, output);
+    }
+
+    private static Edge createTree(Edge edges[], int low, int high) {
+        if (low > high) {
+            return null;
+        }
+        // add midpoint
+        int mid = (low + high) >>> 1;
+        Edge newNode = edges[mid];
+        newNode.size = 1;
+        // add children
+        newNode.left = createTree(edges, low, mid - 1);
+        newNode.right = createTree(edges, mid + 1, high);
+        // pull up max values to this node
+        // and node count
+        if (newNode.left != null) {
+            newNode.maxY = Math.max(newNode.maxY, newNode.left.maxY);
+            newNode.size += newNode.left.size;
+        }
+        if (newNode.right != null) {
+            newNode.maxY = Math.max(newNode.maxY, newNode.right.maxY);
+            newNode.size += newNode.right.size;
+        }
+        return newNode;
+    }
+
+    /**
+     * Object representing an in-memory edge-tree to be serialized
+     */
+    static class Edge implements Comparable<Edge> {
+        final int x1;
+        final int y1;
+        final int x2;
+        final int y2;
+        int minY;
+        int maxY;
+        int size;
+        Edge left;
+        Edge right;
+
+        Edge(int x1, int y1, int x2, int y2, int minY, int maxY) {
+            this.x1 = x1;
+            this.y1 = y1;
+            this.x2 = x2;
+            this.y2 = y2;
+            this.minY = minY;
+            this.maxY = maxY;
+        }
+
+        @Override
+        public int compareTo(Edge other) {
+            int ret = Integer.compare(minY, other.minY);
+            if (ret == 0) {
+                ret = Integer.compare(maxY, other.maxY);
+            }
+            return ret;
+        }
+    }
+}

server/src/main/java/org/elasticsearch/common/io/stream/ByteBufferStreamInput.java

@@ -110,6 +110,14 @@ public long readLong() throws IOException {
         }
     }
 
+    public void position(int newPosition) throws IOException {
+        buffer.position(newPosition);
+    }
+
+    public int position() throws IOException {
+        return buffer.position();
+    }
+
     @Override
     public void reset() throws IOException {
         buffer.reset();