Review fixes: structural augmentation updates, interval semantics, cursor upsert, Verify

CHANGELOG.md

@@ -1,5 +1,51 @@
 # Changelog
 
+## Unreleased
+
+### Fixed
+
+- Splits, merges and rebalances now recompute the restructured node and
+  report the change upward, so augmentations that depend on the shape of
+  the subtree (node counts, heights) stay correct; deleting an absent key
+  that merges nodes is covered. The `Updater` documentation states what
+  an Updater is told.
+- 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.CompareIntervals` is `TieBreak`, consulted only for
+  intervals with equal start keys: overlap searches need intervals ordered
+  by start, and the old field could order them any way.
+- `LowLevelIterator.Config` returns a copy; 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`).

aug/adversarial_test.go

@@ -16,6 +16,7 @@ package aug_test
 
 import (
 	"cmp"
+	"math"
 	"math/rand/v2"
 	"testing"
 
@@ -181,3 +182,78 @@ 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, as the
+// Updater documentation prescribes for such augmentations.
+type nodeCounter[K, V any] struct{}
+
+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.Aug() = 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) 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/config.go

@@ -35,8 +35,9 @@ const MaxDegree = math.MaxInt16 / 2
 // heavier churn should size its own free list with WithFreeList.
 const DefaultFreeListSize = 256
 
-// Config holds a Map's comparison function and Updater. Augmentation code
-// reaches it through LowLevelIterator.Config.
+// Config holds a Map's comparison function and Updater, fixed at
+// construction. Augmentation code reaches a copy through
+// LowLevelIterator.Config.
 type Config[K, V, A any] struct {
 
 	// Updater is used to update the augmentations to the tree.
@@ -47,6 +48,18 @@ type Config[K, V, A any] struct {
 
 // Updater is used to update the augmentation of the node when the subtree
 // changes.
+//
+// The tree tells an Updater about every entry that enters, leaves or is
+// replaced below a node, and recomputes a node (Default) when it is
+// restructured: after a split, merge or rebalance the node whose children
+// changed shape receives Default, and if its augmentation changed the
+// ancestors hear of it through the event that caused the restructuring
+// or through Default. An augmentation that depends only on the entries
+// of the subtree (counts, sums, bounds) therefore needs the incremental
+// cases only; one that depends on the shape of the subtree (node counts,
+// heights) should recompute from its children on every event.
+//
+// An Updater may implement Equaler to give Verify its notion of equality.
 type Updater[K, V, A any] interface {
 
 	// Update should update the augmentation of the passed node, optionally

aug/cursor.go

@@ -199,13 +199,16 @@ func (c *Cursor[K, V, A]) Upsert(k K, v V) (replacedV V, replaced bool) {
 			n = c.node
 			c.pos = int16(i)
 			if found {
+				// Replace the key as well as the value: it compares equal
+				// but need not be identical.
+				prevK := n.keys[i]
 				replacedV = n.values[i]
-				n.values[i] = v
+				n.keys[i], n.values[i] = k, v
 				c.updatePath(UpdateInfo[K, V, A]{
 					Action:        Replacement,
 					RelevantKey:   k,
 					RelevantValue: v,
-					PrevKey:       k,
+					PrevKey:       prevK,
 					PrevValue:     replacedV,
 				})
 				return replacedV, true

aug/cursor_test.go

@@ -467,3 +467,71 @@ func TestCursorWriterWithSnapshotReaders(t *testing.T) {
 		t.Fatal(err)
 	}
 }
+
+// TestCursorComparatorEquivalentKeys uses keys that compare by id only, so
+// an upsert may replace a key with a distinct but equal one, and an
+// augmentation that depends on the key, so a stale key is detectable.
+func TestCursorComparatorEquivalentKeys(t *testing.T) {
+	type wkey struct{ id, weight int }
+	byID := func(a, b wkey) int { return cmp.Compare(a.id, b.id) }
+	weighted := aug.Sum(func(k wkey, v int) int { return k.weight * v })
+	rng := rand.New(rand.NewPCG(21, 22))
+	for _, degree := range []int{2, 16} {
+		m := aug.NewMonoid[wkey, int, int](byID, weighted, aug.WithDegree(degree))
+		type entry struct{ weight, value int }
+		ref := map[int]entry{}
+		c := m.Cursor()
+		check := func() {
+			t.Helper()
+			if err := m.Verify(); err != nil {
+				t.Fatal(err)
+			}
+			want := 0
+			for _, e := range ref {
+				want += e.weight * e.value
+			}
+			if got := m.Total(); got != want {
+				t.Fatalf("Total %d, want %d", got, want)
+			}
+			// The stored key must be the latest one, not just an equal one.
+			for k, v := range m.All() {
+				if e := ref[k.id]; e.weight != k.weight || e.value != v {
+					t.Fatalf("stored {%d %d}:%d, want {%d %d}:%d", k.id, k.weight, v, k.id, e.weight, e.value)
+				}
+			}
+		}
+		for step := range 20000 {
+			id, w, v := rng.IntN(300), 1+rng.IntN(9), rng.IntN(100)
+			switch rng.IntN(5) {
+			case 0:
+				m.Upsert(wkey{id, w}, v)
+				ref[id] = entry{w, v}
+			case 1:
+				c.SeekGE(wkey{id: rng.IntN(300)})
+				c.Upsert(wkey{id, w}, v) // hinted; in place when the leaf allows
+				ref[id] = entry{w, v}
+			case 2:
+				if c.SeekGE(wkey{id: id}) {
+					c.SetValue(v)
+					ref[id] = entry{ref[id].weight, v}
+				}
+			case 3:
+				if c.SeekGE(wkey{id: id}) {
+					c.Delete()
+					delete(ref, id)
+				}
+			case 4:
+				if c.SeekGE(wkey{id: id}) {
+					// Rekey to an equal key with a new weight: in place.
+					old := c.Key()
+					c.Rekey(wkey{id, w})
+					ref[id] = entry{w, ref[old.id].value}
+				}
+			}
+			if step%2000 == 1999 {
+				check()
+			}
+		}
+		check()
+	}
+}

aug/doc.go

@@ -30,6 +30,16 @@
 // writes to the receiver, but the returned clone is independent and may be
 // used from another goroutine immediately.
 //
+// A Map must not be copied by value: a copy shares the root without holding
+// a reference to it, so clearing or writing either corrupts the other. Use
+// the pointer New returns, and Clone for a second handle.
+//
+// Keys, values and augmentations are copied shallowly when a node is
+// copied on write, so a map, its clones and their free-listed nodes may
+// share any references they contain. Treat a key or augmentation that
+// holds a reference as immutable once stored, and have Monoid operations
+// return fresh values rather than modify their arguments.
+//
 // Clear releases the Map's reference to its nodes. Nodes that no other Map
 // references are returned to the free list. A Map that is dropped without
 // Clear leaks nothing to the garbage collector's eyes, but nodes it shared

aug/low_level.go

@@ -12,9 +12,10 @@ func LowLevel[K, V, A any](
 	return it.lowLevel()
 }
 
-// Config returns the Map's config.
-func (i *LowLevelIterator[K, V, A]) Config() *Config[K, V, A] {
-	return &i.r.cfg.Config
+// Config returns a copy of the Map's configuration. A Map's Updater and
+// comparison function are fixed at construction.
+func (i *LowLevelIterator[K, V, A]) Config() Config[K, V, A] {
+	return i.r.cfg.Config
 }
 
 // IncrementPos increments the iterator's position within the current node.

aug/monoid.go

@@ -14,6 +14,8 @@
 
 package aug
 
+import "reflect"
+
 // Monoid describes an augmentation that is the combination of per-entry
 // contributions: the augmentation of a subtree is Combine folded over Of of
 // every entry in it. The zero value of A must be the identity of Combine,
@@ -93,6 +95,9 @@ type countUpdater[K, V any] struct {
 
 func (u *countUpdater[K, V]) monoid() Monoid[K, V, int] { return u.m }
 
+// Equal compares counts.
+func (u *countUpdater[K, V]) Equal(a, b int) bool { return a == b }
+
 func (u *countUpdater[K, V]) Update(n *Node[K, V, int], md UpdateInfo[K, V, int]) bool {
 	a := n.Aug()
 	switch md.Action {
@@ -132,6 +137,15 @@ type monoidUpdater[K, V, A any] struct {
 
 func (u *monoidUpdater[K, V, A]) monoid() Monoid[K, V, A] { return u.m }
 
+// Equal compares aggregates with the Monoid's Equaler, or reflect.DeepEqual
+// without one. Verify uses it.
+func (u *monoidUpdater[K, V, A]) Equal(a, b A) bool {
+	if u.eq != nil {
+		return u.eq.Equal(a, b)
+	}
+	return reflect.DeepEqual(a, b)
+}
+
 func (u *monoidUpdater[K, V, A]) Update(n *Node[K, V, A], md UpdateInfo[K, V, A]) bool {
 	a := n.Aug()
 	prev := *a

aug/monoid_map.go

@@ -88,12 +88,17 @@ func (i *MonoidIterator[K, V, A]) Prefix() A {
 // that prefix. If no entry qualifies the iterator is left past the end and
 // the total is returned.
 //
-// pred must be monotone: if pred(p, x) is false for the aggregate x of a
-// span, it must be false for (p', x') of every entry in the span, where p'
-// is the prefix of that entry. Predicates of the form "prefix combined with
-// x reaches a threshold" have this property, so SeekWhere implements
-// selection by rank (see orderstat) and by cumulative sum. It runs one
-// descent, calling pred and Combine once per child or entry visited.
+// The search descends once, into the first subtree whose aggregate
+// satisfies pred, and never backtracks, so pred must be exact for spans:
+// pred(p, x) for the aggregate x of a span, with p the prefix before the
+// span, must be true if and only if it is true for some entry of the span
+// with that entry's own prefix. Predicates of the form "prefix combined
+// with x reaches a threshold", with a threshold that the running prefix
+// crosses at most once, have this property; they implement selection by
+// rank (see orderstat) and by cumulative sum. A predicate that accepts a
+// span none of whose entries qualifies leaves the iterator past the end
+// without examining later spans. It calls pred and Combine once per child
+// or entry visited.
 func (i *MonoidIterator[K, V, A]) SeekWhere(pred func(prefix, contribution A) bool) A {
 	return i.Iterator.seekWhere(pred)
 }