Merge golang#98 into golang/geo fork

Author: neilaram1

README.md

@@ -1,3 +1,6 @@
+This fork of golang/geo exists to merge https://github.com/golang/geo/pull/98
+and add CellIndex intersection functionality.
+
 # Overview
 
 S2 is a library for spherical geometry that aims to have the same robustness,

go.mod

@@ -1,3 +1,7 @@
-module github.com/golang/geo
+module github.com/MadHive/geo
+
+replace github.com/golang/geo => ./
 
 go 1.12
+
+require github.com/golang/geo v0.0.0-20210211234256-740aa86cb551

r2/rect.go

@@ -18,7 +18,7 @@ import (
 	"fmt"
 	"math"
 
-	"github.com/golang/geo/r1"
+	"github.com/MadHive/geo/r1"
 )
 
 // Point represents a point in ℝ².

r2/rect_test.go

@@ -22,7 +22,7 @@ import (
 	"reflect"
 	"testing"
 
-	"github.com/golang/geo/r1"
+	"github.com/MadHive/geo/r1"
 )
 
 var (

r3/vector.go

@@ -18,7 +18,7 @@ import (
 	"fmt"
 	"math"
 
-	"github.com/golang/geo/s1"
+	"github.com/MadHive/geo/s1"
 )
 
 // Vector represents a point in ℝ³.

s1/example_test.go

@@ -18,7 +18,7 @@ import (
 	"fmt"
 	"math"
 
-	"github.com/golang/geo/s1"
+	"github.com/MadHive/geo/s1"
 )
 
 func ExampleInterval_DirectedHausdorffDistance() {

s2/cap.go

@@ -19,8 +19,8 @@ import (
 	"io"
 	"math"
 
-	"github.com/golang/geo/r1"
-	"github.com/golang/geo/s1"
+	"github.com/MadHive/geo/r1"
+	"github.com/MadHive/geo/s1"
 )
 
 var (

s2/cap_test.go

@@ -18,8 +18,8 @@ import (
 	"math"
 	"testing"
 
-	"github.com/golang/geo/r3"
-	"github.com/golang/geo/s1"
+	"github.com/MadHive/geo/r3"
+	"github.com/MadHive/geo/s1"
 )
 
 const (

s2/cell.go

@@ -18,10 +18,10 @@ import (
 	"io"
 	"math"
 
-	"github.com/golang/geo/r1"
-	"github.com/golang/geo/r2"
-	"github.com/golang/geo/r3"
-	"github.com/golang/geo/s1"
+	"github.com/MadHive/geo/r1"
+	"github.com/MadHive/geo/r2"
+	"github.com/MadHive/geo/r3"
+	"github.com/MadHive/geo/s1"
 )
 
 // Cell is an S2 region object that represents a cell. Unlike CellIDs,

s2/cell_index.go

@@ -34,7 +34,7 @@ type cellIndexNode struct {
 // newCellIndexNode returns a node with the appropriate default values.
 func newCellIndexNode() cellIndexNode {
 	return cellIndexNode{
-		cellID: 0,
+		cellID: CellID(0),
 		label:  cellIndexDoneContents,
 		parent: -1,
 	}
@@ -49,15 +49,52 @@ type rangeNode struct {
 	contents int32  // Contents of this node (an index within the cell tree).
 }
 
+type labels []int32
+
+func (l *labels) Normalize() {
+	encountered := make(map[int32]struct{})
+
+	for i := range *l {
+		encountered[(*l)[i]] = struct{}{}
+	}
+
+	*l = (*l)[:0]
+	for key := range encountered {
+		*l = append(*l, key)
+	}
+}
+
+type cellVisitor func(cellID CellID, label int32) bool
+
 // CellIndexIterator is an iterator that visits the entire set of indexed
 // (CellID, label) pairs in an unspecified order.
 type CellIndexIterator struct {
-	// TODO(roberts): Implement
+	nodes []cellIndexNode
+	pos   int
 }
 
 // NewCellIndexIterator creates an iterator for the given CellIndex.
 func NewCellIndexIterator(index *CellIndex) *CellIndexIterator {
-	return &CellIndexIterator{}
+	return &CellIndexIterator{
+		nodes: index.cellTree,
+		pos:   0,
+	}
+}
+
+func (c *CellIndexIterator) CellID() CellID {
+	return c.nodes[c.pos].cellID
+}
+
+func (c *CellIndexIterator) Label() int32 {
+	return c.nodes[c.pos].label
+}
+
+func (c *CellIndexIterator) Done() bool {
+	return c.pos == len(c.nodes)
+}
+
+func (c *CellIndexIterator) Next() {
+	c.pos++
 }
 
 // CellIndexRangeIterator is an iterator that seeks and iterates over a set of
@@ -202,7 +239,6 @@ func (c *CellIndexRangeIterator) Seek(target CellID) {
 
 	// Nonempty needs to find the next non-empty entry.
 	for c.nonEmpty && c.IsEmpty() && !c.Done() {
-		// c.Next()
 		c.pos++
 	}
 }
@@ -246,7 +282,7 @@ type CellIndexContentsIterator struct {
 func NewCellIndexContentsIterator(index *CellIndex) *CellIndexContentsIterator {
 	it := &CellIndexContentsIterator{
 		cellTree:       index.cellTree,
-		prevStartID:    0,
+		prevStartID:    CellID(0),
 		nodeCutoff:     -1,
 		nextNodeCutoff: -1,
 		node:           cellIndexNode{label: cellIndexDoneContents},
@@ -256,7 +292,7 @@ func NewCellIndexContentsIterator(index *CellIndex) *CellIndexContentsIterator {
 
 // Clear clears all state with respect to which range(s) have been visited.
 func (c *CellIndexContentsIterator) Clear() {
-	c.prevStartID = 0
+	c.prevStartID = CellID(0)
 	c.nodeCutoff = -1
 	c.nextNodeCutoff = -1
 	c.node.label = cellIndexDoneContents
@@ -343,7 +379,7 @@ func (c *CellIndexContentsIterator) StartUnion(r *CellIndexRangeIterator) {
 // There is also a helper method that adds all elements of CellUnion with the
 // same label:
 //
-//     index.AddCellUnion(cellUnion, label)
+//	index.AddCellUnion(cellUnion, label)
 //
 // Note that the index is not dynamic; the contents of the index cannot be
 // changed once it has been built. Adding more after calling Build results in
@@ -353,7 +389,7 @@ func (c *CellIndexContentsIterator) StartUnion(r *CellIndexRangeIterator) {
 // is to use a built-in method such as IntersectingLabels (which returns
 // the labels of all cells that intersect a given target CellUnion):
 //
-//   labels := index.IntersectingLabels(targetUnion);
+//	labels := index.IntersectingLabels(targetUnion);
 //
 // Alternatively, you can use a ClosestCellQuery which computes the cell(s)
 // that are closest to a given target geometry.
@@ -482,7 +518,8 @@ func (c *CellIndex) Build() {
 				c.cellTree = append(c.cellTree, cellIndexNode{
 					cellID: deltas[i].cellID,
 					label:  deltas[i].label,
-					parent: contents})
+					parent: contents,
+				})
 				contents = int32(len(c.cellTree) - 1)
 			} else if deltas[i].cellID == SentinelCellID {
 				contents = c.cellTree[contents].parent
@@ -492,7 +529,48 @@ func (c *CellIndex) Build() {
 	}
 }
 
-// TODO(roberts): Differences from C++
-// IntersectingLabels
-// VisitIntersectingCells
-// CellIndexIterator
+func (c *CellIndex) IntersectingLabels(target CellUnion) []int32 {
+	var labels labels
+	c.VisitIntersectingCells(target, func(cellID CellID, label int32) bool {
+		labels = append(labels, label)
+		return true
+	})
+
+	labels.Normalize()
+	return labels
+}
+
+func (c *CellIndex) VisitIntersectingCells(target CellUnion, visitor cellVisitor) bool {
+	if len(target) == 0 {
+		return true
+	}
+
+	pos := 0
+	contents := NewCellIndexContentsIterator(c)
+	rangeIter := NewCellIndexRangeIterator(c)
+	rangeIter.Begin()
+
+	for ok := true; ok; ok = pos != len(target) {
+		if rangeIter.LimitID() <= target[pos].RangeMin() {
+			rangeIter.Seek(target[pos].RangeMin())
+		}
+
+		for ; rangeIter.StartID() <= target[pos].RangeMax(); rangeIter.Next() {
+			for contents.StartUnion(rangeIter); !contents.Done(); contents.Next() {
+				if !visitor(contents.CellID(), contents.Label()) {
+					return false
+				}
+			}
+		}
+
+		pos++
+		if pos != len(target) && target[pos].RangeMax() < rangeIter.StartID() {
+			pos = target.lowerBound(pos+1, len(target), rangeIter.StartID())
+			if target[pos-1].RangeMax() >= rangeIter.StartID() {
+				pos--
+			}
+		}
+	}
+
+	return true
+}