Review fixes and API pass: correctness fixes, validity-returning navigation, set iterators, Span queries, CommutativeMonoid

CHANGELOG.md

@@ -1,5 +1,89 @@
 # Changelog
 
+## Unreleased
+
+### API (breaking)
+
+- Every positioning method on iterators and cursors (`First`, `Last`,
+  `Next`, `Prev`, the seeks, `SeekNth`) returns whether the iterator is now
+  at an entry. `SeekGE` and `SeekLT` no longer report whether the sought
+  key exists; `SeekExact` does. Code that compiled against `v0.2.0` and
+  used a seek's result as "found" changes meaning: replace it with
+  `SeekExact`.
+- Sets: `Get(probe)` returns the stored item equal to a probe; `Delete`
+  returns the removed item. Set iterators and cursors are their own types
+  (`SetIterator`, `SetCursor`) with `Item` and no `Value`, `SetValue` or
+  phantom value arguments.
+- Cursors: `Upsert` returns the replaced key and value like `Map.Upsert`;
+  `Rekey` returns the entry it displaced at the destination.
+- interval: queries are `Span`s (`HalfOpen`, `Point`) instead of stored
+  intervals; `Overlaps(span)` returns an `OverlapIterator` whose `Next` is
+  always the next overlap and whose `Seek` restarts it for another query,
+  and `Iterator` is a plain iterator with no overlap mode. `Bounds.HasEnd`
+  defaults to "every interval has an end; one whose end is not after its
+  start is a point" instead of the zero-end convention, and the default
+  tie-break treats every representation of a point at a key as the same
+  item. `Bounds.CompareIntervals` is optional again as a single-call
+  comparator that must order by start, which `Verify` checks. `Set` has
+  `Get`, an item-returning `Delete`, `SetIterator`, `SetCursor` and
+  `SetOverlapIterator`.
+- aug: `Monoid`/`Group` are `CommutativeMonoid`/`CommutativeGroup` and
+  require an `Identity` method; `SeekWhere(prefix, contribution)` is
+  `SeekPrefix(inclusivePrefix)` with a predicate that turns true at most
+  once, and returns the prefix and whether an entry was found;
+  `Node.Aug` and `ChildAug` return values and `Node.SetAug` is for
+  Updaters; `MonoidMap` no longer exposes an embedded `Map`; `New` panics
+  on a nil comparison function. The package documents one ownership
+  contract for every collection.
+
+### Fixed
+
+- Splits, merges and rebalances are reported to Updaters that implement
+  the new `Restructurer` interface, so augmentations that depend on the
+  shape of the subtree (node counts, heights) can stay correct; deleting
+  an absent key that merges nodes is covered. Augmentations that depend
+  only on the entries pay nothing. The `Updater` documentation states
+  what an Updater is told.
+- Nodes start from the monoid's `Identity`, fresh and recycled, so a
+  monoid whose identity is not the zero value aggregates correctly.
+- Reference counts are 64-bit; clones dropped without `Clear` can no longer
+  wrap them.
+- An interval whose end is not after its start is a point, so leaf matching
+  and subtree pruning agree on it wherever it sits in the tree.
+- `Cursor.Upsert` replaced only the value when the key compared equal to
+  an existing one; it now replaces the key too and reports the previous
+  key to the Updater, so key-dependent augmentations stay correct.
+- The interval augmentation treated the zero endpoint as its initial
+  bound, which over-estimated bounds for negative endpoints and compared
+  the zero value with comparators that cannot take it (pointer endpoints
+  panicked). The bound now has an explicit unset state.
+- `Next` and `Prev` on an interval iterator end an overlap scan, as the
+  seeks already did, so `NextOverlap` after stepping does not apply stale
+  constraints.
+- `Verify` no longer copies a node's reference count non-atomically while
+  clones may be changing it.
+
+### Changed
+
+- `interval.Bounds.TieBreak` orders intervals with equal start keys; the
+  tree orders by start first. `CompareIntervals` remains for a
+  single-call comparator and must order by start.
+- `LowLevelIterator.Config` returns a copy and `Config.Compare` has a
+  value receiver; a Map's Updater and comparison function are fixed at
+  construction.
+- `Verify` compares augmentations with the Updater's `Equal` when it
+  implements `Equaler`; `MonoidUpdater` forwards the Monoid's.
+- `interval.Cursor`, `interval.FreeList` and `interval.NewFreeList` name
+  the types that were only reachable through inference.
+- The interval package documents its overlap cost as O(log n) plus the
+  ancestors of the k matches, up to O(k log(n/k)) when they are scattered,
+  rather than O(log n + k).
+- `SeekWhere` documents the contract its single descent needs: the
+  predicate must be exact for spans, true for a span exactly when true for
+  some entry of it.
+- The `aug` documentation states that maps must not be copied by value and
+  that keys, values and augmentations are shallow-copied.
+
 ## v0.2.0 (2026-09-29)
 
 A rewrite of the core with the same design (PR #3, from `1cedbd5`).

README.md

@@ -25,7 +25,9 @@ snapshot.Clear() // returns nodes only this tree references to the free list
 
 The stateful `Iterator` and the range-over-func views cost about the same per entry when the loop body does real work; the iterator is the one to use for seeking and stepping, the views for whole-range walks.
 
-A `Cursor` is an iterator that can also mutate the entry it is on and stay valid: `SetValue`, `Delete` (leaves the cursor on the successor), `Rekey` (moves the entry to a new key) and a hinted `Upsert`. Each acts in place when the leaf allows it and otherwise falls back to the top-down algorithm plus a re-seek, so "seek, then move this entry" costs one descent instead of three.
+A `Cursor` is an iterator that can also mutate the entry it is on and stay valid: `SetValue`, `Delete` (leaves the cursor on the successor), `Rekey` (moves the entry to a new key) and a hinted `Upsert`, each returning the displaced entry as the map methods do. Each acts in place when the leaf allows it and otherwise falls back to the top-down algorithm plus a re-seek, so "seek, then move this entry" costs one descent instead of three.
+
+Sets are keyed by a comparison function, so two items that compare equal are the same item; `Set.Get(probe)` returns the stored one. Set iterators and cursors have no phantom values.
 
 Writes to a map must be serialized by the caller; any number of goroutines may read a map, or a clone of it, while no goroutine writes it. See the `aug` package documentation for details.
 
@@ -39,15 +41,15 @@ Against google/btree v1.1.3 at equal degree, point operations (insert, get, dele
 
 ## Interval Trees
 
-The `interval` package provides interval trees for efficiently finding all intervals that overlap a query range. The iterator supports `FirstOverlap()` and `NextOverlap()` methods for querying overlapping intervals.
+The `interval` package provides interval trees for efficiently finding all intervals that overlap a query. Stored intervals are described by `Bounds` (how to read a start and end from your type); queries are `Span`s built with `HalfOpen(start, end)` or `Point(key)`, so asking about a range never requires building a stored object. `Overlaps(span)` returns an `OverlapIterator` whose `Next` always means the next overlap and whose `Seek(span)` reuses it for the next query; `Overlapping(span)` is the range-over-func form.
 
 ## Order-Statistic Trees
 
 The `orderstat` package provides order-statistic trees: `Rank(key)` (the number of keys less than a key, and whether the key exists), `Nth(rank)`, and `Count(lo, hi)`, each in one descent. The iterator adds `Rank()` for its current position and `SeekNth(rank)`.
 
 ## Custom augmentations
 
-Describe an augmentation as a `Monoid` (a per-entry contribution and a commutative, associative `Combine`; add `Uncombine` for O(1) removals) and build the map with `aug.NewMonoid(cmp, m)`. Ready-made pieces: `aug.Count`, `aug.Sum(f)`, and `aug.PairOf(a, b)` to keep several side by side. A `MonoidMap` answers `Total()`, `Prefix(key)` and `Aggregate(lo, hi)`, and its iterators and cursors answer `Prefix()` and `SeekWhere(pred)`, which descends once to the first entry at which a monotone predicate on the running prefix flips (selection by rank, by cumulative sum, and so on). See `Example_customAugmentation` in package `aug`. Augmentations that are not monoids (the interval tree's upper bound) implement `aug.Updater` directly with `aug.New`.
+Describe an augmentation as a `CommutativeMonoid` (an `Identity`, a per-entry contribution and a commutative, associative `Combine`; add `Uncombine` for O(1) removals) and build the map with `aug.NewMonoid(cmp, m)`. Ready-made pieces: `aug.Count`, `aug.Sum(f)`, and `aug.PairOf(a, b)` to keep several side by side. A `MonoidMap` answers `Total()`, `Prefix(key)` and `Aggregate(lo, hi)`, and its iterators and cursors answer `Prefix()` and `SeekPrefix(pred)`, which descends once to the first entry whose inclusive prefix satisfies a predicate that turns from false to true at most once (selection by rank, by cumulative weight, and so on). See `Example_customAugmentation` in package `aug`. Augmentations that are not monoids (the interval tree's upper bound) implement `aug.Updater` directly with `aug.New`.
 
 ## License
 

aug/adversarial_test.go

@@ -16,16 +16,17 @@ package aug_test
 
 import (
 	"cmp"
+	"math"
 	"math/rand/v2"
 	"testing"
 
 	"github.com/ajwerner/btree/aug"
 )
 
-// TestSeekWhereNonMonotone checks that a predicate that breaks the
+// TestSeekPrefixNonMonotone checks that a predicate that breaks the
 // monotonicity contract still leaves the iterator in a consistent state:
 // valid at some entry, or past the end, never panicking or looping.
-func TestSeekWhereNonMonotone(t *testing.T) {
+func TestSeekPrefixNonMonotone(t *testing.T) {
 	rng := rand.New(rand.NewPCG(5, 6))
 	for _, degree := range []int{2, 16} {
 		m := aug.NewMonoid[int, int, int](cmp.Compare[int], aug.Count[int, int]{}, aug.WithDegree(degree))
@@ -35,13 +36,16 @@ func TestSeekWhereNonMonotone(t *testing.T) {
 		it := m.Iterator()
 		for range 2000 {
 			calls := 0
-			prefix := it.SeekWhere(func(int, int) bool {
+			prefix, ok := it.SeekPrefix(func(int) bool {
 				calls++
 				if calls > 100000 {
 					t.Fatal("predicate called too many times")
 				}
 				return rng.IntN(3) == 0
 			})
+			if ok != it.Valid() {
+				t.Fatalf("SeekPrefix reported %v but Valid is %v", ok, it.Valid())
+			}
 			if it.Valid() {
 				if got := it.Prefix(); got != prefix {
 					t.Fatalf("returned prefix %d, Prefix() %d", prefix, got)
@@ -55,21 +59,20 @@ func TestSeekWhereNonMonotone(t *testing.T) {
 				if prefix != 5000 {
 					t.Fatalf("past the end with prefix %d, want the total", prefix)
 				}
-				it.Prev()
-				if !it.Valid() || it.Key() != 4999 {
-					t.Fatalf("Prev after a failed SeekWhere gave valid=%v %d", it.Valid(), it.Key())
+				if !it.Prev() || it.Key() != 4999 {
+					t.Fatalf("Prev after a failed SeekPrefix gave valid=%v %d", it.Valid(), it.Key())
 				}
 			}
 			it.Next()
 			it.Prev()
 		}
 		// A predicate that is always false ends past the end with the total.
-		if p := it.SeekWhere(func(int, int) bool { return false }); p != 5000 || it.Valid() {
-			t.Fatalf("always-false: prefix %d valid %v", p, it.Valid())
+		if p, ok := it.SeekPrefix(func(int) bool { return false }); ok || p != 5000 || it.Valid() {
+			t.Fatalf("always-false: prefix %d ok %v valid %v", p, ok, it.Valid())
 		}
 		// Always true stops at the first entry.
-		if p := it.SeekWhere(func(int, int) bool { return true }); p != 0 || !it.Valid() || it.Key() != 0 {
-			t.Fatalf("always-true: prefix %d valid %v key %d", p, it.Valid(), it.Key())
+		if p, ok := it.SeekPrefix(func(int) bool { return true }); !ok || p != 0 || !it.Valid() || it.Key() != 0 {
+			t.Fatalf("always-true: prefix %d ok %v valid %v key %d", p, ok, it.Valid(), it.Key())
 		}
 	}
 }
@@ -181,3 +184,84 @@ func TestVerifyInterfaceTypes(t *testing.T) {
 		t.Fatal(err)
 	}
 }
+
+// nodeCounter is a shape-dependent augmentation: the number of nodes in
+// the subtree. It recomputes from its children on every event and
+// implements Restructurer, as the Updater documentation prescribes for
+// such augmentations.
+type nodeCounter[K, V any] struct{}
+
+func (u nodeCounter[K, V]) Restructured(n *aug.Node[K, V, int]) bool {
+	return u.Update(n, aug.UpdateInfo[K, V, int]{})
+}
+
+func (nodeCounter[K, V]) Update(n *aug.Node[K, V, int], _ aug.UpdateInfo[K, V, int]) bool {
+	count := 1
+	if !n.IsLeaf() {
+		for i := int16(0); i <= n.Count(); i++ {
+			count += n.ChildAug(i)
+		}
+	}
+	changed := n.Aug() != count
+	n.SetAug(count)
+	return changed
+}
+
+// TestStructuralAugmentation checks that splits, merges and rebalances
+// reach ancestors' augmentations even when no entry below them changed,
+// including the merge caused by deleting a key that is absent.
+func TestStructuralAugmentation(t *testing.T) {
+	// The reviewer's example: degree 2, keys 1..6 leave the root at 3
+	// nodes instead of 4 when the split is not reported.
+	m := aug.New[int, int, int](cmp.Compare[int], nodeCounter[int, int]{}, aug.WithDegree(2))
+	for i := 1; i <= 6; i++ {
+		m.Upsert(i, i)
+	}
+	if err := m.Verify(); err != nil {
+		t.Fatal(err)
+	}
+	for _, degree := range []int{2, 3} {
+		m := aug.New[int, int, int](cmp.Compare[int], nodeCounter[int, int]{}, aug.WithDegree(degree))
+		rng := rand.New(rand.NewPCG(31, 32))
+		for step := range 20000 {
+			k := rng.IntN(500)
+			switch rng.IntN(3) {
+			case 0, 1:
+				m.Upsert(k, step)
+			case 2:
+				m.Delete(rng.IntN(600)) // often absent, which can still merge
+			}
+			if step%500 == 499 {
+				if err := m.Verify(); err != nil {
+					t.Fatalf("degree %d step %d: %v", degree, step, err)
+				}
+			}
+		}
+	}
+}
+
+// maxFloat is a NaN-propagating maximum over non-negative floats whose
+// Equaler treats NaN as equal to itself, as cmp.Compare does, where
+// reflect.DeepEqual would not.
+type maxFloat struct{}
+
+func (maxFloat) Identity() float64           { return 0 }
+func (maxFloat) Of(_ int, v float64) float64 { return v }
+func (maxFloat) Combine(a, b float64) float64 {
+	if math.IsNaN(a) || math.IsNaN(b) {
+		return math.NaN()
+	}
+	return max(a, b)
+}
+func (maxFloat) Equal(a, b float64) bool { return cmp.Compare(a, b) == 0 }
+
+func TestVerifyUsesEqualer(t *testing.T) {
+	m := aug.NewMonoid[int, float64, float64](cmp.Compare[int], maxFloat{}, aug.WithDegree(2))
+	nan := math.NaN()
+	for i := range 50 {
+		m.Upsert(i, nan)
+	}
+	if err := m.Verify(); err != nil {
+		t.Fatal(err)
+	}
+}

aug/aggregate.go

@@ -14,7 +14,7 @@
 
 package aug
 
-func (t *Map[K, V, A]) monoid() Monoid[K, V, A] {
+func (t *Map[K, V, A]) monoid() CommutativeMonoid[K, V, A] {
 	if t.cfg.monoid == nil {
 		panic("aug: operation requires an Updater created by MonoidUpdater")
 	}
@@ -71,16 +71,15 @@ func (c *config[K, V, A]) foldBefore(acc A, n *Node[K, V, A], pos int, withChild
 }
 
 func (t *Map[K, V, A]) total() A {
-	t.monoid()
-	var zero A
+	m := t.monoid()
 	if t.root == nil {
-		return zero
+		return m.Identity()
 	}
 	return t.root.aug
 }
 
 func (t *Map[K, V, A]) prefix(k K) (prefix A, found bool) {
-	t.monoid()
+	prefix = t.monoid().Identity()
 	n := t.root
 	for n != nil {
 		i, found := n.find(&t.cfg, k)
@@ -98,16 +97,15 @@ func (t *Map[K, V, A]) prefix(k K) (prefix A, found bool) {
 
 func (t *Map[K, V, A]) aggregateRange(lo, hi K) A {
 	m := t.monoid()
-	var zero A
 	if t.root == nil || t.cfg.cmp(lo, hi) >= 0 {
-		return zero
+		return m.Identity()
 	}
 	return t.aggregate(m, t.root, &lo, &hi)
 }
 
 // aggregate folds the entries of the subtree at n with keys in [lo, hi),
 // where a nil bound is unconstrained.
-func (t *Map[K, V, A]) aggregate(m Monoid[K, V, A], n *Node[K, V, A], lo, hi *K) A {
+func (t *Map[K, V, A]) aggregate(m CommutativeMonoid[K, V, A], n *Node[K, V, A], lo, hi *K) A {
 	if lo == nil && hi == nil {
 		return n.aug
 	}
@@ -119,7 +117,7 @@ func (t *Map[K, V, A]) aggregate(m Monoid[K, V, A], n *Node[K, V, A], lo, hi *K)
 	if hi != nil {
 		j, _ = n.find(&t.cfg, *hi) // first entry >= hi
 	}
-	var acc A
+	acc := m.Identity()
 	if n.IsLeaf() {
 		return t.cfg.foldEntries(acc, n, i, j)
 	}
@@ -137,9 +135,8 @@ func (t *Map[K, V, A]) aggregate(m Monoid[K, V, A], n *Node[K, V, A], lo, hi *K)
 }
 
 func (i *Iterator[K, V, A]) prefix() A {
-	i.r.monoid()
 	c := &i.r.cfg
-	var p A
+	p := i.r.monoid().Identity()
 	// Ancestor frames contribute what lies left of the child through which
 	// the iterator descended; the current node contributes what lies left
 	// of its position, including the child at the position, whose entries
@@ -154,12 +151,12 @@ func (i *Iterator[K, V, A]) prefix() A {
 	return p
 }
 
-func (i *Iterator[K, V, A]) seekWhere(pred func(prefix, contribution A) bool) A {
+func (i *Iterator[K, V, A]) seekPrefix(pred func(inclusivePrefix A) bool) (A, bool) {
 	m := i.r.monoid()
 	i.Reset()
-	var p A
+	p := m.Identity()
 	if i.node == nil {
-		return p
+		return p, false
 	}
 	ll := i.lowLevel()
 	for {
@@ -170,28 +167,28 @@ func (i *Iterator[K, V, A]) seekWhere(pred func(prefix, contribution A) bool) A
 		for ; pos <= count; pos++ {
 			if !leaf {
 				ca := n.children[pos].aug
-				if pred(p, ca) {
+				if pred(m.Combine(p, ca)) {
 					break
 				}
 				p = m.Combine(p, ca)
 			}
 			if pos < count {
 				oc := m.Of(n.keys[pos], n.values[pos])
-				if pred(p, oc) {
+				if pred(m.Combine(p, oc)) {
 					i.pos = int16(pos)
-					return p
+					return p, true
 				}
 				p = m.Combine(p, oc)
 			}
 		}
 		if pos > count {
-			// Nothing qualifies in this subtree; leave the iterator past
-			// the end and return the total, as documented.
+			// The predicate never became true; leave the iterator past the
+			// end with the total.
 			for i.s.len() > 0 {
 				ll.Ascend()
 			}
 			i.pos = i.node.Count()
-			return i.r.root.aug
+			return i.r.root.aug, false
 		}
 		i.pos = int16(pos)
 		ll.Descend()