From 63ef568cc316d349eb135e25065850942eb65c8d Mon Sep 17 00:00:00 2001 From: Fabien Meurisse Date: Fri, 29 May 2026 09:55:26 +0200 Subject: [PATCH 1/3] add PointIndex and PointIndexIterator Points are stored in a plain sorted slice keyed by leaf CellID. Add and Remove maintain sort order by shifting elements, so each mutation is O(n). This implementation is not a btree and is therefore not suitable for large datasets or frequently mutated indexes; it is designed for the build-once, query-many pattern. --- README.md | 2 +- s2/point_index.go | 225 +++++++++++++++++++++++++++++++ s2/point_index_test.go | 298 +++++++++++++++++++++++++++++++++++++++++ 3 files changed, 524 insertions(+), 1 deletion(-) create mode 100644 s2/point_index.go create mode 100644 s2/point_index_test.go diff --git a/README.md b/README.md index 6abbeac5..47866c07 100644 --- a/README.md +++ b/README.md @@ -154,7 +154,7 @@ S2LaxPolyline | 🟡 S2Loop | ✅ S2PaddedCell | ✅ S2Point | ✅ -S2PointIndex | ❌ +S2PointIndex | ✅ S2PointSpan | ❌ S2PointRegion | ❌ S2PointVector | ✅ diff --git a/s2/point_index.go b/s2/point_index.go new file mode 100644 index 00000000..27ac8d74 --- /dev/null +++ b/s2/point_index.go @@ -0,0 +1,225 @@ +// Copyright 2015 Google Inc. All rights reserved. +// +// Licensed 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 s2 + +import "sort" + +// PointData holds a Point and its associated data. +type PointData[Data comparable] struct { + point Point + data Data +} + +// Point returns the point. +func (pd PointData[Data]) Point() Point { return pd.point } + +// Data returns the associated data. +func (pd PointData[Data]) Data() Data { return pd.data } + +// pointIndexEntry is a single entry in the sorted PointIndex. +type pointIndexEntry[Data comparable] struct { + id CellID + pointData PointData[Data] +} + +// PointIndex maintains an index of points sorted by leaf CellID. Each point +// can optionally store auxiliary data such as an integer or pointer. This can +// be used to map results back to client data structures. +// +// The index supports adding or removing points dynamically, and provides a +// seekable iterator interface for navigating the index. +// +// You can use this class in conjunction with ClosestPointQuery to find the +// closest index points to a given query point. For example: +// +// index := &PointIndex[int]{} +// for i, p := range indexPoints { +// index.Add(p, i) +// } +// +// You can also access the index directly using the iterator interface. For +// example, here is how to iterate through all the points in a given CellID +// target: +// +// it := NewPointIndexIterator(index) +// for it.Seek(target.RangeMin()); !it.Done() && it.CellID() <= target.RangeMax(); it.Next() { +// DoSomething(it.CellID(), it.Point(), it.Data()) +// } +// +// Points can be added or removed from the index at any time by calling Add() +// or Remove(). However when the index is modified, any existing iterator's +// position may refer to the wrong entry; create a new iterator to resume +// traversal safely. +// +// Note: Add and Remove maintain sorted order by shifting elements, so they +// run in O(n) time. This index is suitable for building once and querying +// many times, or for small dynamic datasets. +type PointIndex[Data comparable] struct { + entries []pointIndexEntry[Data] +} + +// NumPoints returns the number of points in the index. +func (p *PointIndex[Data]) NumPoints() int { return len(p.entries) } + +// Add adds the given point with associated data to the index. Invalidates all iterators. +func (p *PointIndex[Data]) Add(point Point, data Data) { + id := cellIDFromPoint(point) + pos := sort.Search(len(p.entries), func(i int) bool { + return p.entries[i].id >= id + }) + entry := pointIndexEntry[Data]{id: id, pointData: PointData[Data]{point: point, data: data}} + p.entries = append(p.entries, pointIndexEntry[Data]{}) + copy(p.entries[pos+1:], p.entries[pos:]) + p.entries[pos] = entry +} + +// Remove removes the given point and data from the index. Returns false if the +// given point was not present. Invalidates all iterators. +func (p *PointIndex[Data]) Remove(point Point, data Data) bool { + id := cellIDFromPoint(point) + pd := PointData[Data]{point: point, data: data} + pos := sort.Search(len(p.entries), func(i int) bool { + return p.entries[i].id >= id + }) + for pos < len(p.entries) && p.entries[pos].id == id { + if p.entries[pos].pointData == pd { + p.entries = append(p.entries[:pos], p.entries[pos+1:]...) + return true + } + pos++ + } + return false +} + +// Clear resets the index to its original empty state. Invalidates all iterators. +func (p *PointIndex[Data]) Clear() { + p.entries = nil +} + +// PointIndexIterator is a seekable iterator for a PointIndex. +// +// The iterator holds a pointer to the index, so the underlying data is always +// live. However, mutations to the index (Add/Remove) may shift entries and +// leave the iterator's position pointing at the wrong entry. Create a new +// iterator after any mutation. +type PointIndexIterator[Data comparable] struct { + index *PointIndex[Data] + position int +} + +// NewPointIndexIterator creates a new iterator for the given PointIndex. +// If the index is non-empty, the iterator is positioned at the first entry. +func NewPointIndexIterator[Data comparable](index *PointIndex[Data]) *PointIndexIterator[Data] { + return &PointIndexIterator[Data]{index: index} +} + +// CellID returns the CellID for the current index entry. +// Requires: !Done() +func (it *PointIndexIterator[Data]) CellID() CellID { + return it.index.entries[it.position].id +} + +// Point returns the point associated with the current index entry. +// Requires: !Done() +func (it *PointIndexIterator[Data]) Point() Point { + return it.index.entries[it.position].pointData.point +} + +// Data returns the data associated with the current index entry. +// Requires: !Done() +func (it *PointIndexIterator[Data]) Data() Data { + return it.index.entries[it.position].pointData.data +} + +// PointData returns the (Point, Data) pair for the current index entry. +// Requires: !Done() +func (it *PointIndexIterator[Data]) PointData() PointData[Data] { + return it.index.entries[it.position].pointData +} + +// Done reports if the iterator is positioned past the last index entry. +func (it *PointIndexIterator[Data]) Done() bool { + return it.position >= len(it.index.entries) +} + +// Begin positions the iterator at the first index entry (if any). +func (it *PointIndexIterator[Data]) Begin() { + it.position = 0 +} + +// Finish positions the iterator so that Done() is true. +func (it *PointIndexIterator[Data]) Finish() { + it.position = len(it.index.entries) +} + +// Next advances the iterator to the next index entry. +// Requires: !Done() +func (it *PointIndexIterator[Data]) Next() { + it.position++ +} + +// Prev positions the iterator at the previous entry and reports whether the +// iterator was not already positioned at the beginning. +func (it *PointIndexIterator[Data]) Prev() bool { + if it.position == 0 { + return false + } + it.position-- + return true +} + +// Seek positions the iterator at the first entry with CellID() >= target, or +// at the end of the index if no such entry exists. +func (it *PointIndexIterator[Data]) Seek(target CellID) { + it.position = sort.Search(len(it.index.entries), func(i int) bool { + return it.index.entries[i].id >= target + }) +} + +// LocatePoint positions the iterator at the entry for the cell containing the +// given point. Returns true if such an entry exists. +func (it *PointIndexIterator[Data]) LocatePoint(target Point) bool { + id := cellIDFromPoint(target) + it.Seek(id) + if !it.Done() && it.CellID().RangeMin() <= id { + return true + } + if it.Prev() && it.CellID().RangeMax() >= id { + return true + } + return false +} + +// LocateCellID positions the iterator given the target CellID. Let T be the +// target CellID. If T is contained by some index cell I (including equality), +// the iterator is positioned at I and Indexed is returned. Otherwise if T +// contains one or more (smaller) index cells, the iterator is positioned at +// the first such cell and Subdivided is returned. Otherwise Disjoint is +// returned and the iterator position is unspecified. +func (it *PointIndexIterator[Data]) LocateCellID(target CellID) CellRelation { + it.Seek(target.RangeMin()) + if !it.Done() { + if it.CellID() >= target && it.CellID().RangeMin() <= target { + return Indexed + } + if it.CellID() <= target.RangeMax() { + return Subdivided + } + } + if it.Prev() && it.CellID().RangeMax() >= target { + return Indexed + } + return Disjoint +} diff --git a/s2/point_index_test.go b/s2/point_index_test.go new file mode 100644 index 00000000..9c1b329b --- /dev/null +++ b/s2/point_index_test.go @@ -0,0 +1,298 @@ +// Copyright 2015 Google Inc. All rights reserved. +// +// Licensed 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 s2 + +import ( + "testing" +) + +// pointIndexTest is a test helper that tracks both the index and expected contents. +type pointIndexTest struct { + t *testing.T + index *PointIndex[int] + contents map[PointData[int]]int // PointData -> occurrence count +} + +func newPointIndexTest(t *testing.T) *pointIndexTest { + return &pointIndexTest{ + t: t, + index: &PointIndex[int]{}, + contents: make(map[PointData[int]]int), + } +} + +func (pt *pointIndexTest) add(point Point, data int) { + pt.index.Add(point, data) + pt.contents[PointData[int]{point: point, data: data}]++ +} + +func (pt *pointIndexTest) remove(point Point, data int) { + pd := PointData[int]{point: point, data: data} + pt.contents[pd]-- + if pt.contents[pd] == 0 { + delete(pt.contents, pd) + } + if !pt.index.Remove(point, data) { + pt.t.Errorf("Remove(%v, %v) returned false, expected true", point, data) + } +} + +func (pt *pointIndexTest) verify() { + pt.verifyContents() + pt.verifyIteratorMethods() +} + +func (pt *pointIndexTest) verifyContents() { + remaining := make(map[PointData[int]]int) + for k, v := range pt.contents { + remaining[k] = v + } + for it := NewPointIndexIterator(pt.index); !it.Done(); it.Next() { + pd := it.PointData() + if got := pd.Point(); got != it.Point() { + pt.t.Errorf("PointData.Point() = %v, want %v", got, it.Point()) + } + if got := pd.Data(); got != it.Data() { + pt.t.Errorf("PointData.Data() = %v, want %v", got, it.Data()) + } + if remaining[pd] <= 0 { + pt.t.Errorf("point_data %v found in index but not in expected contents", pd) + continue + } + remaining[pd]-- + if remaining[pd] == 0 { + delete(remaining, pd) + } + } + if len(remaining) > 0 { + pt.t.Errorf("expected contents not found in index: %v", remaining) + } +} + +func (pt *pointIndexTest) verifyIteratorMethods() { + it := NewPointIndexIterator(pt.index) + if it.Prev() { + pt.t.Error("Prev() returned true on freshly created iterator at position 0") + } + it.Finish() + if !it.Done() { + pt.t.Error("Done() returned false after Finish()") + } + + var prevCellID CellID + minCellID := CellIDFromFace(0).ChildBeginAtLevel(MaxLevel) + + for it.Begin(); !it.Done(); it.Next() { + cellID := it.CellID() + + if got := cellIDFromPoint(it.Point()); got != cellID { + pt.t.Errorf("cellIDFromPoint(it.Point()) = %v, want %v", got, cellID) + } + if cellID < prevCellID { + pt.t.Errorf("iterator not in sorted order: %v < %v", cellID, prevCellID) + } + + it2 := *it + if cellID == prevCellID { + it2.Seek(cellID) + } + + // Verify that seeking to any skipped leaf cell lands at cellID. + if cellID > prevCellID { + for _, skipped := range CellUnionFromRange(minCellID, cellID) { + it2.Seek(skipped) + if it2.Done() || it2.CellID() != cellID { + pt.t.Errorf("Seek(%v): got %v, want %v", skipped, it2.CellID(), cellID) + } + } + } + + // Test Prev, Next, and Seek. + if prevCellID.IsValid() { + it2 = *it + if !it2.Prev() { + pt.t.Error("Prev() returned false, expected true") + } + if it2.CellID() != prevCellID { + pt.t.Errorf("after Prev(), CellID() = %v, want %v", it2.CellID(), prevCellID) + } + it2.Next() + if it2.CellID() != cellID { + pt.t.Errorf("after Next(), CellID() = %v, want %v", it2.CellID(), cellID) + } + it2.Seek(prevCellID) + if it2.CellID() != prevCellID { + pt.t.Errorf("Seek(%v): CellID() = %v, want %v", prevCellID, it2.CellID(), prevCellID) + } + } + + prevCellID = cellID + minCellID = cellID.Next() + } +} + +func TestPointIndexNoPoints(t *testing.T) { + pt := newPointIndexTest(t) + pt.verify() +} + +func TestPointIndexDuplicatePoints(t *testing.T) { + pt := newPointIndexTest(t) + p := PointFromCoords(1, 0, 0) + for range 10 { + pt.add(p, 123) + } + pt.verify() + for range 5 { + pt.remove(p, 123) + } + pt.verify() + + // Remove with wrong data value — point is present but data does not match. + if pt.index.Remove(p, 456) { + t.Error("Remove(p, 456) = true, want false: data 456 was never added") + } + // Remove with a point not in the index at all. + absent := PointFromCoords(0, 1, 0) + if pt.index.Remove(absent, 123) { + t.Error("Remove(absent, 123) = true, want false: point was never added") + } +} + +func TestPointIndexRandomPoints(t *testing.T) { + pt := newPointIndexTest(t) + for range 100 { + pt.add(randomPoint(), randomUniformInt(100)) + } + pt.verify() + + // Remove some points via iterator traversal to a random leaf cell. + for range 10 { + it := NewPointIndexIterator(pt.index) + found := false + for range 100 { + it.Seek(randomCellIDForLevel(MaxLevel)) + if !it.Done() { + found = true + break + } + } + if !found { + t.Fatal("failed to find a non-empty position after 100 seeks") + } + pt.remove(it.Point(), it.Data()) + pt.verify() + } +} + +func TestPointIndexNumPoints(t *testing.T) { + index := &PointIndex[int]{} + if got := index.NumPoints(); got != 0 { + t.Errorf("NumPoints() = %d, want 0 for empty index", got) + } + p := PointFromCoords(1, 0, 0) + for i := range 5 { + index.Add(p, i) + } + if got := index.NumPoints(); got != 5 { + t.Errorf("NumPoints() = %d, want 5 after 5 adds", got) + } + index.Remove(p, 2) + if got := index.NumPoints(); got != 4 { + t.Errorf("NumPoints() = %d, want 4 after one removal", got) + } +} + +func TestPointIndexClear(t *testing.T) { + index := &PointIndex[int]{} + for i := range 10 { + index.Add(randomPoint(), i) + } + index.Clear() + if got := index.NumPoints(); got != 0 { + t.Errorf("NumPoints() = %d, want 0 after Clear", got) + } + it := NewPointIndexIterator(index) + if !it.Done() { + t.Error("iterator not Done() immediately after Clear") + } + // Verify the index is usable after clearing. + p := PointFromCoords(0, 1, 0) + index.Add(p, 99) + if got := index.NumPoints(); got != 1 { + t.Errorf("NumPoints() = %d, want 1 after Add following Clear", got) + } +} + +func TestPointIndexLocatePoint(t *testing.T) { + index := &PointIndex[int]{} + // Three points on distinct S2 faces to guarantee distinct leaf CellIDs. + points := []Point{ + PointFromCoords(1, 0, 0), + PointFromCoords(0, 1, 0), + PointFromCoords(0, 0, 1), + } + for i, p := range points { + index.Add(p, i) + } + + it := NewPointIndexIterator(index) + for i, p := range points { + if !it.LocatePoint(p) { + t.Errorf("LocatePoint(points[%d]) = false, want true", i) + continue + } + if got := it.Data(); got != i { + t.Errorf("LocatePoint(points[%d]): Data() = %d, want %d", i, got, i) + } + } + + absent := PointFromCoords(1, 1, 1) + if it.LocatePoint(absent) { + t.Errorf("LocatePoint(%v) = true, want false for absent point", absent) + } +} + +func TestPointIndexLocateCellID(t *testing.T) { + index := &PointIndex[int]{} + p := PointFromCoords(1, 0, 0) + index.Add(p, 42) + leafID := cellIDFromPoint(p) + + it := NewPointIndexIterator(index) + + // Exact leaf cell in the index → Indexed, iterator at that cell. + if got := it.LocateCellID(leafID); got != Indexed { + t.Errorf("LocateCellID(leaf) = %v, want Indexed", got) + } + if it.CellID() != leafID { + t.Errorf("after LocateCellID(leaf): CellID() = %v, want %v", it.CellID(), leafID) + } + + // Parent cell containing the leaf → Subdivided, iterator at the leaf. + parent := leafID.Parent(MaxLevel - 1) + if got := it.LocateCellID(parent); got != Subdivided { + t.Errorf("LocateCellID(parent) = %v, want Subdivided", got) + } + if it.CellID() != leafID { + t.Errorf("after LocateCellID(parent): CellID() = %v, want %v", it.CellID(), leafID) + } + + // Level-0 cell on a different face → Disjoint. + otherFace := CellIDFromFace(2) + if got := it.LocateCellID(otherFace); got != Disjoint { + t.Errorf("LocateCellID(otherFace) = %v, want Disjoint", got) + } +} From 46e4d45b1f4089a2c5cac553f3d7b2c61a996bb2 Mon Sep 17 00:00:00 2001 From: Fabien Meurisse Date: Wed, 3 Jun 2026 10:50:21 +0200 Subject: [PATCH 2/3] switch PointIndex to use btree for more efficient dynamic updates and lookups --- go.mod | 1 + go.sum | 2 + s2/point_index.go | 267 ++++++++++++++++++++++++++++------------- s2/point_index_test.go | 1 + 4 files changed, 187 insertions(+), 84 deletions(-) diff --git a/go.mod b/go.mod index c20e44e6..3c701629 100644 --- a/go.mod +++ b/go.mod @@ -7,4 +7,5 @@ go 1.23.0 require ( github.com/google/go-cmp v0.7.0 // indirect github.com/google/go-units v0.0.0-20250612230646-eddd77f68220 // indirect + github.com/tidwall/btree v1.8.1 // indirect ) diff --git a/go.sum b/go.sum index 13851409..6e4feeef 100644 --- a/go.sum +++ b/go.sum @@ -2,3 +2,5 @@ github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8= github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU= github.com/google/go-units v0.0.0-20250612230646-eddd77f68220 h1:hM8xVjUr4Iv/iQIx4Jq1xckZkKlXu51Gqku5HlEpQAE= github.com/google/go-units v0.0.0-20250612230646-eddd77f68220/go.mod h1:wBcRMlRM/bVzYk9xtR2hOp3+iWOhEh1FiK8sAzeR9eA= +github.com/tidwall/btree v1.8.1 h1:27ehoXvm5AG/g+1VxLS1SD3vRhp/H7LuEfwNvddEdmA= +github.com/tidwall/btree v1.8.1/go.mod h1:jBbTdUWhSZClZWoDg54VnvV7/54modSOzDN7VXftj1A= diff --git a/s2/point_index.go b/s2/point_index.go index 27ac8d74..13339d0a 100644 --- a/s2/point_index.go +++ b/s2/point_index.go @@ -14,7 +14,7 @@ package s2 -import "sort" +import "github.com/tidwall/btree" // PointData holds a Point and its associated data. type PointData[Data comparable] struct { @@ -28,26 +28,33 @@ func (pd PointData[Data]) Point() Point { return pd.point } // Data returns the associated data. func (pd PointData[Data]) Data() Data { return pd.data } -// pointIndexEntry is a single entry in the sorted PointIndex. -type pointIndexEntry[Data comparable] struct { - id CellID - pointData PointData[Data] -} - -// PointIndex maintains an index of points sorted by leaf CellID. Each point -// can optionally store auxiliary data such as an integer or pointer. This can -// be used to map results back to client data structures. +// PointIndex maintains an index of points sorted by leaf CellID using a B-tree. +// Each point can optionally store auxiliary data such as an integer or pointer. +// This can be used to map results back to client data structures. // -// The index supports adding or removing points dynamically, and provides a +// The index supports adding or removing points dynamically and provides a // seekable iterator interface for navigating the index. // // You can use this class in conjunction with ClosestPointQuery to find the -// closest index points to a given query point. For example: +// closest index points to a given query point. For example, // -// index := &PointIndex[int]{} -// for i, p := range indexPoints { -// index.Add(p, i) -// } +// index := &PointIndex[int]{} +// for i, p := range indexPoints { +// index.Add(p, i) +// } +// TODO(fmeurisse): Implement ClosestPointQuery integration and update example. +// S2ClosestPointQuery query(&index); +// query.mutable_options()->set_max_results(5); +// for (const S2Point& target_point : target_points) { +// S2ClosestPointQueryPointTarget target(target_point); +// for (const auto& result : query.FindClosestPoints(&target)) { +// // The Result class contains the following methods: +// // distance() is the distance to the target. +// // point() is the indexed point. +// // data() is the auxiliary data. +// DoSomething(target_point, result); +// } +// } // // You can also access the index directly using the iterator interface. For // example, here is how to iterate through all the points in a given CellID @@ -59,133 +66,225 @@ type pointIndexEntry[Data comparable] struct { // } // // Points can be added or removed from the index at any time by calling Add() -// or Remove(). However when the index is modified, any existing iterator's -// position may refer to the wrong entry; create a new iterator to resume -// traversal safely. +// or Remove(). However when the index is modified, you must call Init() on +// each iterator before using it again (or simply create a new iterator): +// +// index.Add(newPoint, 123456) +// it.Init(index) +// it.Seek(target.RangeMin()) // -// Note: Add and Remove maintain sorted order by shifting elements, so they -// run in O(n) time. This index is suitable for building once and querying -// many times, or for small dynamic datasets. +// PointIndex is not safe for concurrent use without external synchronization. type PointIndex[Data comparable] struct { - entries []pointIndexEntry[Data] + // tree maps each leaf CellID to the slice of PointData values at that cell. + // Multiple points at the same CellID are stored together in one entry. + tree btree.Map[CellID, []PointData[Data]] + numPoints int } // NumPoints returns the number of points in the index. -func (p *PointIndex[Data]) NumPoints() int { return len(p.entries) } +func (p *PointIndex[Data]) NumPoints() int { return p.numPoints } // Add adds the given point with associated data to the index. Invalidates all iterators. func (p *PointIndex[Data]) Add(point Point, data Data) { id := cellIDFromPoint(point) - pos := sort.Search(len(p.entries), func(i int) bool { - return p.entries[i].id >= id - }) - entry := pointIndexEntry[Data]{id: id, pointData: PointData[Data]{point: point, data: data}} - p.entries = append(p.entries, pointIndexEntry[Data]{}) - copy(p.entries[pos+1:], p.entries[pos:]) - p.entries[pos] = entry + slice, _ := p.tree.Get(id) + p.tree.Set(id, append(slice, PointData[Data]{point: point, data: data})) + p.numPoints++ } -// Remove removes the given point and data from the index. Returns false if the -// given point was not present. Invalidates all iterators. +// Remove removes one occurrence of the given point and data from the index. +// Returns false if no matching entry was found. Invalidates all iterators. func (p *PointIndex[Data]) Remove(point Point, data Data) bool { id := cellIDFromPoint(point) pd := PointData[Data]{point: point, data: data} - pos := sort.Search(len(p.entries), func(i int) bool { - return p.entries[i].id >= id - }) - for pos < len(p.entries) && p.entries[pos].id == id { - if p.entries[pos].pointData == pd { - p.entries = append(p.entries[:pos], p.entries[pos+1:]...) + slice, found := p.tree.Get(id) + if !found { + return false + } + for i, existing := range slice { + if existing == pd { + slice = append(slice[:i], slice[i+1:]...) + if len(slice) == 0 { + p.tree.Delete(id) + } else { + p.tree.Set(id, slice) + } + p.numPoints-- return true } - pos++ } return false } // Clear resets the index to its original empty state. Invalidates all iterators. func (p *PointIndex[Data]) Clear() { - p.entries = nil + p.tree.Clear() + p.numPoints = 0 } // PointIndexIterator is a seekable iterator for a PointIndex. // -// The iterator holds a pointer to the index, so the underlying data is always -// live. However, mutations to the index (Add/Remove) may shift entries and -// leave the iterator's position pointing at the wrong entry. Create a new -// iterator after any mutation. +// Points at the same CellID are yielded consecutively. The iterator is safe +// to copy for save/restore of position: +// +// it2 := *it +// +// After a copy, call Init(), Begin(), or Seek() on the copy before calling +// Next() or Prev() across a CellID boundary — these methods assign a fresh +// internal iterator, decoupling the copy from the original. Calling Next() or +// Prev() across a CellID boundary on a raw copy without a prior repositioning +// call is undefined behaviour. +// +// After any Add or Remove call, call Init() to make the iterator valid again, +// or create a new iterator. type PointIndexIterator[Data comparable] struct { - index *PointIndex[Data] - position int + index *PointIndex[Data] + iter btree.MapIter[CellID, []PointData[Data]] // live cursor on map keys + currentID CellID + currentSlice []PointData[Data] // points slice of the current map entry + sliceIdx int // position within currentSlice + valid bool + atEnd bool // true when positioned logically past the last entry } -// NewPointIndexIterator creates a new iterator for the given PointIndex. -// If the index is non-empty, the iterator is positioned at the first entry. +// NewPointIndexIterator creates a new iterator for the given PointIndex, +// positioned at the first entry (if any). func NewPointIndexIterator[Data comparable](index *PointIndex[Data]) *PointIndexIterator[Data] { - return &PointIndexIterator[Data]{index: index} + var it PointIndexIterator[Data] + it.Init(index) + return &it } -// CellID returns the CellID for the current index entry. -// Requires: !Done() -func (it *PointIndexIterator[Data]) CellID() CellID { - return it.index.entries[it.position].id +// Init (re)initializes the iterator for the given index, positioning it at +// the first entry if any. This may be called multiple times, e.g. to make an +// iterator valid again after the index is modified. +func (it *PointIndexIterator[Data]) Init(index *PointIndex[Data]) { + it.index = index + it.iter = index.tree.Iter() // fresh iter: independent backing array + it.valid = it.iter.First() + it.atEnd = false + it.sliceIdx = 0 + if it.valid { + it.currentID = it.iter.Key() + it.currentSlice = it.iter.Value() + } } -// Point returns the point associated with the current index entry. +// Done reports whether the iterator is positioned past the last entry. +func (it *PointIndexIterator[Data]) Done() bool { return !it.valid } + +// CellID returns the CellID of the current entry. // Requires: !Done() -func (it *PointIndexIterator[Data]) Point() Point { - return it.index.entries[it.position].pointData.point -} +func (it *PointIndexIterator[Data]) CellID() CellID { return it.currentID } -// Data returns the data associated with the current index entry. +// Point returns the point of the current entry. // Requires: !Done() -func (it *PointIndexIterator[Data]) Data() Data { - return it.index.entries[it.position].pointData.data -} +func (it *PointIndexIterator[Data]) Point() Point { return it.currentSlice[it.sliceIdx].point } + +// Data returns the data of the current entry. +// Requires: !Done() +func (it *PointIndexIterator[Data]) Data() Data { return it.currentSlice[it.sliceIdx].data } -// PointData returns the (Point, Data) pair for the current index entry. +// PointData returns the (Point, Data) pair for the current entry. // Requires: !Done() func (it *PointIndexIterator[Data]) PointData() PointData[Data] { - return it.index.entries[it.position].pointData + return it.currentSlice[it.sliceIdx] } -// Done reports if the iterator is positioned past the last index entry. -func (it *PointIndexIterator[Data]) Done() bool { - return it.position >= len(it.index.entries) +// Refresh gives the iterator a fresh internal cursor positioned at the current +// entry without changing the logical position (CellID or slice index). Call +// this after copying an iterator (it2 = *it) and before calling Next() or +// Prev() across a CellID boundary, to decouple the copy's cursor from the +// original. +func (it *PointIndexIterator[Data]) Refresh() { + if !it.valid { + return + } + it.iter = it.index.tree.Iter() + it.iter.Seek(it.currentID) } -// Begin positions the iterator at the first index entry (if any). -func (it *PointIndexIterator[Data]) Begin() { - it.position = 0 -} +// Begin positions the iterator at the first entry (if any). +func (it *PointIndexIterator[Data]) Begin() { it.Init(it.index) } // Finish positions the iterator so that Done() is true. func (it *PointIndexIterator[Data]) Finish() { - it.position = len(it.index.entries) + it.valid = false + it.atEnd = true } -// Next advances the iterator to the next index entry. +// Next advances to the next entry. +// +// Next uses the live internal cursor directly (O(1) per step). After copying +// an iterator (it2 = *it), call it2.Refresh() before the first cross-CellID +// Next() call to decouple the copy's cursor from the original. // Requires: !Done() func (it *PointIndexIterator[Data]) Next() { - it.position++ + // Fast path: advance within the current CellID's group. + if it.sliceIdx+1 < len(it.currentSlice) { + it.sliceIdx++ + return + } + // Slow path: advance the live cursor to the next map entry — O(1). + it.valid = it.iter.Next() + it.atEnd = !it.valid + it.sliceIdx = 0 + if it.valid { + it.currentID = it.iter.Key() + it.currentSlice = it.iter.Value() + } } -// Prev positions the iterator at the previous entry and reports whether the -// iterator was not already positioned at the beginning. +// Prev moves to the previous entry and reports whether the iterator was not +// already at the first entry. If Done() is true (e.g. after Seek past the end), +// Prev navigates to the last entry. +// +// Prev uses the live internal cursor directly (O(1) per step). After copying +// an iterator (it2 = *it), call it2.Refresh() before the first cross-CellID +// Prev() call to decouple the copy's cursor from the original. func (it *PointIndexIterator[Data]) Prev() bool { - if it.position == 0 { + // Fast path: go back within the current CellID's group. + if it.valid && it.sliceIdx > 0 { + it.sliceIdx-- + return true + } + if it.valid { + if it.iter.Prev() { + it.currentID = it.iter.Key() + it.currentSlice = it.iter.Value() + it.sliceIdx = len(it.currentSlice) - 1 + it.atEnd = false + return true + } + // Already at the first entry; restore the cursor so Next() still works. + it.iter.Seek(it.currentID) + return false + } + if it.atEnd { + if it.iter.Last() { + it.currentID = it.iter.Key() + it.currentSlice = it.iter.Value() + it.sliceIdx = len(it.currentSlice) - 1 + it.valid = true + it.atEnd = false + return true + } return false } - it.position-- - return true + return false } -// Seek positions the iterator at the first entry with CellID() >= target, or -// at the end of the index if no such entry exists. +// Seek positions the iterator at the first entry with CellID >= target, +// or at Done if no such entry exists. func (it *PointIndexIterator[Data]) Seek(target CellID) { - it.position = sort.Search(len(it.index.entries), func(i int) bool { - return it.index.entries[i].id >= target - }) + it.iter = it.index.tree.Iter() // fresh iter: decouples from any copy + it.valid = it.iter.Seek(target) + it.atEnd = !it.valid + it.sliceIdx = 0 + if it.valid { + it.currentID = it.iter.Key() + it.currentSlice = it.iter.Value() + } } // LocatePoint positions the iterator at the entry for the cell containing the diff --git a/s2/point_index_test.go b/s2/point_index_test.go index 9c1b329b..963bff21 100644 --- a/s2/point_index_test.go +++ b/s2/point_index_test.go @@ -122,6 +122,7 @@ func (pt *pointIndexTest) verifyIteratorMethods() { // Test Prev, Next, and Seek. if prevCellID.IsValid() { it2 = *it + it2.Refresh() // decouple cursor from original before cross-boundary Prev if !it2.Prev() { pt.t.Error("Prev() returned false, expected true") } From 468bab7cc60402e6cf3e490c58510776e8790d28 Mon Sep 17 00:00:00 2001 From: Fabien Meurisse Date: Wed, 3 Jun 2026 14:39:11 +0200 Subject: [PATCH 3/3] add ClosestPointQuery with support for finding closest points in PointIndex --- README.md | 2 +- s2/closest_point_query.go | 496 +++++++++++++++++++++++++++++++++ s2/closest_point_query_test.go | 281 +++++++++++++++++++ s2/point_index.go | 29 +- 4 files changed, 791 insertions(+), 17 deletions(-) create mode 100644 s2/closest_point_query.go create mode 100644 s2/closest_point_query_test.go diff --git a/README.md b/README.md index 47866c07..bdd8804e 100644 --- a/README.md +++ b/README.md @@ -188,7 +188,7 @@ S2ClosestCell | ❌ S2FurthestCell | ❌ S2ClosestEdge | ✅ S2FurthestEdge | ✅ -S2ClosestPoint | ❌ +S2ClosestPoint | ✅ S2FurthestPoint | ❌ S2ContainsPoint | ✅ S2ContainsVertex | ✅ diff --git a/s2/closest_point_query.go b/s2/closest_point_query.go new file mode 100644 index 00000000..8b9fa0bb --- /dev/null +++ b/s2/closest_point_query.go @@ -0,0 +1,496 @@ +// Copyright 2015 Google Inc. All rights reserved. +// +// Licensed 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 s2 + +import ( + "container/heap" + "sort" + + "github.com/golang/geo/s1" +) + +// minPointsToEnqueue is the minimum number of points in a cell required +// to enqueue it rather than process its contents directly. +const minPointsToEnqueue = 13 + +// closestPointResultHeap is a max-heap of ClosestPointQueryResult values, +// ordered by distance (largest distance at the top). +type closestPointResultHeap[Data comparable] []ClosestPointQueryResult[Data] + +func (h closestPointResultHeap[Data]) Len() int { return len(h) } +func (h closestPointResultHeap[Data]) Less(i, j int) bool { + return h[i].distance > h[j].distance +} +func (h closestPointResultHeap[Data]) Swap(i, j int) { h[i], h[j] = h[j], h[i] } +func (h *closestPointResultHeap[Data]) Push(x any) { + *h = append(*h, x.(ClosestPointQueryResult[Data])) +} +func (h *closestPointResultHeap[Data]) Pop() any { + old := *h + n := len(old) + item := old[n-1] + *h = old[:n-1] + return item +} +func (h closestPointResultHeap[Data]) top() ClosestPointQueryResult[Data] { return h[0] } + +// ClosestPointQueryResult holds one result from ClosestPointQuery. +type ClosestPointQueryResult[Data comparable] struct { + distance s1.ChordAngle + pointData PointData[Data] +} + +// Distance returns the distance from the target to this point. +func (r ClosestPointQueryResult[Data]) Distance() s1.ChordAngle { return r.distance } + +// Point returns the indexed point. +func (r ClosestPointQueryResult[Data]) Point() Point { return r.pointData.Point() } + +// Data returns the client data associated with this point. +func (r ClosestPointQueryResult[Data]) Data() Data { return r.pointData.Data() } + +// IsEmpty reports whether this result is empty (FindClosestPoint found nothing). +func (r ClosestPointQueryResult[Data]) IsEmpty() bool { + return r.distance == s1.InfChordAngle() +} + +// ClosestPointQueryOptions controls the set of points returned by ClosestPointQuery. +// By default all points are returned, so always set MaxResults and/or DistanceLimit. +type ClosestPointQueryOptions struct { + common *queryOptions +} + +// NewClosestPointQueryOptions returns default options for closest-point queries. +func NewClosestPointQueryOptions() *ClosestPointQueryOptions { + return &ClosestPointQueryOptions{common: newQueryOptions(minDistance(0))} +} + +// MaxResults specifies that at most n points should be returned. n must be >= 1. +func (o *ClosestPointQueryOptions) MaxResults(n int) *ClosestPointQueryOptions { + o.common = o.common.MaxResults(n) + return o +} + +// DistanceLimit specifies that only points whose distance to the target is +// strictly less than the limit should be returned. +func (o *ClosestPointQueryOptions) DistanceLimit(limit s1.ChordAngle) *ClosestPointQueryOptions { + o.common = o.common.DistanceLimit(limit) + return o +} + +// InclusiveDistanceLimit is like DistanceLimit but also returns points +// whose distance is exactly equal to the limit. +func (o *ClosestPointQueryOptions) InclusiveDistanceLimit(limit s1.ChordAngle) *ClosestPointQueryOptions { + o.common = o.common.ClosestInclusiveDistanceLimit(limit) + return o +} + +// ConservativeDistanceLimit expands the limit by the maximum distance +// calculation error, ensuring all points whose true distance is <= limit +// are returned (along with some slightly further ones). +func (o *ClosestPointQueryOptions) ConservativeDistanceLimit(limit s1.ChordAngle) *ClosestPointQueryOptions { + o.common = o.common.ClosestConservativeDistanceLimit(limit) + return o +} + +// MaxError specifies that points up to this distance further than the true +// closest may be substituted in the result set, as long as they satisfy +// all other criteria. Only meaningful when MaxResults is also set. +func (o *ClosestPointQueryOptions) MaxError(dist s1.ChordAngle) *ClosestPointQueryOptions { + o.common = o.common.MaxError(dist) + return o +} + +// Region specifies that results must be contained by the given region. +func (o *ClosestPointQueryOptions) Region(region Region) *ClosestPointQueryOptions { + o.common.region = region + return o +} + +// UseBruteForce forces the brute-force algorithm. Useful for testing. +func (o *ClosestPointQueryOptions) UseBruteForce(x bool) *ClosestPointQueryOptions { + o.common = o.common.UseBruteForce(x) + return o +} + +// ClosestPointQuery finds the closest point(s) in a PointIndex to a given +// target (point, edge, cell, or shape index). +// +// Example: +// +// index := &PointIndex[int]{} +// for i, p := range indexPoints { +// index.Add(p, i) +// } +// query := NewClosestPointQuery(index, NewClosestPointQueryOptions().MaxResults(5)) +// target := NewMinDistanceToPointTarget(queryPoint) +// for _, result := range query.FindClosestPoints(target) { +// // result.Distance(), result.Point(), result.Data() +// } +// +// ClosestPointQuery is not safe for concurrent use without external synchronization. +type ClosestPointQuery[Data comparable] struct { + index *PointIndex[Data] + opts *queryOptions + target distanceTarget + + useConservativeCellDistance bool + + // Precomputed covering of the indexed points; cleared on ReInit. + indexCovering []CellID + + // Distance limit, tightened as results are found. + distanceLimit distance + + // Result stores — exactly one is used per query based on opts.maxResults. + resultSingleton ClosestPointQueryResult[Data] + resultVector []ClosestPointQueryResult[Data] + resultSet closestPointResultHeap[Data] + + // Shared iterator for the optimized algorithm; shared across processOrEnqueue calls. + iter PointIndexIterator[Data] + + // Scratch space for direct processing of small cells. Avoids per-call allocation. + tmpPointData [minPointsToEnqueue - 1]PointData[Data] + + // Priority queue for candidate cells. + queue *queryQueue +} + +// NewClosestPointQuery returns a ClosestPointQuery for the given index. +// Pass nil opts to use default options (returns all points). +func NewClosestPointQuery[Data comparable](index *PointIndex[Data], opts *ClosestPointQueryOptions) *ClosestPointQuery[Data] { + if opts == nil { + opts = NewClosestPointQueryOptions() + } + q := &ClosestPointQuery[Data]{queue: newQueryQueue()} + q.Init(index, opts) + return q +} + +// Init (re)initializes the query for the given index and options. +// Must be called (or ReInit called) if the index is modified after this. +func (q *ClosestPointQuery[Data]) Init(index *PointIndex[Data], opts *ClosestPointQueryOptions) { + q.index = index + if opts != nil { + q.opts = opts.common + } + q.ReInit() +} + +// ReInit must be called whenever the underlying index is modified. +func (q *ClosestPointQuery[Data]) ReInit() { + q.iter.Init(q.index) + q.indexCovering = nil +} + +// Options returns the current query options. +func (q *ClosestPointQuery[Data]) Options() *ClosestPointQueryOptions { + return &ClosestPointQueryOptions{common: q.opts} +} + +// FindClosestPoints returns all points satisfying the current options, sorted +// by distance (closest first). This may be called multiple times. +func (q *ClosestPointQuery[Data]) FindClosestPoints(target distanceTarget) []ClosestPointQueryResult[Data] { + return q.findClosestPoints(target, q.opts) +} + +// FindClosestPoint returns the single closest point. If no point satisfies +// the search criteria, returns a result with IsEmpty() == true. +func (q *ClosestPointQuery[Data]) FindClosestPoint(target distanceTarget) ClosestPointQueryResult[Data] { + opts := *q.opts + opts.maxResults = 1 + results := q.findClosestPoints(target, &opts) + if len(results) == 0 { + return ClosestPointQueryResult[Data]{distance: s1.InfChordAngle()} + } + return results[0] +} + +// GetDistance returns the minimum distance to the target. +// Returns InfChordAngle if the index or target is empty. +// Use IsDistanceLess if only comparing against a threshold. +func (q *ClosestPointQuery[Data]) GetDistance(target distanceTarget) s1.ChordAngle { + return q.FindClosestPoint(target).Distance() +} + +// IsDistanceLess reports whether the distance to target is less than limit. +// This is usually faster than GetDistance since the search can stop early. +func (q *ClosestPointQuery[Data]) IsDistanceLess(target distanceTarget, limit s1.ChordAngle) bool { + opts := *q.opts + opts.maxResults = 1 + opts.distanceLimit = limit + opts.maxError = s1.StraightChordAngle + return len(q.findClosestPoints(target, &opts)) > 0 +} + +// IsDistanceLessOrEqual reports whether the distance to target is <= limit. +func (q *ClosestPointQuery[Data]) IsDistanceLessOrEqual(target distanceTarget, limit s1.ChordAngle) bool { + return q.IsDistanceLess(target, limit.Successor()) +} + +// IsConservativeDistanceLessOrEqual reports whether the true distance to the +// target is likely <= limit. It accounts for rounding error: all points whose +// true distance is <= limit are guaranteed to be found. +func (q *ClosestPointQuery[Data]) IsConservativeDistanceLessOrEqual(target distanceTarget, limit s1.ChordAngle) bool { + opts := *q.opts + opts.maxResults = 1 + opts.distanceLimit = limit.Expanded(minUpdateDistanceMaxError(limit)).Successor() + opts.maxError = s1.StraightChordAngle + return len(q.findClosestPoints(target, &opts)) > 0 +} + +func (q *ClosestPointQuery[Data]) findClosestPoints(target distanceTarget, opts *queryOptions) []ClosestPointQueryResult[Data] { + q.findClosestPointsInternal(target, opts) + + if opts.maxResults == 1 { + if q.resultSingleton.IsEmpty() { + return nil + } + return []ClosestPointQueryResult[Data]{q.resultSingleton} + } + + if opts.maxResults == maxQueryResults { + sort.Slice(q.resultVector, func(i, j int) bool { + return q.resultVector[i].distance < q.resultVector[j].distance + }) + results := q.resultVector + q.resultVector = nil + return results + } + + // Drain the max-heap (largest first) then reverse to get ascending order. + results := make([]ClosestPointQueryResult[Data], 0, q.resultSet.Len()) + for q.resultSet.Len() > 0 { + results = append(results, heap.Pop(&q.resultSet).(ClosestPointQueryResult[Data])) + } + for i, j := 0, len(results)-1; i < j; i, j = i+1, j-1 { + results[i], results[j] = results[j], results[i] + } + return results +} + +func (q *ClosestPointQuery[Data]) findClosestPointsInternal(target distanceTarget, opts *queryOptions) { + q.target = target + q.opts = opts + + q.distanceLimit = minDistance(opts.distanceLimit) + q.resultSingleton = ClosestPointQueryResult[Data]{distance: s1.InfChordAngle()} + q.resultVector = nil + q.resultSet = closestPointResultHeap[Data]{} + + if q.distanceLimit == minDistance(0) { + return + } + + targetUsesMaxError := opts.maxError != 0 && target.setMaxError(opts.maxError) + q.useConservativeCellDistance = targetUsesMaxError && + (q.distanceLimit == minDistance(0).infinity() || + minDistance(0).less(q.distanceLimit.sub(minDistance(opts.maxError)))) + + if opts.useBruteForce || q.index.NumPoints() <= target.maxBruteForceIndexSize() { + q.findClosestPointsBruteForce() + } else { + q.findClosestPointsOptimized() + } +} + +func (q *ClosestPointQuery[Data]) findClosestPointsBruteForce() { + for it := NewPointIndexIterator(q.index); !it.Done(); it.Next() { + q.maybeAddResult(it.PointData()) + } +} + +func (q *ClosestPointQuery[Data]) findClosestPointsOptimized() { + q.initQueue() + for q.queue.size() > 0 { + entry := q.queue.pop() + if !entry.distance.less(q.distanceLimit) { + q.queue.reset() + break + } + child := entry.id.ChildBegin() + seek := true + for i := 0; i < 4; i++ { + seek = q.processOrEnqueue(child, seek) + child = child.Next() + } + } +} + +func (q *ClosestPointQuery[Data]) initQueue() { + cb := q.target.capBound() + if cb.IsEmpty() { + return + } + + if q.opts.maxResults == 1 { + // Optimization: seek near the target cap center to get an early upper + // bound on the search radius. The two adjacent index points (in CellID + // order) often yield a tight bound. + q.iter.Seek(cellIDFromPoint(cb.Center())) + if !q.iter.Done() { + q.maybeAddResult(q.iter.PointData()) + } + if q.iter.Prev() { + q.maybeAddResult(q.iter.PointData()) + } + if q.distanceLimit == minDistance(0) { + return + } + } + + if q.indexCovering == nil { + q.initCovering() + } + + initialCells := []CellID(q.indexCovering) + + if q.opts.region != nil { + coverer := &RegionCoverer{MaxCells: 4, LevelMod: 1, MaxLevel: MaxLevel} + regionCover := coverer.Covering(q.opts.region) + initialCells = CellUnionFromIntersection(CellUnion(q.indexCovering), regionCover) + } + + if q.distanceLimit != minDistance(0).infinity() { + coverer := &RegionCoverer{MaxCells: 4, LevelMod: 1, MaxLevel: MaxLevel} + radius := cb.Radius() + q.distanceLimit.chordAngleBound().Angle() + searchCap := CapFromCenterAngle(cb.Center(), radius) + maxDistCover := coverer.FastCovering(searchCap) + initialCells = CellUnionFromIntersection(CellUnion(initialCells), maxDistCover) + } + + q.iter.Begin() + for _, id := range initialCells { + if q.iter.Done() { + break + } + q.processOrEnqueue(id, id.RangeMin() > q.iter.CellID()) + } +} + +func (q *ClosestPointQuery[Data]) initCovering() { + // Compute a minimal covering (at most 6 cells) of all indexed points. + // See the equivalent method in EdgeQuery for a detailed explanation. + q.indexCovering = make([]CellID, 0, 6) + q.iter.Finish() + if !q.iter.Prev() { + return // Empty index. + } + indexLastID := q.iter.CellID() + q.iter.Begin() + if q.iter.CellID() != indexLastID { + level, ok := q.iter.CellID().CommonAncestorLevel(indexLastID) + if !ok { + level = 0 + } else { + level++ + } + lastID := indexLastID.Parent(level) + for id := q.iter.CellID().Parent(level); id != lastID; id = id.Next() { + if id.RangeMax() < q.iter.CellID() { + continue + } + cellFirstID := q.iter.CellID() + q.iter.Seek(id.RangeMax().Next()) + q.iter.Prev() + cellLastID := q.iter.CellID() + q.iter.Next() + q.addInitialRange(cellFirstID, cellLastID) + } + } + q.addInitialRange(q.iter.CellID(), indexLastID) +} + +// addInitialRange appends to indexCovering the lowest common ancestor of firstID and lastID. +func (q *ClosestPointQuery[Data]) addInitialRange(firstID, lastID CellID) { + level, _ := firstID.CommonAncestorLevel(lastID) + q.indexCovering = append(q.indexCovering, firstID.Parent(level)) +} + +func (q *ClosestPointQuery[Data]) maybeAddResult(pd PointData[Data]) { + dist := q.distanceLimit + updated, ok := q.target.updateDistanceToPoint(pd.Point(), dist) + if !ok { + return + } + if q.opts.region != nil && !q.opts.region.ContainsPoint(pd.Point()) { + return + } + result := ClosestPointQueryResult[Data]{ + distance: updated.chordAngle(), + pointData: pd, + } + switch { + case q.opts.maxResults == 1: + q.resultSingleton = result + q.distanceLimit = updated.sub(minDistance(q.opts.maxError)) + case q.opts.maxResults == maxQueryResults: + q.resultVector = append(q.resultVector, result) + default: + if q.resultSet.Len() >= q.opts.maxResults { + heap.Pop(&q.resultSet) + } + heap.Push(&q.resultSet, result) + if q.resultSet.Len() >= q.opts.maxResults { + q.distanceLimit = minDistance(q.resultSet.top().distance).sub(minDistance(q.opts.maxError)) + } + } +} + +// processOrEnqueue either processes the points in id directly (if few enough) +// or enqueues id for later subdivision. +// +// If seek is false, q.iter must already be positioned at the first indexed +// point within or after id. Returns true if the cell was enqueued (caller +// must seek for the next sibling), false if it was processed (q.iter is now +// positioned at the next cell in CellID order). +func (q *ClosestPointQuery[Data]) processOrEnqueue(id CellID, seek bool) bool { + if seek { + q.iter.Seek(id.RangeMin()) + } + if id.IsLeaf() { + for !q.iter.Done() && q.iter.CellID() == id { + q.maybeAddResult(q.iter.PointData()) + q.iter.Next() + } + return false + } + last := id.RangeMax() + numPoints := 0 + for !q.iter.Done() && q.iter.CellID() <= last { + if numPoints == minPointsToEnqueue-1 { + // Cell has at least minPointsToEnqueue points; enqueue for subdivision. + cell := CellFromCellID(id) + dist := q.distanceLimit + if updated, ok := q.target.updateDistanceToCell(cell, dist); ok { + if q.opts.region == nil || q.opts.region.IntersectsCell(cell) { + if q.useConservativeCellDistance { + updated = updated.sub(minDistance(q.opts.maxError)) + } + q.queue.push(&queryQueueEntry{distance: updated, id: id}) + } + } + return true + } + q.tmpPointData[numPoints] = q.iter.PointData() + numPoints++ + q.iter.Next() + } + for i := 0; i < numPoints; i++ { + q.maybeAddResult(q.tmpPointData[i]) + } + return false +} diff --git a/s2/closest_point_query_test.go b/s2/closest_point_query_test.go new file mode 100644 index 00000000..c6da12ee --- /dev/null +++ b/s2/closest_point_query_test.go @@ -0,0 +1,281 @@ +// Copyright 2015 Google Inc. All rights reserved. +// +// Licensed 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 s2 + +import ( + "math" + "testing" + + "github.com/golang/geo/s1" +) + +func TestClosestPointQueryNoPoints(t *testing.T) { + index := &PointIndex[int]{} + query := NewClosestPointQuery(index, nil) + target := NewMinDistanceToPointTarget(PointFromCoords(1, 0, 0)) + if got := query.FindClosestPoints(target); len(got) != 0 { + t.Errorf("FindClosestPoints on empty index: got %d results, want 0", len(got)) + } + if got := query.FindClosestPoint(target); !got.IsEmpty() { + t.Error("FindClosestPoint on empty index should be empty") + } + if got := query.GetDistance(target); got != s1.InfChordAngle() { + t.Errorf("GetDistance on empty index = %v, want InfChordAngle", got) + } +} + +func TestClosestPointQueryManyDuplicatePoints(t *testing.T) { + const numPoints = 10000 + p := PointFromCoords(1, 0, 0) + index := &PointIndex[int]{} + for i := range numPoints { + index.Add(p, i) + } + query := NewClosestPointQuery(index, nil) + target := NewMinDistanceToPointTarget(p) + results := query.FindClosestPoints(target) + if got := len(results); got != numPoints { + t.Errorf("FindClosestPoints on %d duplicates: got %d results, want %d", numPoints, got, numPoints) + } + for i, r := range results { + if r.IsEmpty() { + t.Errorf("result[%d].IsEmpty() = true", i) + } + } +} + +func TestClosestPointQueryEmptyTarget(t *testing.T) { + index := &PointIndex[int]{} + for i := range 1000 { + index.Add(randomPoint(), i) + } + query := NewClosestPointQuery(index, NewClosestPointQueryOptions(). + DistanceLimit(s1.ChordAngleFromAngle(1e-5*s1.Radian))) + emptyTarget := NewMinDistanceToShapeIndexTarget(NewShapeIndex()) + if got := len(query.FindClosestPoints(emptyTarget)); got != 0 { + t.Errorf("FindClosestPoints with empty target: got %d results, want 0", got) + } +} + +// testClosestPointQuery verifies that brute-force and optimized algorithms +// return equivalent results for the given target and query options. +func testClosestPointQuery[Data comparable](t *testing.T, target distanceTarget, query *ClosestPointQuery[Data]) { + t.Helper() + + query.opts.useBruteForce = true + expected := query.FindClosestPoints(target) + query.opts.useBruteForce = false + actual := query.FindClosestPoints(target) + + maxResults := query.opts.maxResults + maxErr := query.opts.maxError + distLimit := query.opts.distanceLimit + + if len(actual) > maxResults { + t.Errorf("got %d results, want <= %d", len(actual), maxResults) + } + + // If no distance limit and no region, the count must match exactly. + if distLimit == s1.InfChordAngle() && query.opts.region == nil { + want := min(maxResults, query.index.NumPoints()) + if len(actual) != want { + t.Errorf("got %d results, want %d (maxResults=%d, numPoints=%d)", len(actual), want, maxResults, query.index.NumPoints()) + } + } + + // All returned results must have distance < distLimit. + for _, r := range actual { + if r.Distance() >= distLimit { + t.Errorf("result distance %v >= limit %v", r.Distance(), distLimit) + } + } + + // The brute-force minimum distance must be within max_error of the + // optimized minimum distance. + if len(expected) > 0 && len(actual) > 0 { + minExpected := expected[0].Distance() + minActual := actual[0].Distance() + if minActual > minExpected+maxErr { + t.Errorf("optimized min dist %v > brute force %v + maxError %v", minActual, minExpected, maxErr) + } + } +} + +func TestClosestPointQueryBruteForceVsOptimized(t *testing.T) { + const ( + numIndexes = 10 + numPoints = 100 + numQueries = 50 + testCapKm = 10.0 + earthRadius = 6371.0 // km + ) + capAngle := s1.Angle(testCapKm/earthRadius) * s1.Radian + + for range numIndexes { + center := randomPoint() + indexCap := CapFromCenterAngle(center, capAngle) + index := &PointIndex[int]{} + for i := range numPoints { + p := samplePointFromCap(indexCap) + index.Add(p, i) + } + + for range numQueries { + queryRadius := 2 * capAngle + queryCap := CapFromCenterAngle(center, queryRadius) + query := NewClosestPointQuery(index, nil) + + // Vary the options. + if randomUniformInt(5) != 0 { + query.opts.maxResults = 1 + randomUniformInt(10) + } + if randomUniformInt(3) != 0 { + frac := randomUniformFloat64(0, 1) + query.opts.distanceLimit = s1.ChordAngleFromAngle(s1.Angle(frac) * queryRadius) + } + if randomUniformInt(2) != 0 { + maxErrFrac := 1e-4 + math.Exp(randomUniformFloat64(0, 1)*math.Log(1.0)) + query.opts.maxError = s1.ChordAngleFromAngle(s1.Angle(maxErrFrac) * queryRadius) + } + + targetType := randomUniformInt(3) + switch targetType { + case 0: + p := samplePointFromCap(queryCap) + testClosestPointQuery(t, NewMinDistanceToPointTarget(p), query) + case 1: + a := samplePointFromCap(queryCap) + bCap := CapFromCenterAngle(a, s1.Angle(1e-4)*queryRadius) + b := samplePointFromCap(bCap) + testClosestPointQuery(t, NewMinDistanceToEdgeTarget(Edge{a, b}), query) + case 2: + minLevel := MaxLevel - 4 + level := minLevel + randomUniformInt(5) + cellID := cellIDFromPoint(samplePointFromCap(queryCap)).Parent(level) + testClosestPointQuery(t, NewMinDistanceToCellTarget(CellFromCellID(cellID)), query) + } + } + } +} + +func TestClosestPointQueryFindClosestPoint(t *testing.T) { + p0 := parsePoint("0:0") + p1 := parsePoint("1:0") + p2 := parsePoint("2:0") + + index := &PointIndex[int]{} + index.Add(p0, 0) + index.Add(p1, 1) + index.Add(p2, 2) + + query := NewClosestPointQuery(index, nil) + target := NewMinDistanceToPointTarget(parsePoint("1.1:0")) + + result := query.FindClosestPoint(target) + if result.IsEmpty() { + t.Fatal("FindClosestPoint returned empty result") + } + if result.Data() != 1 { + t.Errorf("closest point data = %d, want 1", result.Data()) + } + if result.Point() != p1 { + t.Errorf("closest point = %v, want %v", result.Point(), p1) + } +} + +func TestClosestPointQueryDistanceLimitAndMaxResults(t *testing.T) { + index := &PointIndex[int]{} + for i := range 100 { + index.Add(randomPoint(), i) + } + + query := NewClosestPointQuery(index, NewClosestPointQueryOptions(). + MaxResults(5). + DistanceLimit(s1.ChordAngleFromAngle(s1.InfAngle()))) + target := NewMinDistanceToPointTarget(randomPoint()) + + results := query.FindClosestPoints(target) + if len(results) > 5 { + t.Errorf("got %d results, want <= 5", len(results)) + } + // Verify results are sorted by distance. + for i := 1; i < len(results); i++ { + if results[i].Distance() < results[i-1].Distance() { + t.Errorf("results not sorted: results[%d].Distance=%v < results[%d].Distance=%v", + i, results[i].Distance(), i-1, results[i-1].Distance()) + } + } +} + +func TestClosestPointQueryIsDistanceLess(t *testing.T) { + p0 := parsePoint("23:12") + p1 := parsePoint("47:11") + + index := &PointIndex[int]{} + index.Add(p0, 0) + + query := NewClosestPointQuery(index, nil) + target := NewMinDistanceToPointTarget(p0) + + // Distance to p0 is zero. + zeroAngle := s1.ChordAngle(0) + if query.IsDistanceLess(target, zeroAngle) { + t.Error("IsDistanceLess(p0, 0): want false for distance 0") + } + if !query.IsDistanceLessOrEqual(target, zeroAngle) { + t.Error("IsDistanceLessOrEqual(p0, 0): want true") + } + if !query.IsConservativeDistanceLessOrEqual(target, zeroAngle) { + t.Error("IsConservativeDistanceLessOrEqual(p0, 0): want true") + } + + // Distance to p1 is positive. + target1 := NewMinDistanceToPointTarget(p1) + d01 := ChordAngleBetweenPoints(p0, p1) + if !query.IsDistanceLess(target1, d01.Successor()) { + t.Error("IsDistanceLess(p1, d01.Successor): want true") + } + if query.IsDistanceLess(target1, d01) { + t.Error("IsDistanceLess(p1, d01): want false (equal, not less)") + } + if !query.IsDistanceLessOrEqual(target1, d01) { + t.Error("IsDistanceLessOrEqual(p1, d01): want true") + } + if !query.IsConservativeDistanceLessOrEqual(target1, d01) { + t.Error("IsConservativeDistanceLessOrEqual(p1, d01): want true") + } +} + +func TestClosestPointQueryReInit(t *testing.T) { + p := PointFromCoords(1, 0, 0) + index := &PointIndex[int]{} + query := NewClosestPointQuery(index, nil) + target := NewMinDistanceToPointTarget(p) + + if got := len(query.FindClosestPoints(target)); got != 0 { + t.Errorf("empty index: got %d results, want 0", got) + } + + index.Add(p, 42) + query.ReInit() + + results := query.FindClosestPoints(target) + if got := len(results); got != 1 { + t.Fatalf("after ReInit: got %d results, want 1", got) + } + if results[0].Data() != 42 { + t.Errorf("result.Data() = %d, want 42", results[0].Data()) + } +} diff --git a/s2/point_index.go b/s2/point_index.go index 13339d0a..3b725b2c 100644 --- a/s2/point_index.go +++ b/s2/point_index.go @@ -38,23 +38,20 @@ func (pd PointData[Data]) Data() Data { return pd.data } // You can use this class in conjunction with ClosestPointQuery to find the // closest index points to a given query point. For example, // -// index := &PointIndex[int]{} -// for i, p := range indexPoints { -// index.Add(p, i) -// } -// TODO(fmeurisse): Implement ClosestPointQuery integration and update example. -// S2ClosestPointQuery query(&index); -// query.mutable_options()->set_max_results(5); -// for (const S2Point& target_point : target_points) { -// S2ClosestPointQueryPointTarget target(target_point); -// for (const auto& result : query.FindClosestPoints(&target)) { -// // The Result class contains the following methods: -// // distance() is the distance to the target. -// // point() is the indexed point. -// // data() is the auxiliary data. -// DoSomething(target_point, result); +// index := &PointIndex[int]{} +// for i, p := range indexPoints { +// index.Add(p, i) +// } +// query := NewClosestPointQuery(index, NewClosestPointQueryOptions().MaxResults(5)) +// for _, targetPoint := range targetPoints { +// target := NewMinDistanceToPointTarget(targetPoint) +// for _, result := range query.FindClosestPoints(target) { +// // result.Distance() is the distance to the target. +// // result.Point() is the indexed point. +// // result.Data() is the auxiliary data. +// DoSomething(targetPoint, result) // } -// } +// } // // You can also access the index directly using the iterator interface. For // example, here is how to iterate through all the points in a given CellID