From 7747b7a63468e15e35fc244aab552ea0f1f2536a Mon Sep 17 00:00:00 2001 From: Andrew Werner Date: Tue, 29 Sep 2026 07:18:10 -0400 Subject: [PATCH 1/7] Fix cursor upsert of equal keys, interval zero bound, scan state on Next/Prev, Verify ref copy - Cursor.Upsert's in-place path kept the old key when the new one compared equal and reported the new key as PrevKey; it now replaces the key and reports the real previous one. Tested with keys compared by id and an augmentation that depends on the key. - The interval augmentation started from the zero endpoint as its bound: wrong for negative endpoints and a panic for comparators that cannot take the zero value. subtreeBound now has an explicit unset state and findUpperBound reports emptiness. Tested with pointer endpoints and negative spans in the property test. - interval.Iterator.Next and Prev end an overlap scan, as the seeks do. - Verify copies only the fields the Updater reads instead of the whole node, whose reference count another goroutine may be changing atomically. A race test verifies a snapshot while the original is cloned, written and cleared. - SeekWhere's documented contract now matches its single-descent algorithm: the predicate must be exact for spans. The aug package doc states that maps must not be copied by value and that keys, values and augmentations are shallow-copied. Each new test fails on v0.2.0. Co-Authored-By: Claude Fable 5.1 --- CHANGELOG.md | 25 +++++++++++++ aug/cursor.go | 7 ++-- aug/cursor_test.go | 68 +++++++++++++++++++++++++++++++++++ aug/doc.go | 10 ++++++ aug/monoid_map.go | 17 +++++---- aug/retention_test.go | 37 +++++++++++++++++++ aug/verify.go | 8 +++-- interval/aug.go | 38 +++++++++++--------- interval/iterator.go | 12 +++++++ interval/property_test.go | 75 ++++++++++++++++++++++++++++++++++++++- 10 files changed, 269 insertions(+), 28 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index a97246d..a3ff250 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,30 @@ # Changelog +## Unreleased + +### Fixed + +- `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 + +- `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`). diff --git a/aug/cursor.go b/aug/cursor.go index 06b5071..5339014 100644 --- a/aug/cursor.go +++ b/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 diff --git a/aug/cursor_test.go b/aug/cursor_test.go index a71de25..80be0c1 100644 --- a/aug/cursor_test.go +++ b/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() + } +} diff --git a/aug/doc.go b/aug/doc.go index 170969f..6c20d33 100644 --- a/aug/doc.go +++ b/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 diff --git a/aug/monoid_map.go b/aug/monoid_map.go index a9e54db..b7d7d30 100644 --- a/aug/monoid_map.go +++ b/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) } diff --git a/aug/retention_test.go b/aug/retention_test.go index 28c74b3..15363fd 100644 --- a/aug/retention_test.go +++ b/aug/retention_test.go @@ -15,6 +15,7 @@ package aug_test import ( + "cmp" "fmt" "math/rand/v2" "runtime" @@ -170,3 +171,39 @@ func TestPointerTypesUnderGCPressure(t *testing.T) { } } } + +// TestVerifyWhileCloningAndClearing runs Verify on a snapshot while the +// original is cloned, written and cleared from another goroutine, which +// changes the reference counts of nodes the snapshot shares. +func TestVerifyWhileCloningAndClearing(t *testing.T) { + m := aug.NewMonoid[int, int, int](cmp.Compare[int], aug.Count[int, int]{}, aug.WithDegree(4)) + for i := range 5000 { + m.Upsert(i, i) + } + snap := m.Clone() + done := make(chan struct{}) + go func() { + defer close(done) + for i := range 300 { + c := m.Clone() + c.Upsert(i, -i) + c.Delete(i + 1) + c.Clear() + } + }() + for { + select { + case <-done: + if err := snap.Verify(); err != nil { + t.Fatal(err) + } + snap.Clear() + m.Clear() + return + default: + if err := snap.Verify(); err != nil { + t.Fatal(err) + } + } + } +} diff --git a/aug/verify.go b/aug/verify.go index 3b1dfe6..c08c8e5 100644 --- a/aug/verify.go +++ b/aug/verify.go @@ -123,9 +123,11 @@ func (v *verifier[K, V, A]) node(n *Node[K, V, A], depth int, isRoot bool, lo, h } } if v.cfg.Updater != nil { - // Recompute on a shallow copy so that nodes shared with snapshots - // being read elsewhere are never written. - tmp := *n + // Recompute on a copy so that nodes shared with snapshots being + // read elsewhere are never written. Copy only what the Updater + // reads: the reference count may be changing under an atomic in + // another goroutine. + tmp := Node[K, V, A]{aug: n.aug, keys: n.keys, values: n.values, children: n.children} v.cfg.Updater.Update(&tmp, UpdateInfo[K, V, A]{}) if !reflect.DeepEqual(n.aug, tmp.aug) { return fmt.Errorf("node at depth %d has stale augmentation %v, recomputed %v", depth, n.aug, tmp.aug) diff --git a/interval/aug.go b/interval/aug.go index 2ca90c6..67a6c64 100644 --- a/interval/aug.go +++ b/interval/aug.go @@ -19,6 +19,7 @@ import "github.com/ajwerner/btree/aug" type subtreeBound[K any] struct { keyBound[K] + set bool // false until the first entry; the zero K is not a bound } type updater[I, K, V any] struct { @@ -40,8 +41,8 @@ func (u *updater[I, K, V]) Update( up = child.keyBound } } - if a.compare(u.cmp, up) < 0 { - a.keyBound = up + if !a.set || a.compare(u.cmp, up) < 0 { + a.keyBound, a.set = up, true return true } return false @@ -52,23 +53,23 @@ func (u *updater[I, K, V]) Update( up = child.keyBound } } - if a.compare(u.cmp, up) == 0 { - a.keyBound = u.findUpperBound(n) - return a.compare(u.cmp, up) != 0 + if a.set && a.compare(u.cmp, up) == 0 { + a.keyBound, a.set = u.findUpperBound(n) + return !a.set || a.compare(u.cmp, up) != 0 } return false case aug.Split: - if a.compare(u.cmp, md.ModifiedOther.keyBound) != 0 && + if a.set && a.compare(u.cmp, md.ModifiedOther.keyBound) != 0 && a.compare(u.cmp, u.upperBound(md.RelevantKey)) != 0 { return false } fallthrough case aug.Default, aug.Replacement: - prev := a.keyBound - a.keyBound = u.findUpperBound(n) - return a.compare(u.cmp, prev) != 0 + prev, prevSet := a.keyBound, a.set + a.keyBound, a.set = u.findUpperBound(n) + return prevSet != a.set || (a.set && a.compare(u.cmp, prev) != 0) default: - panic("") + panic("interval: unknown action") } } @@ -84,8 +85,9 @@ func (up *updater[I, K, V]) upperBound(interval I) keyBound[K] { return keyBound[K]{k: up.end(interval)} } -func (up *updater[I, K, V]) findUpperBound(n *aug.Node[I, V, subtreeBound[K]]) keyBound[K] { - var max keyBound[K] +// findUpperBound recomputes the bound of n's subtree; ok is false for an +// empty node. +func (up *updater[I, K, V]) findUpperBound(n *aug.Node[I, V, subtreeBound[K]]) (max keyBound[K], ok bool) { var setMax bool for i, cnt := int16(0), n.Count(); i < cnt; i++ { ub := up.upperBound(n.Key(i)) @@ -96,13 +98,17 @@ func (up *updater[I, K, V]) findUpperBound(n *aug.Node[I, V, subtreeBound[K]]) k } if !n.IsLeaf() { for i, cnt := int16(0), n.Count(); i <= cnt; i++ { - ub := n.ChildAug(i).keyBound - if max.compare(up.cmp, ub) < 0 { - max = ub + child := n.ChildAug(i) + if !child.set { + continue + } + if !setMax || max.compare(up.cmp, child.keyBound) < 0 { + setMax = true + max = child.keyBound } } } - return max + return max, setMax } func (b keyBound[K]) compare(cmp func(K, K) int, o keyBound[K]) int { diff --git a/interval/iterator.go b/interval/iterator.go index 807e334..e3e8c07 100644 --- a/interval/iterator.go +++ b/interval/iterator.go @@ -133,6 +133,18 @@ func (i *Iterator[I, K, V]) Last() { i.Iterator.Last() } +// Next steps to the next interval, ending any overlap scan. +func (i *Iterator[I, K, V]) Next() { + i.o.reset() + i.Iterator.Next() +} + +// Prev steps to the previous interval, ending any overlap scan. +func (i *Iterator[I, K, V]) Prev() { + i.o.reset() + i.Iterator.Prev() +} + // SeekGE seeks to the first interval greater than or equal to the // provided one, ending any overlap scan. func (i *Iterator[I, K, V]) SeekGE(bounds I) bool { diff --git a/interval/property_test.go b/interval/property_test.go index b62c8a8..ea3969a 100644 --- a/interval/property_test.go +++ b/interval/property_test.go @@ -71,7 +71,7 @@ func compareSpans(a, b span) int { } func randomSpan(rng *rand.Rand) span { - lo := rng.IntN(200) + lo := rng.IntN(300) - 100 // negative endpoints too switch rng.IntN(4) { case 0: return span{lo, lo} @@ -204,3 +204,76 @@ func TestOverlapScanAfterSeek(t *testing.T) { t.Fatalf("scan after seek: %v, want %v", got, want) } } + +// TestOverlapScanEndsOnStep checks that Next and Prev end an overlap scan, +// so NextOverlap afterwards does not apply stale constraints. +func TestOverlapScanEndsOnStep(t *testing.T) { + m := interval.NewSet(spanBounds()) + for _, s := range []span{{0, 1}, {5, 6}, {10, 11}, {15, 16}} { + m.Upsert(s) + } + it := m.Iterator() + it.FirstOverlap(span{10, 16}) + if !it.Valid() || it.Key() != (span{10, 11}) { + t.Fatalf("FirstOverlap at %v", it.Key()) + } + it.Prev() + it.Prev() + if !it.Valid() || it.Key() != (span{0, 1}) { + t.Fatalf("after two Prev at %v", it.Key()) + } + it.NextOverlap() + if it.Valid() { + t.Fatalf("NextOverlap after Prev returned %v", it.Key()) + } + it.FirstOverlap(span{10, 16}) + it.Next() + if !it.Valid() || it.Key() != (span{15, 16}) { + t.Fatalf("Next after FirstOverlap at %v", it.Key()) + } + it.NextOverlap() + if it.Valid() { + t.Fatalf("NextOverlap after Next returned %v", it.Key()) + } +} + +// TestPointerEndpoints uses endpoints whose comparator cannot take the zero +// value, which the augmentation must therefore never compare. +func TestPointerEndpoints(t *testing.T) { + type ps struct{ lo, hi *int } + deref := func(a, b *int) int { return cmp.Compare(*a, *b) } + m := interval.NewSet(interval.Bounds[ps, *int]{ + Compare: deref, + Key: func(s ps) *int { return s.lo }, + End: func(s ps) *int { return s.hi }, + HasEnd: func(s ps) bool { return s.hi != nil }, + }, aug.WithDegree(2)) + mk := func(lo, hi int) ps { + p := ps{lo: &lo} + if hi > lo { + p.hi = &hi + } + return p + } + for i := -50; i < 50; i++ { + m.Upsert(mk(i, i+3)) + m.Upsert(mk(i, i)) // a point + } + if err := m.Verify(); err != nil { + t.Fatal(err) + } + n := 0 + for range m.Overlapping(mk(-10, -8)) { + n++ + } + // Ranges starting at -12..-9 and points -10, -9. + if n != 6 { + t.Fatalf("overlapping %d, want 6", n) + } + for i := -50; i < 50; i += 2 { + m.Delete(mk(i, i+3)) + } + if err := m.Verify(); err != nil { + t.Fatal(err) + } +} From f30ccc2ffe51d578f74820c26ebe1b97f264754c Mon Sep 17 00:00:00 2001 From: Andrew Werner Date: Tue, 29 Sep 2026 07:25:07 -0400 Subject: [PATCH 2/7] Report structural changes to ancestors, 64-bit refcounts, interval points and TieBreak, config copy, Verify equality - A split, merge or rebalance recomputes the restructured node and reports a change upward, so shape-dependent augmentations (node counts, heights) stay correct, including the merge caused by deleting an absent key. Tested with a node-counting Updater; on v0.2.0 the reviewer's degree-2 example leaves the root at 3 instead of 4. - Reference counts are int64. - An interval whose end is not after its start is a point, so leaf matching and subtree pruning agree wherever it sits; reversed intervals are points too. Tested alone in a leaf, after many splits, and reversed. - interval.Bounds.TieBreak replaces CompareIntervals and is consulted only for equal start keys, since overlap searches require start order. Tested with a tie-breaker preferring ends. - LowLevelIterator.Config returns a copy. - Verify compares augmentations with the Updater's Equal when it implements Equaler; MonoidUpdater forwards the Monoid's. Tested with a NaN-propagating maximum. - interval.Cursor, FreeList and NewFreeList are nameable; the package documents its overlap cost honestly. Co-Authored-By: Claude Fable 5.1 --- CHANGELOG.md | 21 ++++++++++ aug/adversarial_test.go | 76 ++++++++++++++++++++++++++++++++++ aug/config.go | 17 +++++++- aug/low_level.go | 7 ++-- aug/monoid.go | 14 +++++++ aug/node.go | 44 ++++++++++++++------ aug/verify.go | 11 ++++- interval/aug.go | 11 ++++- interval/example_test.go | 8 ++-- interval/interval.go | 60 ++++++++++++++++++++------- interval/interval_span_test.go | 10 ++--- interval/interval_test.go | 8 ++-- interval/property_test.go | 65 +++++++++++++++++++++++++++++ 13 files changed, 302 insertions(+), 50 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index a3ff250..dba3b77 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,15 @@ ### 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. @@ -19,6 +28,18 @@ ### 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. diff --git a/aug/adversarial_test.go b/aug/adversarial_test.go index 5c634a2..df2486e 100644 --- a/aug/adversarial_test.go +++ b/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) + } +} diff --git a/aug/config.go b/aug/config.go index c4257a6..b9e10e3 100644 --- a/aug/config.go +++ b/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 diff --git a/aug/low_level.go b/aug/low_level.go index 5420f71..7a4153b 100644 --- a/aug/low_level.go +++ b/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. diff --git a/aug/monoid.go b/aug/monoid.go index 7fa71a3..4bb12d8 100644 --- a/aug/monoid.go +++ b/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 diff --git a/aug/node.go b/aug/node.go index 2ddc5ac..eaacdb2 100644 --- a/aug/node.go +++ b/aug/node.go @@ -28,7 +28,7 @@ import ( // Keys and values are kept in separate slices so that searching a node // touches only keys, which matters when values are large. type Node[K, V, A any] struct { - ref int32 + ref int64 // wide enough that clones dropped without Clear never wrap it aug A keys []K values []V @@ -122,7 +122,7 @@ func mut[K, V, A any]( c *config[K, V, A], n **Node[K, V, A], ) *Node[K, V, A] { - if atomic.LoadInt32(&(*n).ref) == 1 { + if atomic.LoadInt64(&(*n).ref) == 1 { // Exclusive ownership. Can mutate in place. return *n } @@ -140,7 +140,7 @@ func mut[K, V, A any]( // incRef acquires a reference to the node. func (n *Node[K, V, A]) incRef() { - atomic.AddInt32(&n.ref, 1) + atomic.AddInt64(&n.ref, 1) } // decRef releases a reference to the node. If requested, the method @@ -148,7 +148,7 @@ func (n *Node[K, V, A]) incRef() { func (n *Node[K, V, A]) decRef( c *config[K, V, A], recursive bool, ) { - if atomic.AddInt32(&n.ref, -1) > 0 { + if atomic.AddInt64(&n.ref, -1) > 0 { // Other references remain. Can't free. return } @@ -420,15 +420,22 @@ func (n *Node[K, V, A]) insert(c *config[K, V, A], item K, value V) (replacedK K n.insertAt(i, item, value, nil) return replacedK, replacedV, false, n.updateOn(&c.Config, Insertion, item, value, nil) } + // structural records whether splitting a child changed this node's + // augmentation. The entries below this node are unchanged by a split, + // so only an augmentation that depends on the shape of the subtree + // notices; it is reported upward alongside the insertion. + var structural bool if len(n.children[i].keys) >= c.maxEntries { splitK, splitV, splitNode := mut(c, &n.children[i]).split(c, c.maxEntries/2) n.insertAt(i, splitK, splitV, splitNode) + structural = n.update(&c.Config) if cmp := c.cmp(item, n.keys[i]); cmp < 0 { // no change, we want first split node } else if cmp > 0 { i++ // we want second split node } else { - return n.replaceAt(c, i, item, value) + replacedK, replacedV, replaced, changed = n.replaceAt(c, i, item, value) + return replacedK, replacedV, replaced, changed || structural } } replacedK, replacedV, replaced, changed = @@ -440,7 +447,7 @@ func (n *Node[K, V, A]) insert(c *config[K, V, A], item K, value V) (replacedK K changed = n.updateOn(&c.Config, Insertion, item, value, nil) } } - return replacedK, replacedV, replaced, changed + return replacedK, replacedV, replaced, changed || structural } // replaceAt replaces the entry at index i, which has a key equal to item. @@ -461,7 +468,9 @@ func (n *Node[K, V, A]) removeMax(c *config[K, V, A]) (K, V) { // Recurse into max child. i := len(n.keys) if len(n.children[i].keys) <= c.minEntries { - // Child not large enough to remove from. + // Child not large enough to remove from. Ancestors always receive + // the Removal event, so a structural change here needs no separate + // report. n.rebalanceOrMerge(c, i) return n.removeMax(c) // redo } @@ -472,8 +481,10 @@ func (n *Node[K, V, A]) removeMax(c *config[K, V, A]) (K, V) { } // rebalanceOrMerge grows child 'i' to ensure it has sufficient room to remove -// an item from it while keeping it at or above minEntries. -func (n *Node[K, V, A]) rebalanceOrMerge(c *config[K, V, A], i int) { +// an item from it while keeping it at or above minEntries. It reports +// whether this node's augmentation changed as a result; the entries below +// it did not, so only shape-dependent augmentations do. +func (n *Node[K, V, A]) rebalanceOrMerge(c *config[K, V, A], i int) bool { switch { case i > 0 && len(n.children[i-1].keys) > c.minEntries: // Rebalance from left sibling. @@ -587,7 +598,7 @@ func (n *Node[K, V, A]) rebalanceOrMerge(c *config[K, V, A], i int) { child.children = append(child.children, mergeChild.children...) } child.updateOn(&c.Config, Insertion, mergeK, mergeV, mergeChild) - if atomic.LoadInt32(&mergeChild.ref) == 1 { + if atomic.LoadInt64(&mergeChild.ref) == 1 { // We own mergeChild exclusively, so its references to its // children transfer to child. Drop them from mergeChild so // that freeing it does not release them. @@ -604,6 +615,7 @@ func (n *Node[K, V, A]) rebalanceOrMerge(c *config[K, V, A], i int) { mergeChild.decRef(c, true /* recursive */) } } + return n.update(&c.Config) } // remove removes an item from the subtree rooted at this node. Returns the item @@ -622,8 +634,9 @@ func (n *Node[K, V, A]) remove( } if len(n.children[i].keys) <= c.minEntries { // Child not large enough to remove from. - n.rebalanceOrMerge(c, i) - return n.remove(c, item) // redo + structural := n.rebalanceOrMerge(c, i) + outK, outV, found, changed = n.remove(c, item) // redo + return outK, outV, found, changed || structural } child := mut(c, &n.children[i]) if found { @@ -635,7 +648,12 @@ func (n *Node[K, V, A]) remove( // Item is not in this node and child is large enough to remove from. outK, outV, found, changed = child.remove(c, item) if changed { - changed = n.updateOn(&c.Config, Removal, outK, outV, nil) + if found { + changed = n.updateOn(&c.Config, Removal, outK, outV, nil) + } else { + // Nothing was removed but the subtree was restructured. + changed = n.update(&c.Config) + } } return outK, outV, found, changed } diff --git a/aug/verify.go b/aug/verify.go index c08c8e5..8d3f3df 100644 --- a/aug/verify.go +++ b/aug/verify.go @@ -23,6 +23,9 @@ import ( // Verify checks the structural invariants of the tree and returns an error // describing the first violation found. It is intended for tests. Like any // read, it must not run concurrently with a writer of this tree. +// Augmentations are compared with the Updater's Equal method when it +// implements Equaler (MonoidUpdater forwards the Monoid's), otherwise with +// reflect.DeepEqual. func (t *Map[K, V, A]) Verify() error { if t.root == nil { if t.length != 0 { @@ -54,7 +57,7 @@ type verifier[K, V, A any] struct { } func (v *verifier[K, V, A]) node(n *Node[K, V, A], depth int, isRoot bool, lo, hi *K) error { - if r := atomic.LoadInt32(&n.ref); r < 1 { + if r := atomic.LoadInt64(&n.ref); r < 1 { return fmt.Errorf("node at depth %d has ref %d", depth, r) } count := len(n.keys) @@ -129,7 +132,11 @@ func (v *verifier[K, V, A]) node(n *Node[K, V, A], depth int, isRoot bool, lo, h // another goroutine. tmp := Node[K, V, A]{aug: n.aug, keys: n.keys, values: n.values, children: n.children} v.cfg.Updater.Update(&tmp, UpdateInfo[K, V, A]{}) - if !reflect.DeepEqual(n.aug, tmp.aug) { + equal := reflect.DeepEqual(n.aug, tmp.aug) + if eq, ok := v.cfg.Updater.(Equaler[A]); ok { + equal = eq.Equal(n.aug, tmp.aug) + } + if !equal { return fmt.Errorf("node at depth %d has stale augmentation %v, recomputed %v", depth, n.aug, tmp.aug) } } diff --git a/interval/aug.go b/interval/aug.go index 67a6c64..ff8c282 100644 --- a/interval/aug.go +++ b/interval/aug.go @@ -78,11 +78,18 @@ type keyBound[K any] struct { inclusive bool } +// upperBound is the bound on keys an interval covers: its end, exclusive, +// or its start, inclusive, for a point. An interval whose end is not after +// its start is a point, so that leaf matching and subtree pruning agree. func (up *updater[I, K, V]) upperBound(interval I) keyBound[K] { + k := up.key(interval) if !up.hasEnd(interval) { - return keyBound[K]{k: up.key(interval), inclusive: true} + return keyBound[K]{k: k, inclusive: true} } - return keyBound[K]{k: up.end(interval)} + if end := up.end(interval); up.cmp(end, k) > 0 { + return keyBound[K]{k: end} + } + return keyBound[K]{k: k, inclusive: true} } // findUpperBound recomputes the bound of n's subtree; ok is false for an diff --git a/interval/example_test.go b/interval/example_test.go index 097295b..a53e883 100644 --- a/interval/example_test.go +++ b/interval/example_test.go @@ -35,10 +35,10 @@ func (p pair[T]) second() T { return p[1] } func Example() { m := interval.NewSet(interval.Bounds[pair[int], int]{ - Compare: cmp.Compare[int], - Key: pair[int].first, - End: pair[int].second, - CompareIntervals: pair[int].compare, + Compare: cmp.Compare[int], + Key: pair[int].first, + End: pair[int].second, + TieBreak: pair[int].compare, }) for _, p := range []pair[int]{ {1, 2}, {2, 3}, {1, 5}, {0, 6}, {2, 7}, diff --git a/interval/interval.go b/interval/interval.go index 9e1b876..760839a 100644 --- a/interval/interval.go +++ b/interval/interval.go @@ -13,8 +13,16 @@ // permissions and limitations under the License. // Package interval provides copy-on-write ordered maps and sets of -// intervals whose iterators find every interval overlapping a query in -// O(log n + k). +// intervals whose iterators find every interval overlapping a query. +// +// Intervals are ordered by their start key; a query walks the tree +// pruning subtrees whose largest end key precedes the query's start. That +// visits O(log n) nodes plus the ancestors of the k matching intervals, +// which is O(k) when matches are adjacent and up to O(k log(n/k)) when +// they are scattered. +// +// An interval whose end is not after its start (or which has no end, see +// Bounds.HasEnd) is a point containing only its start key. package interval import ( @@ -44,13 +52,17 @@ type Bounds[I, K any] struct { End func(I) K // HasEnd reports whether an interval has an end. It is optional: by - // default an interval whose End is the zero K has none. + // default an interval whose End is the zero K has none. An interval + // without an end, or whose end is not after its start, is a point. HasEnd func(I) bool - // CompareIntervals orders intervals and so defines which intervals a - // Map treats as the same key. It is optional: by default intervals are - // ordered by Key, then points before ranges, then by End. - CompareIntervals func(I, I) int + // TieBreak orders intervals with equal start keys and so defines which + // intervals a Map treats as the same key. Overlap searches rely on + // intervals being ordered by start, so it is consulted only after the + // starts compare equal. It is optional: by default points sort before + // ranges, then ranges by End, and intervals equal by those are the same + // key. + TieBreak func(I, I) int } // Interval is implemented by interval types that expose their own bounds: @@ -81,11 +93,8 @@ func (b Bounds[I, K]) withDefaults() Bounds[I, K] { return b.Compare(b.End(i), zero) != 0 } } - if b.CompareIntervals == nil { - b.CompareIntervals = func(x, y I) int { - if c := b.Compare(b.Key(x), b.Key(y)); c != 0 { - return c - } + if b.TieBreak == nil { + b.TieBreak = func(x, y I) int { xEnd, yEnd := b.HasEnd(x), b.HasEnd(y) switch { case xEnd && yEnd: @@ -102,12 +111,20 @@ func (b Bounds[I, K]) withDefaults() Bounds[I, K] { return b } +// compare orders intervals by start, then by TieBreak. +func (b Bounds[I, K]) compare(x, y I) int { + if c := b.Compare(b.Key(x), b.Key(y)); c != 0 { + return c + } + return b.TieBreak(x, y) +} + // New constructs a Map over intervals described by b. See aug.WithDegree // and aug.WithFreeList for the options. func New[I, K, V any](b Bounds[I, K], opts ...aug.Option) *Map[I, K, V] { b = b.withDefaults() return (*Map[I, K, V])(aug.New[I, V, subtreeBound[K]]( - b.CompareIntervals, + b.compare, &updater[I, K, V]{ cmp: b.Compare, key: b.Key, @@ -178,7 +195,7 @@ func (m *Map[I, K, V]) Iterator() Iterator[I, K, V] { } // Cursor returns a new cursor positioned before the first entry. -func (m *Map[I, K, V]) Cursor() aug.Cursor[I, V, subtreeBound[K]] { return m.a().Cursor() } +func (m *Map[I, K, V]) Cursor() Cursor[I, K, V] { return m.a().Cursor() } // All returns an iterator over every entry in key order. func (m *Map[I, K, V]) All() iter.Seq2[I, V] { return m.a().All() } @@ -198,6 +215,19 @@ func (m *Map[I, K, V]) Min() (k I, v V, ok bool) { return m.a().Min() } // Max returns the entry with the largest key. func (m *Map[I, K, V]) Max() (k I, v V, ok bool) { return m.a().Max() } +// Cursor is a cursor over a Map; see aug.Cursor. +type Cursor[I, K, V any] = aug.Cursor[I, V, subtreeBound[K]] + +// FreeList recycles nodes between interval Maps and Sets with the same +// type parameters. +type FreeList[I, K, V any] = aug.FreeList[I, V, subtreeBound[K]] + +// NewFreeList returns a FreeList that retains up to size nodes. See +// aug.NewFreeList. +func NewFreeList[I, K, V any](size int) FreeList[I, K, V] { + return aug.NewFreeList[I, V, subtreeBound[K]](size) +} + // Set is an ordered set of intervals of type I with bounds of type K whose // iterator provides efficient overlap queries. type Set[I, K any] Map[I, K, struct{}] @@ -262,7 +292,7 @@ func (s *Set[I, K]) Verify() error { return s.m().Verify() } func (s *Set[I, K]) Iterator() Iterator[I, K, struct{}] { return s.m().Iterator() } // Cursor returns a new cursor positioned before the first item. -func (s *Set[I, K]) Cursor() aug.Cursor[I, struct{}, subtreeBound[K]] { return s.m().Cursor() } +func (s *Set[I, K]) Cursor() Cursor[I, K, struct{}] { return s.m().Cursor() } // All returns an iterator over every item in order. func (s *Set[I, K]) All() iter.Seq[I] { return keys(s.m().All()) } diff --git a/interval/interval_span_test.go b/interval/interval_span_test.go index 72c50c0..53cd435 100644 --- a/interval/interval_span_test.go +++ b/interval/interval_span_test.go @@ -196,11 +196,11 @@ type btree = *Map[*latch, Key, struct{}] func makeBTree() btree { return New[*latch, Key, struct{}](Bounds[*latch, Key]{ - Compare: Key.Compare, - Key: func(l *latch) Key { return l.span.key }, - End: func(l *latch) Key { return l.span.endKey }, - HasEnd: func(l *latch) bool { return len(l.span.endKey) > 0 }, - CompareIntervals: compareLatches, + Compare: Key.Compare, + Key: func(l *latch) Key { return l.span.key }, + End: func(l *latch) Key { return l.span.endKey }, + HasEnd: func(l *latch) bool { return len(l.span.endKey) > 0 }, + TieBreak: compareLatches, }) } diff --git a/interval/interval_test.go b/interval/interval_test.go index 5a1ff53..dc7a78f 100644 --- a/interval/interval_test.go +++ b/interval/interval_test.go @@ -62,10 +62,10 @@ func TestIntervalTree(t *testing.T) { } } tree := New[IntInterval, int, struct{}](Bounds[IntInterval, int]{ - Compare: cmp.Compare[int], - Key: IntInterval.Key, - End: IntInterval.End, - CompareIntervals: IntervalCompare[IntInterval](cmp.Compare[int]), + Compare: cmp.Compare[int], + Key: IntInterval.Key, + End: IntInterval.End, + TieBreak: IntervalCompare[IntInterval](cmp.Compare[int]), }) items := []IntInterval{{1, 4}, {2, 5}, {3, 3}, {3, 6}, {4, 7}} for _, item := range items { diff --git a/interval/property_test.go b/interval/property_test.go index ea3969a..608487e 100644 --- a/interval/property_test.go +++ b/interval/property_test.go @@ -277,3 +277,68 @@ func TestPointerEndpoints(t *testing.T) { t.Fatal(err) } } + +// TestEmptyIntervalsArePoints checks that an interval whose end is not +// after its start behaves as a point, wherever it sits in the tree. +func TestEmptyIntervalsArePoints(t *testing.T) { + b := interval.BoundsOf[span](cmp.Compare[int]) + b.HasEnd = func(span) bool { return true } // [3,3) claims an end + m := interval.NewSet(b, aug.WithDegree(2)) + m.Upsert(span{3, 3}) + collect := func(q span) (got []span) { + for s := range m.Overlapping(q) { + got = append(got, s) + } + return got + } + alone := collect(span{3, 4}) + if !slices.Equal(alone, []span{{3, 3}}) { + t.Fatalf("alone in a leaf: %v", alone) + } + if got := collect(span{2, 3}); len(got) != 0 { + t.Fatalf("[2,3) should not cover the point 3: %v", got) + } + // Split the leaf many times over; the answer must not change. + for i := range 200 { + m.Upsert(span{i * 10, i*10 + 1}) + } + if got := collect(span{3, 4}); !slices.Equal(got, []span{{0, 1}, {3, 3}}) && !slices.Equal(got, []span{{3, 3}}) { + t.Fatalf("after splits: %v", got) + } + found := false + for _, s := range collect(span{3, 4}) { + found = found || s == span{3, 3} + } + if !found { + t.Fatal("[3,3) lost after the leaf split") + } + // A reversed interval is also a point at its start. + m.Upsert(span{50, 40}) + if got := collect(span{50, 51}); !slices.Contains(got, span{50, 40}) { + t.Fatalf("reversed interval not found as a point: %v", got) + } + if got := collect(span{41, 49}); slices.Contains(got, span{50, 40}) { + t.Fatalf("reversed interval matched inside its reversed range: %v", got) + } + if err := m.Verify(); err != nil { + t.Fatal(err) + } +} + +// TestTieBreakCannotReorderStarts checks that a tie-breaker preferring +// ends cannot break the start ordering overlap searches rely on. +func TestTieBreakCannotReorderStarts(t *testing.T) { + b := interval.BoundsOf[span](cmp.Compare[int]) + b.TieBreak = func(a, c span) int { return cmp.Compare(a.hi, c.hi) } + m := interval.NewSet(b) + for _, s := range []span{{0, 8}, {10, 11}, {5, 20}} { + m.Upsert(s) + } + var got []span + for s := range m.Overlapping(span{6, 7}) { + got = append(got, s) + } + if want := []span{{0, 8}, {5, 20}}; !slices.Equal(got, want) { + t.Fatalf("got %v, want %v", got, want) + } +} From c4c3858ad7a5101d80d6a1302c7ac7c2063f303c Mon Sep 17 00:00:00 2001 From: Andrew Werner Date: Tue, 29 Sep 2026 07:41:05 -0400 Subject: [PATCH 3/7] API pass: positioning reports validity, real set iterators, Span queries, CommutativeMonoid Iterators and cursors: every positioning method (First, Last, Next, Prev, the seeks, SeekNth) returns whether the iterator is now at an entry, so a range loop is "for ok := it.SeekGE(lo); ok && ...; ok = it.Next()". SeekGE and SeekLT no longer report whether the sought key exists; SeekExact does. Cursor.Upsert returns the replaced key and value like Map.Upsert, Rekey returns the entry it displaced, and Delete documents the deletion loop. Sets: Get(probe) returns the stored item equal to a probe and Delete returns the removed item, since comparator equality defines identity. SetIterator and SetCursor are their own types with Item and no Value, SetValue or phantom value arguments. The btree package documents the comparison contract and NewOrdered's NaN limitation. interval: queries are Spans (HalfOpen, Point) separate from the stored type; Overlaps(span) returns an OverlapIterator whose Next always means the next overlap and Iterator has no overlap mode; Bounds.HasEnd defaults to "every interval has an end, and one whose end is not after its start is a point" instead of the zero-end convention. Empty and reversed query spans overlap nothing. aug: Monoid and Group are CommutativeMonoid and CommutativeGroup with an Identity method, and the zero value is no longer assumed to be the identity. SeekWhere(prefix, contribution) is SeekPrefix(inclusivePrefix) with a predicate that turns true at most once, returning the prefix and whether an entry was found. Node.Aug and ChildAug return values and Node.SetAug is for Updaters, so traversal cannot modify augmentations; LowLevelIterator.Config already returns a copy. MonoidMap is a defined type over Map with no exported embedded field. New panics on a nil comparison function, the zero Map is documented as unusable, and the package states one ownership contract for every collection. Co-Authored-By: Claude Fable 5.1 --- CHANGELOG.md | 29 +++++ README.md | 8 +- aug/adversarial_test.go | 29 ++--- aug/aggregate.go | 37 +++---- aug/config.go | 11 +- aug/cursor.go | 48 +++++---- aug/cursor_test.go | 8 +- aug/doc.go | 19 +++- aug/iterator.go | 127 +++++++++++++--------- aug/low_level.go | 4 +- aug/map.go | 4 +- aug/monoid.go | 155 ++++++++++++++------------ aug/monoid_map.go | 141 +++++++++++++++--------- aug/monoid_test.go | 18 ++-- aug/node.go | 21 ++-- aug/retention_test.go | 2 +- btree.go | 161 ++++++++++++++++++++++++--- btree_example_test.go | 2 +- btree_test.go | 93 +++++++++++++--- docs/design.md | 40 ++++--- interval/aug.go | 1 + interval/example_test.go | 5 +- interval/interval.go | 44 ++++---- interval/interval_span_test.go | 123 +++++++++++---------- interval/interval_test.go | 4 +- interval/iterator.go | 191 +++++++++++++++------------------ interval/property_test.go | 92 +++++++--------- orderstat/bench_test.go | 13 ++- orderstat/example_test.go | 2 +- orderstat/order_stat.go | 189 +++++++++++++++++++++++++++----- orderstat/order_stat_test.go | 10 +- 31 files changed, 1047 insertions(+), 584 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index dba3b77..913ecc0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,35 @@ ## 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 `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. +- 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 now recompute the restructured node and diff --git a/README.md b/README.md index 67f2ed4..771a79e 100644 --- a/README.md +++ b/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,7 +41,7 @@ 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; `Overlapping(span)` is the range-over-func form. ## Order-Statistic Trees @@ -47,7 +49,7 @@ The `orderstat` package provides order-statistic trees: `Rank(key)` (the number ## 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 diff --git a/aug/adversarial_test.go b/aug/adversarial_test.go index df2486e..faccd1e 100644 --- a/aug/adversarial_test.go +++ b/aug/adversarial_test.go @@ -23,10 +23,10 @@ import ( "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)) @@ -36,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) @@ -56,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()) } } } @@ -192,11 +194,11 @@ func (nodeCounter[K, V]) Update(n *aug.Node[K, V, int], _ aug.UpdateInfo[K, V, i count := 1 if !n.IsLeaf() { for i := int16(0); i <= n.Count(); i++ { - count += *n.ChildAug(i) + count += n.ChildAug(i) } } - changed := *n.Aug() != count - *n.Aug() = count + changed := n.Aug() != count + n.SetAug(count) return changed } @@ -238,6 +240,7 @@ func TestStructuralAugmentation(t *testing.T) { // 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) { diff --git a/aug/aggregate.go b/aug/aggregate.go index 5856b4e..2be9a9a 100644 --- a/aug/aggregate.go +++ b/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() diff --git a/aug/config.go b/aug/config.go index b9e10e3..f0d9037 100644 --- a/aug/config.go +++ b/aug/config.go @@ -134,7 +134,7 @@ func (c *Config[K, V, A]) Compare(a, b K) int { return c.cmp(a, b) } type config[K, V, A any] struct { Config[K, V, A] find func(*Node[K, V, A], K) (int, bool) // nil: binary search with cmp - monoid Monoid[K, V, A] // set when Updater came from MonoidUpdater + monoid CommutativeMonoid[K, V, A] // set when Updater came from MonoidUpdater folder Folder[K, V, A] // set when the monoid implements Folder fl FreeList[K, V, A] maxEntries int @@ -197,9 +197,14 @@ func makeConfig[K, V, A any]( } else { fl = NewFreeList[K, V, A](DefaultFreeListSize) } - var m Monoid[K, V, A] + if cmp == nil { + panic("btree: a comparison function is required") + } + var m CommutativeMonoid[K, V, A] var f Folder[K, V, A] - if mu, ok := up.(interface{ monoid() Monoid[K, V, A] }); ok { + if mu, ok := up.(interface { + monoid() CommutativeMonoid[K, V, A] + }); ok { m = mu.monoid() f, _ = m.(Folder[K, V, A]) } diff --git a/aug/cursor.go b/aug/cursor.go index 5339014..63ac60c 100644 --- a/aug/cursor.go +++ b/aug/cursor.go @@ -114,7 +114,17 @@ func (c *Cursor[K, V, A]) SetValue(v V) { // Delete removes the current entry and returns it. The cursor must be // valid. Afterwards the cursor is on the entry that followed the removed -// one, or past the end if there was none. +// one, or past the end if there was none, so a loop that deletes must not +// also call Next for the deleted entry: +// +// for ok := c.First(); ok; { +// if drop(c.Key()) { +// c.Delete() +// ok = c.Valid() +// } else { +// ok = c.Next() +// } +// } func (c *Cursor[K, V, A]) Delete() (K, V) { if !c.Valid() { panic("aug: Delete on an invalid Cursor") @@ -146,12 +156,13 @@ func (c *Cursor[K, V, A]) Delete() (K, V) { return k, v } -// Rekey moves the current entry to key k, keeping its value. The cursor -// must be valid. If another entry has key k it is replaced. Afterwards the -// cursor is on the moved entry. When k still sorts between the entry's -// neighbours the move is in place and the Updater sees a Replacement whose -// key changed; otherwise it sees the Removal and Insertion of the fallback. -func (c *Cursor[K, V, A]) Rekey(k K) { +// Rekey moves the current entry to key k, keeping its value, and returns +// any entry that had key k and was replaced. The cursor must be valid. +// Afterwards the cursor is on the moved entry. When k still sorts between +// the entry's neighbours the move is in place and the Updater sees a +// Replacement whose key changed; otherwise it sees the Removal and +// Insertion of the fallback. +func (c *Cursor[K, V, A]) Rekey(k K) (replacedK K, replacedV V, replaced bool) { if !c.Valid() { panic("aug: Rekey on an invalid Cursor") } @@ -177,19 +188,19 @@ func (c *Cursor[K, V, A]) Rekey(k K) { PrevKey: prev, PrevValue: v, }) - return + return replacedK, replacedV, false } } // Otherwise remove it and re-insert it near the cursor, which needs no // further descent when the new position is in the same leaf. _, v := c.Delete() - c.Upsert(k, v) + return c.Upsert(k, v) } // Upsert inserts or replaces the entry with key k, using the cursor's -// position as a hint. Afterwards the cursor is on that entry. It returns -// the replaced value, if any. -func (c *Cursor[K, V, A]) Upsert(k K, v V) (replacedV V, replaced bool) { +// position as a hint, and returns any replaced entry, as Map.Upsert does. +// Afterwards the cursor is on that entry. +func (c *Cursor[K, V, A]) Upsert(k K, v V) (replacedK K, replacedV V, replaced bool) { t := c.r n := c.node if n != nil && n.IsLeaf() && len(n.keys) < t.cfg.maxEntries { @@ -201,25 +212,24 @@ func (c *Cursor[K, V, A]) Upsert(k K, v V) (replacedV V, replaced bool) { 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] + replacedK, replacedV = n.keys[i], n.values[i] n.keys[i], n.values[i] = k, v c.updatePath(UpdateInfo[K, V, A]{ Action: Replacement, RelevantKey: k, RelevantValue: v, - PrevKey: prevK, + PrevKey: replacedK, PrevValue: replacedV, }) - return replacedV, true + return replacedK, replacedV, true } n.insertAt(i, k, v, nil) t.length++ c.updatePath(UpdateInfo[K, V, A]{Action: Insertion, RelevantKey: k, RelevantValue: v}) - return replacedV, false + return replacedK, replacedV, false } } - _, replacedV, replaced = t.Upsert(k, v) + replacedK, replacedV, replaced = t.Upsert(k, v) c.SeekGE(k) - return replacedV, replaced + return replacedK, replacedV, replaced } diff --git a/aug/cursor_test.go b/aug/cursor_test.go index 80be0c1..edac1f2 100644 --- a/aug/cursor_test.go +++ b/aug/cursor_test.go @@ -167,7 +167,7 @@ func testCursor(t *testing.T, degree int) { k := rng.IntN(1200) v := rng.IntN(1000) pv, had := ref[k] - gotV, replaced := c.Upsert(k, v) + _, gotV, replaced := c.Upsert(k, v) if replaced != had || (had && gotV != pv) { t.Fatalf("step %d: Upsert(%d) = (%d, %v), want (%d, %v)", step, k, gotV, replaced, pv, had) } @@ -511,17 +511,17 @@ func TestCursorComparatorEquivalentKeys(t *testing.T) { 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}) { + if c.SeekExact(wkey{id: id}) { c.SetValue(v) ref[id] = entry{ref[id].weight, v} } case 3: - if c.SeekGE(wkey{id: id}) { + if c.SeekExact(wkey{id: id}) { c.Delete() delete(ref, id) } case 4: - if c.SeekGE(wkey{id: id}) { + if c.SeekExact(wkey{id: id}) { // Rekey to an equal key with a new weight: in place. old := c.Key() c.Rekey(wkey{id, w}) diff --git a/aug/doc.go b/aug/doc.go index 6c20d33..53d68e0 100644 --- a/aug/doc.go +++ b/aug/doc.go @@ -16,7 +16,24 @@ // shared by the btree, orderstat and interval packages and is exported so // that users can define their own augmentations. // -// # Ownership and concurrency +// # Ownership +// +// The same contract holds for every collection in this module: +// +// - New constructs a collection; its zero value is not usable, and it +// must not be copied by value. +// - Clone returns another collection that may be mutated independently +// of the original, in constant time; the two share nodes until one +// writes them. +// - Iterators and cursors borrow the collection they came from and +// observe it as it is when they are positioned; sequence views (All, +// Range, ...) observe it each time they are ranged over. +// - A mutation invalidates every borrowed position on that collection +// except the cursor performing it; positioning again is always safe. +// - Clear empties a collection and releases what it shared; the +// collection stays usable. +// +// # Concurrency // // A Map owns a reference to its root node; nodes are reference counted and // shared between a Map and its clones. A mutation copies each node on the diff --git a/aug/iterator.go b/aug/iterator.go index 0686f32..b4b685b 100644 --- a/aug/iterator.go +++ b/aug/iterator.go @@ -15,11 +15,17 @@ package aug -// Iterator is responsible for search and traversal within a Map. It holds -// the path from the root to its position, so it must not be copied once -// positioned; take a fresh one from Map.Iterator instead. Reset, First, -// Last, the Seek methods and SeekWhere position it from scratch and are safe -// after any mutation of the Map; the other methods are not. +// Iterator is responsible for search and traversal within a Map. Every +// positioning method reports whether the Iterator is now at an entry, so +// that +// +// for ok := it.SeekGE(lo); ok && it.Compare(it.Key(), hi) < 0; ok = it.Next() { +// +// visits a key range. It holds the path from the root to its position, so +// it must not be copied once positioned; take a fresh one from +// Map.Iterator instead. Reset, First, Last and the Seek methods position it +// from scratch and are safe after any mutation of the Map; the other +// methods are not. type Iterator[K, V, A any] struct { r *Map[K, V, A] iterFrame[K, V, A] @@ -42,10 +48,24 @@ func (i *Iterator[K, V, A]) Reset() { i.s.reset() } -// SeekGE seeks to the first key greater than or equal to the provided key -// and reports whether that key is equal to it. If no such key exists the -// Iterator is left past the end. -func (i *Iterator[K, V, A]) SeekGE(key K) (found bool) { +// SeekGE positions the Iterator at the first key greater than or equal to +// key and reports whether there is one. If there is none the Iterator is +// left past the end. +func (i *Iterator[K, V, A]) SeekGE(key K) bool { + i.seekGE(key) + return i.Valid() +} + +// SeekExact positions the Iterator at key and reports whether it exists. +// If it does not, the Iterator is at the first greater key, or past the +// end. +func (i *Iterator[K, V, A]) SeekExact(key K) bool { + return i.seekGE(key) +} + +// seekGE positions the Iterator at the first key >= key and reports +// whether that key is equal to it. +func (i *Iterator[K, V, A]) seekGE(key K) (found bool) { i.Reset() if i.node == nil { return false @@ -67,30 +87,30 @@ func (i *Iterator[K, V, A]) SeekGE(key K) (found bool) { } } -// SeekGT seeks to the first key greater than the provided key and reports -// whether the provided key exists. If no such key exists the Iterator is -// left past the end. -func (i *Iterator[K, V, A]) SeekGT(key K) (found bool) { - if found = i.SeekGE(key); found { - i.Next() +// SeekGT positions the Iterator at the first key greater than key and +// reports whether there is one. If there is none the Iterator is left past +// the end. +func (i *Iterator[K, V, A]) SeekGT(key K) bool { + if i.seekGE(key) { + return i.Next() } - return found + return i.Valid() } -// SeekLE seeks to the last key less than or equal to the provided key and -// reports whether that key is equal to it. If no such key exists the -// Iterator is left before the beginning. -func (i *Iterator[K, V, A]) SeekLE(key K) (found bool) { - if found = i.SeekGE(key); !found { - i.Prev() +// SeekLE positions the Iterator at the last key less than or equal to key +// and reports whether there is one. If there is none the Iterator is left +// before the beginning. +func (i *Iterator[K, V, A]) SeekLE(key K) bool { + if !i.seekGE(key) { + return i.Prev() } - return found + return true } -// SeekLT seeks to the last key less than the provided key and reports -// whether the provided key exists. If no such key exists the Iterator is -// left before the beginning. -func (i *Iterator[K, V, A]) SeekLT(key K) (found bool) { +// SeekLT positions the Iterator at the last key less than key and reports +// whether there is one. If there is none the Iterator is left before the +// beginning. +func (i *Iterator[K, V, A]) SeekLT(key K) bool { i.Reset() if i.node == nil { return false @@ -100,18 +120,18 @@ func (i *Iterator[K, V, A]) SeekLT(key K) (found bool) { pos, found := i.node.find(&i.r.cfg, key) i.pos = int16(pos) if found || i.node.IsLeaf() { - i.Prev() - return found + return i.Prev() } ll.Descend() } } -// First seeks to the first key in the Map. -func (i *Iterator[K, V, A]) First() { +// First positions the Iterator at the first key and reports whether there +// is one. +func (i *Iterator[K, V, A]) First() bool { i.Reset() if i.node == nil { - return + return false } ll := i.lowLevel() i.pos = 0 // Reset leaves -1; Descend follows children[pos] @@ -119,13 +139,15 @@ func (i *Iterator[K, V, A]) First() { ll.Descend() } i.pos = 0 + return i.Valid() } -// Last seeks to the last key in the Map. -func (i *Iterator[K, V, A]) Last() { +// Last positions the Iterator at the last key and reports whether there is +// one. +func (i *Iterator[K, V, A]) Last() bool { i.Reset() if i.node == nil { - return + return false } ll := i.lowLevel() for !i.node.IsLeaf() { @@ -133,15 +155,16 @@ func (i *Iterator[K, V, A]) Last() { ll.Descend() } i.pos = i.node.Count() - 1 + return i.Valid() } -// Next positions the Iterator to the key immediately following its current -// position. If the Iterator is positioned before the first key (as after -// Reset), Next positions it at the first key. If the Iterator is already -// past the last key, Next leaves it there. -func (i *Iterator[K, V, A]) Next() { +// Next positions the Iterator at the key following its current position +// and reports whether there is one. If the Iterator is before the first +// key (as after Reset), Next positions it at the first key. If it is +// already past the last key, Next leaves it there. +func (i *Iterator[K, V, A]) Next() bool { if i.node == nil { - return + return false } ll := i.lowLevel() if i.node.IsLeaf() { @@ -149,11 +172,11 @@ func (i *Iterator[K, V, A]) Next() { i.pos++ } i.settle() - return + return i.Valid() } if i.pos >= i.node.Count() { // Past the end; stay there. - return + return false } i.pos++ ll.Descend() @@ -162,6 +185,7 @@ func (i *Iterator[K, V, A]) Next() { ll.Descend() } i.pos = 0 + return true } // settle moves an iterator whose position ran off the end of a leaf up to @@ -174,13 +198,13 @@ func (i *Iterator[K, V, A]) settle() { } } -// Prev positions the Iterator to the key immediately preceding its current -// position. If the Iterator is positioned past the last key, Prev positions -// it at the last key. If the Iterator is already before the first key (as -// after Reset), Prev leaves it there. -func (i *Iterator[K, V, A]) Prev() { +// Prev positions the Iterator at the key preceding its current position +// and reports whether there is one. If the Iterator is past the last key, +// Prev positions it at the last key. If it is already before the first key +// (as after Reset), Prev leaves it there. +func (i *Iterator[K, V, A]) Prev() bool { if i.node == nil { - return + return false } ll := i.lowLevel() if i.node.IsLeaf() { @@ -191,11 +215,11 @@ func (i *Iterator[K, V, A]) Prev() { ll.Ascend() i.pos-- } - return + return i.Valid() } if i.pos < 0 { // Before the beginning; stay there. - return + return false } ll.Descend() for !i.node.IsLeaf() { @@ -203,6 +227,7 @@ func (i *Iterator[K, V, A]) Prev() { ll.Descend() } i.pos = i.node.Count() - 1 + return true } // Valid returns whether the Iterator is positioned at a valid position. diff --git a/aug/low_level.go b/aug/low_level.go index 7a4153b..7e5983a 100644 --- a/aug/low_level.go +++ b/aug/low_level.go @@ -61,8 +61,8 @@ func (i *LowLevelIterator[K, V, A]) Frame(d int) (*Node[K, V, A], int16) { // ChildAug returns the augmentation of the child node at the current position. // It is illegal to call if this is a leaf node or there is no child // node at the current position. -func (i *LowLevelIterator[K, V, A]) ChildAug() *A { - return &i.node.children[i.pos].aug +func (i *LowLevelIterator[K, V, A]) ChildAug() A { + return i.node.children[i.pos].aug } // Descend pushes the current position into the iterators stack and diff --git a/aug/map.go b/aug/map.go index c976083..8aa0783 100644 --- a/aug/map.go +++ b/aug/map.go @@ -22,7 +22,9 @@ import ( // Map is an augmented copy-on-write B-tree map from K to V. Each node // carries an augmentation of type A maintained by the Updater given to New. -// See the package documentation for the ownership and concurrency rules. +// The zero Map is not usable; construct one with New or NewOrdered, and do +// not copy it by value. See the package documentation for the ownership +// and concurrency rules. type Map[K, V, A any] struct { root *Node[K, V, A] length int diff --git a/aug/monoid.go b/aug/monoid.go index 4bb12d8..15ebcc7 100644 --- a/aug/monoid.go +++ b/aug/monoid.go @@ -16,19 +16,22 @@ 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, -// and Combine must be associative and commutative: the Updater folds -// contributions in whatever order the tree operations produce them (a new -// entry is combined onto the right of a node's aggregate wherever it -// sits). Counts, sums, minima, maxima and bounds all qualify; an -// order-dependent aggregate such as concatenation needs a hand-written -// Updater that recomputes nodes. +// CommutativeMonoid 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, starting from Identity. Combine must be +// associative and commutative: the Updater folds contributions in whatever +// order the tree operations produce them. Counts, sums, minima, maxima and +// bounds all qualify; an order-dependent aggregate such as concatenation +// needs a hand-written Updater that recomputes nodes. Operations must not +// modify their arguments. // -// A Map whose Updater comes from MonoidUpdater supports Prefix, Aggregate, -// Total, Iterator.Prefix and Iterator.SeekWhere. -type Monoid[K, V, A any] interface { +// A MonoidMap (see NewMonoid) supports Total, Prefix, Aggregate and, on its +// iterators, Prefix and SeekPrefix. +type CommutativeMonoid[K, V, A any] interface { + + // Identity returns the aggregate of no entries: Combine(Identity(), x) + // == x for every x. + Identity() A // Of returns the contribution of a single entry. Of(k K, v V) A @@ -37,25 +40,26 @@ type Monoid[K, V, A any] interface { Combine(a, b A) A } -// Group is a Monoid whose contributions can be removed again. It lets the -// Updater maintain augmentations in O(1) on removal, split and replacement -// instead of recomputing the node. -type Group[K, V, A any] interface { - Monoid[K, V, A] +// CommutativeGroup is a CommutativeMonoid whose contributions can be +// removed again. It lets the Updater maintain augmentations in O(1) on +// removal, split and replacement instead of recomputing the node. +type CommutativeGroup[K, V, A any] interface { + CommutativeMonoid[K, V, A] // Uncombine returns the aggregate a without the contribution b, i.e. // Uncombine(Combine(x, b), b) == x for every x. Uncombine(a, b A) A } -// Equaler may optionally be implemented by a Monoid to let the Updater -// stop propagating an update up the tree when a node's augmentation did not -// change. Without it every update propagates to the root. +// Equaler may optionally be implemented by a CommutativeMonoid to let the +// Updater stop propagating an update up the tree when a node's +// augmentation did not change, and by any Updater to give Verify its +// notion of equality. Without it every update propagates to the root. type Equaler[A any] interface { Equal(a, b A) bool } -// Folder may optionally be implemented by a Monoid to fold spans of a node +// Folder may optionally be implemented by a CommutativeMonoid to fold spans of a node // in one call instead of one Of and Combine call per entry or child. Prefix, // Aggregate and Iterator.Prefix use it when present. type Folder[K, V, A any] interface { @@ -71,17 +75,17 @@ type Folder[K, V, A any] interface { } // MonoidUpdater returns an Updater maintaining the augmentation described -// by m. If m is also a Group, removals, splits and replacements are applied -// incrementally; otherwise they recompute the node from its entries and -// children in O(degree). -func MonoidUpdater[K, V, A any](m Monoid[K, V, A]) Updater[K, V, A] { +// by m. If m is also a CommutativeGroup, removals, splits and replacements +// are applied incrementally; otherwise they recompute the node from its +// entries and children in O(degree). +func MonoidUpdater[K, V, A any](m CommutativeMonoid[K, V, A]) Updater[K, V, A] { if c, ok := any(m).(Count[K, V]); ok { // Count is common enough to deserve an Updater with no interface // calls; the conversion below is a no-op since A is int. return any(&countUpdater[K, V]{m: c}).(Updater[K, V, A]) } u := &monoidUpdater[K, V, A]{m: m} - u.g, _ = m.(Group[K, V, A]) + u.g, _ = m.(CommutativeGroup[K, V, A]) u.eq, _ = m.(Equaler[A]) return u } @@ -93,49 +97,49 @@ type countUpdater[K, V any] struct { m Count[K, V] } -func (u *countUpdater[K, V]) monoid() Monoid[K, V, int] { return u.m } +func (u *countUpdater[K, V]) monoid() CommutativeMonoid[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() + a := n.aug switch md.Action { case Insertion: - *a++ + a++ if md.ModifiedOther != nil { - *a += *md.ModifiedOther + a += *md.ModifiedOther } - return true case Removal: - *a-- + a-- if md.ModifiedOther != nil { - *a -= *md.ModifiedOther + a -= *md.ModifiedOther } - return true case Split: - *a -= 1 + *md.ModifiedOther - return true + a -= 1 + *md.ModifiedOther case Replacement: return false default: - prev := *a count := len(n.keys) for _, c := range n.children { count += c.aug } - *a = count - return prev != count + if count == a { + return false + } + a = count } + n.aug = a + return true } type monoidUpdater[K, V, A any] struct { - m Monoid[K, V, A] - g Group[K, V, A] + m CommutativeMonoid[K, V, A] + g CommutativeGroup[K, V, A] eq Equaler[A] } -func (u *monoidUpdater[K, V, A]) monoid() Monoid[K, V, A] { return u.m } +func (u *monoidUpdater[K, V, A]) monoid() CommutativeMonoid[K, V, A] { return u.m } // Equal compares aggregates with the Monoid's Equaler, or reflect.DeepEqual // without one. Verify uses it. @@ -147,48 +151,49 @@ func (u *monoidUpdater[K, V, A]) Equal(a, b A) bool { } func (u *monoidUpdater[K, V, A]) Update(n *Node[K, V, A], md UpdateInfo[K, V, A]) bool { - a := n.Aug() - prev := *a + prev := n.aug + a := prev switch md.Action { case Insertion: - *a = u.m.Combine(*a, u.m.Of(md.RelevantKey, md.RelevantValue)) + a = u.m.Combine(a, u.m.Of(md.RelevantKey, md.RelevantValue)) if md.ModifiedOther != nil { - *a = u.m.Combine(*a, *md.ModifiedOther) + a = u.m.Combine(a, *md.ModifiedOther) } case Removal: if u.g == nil { - *a = u.recompute(n) + a = u.recompute(n) break } - *a = u.g.Uncombine(*a, u.m.Of(md.RelevantKey, md.RelevantValue)) + a = u.g.Uncombine(a, u.m.Of(md.RelevantKey, md.RelevantValue)) if md.ModifiedOther != nil { - *a = u.g.Uncombine(*a, *md.ModifiedOther) + a = u.g.Uncombine(a, *md.ModifiedOther) } case Split: if u.g == nil { - *a = u.recompute(n) + a = u.recompute(n) break } - *a = u.g.Uncombine(*a, u.m.Of(md.RelevantKey, md.RelevantValue)) - *a = u.g.Uncombine(*a, *md.ModifiedOther) + a = u.g.Uncombine(a, u.m.Of(md.RelevantKey, md.RelevantValue)) + a = u.g.Uncombine(a, *md.ModifiedOther) case Replacement: if u.g == nil { - *a = u.recompute(n) + a = u.recompute(n) break } - *a = u.g.Uncombine(*a, u.m.Of(md.PrevKey, md.PrevValue)) - *a = u.m.Combine(*a, u.m.Of(md.RelevantKey, md.RelevantValue)) + a = u.g.Uncombine(a, u.m.Of(md.PrevKey, md.PrevValue)) + a = u.m.Combine(a, u.m.Of(md.RelevantKey, md.RelevantValue)) default: - *a = u.recompute(n) + a = u.recompute(n) } + n.aug = a if u.eq != nil { - return !u.eq.Equal(prev, *a) + return !u.eq.Equal(prev, a) } return true } func (u *monoidUpdater[K, V, A]) recompute(n *Node[K, V, A]) A { - var acc A + acc := u.m.Identity() leaf := n.IsLeaf() for i, k := range n.keys { if !leaf { @@ -202,10 +207,12 @@ func (u *monoidUpdater[K, V, A]) recompute(n *Node[K, V, A]) A { return acc } -// Count is a Group counting entries; its aggregate over a span is the -// number of entries in it. It is the augmentation of the orderstat package. +// Count is a CommutativeGroup counting entries; its aggregate over a span +// is the number of entries in it. It is the augmentation of the orderstat +// package. type Count[K, V any] struct{} +func (Count[K, V]) Identity() int { return 0 } func (Count[K, V]) Of(K, V) int { return 1 } func (Count[K, V]) Combine(a, b int) int { return a + b } func (Count[K, V]) Uncombine(a, b int) int { return a - b } @@ -230,8 +237,8 @@ type Integer interface { ~uint | ~uint8 | ~uint16 | ~uint32 | ~uint64 | ~uintptr } -// Sum returns a Group summing f over entries. -func Sum[K, V any, N Integer](f func(K, V) N) Group[K, V, N] { +// Sum returns a CommutativeGroup summing f over entries. +func Sum[K, V any, N Integer](f func(K, V) N) CommutativeGroup[K, V, N] { return sum[K, V, N]{f: f} } @@ -239,6 +246,7 @@ type sum[K, V any, N Integer] struct { f func(K, V) N } +func (s sum[K, V, N]) Identity() N { return 0 } func (s sum[K, V, N]) Of(k K, v V) N { return s.f(k, v) } func (s sum[K, V, N]) Combine(a, b N) N { return a + b } func (s sum[K, V, N]) Uncombine(a, b N) N { return a - b } @@ -266,12 +274,13 @@ type Pair[A, B any] struct { B B } -// PairOf returns a Monoid maintaining a and b together. The result is a -// Group if both a and b are Groups, and an Equaler if both are Equalers. -func PairOf[K, V, A, B any](a Monoid[K, V, A], b Monoid[K, V, B]) Monoid[K, V, Pair[A, B]] { +// PairOf returns a CommutativeMonoid maintaining a and b together. The +// result is a CommutativeGroup if both a and b are, and an Equaler if both +// are Equalers. +func PairOf[K, V, A, B any](a CommutativeMonoid[K, V, A], b CommutativeMonoid[K, V, B]) CommutativeMonoid[K, V, Pair[A, B]] { p := pair[K, V, A, B]{a: a, b: b} - ga, aOK := a.(Group[K, V, A]) - gb, bOK := b.(Group[K, V, B]) + ga, aOK := a.(CommutativeGroup[K, V, A]) + gb, bOK := b.(CommutativeGroup[K, V, B]) ea, aEq := a.(Equaler[A]) eb, bEq := b.(Equaler[B]) switch { @@ -287,8 +296,12 @@ func PairOf[K, V, A, B any](a Monoid[K, V, A], b Monoid[K, V, B]) Monoid[K, V, P } type pair[K, V, A, B any] struct { - a Monoid[K, V, A] - b Monoid[K, V, B] + a CommutativeMonoid[K, V, A] + b CommutativeMonoid[K, V, B] +} + +func (p pair[K, V, A, B]) Identity() Pair[A, B] { + return Pair[A, B]{A: p.a.Identity(), B: p.b.Identity()} } func (p pair[K, V, A, B]) Of(k K, v V) Pair[A, B] { @@ -301,8 +314,8 @@ func (p pair[K, V, A, B]) Combine(x, y Pair[A, B]) Pair[A, B] { type pairGroup[K, V, A, B any] struct { pair[K, V, A, B] - ga Group[K, V, A] - gb Group[K, V, B] + ga CommutativeGroup[K, V, A] + gb CommutativeGroup[K, V, B] } func (p pairGroup[K, V, A, B]) Uncombine(x, y Pair[A, B]) Pair[A, B] { diff --git a/aug/monoid_map.go b/aug/monoid_map.go index b7d7d30..c1cdf00 100644 --- a/aug/monoid_map.go +++ b/aug/monoid_map.go @@ -14,61 +14,108 @@ package aug -import "cmp" - -// MonoidMap is a Map whose augmentation is a Monoid, which makes aggregate -// queries available: Total, Prefix and Aggregate on the map, Prefix and -// SeekWhere on its iterators and cursors. It embeds Map and so offers -// everything a Map does. -type MonoidMap[K, V, A any] struct { - Map[K, V, A] -} +import ( + "cmp" + "iter" +) + +// MonoidMap is a Map whose augmentation is a CommutativeMonoid, which +// makes aggregate queries available: Total, Prefix and Aggregate on the +// map, Prefix and SeekPrefix on its iterators and cursors. It has the same +// representation as Map; the conversion between the two is free. +type MonoidMap[K, V, A any] Map[K, V, A] // NewMonoid constructs a MonoidMap maintaining the augmentation described // by m (see MonoidUpdater). See WithDegree and WithFreeList for the // options. -func NewMonoid[K, V, A any](cmp func(K, K) int, m Monoid[K, V, A], opts ...Option) *MonoidMap[K, V, A] { - return &MonoidMap[K, V, A]{Map: Map[K, V, A]{cfg: makeConfig(cmp, MonoidUpdater(m), opts)}} +func NewMonoid[K, V, A any](cmp func(K, K) int, m CommutativeMonoid[K, V, A], opts ...Option) *MonoidMap[K, V, A] { + return (*MonoidMap[K, V, A])(New(cmp, MonoidUpdater(m), opts...)) } // NewOrderedMonoid is NewMonoid for keys that support <; see NewOrdered. -func NewOrderedMonoid[K cmp.Ordered, V, A any](m Monoid[K, V, A], opts ...Option) *MonoidMap[K, V, A] { - opts = append(opts, withFind(findOrdered[K, V, A])) - return &MonoidMap[K, V, A]{Map: Map[K, V, A]{cfg: makeConfig(cmp.Compare[K], MonoidUpdater(m), opts)}} +func NewOrderedMonoid[K cmp.Ordered, V, A any](m CommutativeMonoid[K, V, A], opts ...Option) *MonoidMap[K, V, A] { + return (*MonoidMap[K, V, A])(NewOrdered(MonoidUpdater(m), opts...)) } -// Clone clones the MonoidMap, lazily. It does so in constant time. -func (t *MonoidMap[K, V, A]) Clone() *MonoidMap[K, V, A] { - return &MonoidMap[K, V, A]{Map: t.Map.clone()} -} +func (t *MonoidMap[K, V, A]) a() *Map[K, V, A] { return (*Map[K, V, A])(t) } // Total returns the aggregate over every entry in the map. -func (t *MonoidMap[K, V, A]) Total() A { - return t.Map.total() -} +func (t *MonoidMap[K, V, A]) Total() A { return t.a().total() } // Prefix returns the aggregate over every entry whose key is less than k, // and whether an entry with key k exists. It runs one descent. -func (t *MonoidMap[K, V, A]) Prefix(k K) (prefix A, found bool) { - return t.Map.prefix(k) -} +func (t *MonoidMap[K, V, A]) Prefix(k K) (prefix A, found bool) { return t.a().prefix(k) } // Aggregate returns the aggregate over every entry whose key is in // [lo, hi). It runs in O(degree * height). -func (t *MonoidMap[K, V, A]) Aggregate(lo, hi K) A { - return t.Map.aggregateRange(lo, hi) +func (t *MonoidMap[K, V, A]) Aggregate(lo, hi K) A { return t.a().aggregateRange(lo, hi) } + +// Clear removes all entries; see Map.Clear. +func (t *MonoidMap[K, V, A]) Clear() { t.a().Clear() } + +// Clone clones the map, lazily. It does so in constant time. +func (t *MonoidMap[K, V, A]) Clone() *MonoidMap[K, V, A] { + return (*MonoidMap[K, V, A])(t.a().Clone()) +} + +// Delete removes the entry with key k, returning it. +func (t *MonoidMap[K, V, A]) Delete(k K) (removedK K, v V, found bool) { return t.a().Delete(k) } + +// Upsert inserts or replaces the entry with key k, returning any replaced +// entry. +func (t *MonoidMap[K, V, A]) Upsert(k K, v V) (replacedK K, replacedV V, replaced bool) { + return t.a().Upsert(k, v) } -// Iterator returns a new MonoidIterator. +// Get returns the value for k, if any. +func (t *MonoidMap[K, V, A]) Get(k K) (v V, ok bool) { return t.a().Get(k) } + +// Len returns the number of entries. +func (t *MonoidMap[K, V, A]) Len() int { return t.a().Len() } + +// Height returns the height of the tree. +func (t *MonoidMap[K, V, A]) Height() int { return t.a().Height() } + +// Degree returns the degree of the tree. +func (t *MonoidMap[K, V, A]) Degree() int { return t.a().Degree() } + +// Compare compares two keys with the map's comparison function. +func (t *MonoidMap[K, V, A]) Compare(a, b K) int { return t.a().Compare(a, b) } + +// String renders the tree in a Newick-like format. +func (t *MonoidMap[K, V, A]) String() string { return t.a().String() } + +// Verify checks the tree's invariants; see Map.Verify. +func (t *MonoidMap[K, V, A]) Verify() error { return t.a().Verify() } + +// Iterator returns a new MonoidIterator positioned before the first entry. func (t *MonoidMap[K, V, A]) Iterator() MonoidIterator[K, V, A] { - return MonoidIterator[K, V, A]{Iterator: t.Map.Iterator()} + return MonoidIterator[K, V, A]{Iterator: t.a().Iterator()} } // Cursor returns a new MonoidCursor positioned before the first entry. func (t *MonoidMap[K, V, A]) Cursor() MonoidCursor[K, V, A] { - return MonoidCursor[K, V, A]{Cursor: t.Map.Cursor()} + return MonoidCursor[K, V, A]{Cursor: t.a().Cursor()} } +// All returns an iterator over every entry in key order. +func (t *MonoidMap[K, V, A]) All() iter.Seq2[K, V] { return t.a().All() } + +// Backward returns an iterator over every entry in reverse key order. +func (t *MonoidMap[K, V, A]) Backward() iter.Seq2[K, V] { return t.a().Backward() } + +// Range returns an iterator over the entries with keys in [lo, hi). +func (t *MonoidMap[K, V, A]) Range(lo, hi K) iter.Seq2[K, V] { return t.a().Range(lo, hi) } + +// From returns an iterator over the entries with keys >= lo. +func (t *MonoidMap[K, V, A]) From(lo K) iter.Seq2[K, V] { return t.a().From(lo) } + +// Min returns the entry with the smallest key. +func (t *MonoidMap[K, V, A]) Min() (k K, v V, ok bool) { return t.a().Min() } + +// Max returns the entry with the largest key. +func (t *MonoidMap[K, V, A]) Max() (k K, v V, ok bool) { return t.a().Max() } + // MonoidIterator is an Iterator over a MonoidMap. type MonoidIterator[K, V, A any] struct { Iterator[K, V, A] @@ -82,25 +129,19 @@ func (i *MonoidIterator[K, V, A]) Prefix() A { return i.Iterator.prefix() } -// SeekWhere positions the iterator at the first entry, in key order, for -// which pred(prefix, contribution) is true, where prefix is the aggregate -// over every entry before it and contribution is Of that entry. It returns -// that prefix. If no entry qualifies the iterator is left past the end and -// the total is returned. +// SeekPrefix positions the iterator at the first entry, in key order, +// whose inclusive prefix satisfies pred, where an entry's inclusive prefix +// is the aggregate over it and every entry before it. It returns the +// aggregate over the entries before that one and whether one was found; +// when none was, the iterator is past the end and the total is returned. // -// 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) +// pred must change from false to true at most once as the prefix grows: +// "the running count exceeds n" or "the running sum reaches s" qualify, so +// SeekPrefix implements selection by rank (see orderstat) and by +// cumulative weight in one descent, calling pred and Combine once per +// child or entry visited. Other predicates are not supported. +func (i *MonoidIterator[K, V, A]) SeekPrefix(pred func(inclusivePrefix A) bool) (prefix A, ok bool) { + return i.Iterator.seekPrefix(pred) } // MonoidCursor is a Cursor over a MonoidMap. @@ -114,7 +155,7 @@ func (c *MonoidCursor[K, V, A]) Prefix() A { return c.Cursor.Iterator.prefix() } -// SeekWhere positions the cursor as MonoidIterator.SeekWhere does. -func (c *MonoidCursor[K, V, A]) SeekWhere(pred func(prefix, contribution A) bool) A { - return c.Cursor.Iterator.seekWhere(pred) +// SeekPrefix positions the cursor as MonoidIterator.SeekPrefix does. +func (c *MonoidCursor[K, V, A]) SeekPrefix(pred func(inclusivePrefix A) bool) (prefix A, ok bool) { + return c.Cursor.Iterator.seekPrefix(pred) } diff --git a/aug/monoid_test.go b/aug/monoid_test.go index 191267e..505f4d4 100644 --- a/aug/monoid_test.go +++ b/aug/monoid_test.go @@ -28,6 +28,7 @@ import ( // maxOf is a Monoid (not a Group) tracking the largest value. type maxOf struct{} +func (maxOf) Identity() int { return 0 } func (maxOf) Of(_ int, v int) int { return v } func (maxOf) Combine(a, b int) int { return max(a, b) } func (maxOf) Equal(a, b int) bool { return a == b } @@ -111,12 +112,15 @@ func testMonoidAggregates(t *testing.T, degree int) { if got := it.Prefix(); got != (stats{}) { t.Fatalf("Prefix after Reset = %v", got) } - // SeekWhere by cumulative sum: first entry at which the running sum + // SeekPrefix by cumulative sum: first entry at which the running sum // of values reaches a threshold. total := statsOf(keys, ref, -1<<31, 1<<31).B.A for range 10 { threshold := rng.IntN(total + 2) - prefix := it.SeekWhere(func(p, c stats) bool { return p.B.A+c.B.A >= threshold }) + prefix, ok := it.SeekPrefix(func(p stats) bool { return p.B.A >= threshold }) + if ok != it.Valid() { + t.Fatalf("SeekPrefix reported %v but Valid is %v", ok, it.Valid()) + } running := 0 wantIdx := -1 for j, k := range keys { @@ -128,18 +132,18 @@ func testMonoidAggregates(t *testing.T, degree int) { } if wantIdx == -1 { if it.Valid() { - t.Fatalf("SeekWhere(sum >= %d) valid at %d, want past end", threshold, it.Key()) + t.Fatalf("SeekPrefix(sum >= %d) valid at %d, want past end", threshold, it.Key()) } if prefix != statsOf(keys, ref, -1<<31, 1<<31) { - t.Fatalf("SeekWhere past end returned prefix %v", prefix) + t.Fatalf("SeekPrefix past end returned prefix %v", prefix) } continue } if !it.Valid() || it.Key() != keys[wantIdx] { - t.Fatalf("SeekWhere(sum >= %d) = %v %d, want %d", threshold, it.Valid(), it.Key(), keys[wantIdx]) + t.Fatalf("SeekPrefix(sum >= %d) = %v %d, want %d", threshold, it.Valid(), it.Key(), keys[wantIdx]) } if want := statsOf(keys[:wantIdx], ref, -1<<31, 1<<31); prefix != want { - t.Fatalf("SeekWhere(sum >= %d) prefix %v, want %v", threshold, prefix, want) + t.Fatalf("SeekPrefix(sum >= %d) prefix %v, want %v", threshold, prefix, want) } } } @@ -202,7 +206,7 @@ func Example_customAugmentation() { fmt.Println("pool 2:", pool2.A, "hosts,", pool2.B, "load") it := m.Iterator() - it.SeekWhere(func(prefix, c load) bool { return prefix.B+c.B >= 200 }) + it.SeekPrefix(func(p load) bool { return p.B >= 200 }) fmt.Println("cumulative load reaches 200 at", it.Key()) // Output: // pool 2: 3 hosts, 150 load diff --git a/aug/node.go b/aug/node.go index eaacdb2..cf89b9d 100644 --- a/aug/node.go +++ b/aug/node.go @@ -35,9 +35,15 @@ type Node[K, V, A any] struct { children []*Node[K, V, A] // empty for leaves } -// Aug returns a pointer to the node's augmentation. -func (n *Node[K, V, A]) Aug() *A { - return &n.aug +// Aug returns the node's augmentation. +func (n *Node[K, V, A]) Aug() A { + return n.aug +} + +// SetAug replaces the node's augmentation. Only an Updater, in its Update +// method, may call it. +func (n *Node[K, V, A]) SetAug(a A) { + n.aug = a } // IsLeaf returns true if the node is a leaf. @@ -61,12 +67,9 @@ func (n *Node[K, V, A]) Value(i int16) V { } // ChildAug returns the augmentation of the child at position i, which must -// be in [0, Count()] for a non-leaf node. It returns nil for a leaf. -func (n *Node[K, V, A]) ChildAug(i int16) *A { - if int(i) < len(n.children) { - return &n.children[i].aug - } - return nil +// be in [0, Count()]. It is illegal to call on a leaf. +func (n *Node[K, V, A]) ChildAug(i int16) A { + return n.children[i].aug } func (c *config[K, V, A]) getNode() *Node[K, V, A] { diff --git a/aug/retention_test.go b/aug/retention_test.go index 15363fd..b932b1b 100644 --- a/aug/retention_test.go +++ b/aug/retention_test.go @@ -140,7 +140,7 @@ func TestPointerTypesUnderGCPressure(t *testing.T) { m.Delete(k) delete(ref, k) case 3: - if c.SeekGE(k) { + if c.SeekExact(k) { c.Delete() delete(ref, k) } diff --git a/btree.go b/btree.go index 97a2dc4..7ad83b7 100644 --- a/btree.go +++ b/btree.go @@ -15,6 +15,20 @@ // Package btree provides copy-on-write ordered maps and sets backed by a // B-tree. See package aug for the ownership and concurrency rules, and the // orderstat and interval packages for augmented variants. +// +// # Comparison +// +// A map or set is ordered, and its entries identified, by a comparison +// function: cmp(a, b) is negative, zero or positive as a sorts before, with +// or after b, and must be a total order (in particular, transitive). Two +// keys that compare equal are the same key, so an Upsert of a key equal to +// a stored one replaces it, and Set.Get returns the stored item for an +// equal probe. Data that the comparison reads must not change while the +// key is stored. +// +// NewOrdered and NewOrderedSet order keys with <, as cmp.Compare does +// except for floating-point NaN, which they cannot place; use New with +// cmp.Compare for keys that may be NaN. package btree import ( @@ -134,16 +148,26 @@ func (s *Set[T]) Upsert(item T) (replaced T, overwrote bool) { return replaced, overwrote } -// Delete removes item, reporting whether it was present. -func (s *Set[T]) Delete(item T) (removed bool) { - _, _, removed = s.m().Delete(item) - return removed +// Delete removes the item equal to item and returns it. +func (s *Set[T]) Delete(item T) (removed T, ok bool) { + removed, _, ok = s.m().Delete(item) + return removed, ok } -// Contains reports whether item is in the set. +// Contains reports whether an item equal to item is in the set. func (s *Set[T]) Contains(item T) bool { - _, ok := s.m().Get(item) - return ok + _, found := s.Get(item) + return found +} + +// Get returns the item in the set equal to probe, if any. Items that +// compare equal may differ in fields the comparison ignores. +func (s *Set[T]) Get(probe T) (item T, found bool) { + it := s.m().a().Iterator() + if it.SeekExact(probe) { + return it.Key(), true + } + return item, false } // Len returns the number of items. @@ -165,10 +189,10 @@ func (s *Set[T]) String() string { return s.m().String() } func (s *Set[T]) Verify() error { return s.m().Verify() } // Iterator returns a new iterator positioned before the first item. -func (s *Set[T]) Iterator() SetIterator[T] { return s.m().Iterator() } +func (s *Set[T]) Iterator() SetIterator[T] { return SetIterator[T]{it: s.m().Iterator()} } // Cursor returns a new cursor positioned before the first item. -func (s *Set[T]) Cursor() SetCursor[T] { return s.m().Cursor() } +func (s *Set[T]) Cursor() SetCursor[T] { return SetCursor[T]{c: s.m().Cursor()} } // All returns an iterator over every item in order. func (s *Set[T]) All() iter.Seq[T] { return keys(s.m().All()) } @@ -231,12 +255,121 @@ func WithFreeList[K, V any](fl FreeList[K, V]) Option { // MapIterator is an iterator for a Map. type MapIterator[K, V any] = aug.Iterator[K, V, struct{}] -// SetIterator is an iterator for a Set. -type SetIterator[T any] = MapIterator[T, struct{}] - // MapCursor is a cursor for a Map: an iterator that can mutate the entry it // is positioned on. See aug.Cursor. type MapCursor[K, V any] = aug.Cursor[K, V, struct{}] -// SetCursor is a cursor for a Set. See aug.Cursor. -type SetCursor[T any] = MapCursor[T, struct{}] +// SetIterator iterates over a set. Every positioning method reports whether +// the iterator is now at an item. +type SetIterator[T any] struct { + it MapIterator[T, struct{}] +} + +// Reset marks the iterator invalid, before the first item. +func (i *SetIterator[T]) Reset() { i.it.Reset() } + +// First positions the iterator at the smallest item. +func (i *SetIterator[T]) First() bool { return i.it.First() } + +// Last positions the iterator at the largest item. +func (i *SetIterator[T]) Last() bool { return i.it.Last() } + +// Next positions the iterator at the following item. +func (i *SetIterator[T]) Next() bool { return i.it.Next() } + +// Prev positions the iterator at the preceding item. +func (i *SetIterator[T]) Prev() bool { return i.it.Prev() } + +// SeekGE positions the iterator at the first item >= item. +func (i *SetIterator[T]) SeekGE(item T) bool { return i.it.SeekGE(item) } + +// SeekGT positions the iterator at the first item > item. +func (i *SetIterator[T]) SeekGT(item T) bool { return i.it.SeekGT(item) } + +// SeekLE positions the iterator at the last item <= item. +func (i *SetIterator[T]) SeekLE(item T) bool { return i.it.SeekLE(item) } + +// SeekLT positions the iterator at the last item < item. +func (i *SetIterator[T]) SeekLT(item T) bool { return i.it.SeekLT(item) } + +// SeekExact positions the iterator at item and reports whether it is in +// the set; if not, the iterator is at the first greater item. +func (i *SetIterator[T]) SeekExact(item T) bool { return i.it.SeekExact(item) } + +// Valid reports whether the iterator is at an item. +func (i *SetIterator[T]) Valid() bool { return i.it.Valid() } + +// Item returns the item at the iterator's position, which must be valid. +func (i *SetIterator[T]) Item() T { return i.it.Key() } + +// Compare compares two items with the set's comparison function. +func (i *SetIterator[T]) Compare(a, b T) int { return i.it.Compare(a, b) } + +// SetCursor is a SetIterator that can also mutate the set at its position; +// see aug.Cursor for the rules. +type SetCursor[T any] struct { + c MapCursor[T, struct{}] +} + +// Reset marks the cursor invalid, before the first item. +func (c *SetCursor[T]) Reset() { c.c.Reset() } + +// First positions the cursor at the smallest item. +func (c *SetCursor[T]) First() bool { return c.c.First() } + +// Last positions the cursor at the largest item. +func (c *SetCursor[T]) Last() bool { return c.c.Last() } + +// Next positions the cursor at the following item. +func (c *SetCursor[T]) Next() bool { return c.c.Next() } + +// Prev positions the cursor at the preceding item. +func (c *SetCursor[T]) Prev() bool { return c.c.Prev() } + +// SeekGE positions the cursor at the first item >= item. +func (c *SetCursor[T]) SeekGE(item T) bool { return c.c.SeekGE(item) } + +// SeekGT positions the cursor at the first item > item. +func (c *SetCursor[T]) SeekGT(item T) bool { return c.c.SeekGT(item) } + +// SeekLE positions the cursor at the last item <= item. +func (c *SetCursor[T]) SeekLE(item T) bool { return c.c.SeekLE(item) } + +// SeekLT positions the cursor at the last item < item. +func (c *SetCursor[T]) SeekLT(item T) bool { return c.c.SeekLT(item) } + +// SeekExact positions the cursor at item and reports whether it is in the +// set; if not, the cursor is at the first greater item. +func (c *SetCursor[T]) SeekExact(item T) bool { return c.c.SeekExact(item) } + +// Valid reports whether the cursor is at an item. +func (c *SetCursor[T]) Valid() bool { return c.c.Valid() } + +// Item returns the item at the cursor's position, which must be valid. +func (c *SetCursor[T]) Item() T { return c.c.Key() } + +// Compare compares two items with the set's comparison function. +func (c *SetCursor[T]) Compare(a, b T) int { return c.c.Compare(a, b) } + +// Delete removes the current item and returns it, leaving the cursor on +// the following item or past the end; see aug.Cursor.Delete. +func (c *SetCursor[T]) Delete() T { + item, _ := c.c.Delete() + return item +} + +// Rekey replaces the current item with item, which may sort elsewhere, +// and returns any other item equal to it that was displaced. Afterwards +// the cursor is on item. +func (c *SetCursor[T]) Rekey(item T) (displaced T, ok bool) { + displaced, _, ok = c.c.Rekey(item) + return displaced, ok +} + +// Upsert inserts item, or replaces the equal item already present and +// returns it, using the cursor's position as a hint. Afterwards the cursor +// is on item. +func (c *SetCursor[T]) Upsert(item T) (replaced T, ok bool) { + replaced, _, ok = c.c.Upsert(item, struct{}{}) + return replaced, ok +} diff --git a/btree_example_test.go b/btree_example_test.go index c1278ab..46d4f34 100644 --- a/btree_example_test.go +++ b/btree_example_test.go @@ -61,7 +61,7 @@ func ExampleMapCursor() { c := fleet.Cursor() for range 3 { c.First() - h := c.Key() + h := c.Item() h.load += 2 c.Rekey(h) fmt.Println("charged", h.name, "to", h.load) diff --git a/btree_test.go b/btree_test.go index dae1f9d..0d7c044 100644 --- a/btree_test.go +++ b/btree_test.go @@ -33,7 +33,7 @@ func TestBTree(t *testing.T) { it.First() expected := []int{1, 2, 12} for _, exp := range expected { - if got := it.Key(); got != exp { + if got := it.Item(); got != exp { t.Fatalf("expected %d, got %d", exp, got) } it.Next() @@ -103,9 +103,10 @@ func checkMap(t *testing.T, tag string, m *Map[int, int], r *model, rng *rand.Ra t.Fatalf("%s: Get(%d) = (%d, %v), want (%d, %v)", tag, k, v, ok, rv, rok) } idx, exists := slices.BinarySearch(keys, k) - if found := it.SeekGE(k); found != exists { - t.Fatalf("%s: SeekGE(%d) found=%v, want %v", tag, k, found, exists) + if found := it.SeekExact(k); found != exists { + t.Fatalf("%s: SeekExact(%d) found=%v, want %v", tag, k, found, exists) } + it.SeekGE(k) if idx < len(keys) { if !it.Valid() || it.Key() != keys[idx] { t.Fatalf("%s: SeekGE(%d) valid=%v cur=%d, want %d", tag, k, it.Valid(), it.Key(), keys[idx]) @@ -216,8 +217,8 @@ func testIteratorEdges(t *testing.T, degree, n int) { } return } - if !it.Valid() || it.Key() != 0 { - t.Fatalf("n=%d: Next after reset gives %v %d", n, it.Valid(), it.Key()) + if !it.Valid() || it.Item() != 0 { + t.Fatalf("n=%d: Next after reset gives %v %d", n, it.Valid(), it.Item()) } // Next past the end stays invalid; Prev then yields the last key. it.Last() @@ -227,8 +228,8 @@ func testIteratorEdges(t *testing.T, degree, n int) { t.Fatalf("n=%d: Next past end is valid", n) } it.Prev() - if !it.Valid() || it.Key() != n-1 { - t.Fatalf("n=%d: Prev after end gives %v %d", n, it.Valid(), it.Key()) + if !it.Valid() || it.Item() != n-1 { + t.Fatalf("n=%d: Prev after end gives %v %d", n, it.Valid(), it.Item()) } // Prev past the beginning stays invalid; Next then yields the first. it.First() @@ -238,8 +239,8 @@ func testIteratorEdges(t *testing.T, degree, n int) { t.Fatalf("n=%d: Prev past beginning is valid", n) } it.Next() - if !it.Valid() || it.Key() != 0 { - t.Fatalf("n=%d: Next after beginning gives %v %d", n, it.Valid(), it.Key()) + if !it.Valid() || it.Item() != 0 { + t.Fatalf("n=%d: Next after beginning gives %v %d", n, it.Valid(), it.Item()) } // SeekGE past the end then Prev yields the last key. it.SeekGE(n + 5) @@ -247,8 +248,8 @@ func testIteratorEdges(t *testing.T, degree, n int) { t.Fatalf("n=%d: SeekGE past end is valid", n) } it.Prev() - if !it.Valid() || it.Key() != n-1 { - t.Fatalf("n=%d: Prev after SeekGE past end gives %v %d", n, it.Valid(), it.Key()) + if !it.Valid() || it.Item() != n-1 { + t.Fatalf("n=%d: Prev after SeekGE past end gives %v %d", n, it.Valid(), it.Item()) } // SeekLT before the beginning then Next yields the first key. it.SeekLT(0) @@ -256,8 +257,8 @@ func testIteratorEdges(t *testing.T, degree, n int) { t.Fatalf("n=%d: SeekLT before beginning is valid", n) } it.Next() - if !it.Valid() || it.Key() != 0 { - t.Fatalf("n=%d: Next after SeekLT before beginning gives %v %d", n, it.Valid(), it.Key()) + if !it.Valid() || it.Item() != 0 { + t.Fatalf("n=%d: Next after SeekLT before beginning gives %v %d", n, it.Valid(), it.Item()) } } } @@ -351,3 +352,69 @@ func FuzzMap(f *testing.F) { } }) } + +// TestPositioningReportsValidity checks that every positioning method +// reports whether the iterator is now at an entry, and that SeekExact is +// the way to ask whether a key exists. +func TestPositioningReportsValidity(t *testing.T) { + s := NewOrderedSet[int]() + s.Upsert(10) + it := s.Iterator() + for _, tc := range []struct { + name string + call func() bool + }{ + {"SeekGT(10)", func() bool { return it.SeekGT(10) }}, + {"SeekLT(10)", func() bool { return it.SeekLT(10) }}, + {"SeekExact(9)", func() bool { return it.SeekExact(9) }}, + } { + if got := tc.call(); got || (it.Valid() && tc.name != "SeekExact(9)") { + t.Fatalf("%s reported %v, valid %v; want false", tc.name, got, it.Valid()) + } + } + for _, tc := range []struct { + name string + call func() bool + }{ + {"SeekGT(9)", func() bool { return it.SeekGT(9) }}, + {"SeekGE(10)", func() bool { return it.SeekGE(10) }}, + {"SeekLE(11)", func() bool { return it.SeekLE(11) }}, + {"SeekExact(10)", func() bool { return it.SeekExact(10) }}, + {"First", it.First}, + {"Last", it.Last}, + } { + if got := tc.call(); !got || !it.Valid() || it.Item() != 10 { + t.Fatalf("%s reported %v, valid %v; want true, true", tc.name, got, it.Valid()) + } + } + if it.Next() || it.Valid() { + t.Fatal("Next past the only entry reported true") + } + if !it.Prev() || it.Item() != 10 { + t.Fatal("Prev from past the end did not report the entry") + } + // The range-loop idiom. + m := NewOrdered[int, string]() + for i := range 10 { + m.Upsert(i, "") + } + mi := m.Iterator() + n := 0 + for ok := mi.SeekGE(3); ok && mi.Key() < 7; ok = mi.Next() { + n++ + } + if n != 4 { + t.Fatalf("range loop visited %d, want 4", n) + } + // Set.Get returns the stored representative of an equal probe. + type rec struct{ id, payload int } + byID := func(a, b rec) int { return cmp.Compare(a.id, b.id) } + rs := NewSet(byID) + rs.Upsert(rec{1, 100}) + if got, ok := rs.Get(rec{1, 0}); !ok || got.payload != 100 { + t.Fatalf("Get = %v %v", got, ok) + } + if removed, ok := rs.Delete(rec{1, 0}); !ok || removed.payload != 100 { + t.Fatalf("Delete = %v %v", removed, ok) + } +} diff --git a/docs/design.md b/docs/design.md index 648b2ff..cf52902 100644 --- a/docs/design.md +++ b/docs/design.md @@ -18,9 +18,10 @@ fix the augmentation and add the queries it enables. Each wrapper is a defined type over the `aug` type (`type Map[K, V any] aug.Map[K, V, struct{}]`) with explicit forwarding methods, so it has the same memory representation, the conversion is free, `New` and `Clone` allocate once, -and a method call is a direct call. Embedding a pointer would have meant a -second allocation and an indirection per call. `aug.MonoidMap` likewise -embeds `Map` by value. +a method call is a direct call, and nothing of the core is reachable +through an exported field. `aug.MonoidMap` is a defined type over +`aug.Map` the same way. `Node.Aug` and `ChildAug` return values; +`Node.SetAug` is for Updaters. ## Nodes @@ -98,9 +99,9 @@ neighbours), and `Default` (recompute from scratch). The return value says wheth whether ancestors need updating; the write paths stop propagating as soon as it is false. -Most augmentations are commutative monoids: an aggregate is `Combine` -folded over `Of` of every entry, in whatever order the operations produce -them. `MonoidUpdater` implements `Updater` for any `Monoid`, in +Most augmentations are commutative monoids (`CommutativeMonoid`: an +`Identity`, `Of` and `Combine`): an aggregate is `Combine` folded over `Of` +of every entry, in whatever order the operations produce them. `MonoidUpdater` implements `Updater` for any `Monoid`, in O(1) per level for `Group`s (which can `Uncombine`) and by recomputing the node otherwise. An optional `Equaler` supplies the changed check and an optional `Folder` folds a span of a node in one call, which matters because @@ -111,16 +112,21 @@ building blocks; orderstat is `Count`. queries, so a plain map cannot expose methods that would panic: `Total`, `Prefix(key)` (aggregate of all smaller keys, one descent) and `Aggregate(lo, hi)` (two boundary paths) on the map, and `Prefix()` at the -position and `SeekWhere(pred)`, one descent to the first entry where a -monotone predicate on the running prefix flips, on its iterators and -cursors. `SeekWhere` pays a predicate and a `Combine` per child -visited; orderstat's `SeekNth` is a native descent by counts for that -reason. +position and `SeekPrefix(pred)`, one descent to the first entry whose +inclusive prefix satisfies a predicate that turns from false to true at +most once, on its iterators and cursors. `SeekPrefix` pays a predicate and +a `Combine` per child visited; orderstat's `SeekNth` is a native descent +by counts for that reason. The contract is deliberately narrow: the +descent never backtracks, so a predicate that accepts a subtree must be +satisfied by some entry in it. The interval bound is a monoid without an inverse whose useful queries are not prefix-shaped, so `interval` implements `Updater` directly and its -overlap scan works on the `LowLevelIterator` (`Descend`, `Ascend`, -`ChildAug`, `Frame`). +overlap scan, `OverlapIterator`, works on the `LowLevelIterator` +(`Descend`, `Ascend`, `ChildAug`, `Frame`). Queries are `Span`s +(`HalfOpen`, `Point`) rather than stored intervals, and the tree orders +intervals by start with `Bounds.TieBreak` deciding only among equal +starts, because the scan relies on start order. Prefixes are computed lazily by walking the iterator's frames. An eager per-frame prefix would make `Rank()` O(1) after `Next` at the cost of a @@ -132,8 +138,12 @@ per-frame prefix would make `Rank()` O(1) after `Next` at the cost of a `Iterator` is a stateful cursor with an inline stack of ten frames (node, position); it spills to the heap only for trees deeper than that, which means degree 2 or 3 at large sizes. `Valid`, `Key`, `Value`, `Next`, `Prev` -and the seeks are the whole surface; `Next` past the end and `Prev` before -the beginning stay invalid, and each steps back in from the other end. +and the seeks are the whole surface. Every positioning method returns +whether the iterator is now at an entry, so the answer to "is the key +present" is `SeekExact`, never the result of `SeekGE`; `Next` past the end +and `Prev` before the beginning stay invalid, and each steps back in from +the other end. Set iterators and cursors are their own types with no +`Value`. `Cursor` embeds `Iterator` and adds writes at the position. The first write pins the path: `mut` on every frame from the root down, updating the child diff --git a/interval/aug.go b/interval/aug.go index ff8c282..d84d225 100644 --- a/interval/aug.go +++ b/interval/aug.go @@ -33,6 +33,7 @@ func (u *updater[I, K, V]) Update( md aug.UpdateInfo[I, V, subtreeBound[K]], ) (updated bool) { a := n.Aug() + defer func() { n.SetAug(a) }() // a is modified below; store it on the way out switch md.Action { case aug.Insertion: up := u.upperBound(md.RelevantKey) diff --git a/interval/example_test.go b/interval/example_test.go index a53e883..1196592 100644 --- a/interval/example_test.go +++ b/interval/example_test.go @@ -45,9 +45,8 @@ func Example() { } { m.Upsert(p) } - it := m.Iterator() - for it.FirstOverlap(pair[int]{4, 5}); it.Valid(); it.NextOverlap() { - fmt.Println(it.Key()) + for p := range m.Overlapping(interval.HalfOpen(4, 5)) { + fmt.Println(p) } // Output: // [0 6] diff --git a/interval/interval.go b/interval/interval.go index 760839a..a79c981 100644 --- a/interval/interval.go +++ b/interval/interval.go @@ -21,8 +21,9 @@ // which is O(k) when matches are adjacent and up to O(k log(n/k)) when // they are scattered. // -// An interval whose end is not after its start (or which has no end, see -// Bounds.HasEnd) is a point containing only its start key. +// A stored interval whose end is not after its start (or which has no +// end, see Bounds.HasEnd) is a point containing only its start key. +// Queries are Spans: HalfOpen(start, end) or Point(key). package interval import ( @@ -52,8 +53,9 @@ type Bounds[I, K any] struct { End func(I) K // HasEnd reports whether an interval has an end. It is optional: by - // default an interval whose End is the zero K has none. An interval - // without an end, or whose end is not after its start, is a point. + // default every interval has one, and an interval whose End is not + // after its Key is a point containing only its Key. A type that marks + // points some other way (a nil end, a flag) provides HasEnd. HasEnd func(I) bool // TieBreak orders intervals with equal start keys and so defines which @@ -88,10 +90,7 @@ func (b Bounds[I, K]) withDefaults() Bounds[I, K] { panic("interval: Bounds.Compare, Key and End are required") } if b.HasEnd == nil { - b.HasEnd = func(i I) bool { - var zero K - return b.Compare(b.End(i), zero) != 0 - } + b.HasEnd = func(I) bool { return true } } if b.TieBreak == nil { b.TieBreak = func(x, y I) int { @@ -139,12 +138,17 @@ func (m *Map[I, K, V]) a() *aug.Map[I, V, subtreeBound[K]] { return (*aug.Map[I, V, subtreeBound[K]])(m) } +// Overlaps returns an OverlapIterator positioned at the first entry whose +// interval overlaps span. +func (m *Map[I, K, V]) Overlaps(span Span[K]) OverlapIterator[I, K, V] { + return newOverlapIterator[I, K, V](m.a(), span) +} + // Overlapping returns an iterator over the entries whose intervals overlap -// bounds, in order of their start keys. -func (m *Map[I, K, V]) Overlapping(bounds I) iter.Seq2[I, V] { +// span, in order of their start keys. +func (m *Map[I, K, V]) Overlapping(span Span[K]) iter.Seq2[I, V] { return func(yield func(I, V) bool) { - it := m.Iterator() - for it.FirstOverlap(bounds); it.Valid(); it.NextOverlap() { + for it := m.Overlaps(span); it.Valid(); it.Next() { if !yield(it.Key(), it.Value()) { return } @@ -190,9 +194,7 @@ func (m *Map[I, K, V]) String() string { return m.a().String() } func (m *Map[I, K, V]) Verify() error { return m.a().Verify() } // Iterator returns a new iterator positioned before the first entry. -func (m *Map[I, K, V]) Iterator() Iterator[I, K, V] { - return Iterator[I, K, V]{Iterator: m.a().Iterator()} -} +func (m *Map[I, K, V]) Iterator() Iterator[I, K, V] { return m.a().Iterator() } // Cursor returns a new cursor positioned before the first entry. func (m *Map[I, K, V]) Cursor() Cursor[I, K, V] { return m.a().Cursor() } @@ -239,10 +241,16 @@ func NewSet[I, K any](b Bounds[I, K], opts ...aug.Option) *Set[I, K] { func (s *Set[I, K]) m() *Map[I, K, struct{}] { return (*Map[I, K, struct{}])(s) } -// Overlapping returns an iterator over the items that overlap bounds, in +// Overlaps returns an OverlapIterator positioned at the first item that +// overlaps span. +func (s *Set[I, K]) Overlaps(span Span[K]) OverlapIterator[I, K, struct{}] { + return s.m().Overlaps(span) +} + +// Overlapping returns an iterator over the items that overlap span, in // order of their start keys. -func (s *Set[I, K]) Overlapping(bounds I) iter.Seq[I] { - return keys(s.m().Overlapping(bounds)) +func (s *Set[I, K]) Overlapping(span Span[K]) iter.Seq[I] { + return keys(s.m().Overlapping(span)) } // Clear removes all items, returning nodes no other set references to the diff --git a/interval/interval_span_test.go b/interval/interval_span_test.go index 53cd435..84a5cb8 100644 --- a/interval/interval_span_test.go +++ b/interval/interval_span_test.go @@ -31,7 +31,7 @@ const maxEntries = 2*aug.DefaultDegree - 1 type Key []byte -type Span struct { +type keySpan struct { key, endKey Key } @@ -63,30 +63,39 @@ func BytesNext(b []byte) []byte { return bn } -func span(i int) Span { +func span(i int) keySpan { switch i % 10 { case 0: - return Span{key: key(i)} + return keySpan{key: key(i)} case 1: - return Span{key: key(i), endKey: key(i).Next()} + return keySpan{key: key(i), endKey: key(i).Next()} case 2: - return Span{key: key(i), endKey: key(i + 64)} + return keySpan{key: key(i), endKey: key(i + 64)} default: - return Span{key: key(i), endKey: key(i + 4)} + return keySpan{key: key(i), endKey: key(i + 4)} } } -func spanWithEnd(start, end int) Span { +// query is the Span asked about for a key span: a point when it has no +// end key. +func query(sp keySpan) Span[Key] { + if len(sp.endKey) == 0 { + return Point(sp.key) + } + return HalfOpen(sp.key, sp.endKey) +} + +func spanWithEnd(start, end int) keySpan { if start < end { - return Span{key: key(start), endKey: key(end)} + return keySpan{key: key(start), endKey: key(end)} } else if start == end { - return Span{key: key(start)} + return keySpan{key: key(start)} } else { panic("illegal span") } } -func spanWithMemo(i int, memo map[int]Span) Span { +func spanWithMemo(i int, memo map[int]keySpan) keySpan { if s, ok := memo[i]; ok { return s } @@ -102,7 +111,7 @@ func newRNG(tb testing.TB) *rand.Rand { return rand.New(rand.NewPCG(seed, 0)) } -func randomSpan(rng *rand.Rand, n int) Span { +func randomSpan(rng *rand.Rand, n int) keySpan { start := rng.IntN(n) end := rng.IntN(n + 1) if end < start { @@ -111,12 +120,12 @@ func randomSpan(rng *rand.Rand, n int) Span { return spanWithEnd(start, end) } -func newLatch(s Span) *latch { +func newLatch(s keySpan) *latch { return &latch{span: s} } type latch struct { - span Span + span keySpan id int } @@ -131,7 +140,7 @@ func (l *latch) End() Key { return l.span.endKey } -func (sp Span) Equal(other Span) bool { +func (sp keySpan) Equal(other keySpan) bool { if bytes.Equal(sp.key, other.key) { return bytes.Equal(sp.endKey, other.endKey) } @@ -140,7 +149,7 @@ func (sp Span) Equal(other Span) bool { type iterator = Iterator[*latch, Key, struct{}] -func checkIter(t *testing.T, it iterator, start, end int, spanMemo map[int]Span) { +func checkIter(t *testing.T, tr btree, it iterator, start, end int, spanMemo map[int]keySpan) { i := start for it.First(); it.Valid(); it.Next() { la := it.Key() @@ -166,9 +175,8 @@ func checkIter(t *testing.T, it iterator, start, end int, spanMemo map[int]Span) t.Fatalf("expected %d, but at %d: %+v", start, i, it) } - all := newLatch(spanWithEnd(start, end)) - for it.FirstOverlap(all); it.Valid(); it.NextOverlap() { - la := it.Key() + for ov := tr.Overlaps(query(spanWithEnd(start, end))); ov.Valid(); ov.Next() { + la := ov.Key() expected := spanWithMemo(i, spanMemo) if !expected.Equal(la.span) { t.Fatalf("expected %s, but found %s", expected, la.span) @@ -206,7 +214,7 @@ func makeBTree() btree { func TestBTree(t *testing.T) { tr := makeBTree() - spanMemo := make(map[int]Span) + spanMemo := make(map[int]keySpan) // With degree == 16 (max-items/node == 31) we need 513 items in order for // there to be 3 levels in the tree. The count here is comfortably above @@ -220,7 +228,7 @@ func TestBTree(t *testing.T) { if e := i + 1; e != tr.Len() { t.Fatalf("expected length %d, but found %d", e, tr.Len()) } - checkIter(t, tr.Iterator(), 0, i+1, spanMemo) + checkIter(t, tr, tr.Iterator(), 0, i+1, spanMemo) } // Delete keys in sorted order. @@ -230,7 +238,7 @@ func TestBTree(t *testing.T) { if e := count - (i + 1); e != tr.Len() { t.Fatalf("expected length %d, but found %d", e, tr.Len()) } - checkIter(t, tr.Iterator(), i+1, count, spanMemo) + checkIter(t, tr, tr.Iterator(), i+1, count, spanMemo) } // Add keys in reverse sorted order. @@ -240,7 +248,7 @@ func TestBTree(t *testing.T) { if e := i + 1; e != tr.Len() { t.Fatalf("expected length %d, but found %d", e, tr.Len()) } - checkIter(t, tr.Iterator(), count-i, count+1, spanMemo) + checkIter(t, tr, tr.Iterator(), count-i, count+1, spanMemo) } // Delete keys in reverse sorted order. @@ -250,7 +258,7 @@ func TestBTree(t *testing.T) { if e := count - (i + 1); e != tr.Len() { t.Fatalf("expected length %d, but found %d", e, tr.Len()) } - checkIter(t, tr.Iterator(), 1, count-i, spanMemo) + checkIter(t, tr, tr.Iterator(), 1, count-i, spanMemo) } } @@ -306,9 +314,9 @@ func TestBTreeSeekOverlap(t *testing.T) { } // Iterate over overlaps with a point scan. - it := tr.Iterator() + var ov OverlapIterator[*latch, Key, struct{}] for i := range count + size { - it.FirstOverlap(newLatch(spanWithEnd(i, i))) + ov = tr.Overlaps(query(spanWithEnd(i, i))) for j := range size + 1 { expStart := i - size + j if expStart < 0 { @@ -318,30 +326,29 @@ func TestBTreeSeekOverlap(t *testing.T) { continue } - if !it.Valid() { + if !ov.Valid() { t.Fatalf("%d/%d: expected valid iterator", i, j) } - la := it.Key() + la := ov.Key() expected := spanWithEnd(expStart, expStart+size+1) if !expected.Equal(la.span) { t.Fatalf("%d: expected %s, but found %s", i, expected, la.span) } - it.NextOverlap() + ov.Next() } - if it.Valid() { - t.Fatalf("%d: expected invalid iterator %v", i, it.Key()) + if ov.Valid() { + t.Fatalf("%d: expected invalid iterator %v", i, ov.Key()) } } - it.FirstOverlap(newLatch(span(count + size + 1))) - if it.Valid() { + ov = tr.Overlaps(query(span(count + size + 1))) + if ov.Valid() { t.Fatalf("expected invalid iterator") } // Iterate over overlaps with a range scan. - it = tr.Iterator() for i := range count + size { - it.FirstOverlap(newLatch(spanWithEnd(i, i+size+1))) + ov = tr.Overlaps(query(spanWithEnd(i, i+size+1))) for j := range 2*size + 1 { expStart := i - size + j if expStart < 0 { @@ -351,23 +358,23 @@ func TestBTreeSeekOverlap(t *testing.T) { continue } - if !it.Valid() { + if !ov.Valid() { t.Fatalf("%d/%d: expected valid iterator", i, j) } - la := it.Key() + la := ov.Key() expected := spanWithEnd(expStart, expStart+size+1) if !expected.Equal(la.span) { t.Fatalf("%d: expected %s, but found %s", i, expected, la.span) } - it.NextOverlap() + ov.Next() } - if it.Valid() { - t.Fatalf("%d: expected invalid iterator %v", i, it.Key()) + if ov.Valid() { + t.Fatalf("%d: expected invalid iterator %v", i, ov.Key()) } } - it.FirstOverlap(newLatch(span(count + size + 1))) - if it.Valid() { + ov = tr.Overlaps(query(span(count + size + 1))) + if ov.Valid() { t.Fatalf("expected invalid iterator") } } @@ -415,11 +422,8 @@ func TestBTreeSeekOverlapRandom(t *testing.T) { } } - it := tr.Iterator() - it.FirstOverlap(scanLa) - for it.Valid() { - found = append(found, it.Key()) - it.NextOverlap() + for ov := tr.Overlaps(query(scanLa.span)); ov.Valid(); ov.Next() { + found = append(found, ov.Key()) } if len(exp) != len(found) { @@ -643,7 +647,7 @@ func BenchmarkBTreeMakeIter(b *testing.B) { func BenchmarkBTreeIterSeekGE(b *testing.B) { forBenchmarkSizes(b, func(b *testing.B, count int) { - var spans []Span + var spans []keySpan tr := makeBTree() for i := range count { @@ -675,7 +679,7 @@ func BenchmarkBTreeIterSeekGE(b *testing.B) { func BenchmarkBTreeIterSeekLT(b *testing.B) { forBenchmarkSizes(b, func(b *testing.B, count int) { - var spans []Span + var spans []keySpan tr := makeBTree() for i := range count { @@ -715,7 +719,7 @@ func BenchmarkBTreeIterSeekLT(b *testing.B) { func BenchmarkBTreeIterFirstOverlap(b *testing.B) { forBenchmarkSizes(b, func(b *testing.B, count int) { - var spans []Span + var spans []keySpan var latches []*latch tr := makeBTree() @@ -728,14 +732,12 @@ func BenchmarkBTreeIterFirstOverlap(b *testing.B) { } rng := newRNG(b) - it := tr.Iterator() b.ResetTimer() for i := 0; i < b.N; i++ { j := rng.IntN(len(spans)) s := spans[j] - la := latches[j] - it.FirstOverlap(la) + it := tr.Overlaps(query(s)) if testing.Verbose() { if !it.Valid() { b.Fatal("expected to find key") @@ -798,14 +800,14 @@ func BenchmarkBTreeIterNextOverlap(b *testing.B) { tr.Upsert(la, struct{}{}) } - allCmd := newLatch(spanWithEnd(0, count+1)) - it := tr.Iterator() + all := query(spanWithEnd(0, count+1)) + ov := tr.Overlaps(all) b.ResetTimer() for i := 0; i < b.N; i++ { - if !it.Valid() { - it.FirstOverlap(allCmd) + if !ov.Valid() { + ov = tr.Overlaps(all) } - it.NextOverlap() + ov.Next() } } @@ -819,14 +821,9 @@ func BenchmarkBTreeIterOverlapScan(b *testing.B) { tr.Upsert(newLatch(spanWithEnd(i, i+size+1)), struct{}{}) } - la := new(latch) b.ResetTimer() for i := 0; i < b.N; i++ { - la.span = randomSpan(rng, count) - it := tr.Iterator() - it.FirstOverlap(la) - for it.Valid() { - it.NextOverlap() + for ov := tr.Overlaps(query(randomSpan(rng, count))); ov.Valid(); ov.Next() { } } } diff --git a/interval/interval_test.go b/interval/interval_test.go index dc7a78f..197fc0e 100644 --- a/interval/interval_test.go +++ b/interval/interval_test.go @@ -92,8 +92,8 @@ func TestIntervalTree(t *testing.T) { }, } { var res []IntInterval - for iter.FirstOverlap(tc.q); iter.Valid(); iter.NextOverlap() { - res = append(res, iter.Key()) + for it := tree.Overlaps(HalfOpen(tc.q[0], tc.q[1])); it.Valid(); it.Next() { + res = append(res, it.Key()) } if !slices.Equal(tc.res, res) { t.Fatalf("expected %v, got %v", tc.res, res) diff --git a/interval/iterator.go b/interval/iterator.go index e3e8c07..8f15ec9 100644 --- a/interval/iterator.go +++ b/interval/iterator.go @@ -21,10 +21,48 @@ import ( "github.com/ajwerner/btree/aug" ) -type Iterator[I, K, V any] struct { - aug.Iterator[I, V, subtreeBound[K]] +// Span is a query: the keys a search asks about. It is separate from the +// stored interval type so that asking about a range does not require +// building a stored object. +type Span[K any] struct { + start, end K + point bool +} + +// HalfOpen returns the span of keys in [start, end). A span whose end is +// not after its start is empty and overlaps nothing. +func HalfOpen[K any](start, end K) Span[K] { + return Span[K]{start: start, end: end} +} + +// Point returns the span containing only key. +func Point[K any](key K) Span[K] { + return Span[K]{start: key, point: true} +} + +// upperBound is the span's bound, in the augmentation's terms. +func (s Span[K]) upperBound() keyBound[K] { + if s.point { + return keyBound[K]{k: s.start, inclusive: true} + } + return keyBound[K]{k: s.end} +} - o overlapScan[I, K, V] +// empty reports whether the span covers no key. +func (s Span[K]) empty(cmp func(K, K) int) bool { + return !s.point && cmp(s.end, s.start) <= 0 +} + +// Iterator is an iterator over a Map in key order; see aug.Iterator. +type Iterator[I, K, V any] = aug.Iterator[I, V, subtreeBound[K]] + +// OverlapIterator visits the entries whose intervals overlap a Span, in +// order of their start keys. Map.Overlaps returns one positioned at the +// first overlap; Next moves to the following one. It borrows the map like +// an Iterator does. +type OverlapIterator[I, K, V any] struct { + it aug.Iterator[I, V, subtreeBound[K]] + o overlapScan[I, K, V] } // An overlap scan is a scan over all intervals that overlap with the provided @@ -72,8 +110,7 @@ type Iterator[I, K, V any] struct { // It does so because the interval at this position is the first interval with a // start key larger than the search range's end key. type overlapScan[I, K, V any] struct { - bounds I - set bool + span Span[K] // The "soft" lower-bound constraint. constrMinN *aug.Node[I, V, subtreeBound[K]] @@ -85,114 +122,53 @@ type overlapScan[I, K, V any] struct { constrMaxPos int16 } -func (o *overlapScan[I, K, V]) reset() { - *o = overlapScan[I, K, V]{} -} - -func (o *overlapScan[I, K, V]) empty() bool { - return !o.set -} - -// FirstOverlap seeks to the first interval in the tree that overlaps with the -// provided search interval. -func (i *Iterator[I, K, V]) FirstOverlap(bounds I) { - i.Reset() - it := lowLevel(i) - it.IncrementPos() - if !i.Valid() { - return +func newOverlapIterator[I, K, V any](m *aug.Map[I, V, subtreeBound[K]], span Span[K]) OverlapIterator[I, K, V] { + i := OverlapIterator[I, K, V]{it: m.Iterator(), o: overlapScan[I, K, V]{span: span}} + ll := aug.LowLevel(&i.it) + cfg := i.cfg() + if span.empty(cfg.cmp) { + return i + } + ll.IncrementPos() + if !i.it.Valid() { + return i } - i.o = overlapScan[I, K, V]{bounds: bounds, set: true} i.constrainMinSearchBounds() i.constrainMaxSearchBounds() i.findNextOverlap() + return i } -func lowLevel[I, K, V any]( - it *Iterator[I, K, V], -) *aug.LowLevelIterator[I, V, subtreeBound[K]] { - return aug.LowLevel(&it.Iterator) +func (i *OverlapIterator[I, K, V]) cfg() *updater[I, K, V] { + return aug.LowLevel(&i.it).Config().Updater.(*updater[I, K, V]) } -// Reset marks the iterator as invalid and clears any state, including an -// overlap scan in progress. -func (i *Iterator[I, K, V]) Reset() { - i.o.reset() - i.Iterator.Reset() -} - -// First seeks to the first interval, ending any overlap scan. -func (i *Iterator[I, K, V]) First() { - i.o.reset() - i.Iterator.First() -} - -// Last seeks to the last interval, ending any overlap scan. -func (i *Iterator[I, K, V]) Last() { - i.o.reset() - i.Iterator.Last() -} - -// Next steps to the next interval, ending any overlap scan. -func (i *Iterator[I, K, V]) Next() { - i.o.reset() - i.Iterator.Next() -} - -// Prev steps to the previous interval, ending any overlap scan. -func (i *Iterator[I, K, V]) Prev() { - i.o.reset() - i.Iterator.Prev() -} - -// SeekGE seeks to the first interval greater than or equal to the -// provided one, ending any overlap scan. -func (i *Iterator[I, K, V]) SeekGE(bounds I) bool { - i.o.reset() - return i.Iterator.SeekGE(bounds) -} - -// SeekGT seeks to the first interval greater than the provided one, ending -// any overlap scan. -func (i *Iterator[I, K, V]) SeekGT(bounds I) bool { - i.o.reset() - return i.Iterator.SeekGT(bounds) -} +// Valid reports whether the iterator is at an overlapping entry. +func (i *OverlapIterator[I, K, V]) Valid() bool { return i.it.Valid() } -// SeekLE seeks to the last interval less than or equal to the provided one, -// ending any overlap scan. -func (i *Iterator[I, K, V]) SeekLE(bounds I) bool { - i.o.reset() - return i.Iterator.SeekLE(bounds) -} +// Key returns the interval at the iterator's position, which must be +// valid. +func (i *OverlapIterator[I, K, V]) Key() I { return i.it.Key() } -// SeekLT seeks to the last interval less than the provided one, ending any -// overlap scan. -func (i *Iterator[I, K, V]) SeekLT(bounds I) bool { - i.o.reset() - return i.Iterator.SeekLT(bounds) -} +// Value returns the value at the iterator's position, which must be valid. +func (i *OverlapIterator[I, K, V]) Value() V { return i.it.Value() } -// NextOverlap positions the iterator to the interval immediately following -// its current position that overlaps with the search interval. -func (i *Iterator[I, K, V]) NextOverlap() { - if !i.Valid() { - return - } - if i.o.empty() { - // Invalid. Mixed overlap scan with non-overlap scan. - i.Reset() - return +// Next positions the iterator at the following overlapping entry and +// reports whether there is one. +func (i *OverlapIterator[I, K, V]) Next() bool { + if !i.it.Valid() { + return false } - lowLevel(i).IncrementPos() + aug.LowLevel(&i.it).IncrementPos() i.findNextOverlap() + return i.it.Valid() } -func (i *Iterator[I, K, V]) constrainMinSearchBounds() { - ll := lowLevel(i) - cfg := ll.Config().Updater.(*updater[I, K, V]) +func (i *OverlapIterator[I, K, V]) constrainMinSearchBounds() { + ll := aug.LowLevel(&i.it) + cfg := i.cfg() cmp := cfg.cmp - k := cfg.key(i.o.bounds) + k := i.o.span.start n := ll.Node() j := sort.Search(int(n.Count()), func(j int) bool { return cmp(k, cfg.key(n.Key(int16(j)))) <= 0 @@ -201,11 +177,11 @@ func (i *Iterator[I, K, V]) constrainMinSearchBounds() { i.o.constrMinPos = int16(j) } -func (i *Iterator[I, K, V]) constrainMaxSearchBounds() { - ll := lowLevel(i) - cfg := ll.Config().Updater.(*updater[I, K, V]) +func (i *OverlapIterator[I, K, V]) constrainMaxSearchBounds() { + ll := aug.LowLevel(&i.it) + cfg := i.cfg() cmp := cfg.cmp - up := cfg.upperBound(i.o.bounds) + up := i.o.span.upperBound() n := ll.Node() j := sort.Search(int(n.Count()), func(j int) bool { return !up.contains(cmp, cfg.key(n.Key(int16(j)))) @@ -214,17 +190,18 @@ func (i *Iterator[I, K, V]) constrainMaxSearchBounds() { i.o.constrMaxPos = int16(j) } -func (i *Iterator[I, K, V]) findNextOverlap() { - ll := lowLevel(i) - cfg := ll.Config().Updater.(*updater[I, K, V]) +func (i *OverlapIterator[I, K, V]) findNextOverlap() { + ll := aug.LowLevel(&i.it) + cfg := i.cfg() cmp := cfg.cmp + start := i.o.span.start for { if ll.Pos() > ll.Node().Count() { // Iterate up tree. ll.Ascend() } else if !ll.Node().IsLeaf() { // Iterate down tree. - if i.o.constrMinReached || ll.ChildAug().contains(cmp, cfg.key(i.o.bounds)) { + if i.o.constrMinReached || ll.ChildAug().contains(cmp, start) { par := ll.Node() pos := ll.Pos() ll.Descend() @@ -243,7 +220,7 @@ func (i *Iterator[I, K, V]) findNextOverlap() { // Check search bounds. if ll.Node() == i.o.constrMaxN && ll.Pos() == i.o.constrMaxPos { // Invalid. Past possible overlaps. - i.Reset() + i.it.Reset() return } if ll.Node() == i.o.constrMinN && ll.Pos() == i.o.constrMinPos { @@ -260,7 +237,7 @@ func (i *Iterator[I, K, V]) findNextOverlap() { // span's start key. return } - if cfg.upperBound(i.Key()).contains(cmp, cfg.key(i.o.bounds)) { + if cfg.upperBound(i.it.Key()).contains(cmp, start) { return } } diff --git a/interval/property_test.go b/interval/property_test.go index 608487e..92421ca 100644 --- a/interval/property_test.go +++ b/interval/property_test.go @@ -38,6 +38,14 @@ func (s span) String() string { return fmt.Sprintf("[%d,%d)", s.lo, s.hi) } +// toSpan is the query for a test span: a point when hi is not after lo. +func toSpan(s span) interval.Span[int] { + if s.hi > s.lo { + return interval.HalfOpen(s.lo, s.hi) + } + return interval.Point(s.lo) +} + func spanBounds() interval.Bounds[span, int] { b := interval.BoundsOf[span](cmp.Compare[int]) b.HasEnd = func(s span) bool { return s.hi > s.lo } @@ -119,7 +127,7 @@ func testOverlapProperty(t *testing.T, degree int) { } slices.SortFunc(want, compareSpans) var got []span - for s := range m.Overlapping(q) { + for s := range m.Overlapping(toSpan(q)) { got = append(got, s) } if !slices.Equal(got, want) { @@ -173,67 +181,38 @@ func testOverlapProperty(t *testing.T, degree int) { check(m, ref) } -// TestOverlapScanAfterSeek checks that a plain seek ends an overlap scan -// in progress, so that a later NextOverlap does not continue with stale -// constraints. -func TestOverlapScanAfterSeek(t *testing.T) { - m := interval.NewSet(spanBounds()) - for i := 0; i < 100; i += 2 { - m.Upsert(span{i, i + 3}) - } - it := m.Iterator() - it.FirstOverlap(span{10, 12}) - if !it.Valid() || it.Key() != (span{8, 11}) { - t.Fatalf("FirstOverlap at %v", it.Key()) - } - it.SeekGE(span{50, 53}) - if !it.Valid() || it.Key() != (span{50, 53}) { - t.Fatalf("SeekGE at %v", it.Key()) - } - // With the scan ended, NextOverlap invalidates rather than scanning. - it.NextOverlap() - if it.Valid() { - t.Fatalf("NextOverlap after SeekGE is valid at %v", it.Key()) - } - // And a fresh scan works. - var got []span - for it.FirstOverlap(span{50, 51}); it.Valid(); it.NextOverlap() { - got = append(got, it.Key()) - } - if want := []span{{48, 51}, {50, 53}}; !slices.Equal(got, want) { - t.Fatalf("scan after seek: %v, want %v", got, want) - } -} - -// TestOverlapScanEndsOnStep checks that Next and Prev end an overlap scan, -// so NextOverlap afterwards does not apply stale constraints. -func TestOverlapScanEndsOnStep(t *testing.T) { +// TestOverlapIteratorIsIndependent checks that an OverlapIterator and a +// plain Iterator on the same map do not disturb each other, and that an +// empty or reversed query span overlaps nothing. +func TestOverlapIteratorIsIndependent(t *testing.T) { m := interval.NewSet(spanBounds()) for _, s := range []span{{0, 1}, {5, 6}, {10, 11}, {15, 16}} { m.Upsert(s) } - it := m.Iterator() - it.FirstOverlap(span{10, 16}) - if !it.Valid() || it.Key() != (span{10, 11}) { - t.Fatalf("FirstOverlap at %v", it.Key()) + ov := m.Overlaps(interval.HalfOpen(10, 16)) + if !ov.Valid() || ov.Key() != (span{10, 11}) { + t.Fatalf("Overlaps at %v", ov.Key()) } + it := m.Iterator() + it.SeekGE(span{5, 6}) it.Prev() - it.Prev() - if !it.Valid() || it.Key() != (span{0, 1}) { - t.Fatalf("after two Prev at %v", it.Key()) + it.Next() + if !ov.Next() || ov.Key() != (span{15, 16}) { + t.Fatalf("Next after plain iteration at valid=%v %v", ov.Valid(), ov.Key()) } - it.NextOverlap() - if it.Valid() { - t.Fatalf("NextOverlap after Prev returned %v", it.Key()) + if ov.Next() || ov.Valid() { + t.Fatalf("Next past the last overlap is valid at %v", ov.Key()) } - it.FirstOverlap(span{10, 16}) - it.Next() - if !it.Valid() || it.Key() != (span{15, 16}) { - t.Fatalf("Next after FirstOverlap at %v", it.Key()) + if !it.Valid() || it.Key() != (span{5, 6}) { + t.Fatalf("plain iterator disturbed: valid=%v %v", it.Valid(), it.Key()) + } + for _, q := range []interval.Span[int]{interval.HalfOpen(5, 5), interval.HalfOpen(6, 5)} { + if e := m.Overlaps(q); e.Valid() { + t.Fatalf("empty span %v overlaps %v", q, e.Key()) + } } - it.NextOverlap() - if it.Valid() { - t.Fatalf("NextOverlap after Next returned %v", it.Key()) + if e := m.Overlaps(interval.Point(5)); !e.Valid() || e.Key() != (span{5, 6}) { + t.Fatalf("Point(5) = valid %v %v", e.Valid(), e.Key()) } } @@ -263,7 +242,8 @@ func TestPointerEndpoints(t *testing.T) { t.Fatal(err) } n := 0 - for range m.Overlapping(mk(-10, -8)) { + lo, hi := -10, -8 + for range m.Overlapping(interval.HalfOpen(&lo, &hi)) { n++ } // Ranges starting at -12..-9 and points -10, -9. @@ -286,7 +266,7 @@ func TestEmptyIntervalsArePoints(t *testing.T) { m := interval.NewSet(b, aug.WithDegree(2)) m.Upsert(span{3, 3}) collect := func(q span) (got []span) { - for s := range m.Overlapping(q) { + for s := range m.Overlapping(toSpan(q)) { got = append(got, s) } return got @@ -335,7 +315,7 @@ func TestTieBreakCannotReorderStarts(t *testing.T) { m.Upsert(s) } var got []span - for s := range m.Overlapping(span{6, 7}) { + for s := range m.Overlapping(interval.HalfOpen(6, 7)) { got = append(got, s) } if want := []span{{0, 8}, {5, 20}}; !slices.Equal(got, want) { diff --git a/orderstat/bench_test.go b/orderstat/bench_test.go index 1e5c2f1..8f0763c 100644 --- a/orderstat/bench_test.go +++ b/orderstat/bench_test.go @@ -73,15 +73,20 @@ func BenchmarkSeekNth(b *testing.B) { } } -func BenchmarkSeekNthViaSeekWhere(b *testing.B) { +func BenchmarkSeekNthViaSeekPrefix(b *testing.B) { for _, n := range []int{10_000, 1_000_000} { - s, p := benchSet(n) + m := New[int, struct{}](cmp.Compare[int]) + rng := rand.New(rand.NewPCG(1, 2)) + p := rng.Perm(n) + for _, k := range p { + m.Upsert(k, struct{}{}) + } b.Run(fmt.Sprintf("n=%d", n), func(b *testing.B) { - it := s.Iterator() + it := m.Iterator() i := 0 for b.Loop() { nth := p[i%n] - it.SeekWhere(func(prefix, contribution int) bool { return prefix+contribution > nth }) + it.SeekPrefix(func(inclusive int) bool { return inclusive > nth }) i++ } }) diff --git a/orderstat/example_test.go b/orderstat/example_test.go index a2a3ce0..1dd92c1 100644 --- a/orderstat/example_test.go +++ b/orderstat/example_test.go @@ -48,7 +48,7 @@ func ExampleIterator_SeekNth() { } it := s.Iterator() for it.SeekNth(1); it.Valid(); it.Next() { - fmt.Println(it.Rank(), it.Key()) + fmt.Println(it.Rank(), it.Item()) } // Output: // 1 fig diff --git a/orderstat/order_stat.go b/orderstat/order_stat.go index e0fe8dc..eb0c257 100644 --- a/orderstat/order_stat.go +++ b/orderstat/order_stat.go @@ -55,8 +55,7 @@ func (m *Map[K, V]) Count(lo, hi K) int { return m.a().Aggregate(lo, hi) } // Nth returns the entry with rank n, i.e. with exactly n entries before it. func (m *Map[K, V]) Nth(n int) (k K, v V, ok bool) { it := m.Iterator() - it.SeekNth(n) - if !it.Valid() { + if !it.SeekNth(n) { return k, v, false } return it.Key(), it.Value(), true @@ -169,16 +168,26 @@ func (s *Set[T]) Upsert(item T) (replaced T, overwrote bool) { return replaced, overwrote } -// Delete removes item, reporting whether it was present. -func (s *Set[T]) Delete(item T) (removed bool) { - _, _, removed = s.m().Delete(item) - return removed +// Delete removes the item equal to item and returns it. +func (s *Set[T]) Delete(item T) (removed T, ok bool) { + removed, _, ok = s.m().Delete(item) + return removed, ok } -// Contains reports whether item is in the set. +// Contains reports whether an item equal to item is in the set. func (s *Set[T]) Contains(item T) bool { - _, ok := s.m().Get(item) - return ok + _, found := s.Get(item) + return found +} + +// Get returns the item in the set equal to probe, if any. Items that +// compare equal may differ in fields the comparison ignores. +func (s *Set[T]) Get(probe T) (item T, found bool) { + it := s.m().a().Iterator() + if it.SeekExact(probe) { + return it.Key(), true + } + return item, false } // Len returns the number of items. @@ -200,10 +209,10 @@ func (s *Set[T]) String() string { return s.m().String() } func (s *Set[T]) Verify() error { return s.m().Verify() } // Iterator returns a new iterator positioned before the first item. -func (s *Set[T]) Iterator() Iterator[T, struct{}] { return s.m().Iterator() } +func (s *Set[T]) Iterator() SetIterator[T] { return SetIterator[T]{it: s.m().Iterator()} } // Cursor returns a new cursor positioned before the first item. -func (s *Set[T]) Cursor() Cursor[T, struct{}] { return s.m().Cursor() } +func (s *Set[T]) Cursor() SetCursor[T] { return SetCursor[T]{c: s.m().Cursor()} } // All returns an iterator over every item in order. func (s *Set[T]) All() iter.Seq[T] { return keys(s.m().All()) } @@ -262,41 +271,40 @@ func (it *Iterator[K, V]) Rank() int { return it.Prefix() } -// SeekNth seeks the iterator to the entry with rank nth, i.e. with exactly -// nth entries before it. If nth is out of range the iterator is left -// invalid. -func (it *Iterator[K, V]) SeekNth(nth int) { - seekNth(aug.LowLevel(&it.Iterator), nth) +// SeekNth positions the iterator at the entry with rank nth, i.e. with +// exactly nth entries before it, and reports whether there is one. +func (it *Iterator[K, V]) SeekNth(nth int) bool { + return seekNth(aug.LowLevel(&it.Iterator), nth) } -// seekNth descends by subtree counts. It is what SeekWhere does with a +// seekNth descends by subtree counts. It is what SeekPrefix does with a // rank predicate, without the calls per child. -func seekNth[K, V any](ll *aug.LowLevelIterator[K, V, int], nth int) { +func seekNth[K, V any](ll *aug.LowLevelIterator[K, V, int], nth int) bool { it := (*aug.Iterator[K, V, int])(ll) it.Reset() n := ll.Node() - if n == nil || nth < 0 || nth >= *n.Aug() { + if n == nil || nth < 0 || nth >= n.Aug() { if n != nil && nth >= 0 { ll.SetPos(n.Count()) } - return + return false } for { n = ll.Node() if n.IsLeaf() { ll.SetPos(int16(nth)) - return + return true } pos := int16(0) for ; ; pos++ { - c := *n.ChildAug(pos) + c := n.ChildAug(pos) if nth < c { break } nth -= c if nth == 0 { ll.SetPos(pos) - return + return true } nth-- } @@ -316,8 +324,135 @@ func (c *Cursor[K, V]) Rank() int { return c.Prefix() } -// SeekNth seeks the cursor to the entry with rank nth. If nth is out of -// range the cursor is left invalid. -func (c *Cursor[K, V]) SeekNth(nth int) { - seekNth(aug.LowLevel(&c.Iterator), nth) +// SeekNth positions the cursor at the entry with rank nth and reports +// whether there is one. +func (c *Cursor[K, V]) SeekNth(nth int) bool { + return seekNth(aug.LowLevel(&c.Iterator), nth) +} + +// SetIterator iterates over a set. Every positioning method reports whether +// the iterator is now at an item. +type SetIterator[T any] struct { + it Iterator[T, struct{}] +} + +// Reset marks the iterator invalid, before the first item. +func (i *SetIterator[T]) Reset() { i.it.Reset() } + +// First positions the iterator at the smallest item. +func (i *SetIterator[T]) First() bool { return i.it.First() } + +// Last positions the iterator at the largest item. +func (i *SetIterator[T]) Last() bool { return i.it.Last() } + +// Next positions the iterator at the following item. +func (i *SetIterator[T]) Next() bool { return i.it.Next() } + +// Prev positions the iterator at the preceding item. +func (i *SetIterator[T]) Prev() bool { return i.it.Prev() } + +// SeekGE positions the iterator at the first item >= item. +func (i *SetIterator[T]) SeekGE(item T) bool { return i.it.SeekGE(item) } + +// SeekGT positions the iterator at the first item > item. +func (i *SetIterator[T]) SeekGT(item T) bool { return i.it.SeekGT(item) } + +// SeekLE positions the iterator at the last item <= item. +func (i *SetIterator[T]) SeekLE(item T) bool { return i.it.SeekLE(item) } + +// SeekLT positions the iterator at the last item < item. +func (i *SetIterator[T]) SeekLT(item T) bool { return i.it.SeekLT(item) } + +// SeekExact positions the iterator at item and reports whether it is in +// the set; if not, the iterator is at the first greater item. +func (i *SetIterator[T]) SeekExact(item T) bool { return i.it.SeekExact(item) } + +// Valid reports whether the iterator is at an item. +func (i *SetIterator[T]) Valid() bool { return i.it.Valid() } + +// Item returns the item at the iterator's position, which must be valid. +func (i *SetIterator[T]) Item() T { return i.it.Key() } + +// Compare compares two items with the set's comparison function. +func (i *SetIterator[T]) Compare(a, b T) int { return i.it.Compare(a, b) } + +// Rank returns the number of items before the iterator's position. +func (i *SetIterator[T]) Rank() int { return i.it.Rank() } + +// SeekNth positions the iterator at the item with rank nth. +func (i *SetIterator[T]) SeekNth(nth int) bool { return i.it.SeekNth(nth) } + +// SetCursor is a SetIterator that can also mutate the set at its position; +// see aug.Cursor for the rules. +type SetCursor[T any] struct { + c Cursor[T, struct{}] +} + +// Reset marks the cursor invalid, before the first item. +func (c *SetCursor[T]) Reset() { c.c.Reset() } + +// First positions the cursor at the smallest item. +func (c *SetCursor[T]) First() bool { return c.c.First() } + +// Last positions the cursor at the largest item. +func (c *SetCursor[T]) Last() bool { return c.c.Last() } + +// Next positions the cursor at the following item. +func (c *SetCursor[T]) Next() bool { return c.c.Next() } + +// Prev positions the cursor at the preceding item. +func (c *SetCursor[T]) Prev() bool { return c.c.Prev() } + +// SeekGE positions the cursor at the first item >= item. +func (c *SetCursor[T]) SeekGE(item T) bool { return c.c.SeekGE(item) } + +// SeekGT positions the cursor at the first item > item. +func (c *SetCursor[T]) SeekGT(item T) bool { return c.c.SeekGT(item) } + +// SeekLE positions the cursor at the last item <= item. +func (c *SetCursor[T]) SeekLE(item T) bool { return c.c.SeekLE(item) } + +// SeekLT positions the cursor at the last item < item. +func (c *SetCursor[T]) SeekLT(item T) bool { return c.c.SeekLT(item) } + +// SeekExact positions the cursor at item and reports whether it is in the +// set; if not, the cursor is at the first greater item. +func (c *SetCursor[T]) SeekExact(item T) bool { return c.c.SeekExact(item) } + +// Valid reports whether the cursor is at an item. +func (c *SetCursor[T]) Valid() bool { return c.c.Valid() } + +// Item returns the item at the cursor's position, which must be valid. +func (c *SetCursor[T]) Item() T { return c.c.Key() } + +// Compare compares two items with the set's comparison function. +func (c *SetCursor[T]) Compare(a, b T) int { return c.c.Compare(a, b) } + +// Delete removes the current item and returns it, leaving the cursor on +// the following item or past the end; see aug.Cursor.Delete. +func (c *SetCursor[T]) Delete() T { + item, _ := c.c.Delete() + return item } + +// Rekey replaces the current item with item, which may sort elsewhere, +// and returns any other item equal to it that was displaced. Afterwards +// the cursor is on item. +func (c *SetCursor[T]) Rekey(item T) (displaced T, ok bool) { + displaced, _, ok = c.c.Rekey(item) + return displaced, ok +} + +// Upsert inserts item, or replaces the equal item already present and +// returns it, using the cursor's position as a hint. Afterwards the cursor +// is on item. +func (c *SetCursor[T]) Upsert(item T) (replaced T, ok bool) { + replaced, _, ok = c.c.Upsert(item, struct{}{}) + return replaced, ok +} + +// Rank returns the number of items before the cursor's position. +func (c *SetCursor[T]) Rank() int { return c.c.Rank() } + +// SeekNth positions the cursor at the item with rank nth. +func (c *SetCursor[T]) SeekNth(nth int) bool { return c.c.SeekNth(nth) } diff --git a/orderstat/order_stat_test.go b/orderstat/order_stat_test.go index 1fbbb9d..e81cdda 100644 --- a/orderstat/order_stat_test.go +++ b/orderstat/order_stat_test.go @@ -74,10 +74,10 @@ func TestOrderStatNth(t *testing.T) { iter := tree.Iterator() for _, idx := range perm { iter.SeekNth(idx) - requireEqual(t, items[idx], iter.Key()) + requireEqual(t, items[idx], iter.Item()) for i := idx + 1; i < N; i++ { iter.Next() - requireEqual(t, items[i], iter.Key()) + requireEqual(t, items[i], iter.Item()) } requireEqual(t, true, iter.Valid()) iter.Next() @@ -98,7 +98,7 @@ func Example_blog() { fmt.Println(s.Len()) it := s.Iterator() it.SeekNth(90) - fmt.Println(it.Key()) + fmt.Println(it.Item()) // Output: // 100 @@ -154,8 +154,8 @@ func testRankAndSeekNth(t *testing.T, degree int) { idx = rng.IntN(n) } it.SeekNth(idx) - if !it.Valid() || it.Key() != keys[idx] { - t.Fatalf("n=%d h=%d: SeekNth(%d) = %v %d, want %d", n, height, idx, it.Valid(), it.Key(), keys[idx]) + if !it.Valid() || it.Item() != keys[idx] { + t.Fatalf("n=%d h=%d: SeekNth(%d) = %v %d, want %d", n, height, idx, it.Valid(), it.Item(), keys[idx]) } if r := it.Rank(); r != idx { t.Fatalf("n=%d h=%d: Rank after SeekNth(%d) = %d", n, height, idx, r) From 2bbaae655f36b08681a9b4e92e69e7dd21c0a6c8 Mon Sep 17 00:00:00 2001 From: Andrew Werner Date: Tue, 29 Sep 2026 08:30:05 -0400 Subject: [PATCH 4/7] interval: drop an unused slice in a benchmark (staticcheck SA4010) Co-Authored-By: Claude Fable 5.1 --- interval/interval_span_test.go | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/interval/interval_span_test.go b/interval/interval_span_test.go index 84a5cb8..b134817 100644 --- a/interval/interval_span_test.go +++ b/interval/interval_span_test.go @@ -720,15 +720,12 @@ func BenchmarkBTreeIterSeekLT(b *testing.B) { func BenchmarkBTreeIterFirstOverlap(b *testing.B) { forBenchmarkSizes(b, func(b *testing.B, count int) { var spans []keySpan - var latches []*latch tr := makeBTree() for i := range count { s := spanWithEnd(i, i+1) spans = append(spans, s) - la := newLatch(s) - latches = append(latches, la) - tr.Upsert(la, struct{}{}) + tr.Upsert(newLatch(s), struct{}{}) } rng := newRNG(b) From 8925e00186e86876a85da02380a43780d2a19022 Mon Sep 17 00:00:00 2001 From: Andrew Werner Date: Tue, 29 Sep 2026 08:54:46 -0400 Subject: [PATCH 5/7] Restructurer opt-in, identity-initialised nodes, interval set API, reusable overlap iterator - Shape-dependent augmentations opt into structural events by implementing Restructurer; the tree no longer recomputes every node after a split, merge or rebalance for augmentations that depend only on the entries, which had made the interval bound recompute run on every merge (Delete +30-70% in the interval benchmarks). - Nodes start from the monoid's Identity, fresh and recycled, so a monoid whose identity is not the zero value (a product) aggregates correctly. - interval.Set gets Get, an item-returning Delete, SetIterator, SetCursor and SetOverlapIterator with Item and no values. The default TieBreak treats every representation of a point at the same key as the same item. Bounds.CompareIntervals is back as an optional single-call comparator that must order by start, which Verify now checks. - OverlapIterator.Seek restarts a scan on an existing iterator, so a query no longer costs a fresh 300-byte iterator; the ported benchmarks use it. - Config.Compare has a value receiver; the interval updater stores its bound without a deferred closure; Next and Prev return their validity from the position they already know; Map.Iterator's doc says to reposition rather than recreate after a mutation. Co-Authored-By: Claude Fable 5.1 --- aug/adversarial_test.go | 9 +- aug/config.go | 29 +++-- aug/iterator.go | 4 +- aug/map.go | 6 +- aug/monoid_test.go | 46 +++++++ aug/node.go | 17 ++- interval/aug.go | 13 +- interval/interval.go | 232 +++++++++++++++++++++++++++++---- interval/interval_span_test.go | 6 +- interval/iterator.go | 18 ++- interval/property_test.go | 61 +++++++-- 11 files changed, 371 insertions(+), 70 deletions(-) diff --git a/aug/adversarial_test.go b/aug/adversarial_test.go index faccd1e..2c345c5 100644 --- a/aug/adversarial_test.go +++ b/aug/adversarial_test.go @@ -186,10 +186,15 @@ func TestVerifyInterfaceTypes(t *testing.T) { } // 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. +// 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() { diff --git a/aug/config.go b/aug/config.go index f0d9037..c6e3819 100644 --- a/aug/config.go +++ b/aug/config.go @@ -50,14 +50,11 @@ type Config[K, V, A any] struct { // 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. +// replaced below a node, which is all an augmentation that depends only +// on the entries of the subtree (counts, sums, bounds) needs. An +// augmentation that depends on the shape of the subtree (node counts, +// heights) must also implement Restructurer to hear about splits, merges +// and rebalances, and 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 { @@ -69,6 +66,17 @@ type Updater[K, V, A any] interface { Update(*Node[K, V, A], UpdateInfo[K, V, A]) (changed bool) } +// Restructurer may be implemented by an Updater whose augmentation depends +// on the shape of the subtree rather than only on its entries. Restructured +// is called on a node after a child of it was split, or after two of its +// children were merged or rebalanced, with the entries below it unchanged; +// it should recompute the node's augmentation and report whether it +// changed, in which case the ancestors are told through the event that +// caused the change or through Restructured. +type Restructurer[K, V, A any] interface { + Restructured(*Node[K, V, A]) (changed bool) +} + // UpdateInfo is used to describe the update operation. type UpdateInfo[K, V, A any] struct { @@ -129,11 +137,12 @@ const ( ) // Compare compares two values using the same comparison function as the Map. -func (c *Config[K, V, A]) Compare(a, b K) int { return c.cmp(a, b) } +func (c Config[K, V, A]) Compare(a, b K) int { return c.cmp(a, b) } type config[K, V, A any] struct { Config[K, V, A] find func(*Node[K, V, A], K) (int, bool) // nil: binary search with cmp + restr Restructurer[K, V, A] // set when Updater implements it monoid CommutativeMonoid[K, V, A] // set when Updater came from MonoidUpdater folder Folder[K, V, A] // set when the monoid implements Folder fl FreeList[K, V, A] @@ -212,9 +221,11 @@ func makeConfig[K, V, A any]( if o.find != nil { find, _ = o.find.(func(*Node[K, V, A], K) (int, bool)) } + restr, _ := up.(Restructurer[K, V, A]) return config[K, V, A]{ Config: Config[K, V, A]{Updater: up, cmp: cmp}, find: find, + restr: restr, monoid: m, folder: f, fl: fl, diff --git a/aug/iterator.go b/aug/iterator.go index b4b685b..d3a53de 100644 --- a/aug/iterator.go +++ b/aug/iterator.go @@ -172,7 +172,7 @@ func (i *Iterator[K, V, A]) Next() bool { i.pos++ } i.settle() - return i.Valid() + return i.pos < i.node.Count() } if i.pos >= i.node.Count() { // Past the end; stay there. @@ -215,7 +215,7 @@ func (i *Iterator[K, V, A]) Prev() bool { ll.Ascend() i.pos-- } - return i.Valid() + return i.pos >= 0 } if i.pos < 0 { // Before the beginning; stay there. diff --git a/aug/map.go b/aug/map.go index 8aa0783..0ccd4a4 100644 --- a/aug/map.go +++ b/aug/map.go @@ -144,9 +144,9 @@ func (t *Map[K, V, A]) Upsert(item K, value V) (replacedK K, replacedV V, replac return replacedK, replacedV, replaced } -// Iterator returns a new Iterator object. It is not safe to continue using an -// Iterator after modifications are made to the tree. If modifications are made, -// create a new Iterator. +// Iterator returns a new Iterator positioned before the first entry. A +// mutation of the Map invalidates the Iterator's position; reposition it +// (Reset, First, a Seek) before using it again. func (t *Map[K, V, A]) Iterator() Iterator[K, V, A] { it := Iterator[K, V, A]{r: t} it.Reset() diff --git a/aug/monoid_test.go b/aug/monoid_test.go index 505f4d4..3515d56 100644 --- a/aug/monoid_test.go +++ b/aug/monoid_test.go @@ -212,3 +212,49 @@ func Example_customAugmentation() { // pool 2: 3 hosts, 150 load // cumulative load reaches 200 at {1 4} } + +// product is a CommutativeMonoid whose identity is 1, not the zero value. +type product struct{} + +func (product) Identity() int { return 1 } +func (product) Of(_ int, v int) int { return v } +func (product) Combine(a, b int) int { return a * b } +func (product) Uncombine(a, b int) int { return a / b } + +// TestNonZeroIdentity checks that fresh and recycled nodes start from the +// monoid's identity rather than the zero value. +func TestNonZeroIdentity(t *testing.T) { + m := aug.NewMonoid[int, int, int](cmp.Compare[int], product{}, aug.WithDegree(2)) + if m.Total() != 1 { + t.Fatalf("empty Total %d, want the identity", m.Total()) + } + m.Upsert(1, 7) + if m.Total() != 7 || m.Aggregate(0, 10) != 7 { + t.Fatalf("Total %d, Aggregate %d after inserting 7", m.Total(), m.Aggregate(0, 10)) + } + want := 7 + for i := 2; i <= 12; i++ { + m.Upsert(i, 2) + want *= 2 + } + if m.Total() != want { + t.Fatalf("Total %d, want %d after splits", m.Total(), want) + } + if err := m.Verify(); err != nil { + t.Fatal(err) + } + // Clear recycles nodes into the free list; they must start from the + // identity again. + m.Clear() + m.Upsert(1, 5) + m.Upsert(2, 3) + if m.Total() != 15 { + t.Fatalf("Total %d after Clear and reinsert, want 15", m.Total()) + } + if p, _ := m.Prefix(2); p != 5 { + t.Fatalf("Prefix(2) = %d, want 5", p) + } + if err := m.Verify(); err != nil { + t.Fatal(err) + } +} diff --git a/aug/node.go b/aug/node.go index cf89b9d..fa53379 100644 --- a/aug/node.go +++ b/aug/node.go @@ -81,6 +81,9 @@ func (c *config[K, V, A]) getNode() *Node[K, V, A] { n.keys = make([]K, 0, c.maxEntries) n.values = make([]V, 0, c.maxEntries) } + if c.monoid != nil { + n.aug = c.monoid.Identity() + } n.ref = 1 return n } @@ -381,6 +384,14 @@ func (n *Node[K, V, A]) update(cfg *Config[K, V, A]) bool { return cfg.Updater.Update(n, UpdateInfo[K, V, A]{}) } +// restructured tells a Restructurer that the shape below n changed. +func (n *Node[K, V, A]) restructured(c *config[K, V, A]) bool { + if c.restr == nil { + return false + } + return c.restr.Restructured(n) +} + func (n *Node[K, V, A]) updateOn(cfg *Config[K, V, A], action Action, k K, v V, affected *Node[K, V, A]) bool { if cfg.Updater == nil { return false @@ -431,7 +442,7 @@ func (n *Node[K, V, A]) insert(c *config[K, V, A], item K, value V) (replacedK K if len(n.children[i].keys) >= c.maxEntries { splitK, splitV, splitNode := mut(c, &n.children[i]).split(c, c.maxEntries/2) n.insertAt(i, splitK, splitV, splitNode) - structural = n.update(&c.Config) + structural = n.restructured(c) if cmp := c.cmp(item, n.keys[i]); cmp < 0 { // no change, we want first split node } else if cmp > 0 { @@ -618,7 +629,7 @@ func (n *Node[K, V, A]) rebalanceOrMerge(c *config[K, V, A], i int) bool { mergeChild.decRef(c, true /* recursive */) } } - return n.update(&c.Config) + return n.restructured(c) } // remove removes an item from the subtree rooted at this node. Returns the item @@ -655,7 +666,7 @@ func (n *Node[K, V, A]) remove( changed = n.updateOn(&c.Config, Removal, outK, outV, nil) } else { // Nothing was removed but the subtree was restructured. - changed = n.update(&c.Config) + changed = n.restructured(c) } } return outK, outV, found, changed diff --git a/interval/aug.go b/interval/aug.go index d84d225..96a8915 100644 --- a/interval/aug.go +++ b/interval/aug.go @@ -33,7 +33,6 @@ func (u *updater[I, K, V]) Update( md aug.UpdateInfo[I, V, subtreeBound[K]], ) (updated bool) { a := n.Aug() - defer func() { n.SetAug(a) }() // a is modified below; store it on the way out switch md.Action { case aug.Insertion: up := u.upperBound(md.RelevantKey) @@ -44,6 +43,7 @@ func (u *updater[I, K, V]) Update( } if !a.set || a.compare(u.cmp, up) < 0 { a.keyBound, a.set = up, true + n.SetAug(a) return true } return false @@ -56,6 +56,7 @@ func (u *updater[I, K, V]) Update( } if a.set && a.compare(u.cmp, up) == 0 { a.keyBound, a.set = u.findUpperBound(n) + n.SetAug(a) return !a.set || a.compare(u.cmp, up) != 0 } return false @@ -68,6 +69,7 @@ func (u *updater[I, K, V]) Update( case aug.Default, aug.Replacement: prev, prevSet := a.keyBound, a.set a.keyBound, a.set = u.findUpperBound(n) + n.SetAug(a) return prevSet != a.set || (a.set && a.compare(u.cmp, prev) != 0) default: panic("interval: unknown action") @@ -84,11 +86,10 @@ type keyBound[K any] struct { // its start is a point, so that leaf matching and subtree pruning agree. func (up *updater[I, K, V]) upperBound(interval I) keyBound[K] { k := up.key(interval) - if !up.hasEnd(interval) { - return keyBound[K]{k: k, inclusive: true} - } - if end := up.end(interval); up.cmp(end, k) > 0 { - return keyBound[K]{k: end} + if up.hasEnd(interval) { + if end := up.end(interval); up.cmp(end, k) > 0 { + return keyBound[K]{k: end} + } } return keyBound[K]{k: k, inclusive: true} } diff --git a/interval/interval.go b/interval/interval.go index a79c981..296e386 100644 --- a/interval/interval.go +++ b/interval/interval.go @@ -27,6 +27,7 @@ package interval import ( + "fmt" "iter" "github.com/ajwerner/btree/aug" @@ -63,8 +64,14 @@ type Bounds[I, K any] struct { // intervals being ordered by start, so it is consulted only after the // starts compare equal. It is optional: by default points sort before // ranges, then ranges by End, and intervals equal by those are the same - // key. + // key; an interval whose end is not after its start is a point here as + // everywhere else. TieBreak func(I, I) int + + // CompareIntervals, when set, replaces the Key-then-TieBreak comparison + // with a single call, which saves two indirect calls per probe. It must + // order intervals by Key first; Verify reports a tree that is not. + CompareIntervals func(I, I) int } // Interval is implemented by interval types that expose their own bounds: @@ -94,28 +101,34 @@ func (b Bounds[I, K]) withDefaults() Bounds[I, K] { } if b.TieBreak == nil { b.TieBreak = func(x, y I) int { - xEnd, yEnd := b.HasEnd(x), b.HasEnd(y) + xPoint, yPoint := b.isPoint(x), b.isPoint(y) switch { - case xEnd && yEnd: + case !xPoint && !yPoint: return b.Compare(b.End(x), b.End(y)) - case xEnd: - return 1 - case yEnd: + case xPoint && yPoint: + return 0 + case xPoint: return -1 default: - return 0 + return 1 + } + } + } + if b.CompareIntervals == nil { + b.CompareIntervals = func(x, y I) int { + if c := b.Compare(b.Key(x), b.Key(y)); c != 0 { + return c } + return b.TieBreak(x, y) } } return b } -// compare orders intervals by start, then by TieBreak. -func (b Bounds[I, K]) compare(x, y I) int { - if c := b.Compare(b.Key(x), b.Key(y)); c != 0 { - return c - } - return b.TieBreak(x, y) +// isPoint reports whether an interval covers only its start key: it has no +// end, or its end is not after its start. +func (b Bounds[I, K]) isPoint(i I) bool { + return !b.HasEnd(i) || b.Compare(b.End(i), b.Key(i)) <= 0 } // New constructs a Map over intervals described by b. See aug.WithDegree @@ -123,7 +136,7 @@ func (b Bounds[I, K]) compare(x, y I) int { func New[I, K, V any](b Bounds[I, K], opts ...aug.Option) *Map[I, K, V] { b = b.withDefaults() return (*Map[I, K, V])(aug.New[I, V, subtreeBound[K]]( - b.compare, + b.CompareIntervals, &updater[I, K, V]{ cmp: b.Compare, key: b.Key, @@ -190,8 +203,26 @@ func (m *Map[I, K, V]) Compare(a, b I) int { return m.a().Compare(a, b) } // String renders the tree in a Newick-like format. func (m *Map[I, K, V]) String() string { return m.a().String() } -// Verify checks the tree's invariants; see aug.Map.Verify. -func (m *Map[I, K, V]) Verify() error { return m.a().Verify() } +// Verify checks the tree's invariants (see aug.Map.Verify) and that the +// intervals are ordered by start key, which a custom CompareIntervals must +// guarantee. +func (m *Map[I, K, V]) Verify() error { + if err := m.a().Verify(); err != nil { + return err + } + it := m.a().Iterator() + u := aug.LowLevel(&it).Config().Updater.(*updater[I, K, V]) + var prev K + first := true + for i := range m.a().All() { + k := u.key(i) + if !first && u.cmp(prev, k) > 0 { + return fmt.Errorf("interval: CompareIntervals does not order by start key") + } + prev, first = k, false + } + return nil +} // Iterator returns a new iterator positioned before the first entry. func (m *Map[I, K, V]) Iterator() Iterator[I, K, V] { return m.a().Iterator() } @@ -241,12 +272,32 @@ func NewSet[I, K any](b Bounds[I, K], opts ...aug.Option) *Set[I, K] { func (s *Set[I, K]) m() *Map[I, K, struct{}] { return (*Map[I, K, struct{}])(s) } -// Overlaps returns an OverlapIterator positioned at the first item that +// Overlaps returns a SetOverlapIterator positioned at the first item that // overlaps span. -func (s *Set[I, K]) Overlaps(span Span[K]) OverlapIterator[I, K, struct{}] { - return s.m().Overlaps(span) +func (s *Set[I, K]) Overlaps(span Span[K]) SetOverlapIterator[I, K] { + return SetOverlapIterator[I, K]{o: s.m().Overlaps(span)} +} + +// SetOverlapIterator visits the items of a Set that overlap a Span, in +// order of their start keys. See OverlapIterator. +type SetOverlapIterator[I, K any] struct { + o OverlapIterator[I, K, struct{}] } +// Seek restarts the iterator at the first item overlapping span and +// reports whether there is one. +func (i *SetOverlapIterator[I, K]) Seek(span Span[K]) bool { return i.o.Seek(span) } + +// Valid reports whether the iterator is at an overlapping item. +func (i *SetOverlapIterator[I, K]) Valid() bool { return i.o.Valid() } + +// Item returns the item at the iterator's position, which must be valid. +func (i *SetOverlapIterator[I, K]) Item() I { return i.o.Key() } + +// Next positions the iterator at the following overlapping item and +// reports whether there is one. +func (i *SetOverlapIterator[I, K]) Next() bool { return i.o.Next() } + // Overlapping returns an iterator over the items that overlap span, in // order of their start keys. func (s *Set[I, K]) Overlapping(span Span[K]) iter.Seq[I] { @@ -266,16 +317,26 @@ func (s *Set[I, K]) Upsert(item I) (replaced I, overwrote bool) { return replaced, overwrote } -// Delete removes item, reporting whether it was present. -func (s *Set[I, K]) Delete(item I) (removed bool) { - _, _, removed = s.m().Delete(item) - return removed +// Delete removes the item equal to item and returns it. +func (s *Set[I, K]) Delete(item I) (removed I, ok bool) { + removed, _, ok = s.m().Delete(item) + return removed, ok } -// Contains reports whether item is in the set. +// Contains reports whether an item equal to item is in the set. func (s *Set[I, K]) Contains(item I) bool { - _, ok := s.m().Get(item) - return ok + _, found := s.Get(item) + return found +} + +// Get returns the item in the set equal to probe, if any. Items that +// compare equal may differ in fields the comparison ignores. +func (s *Set[I, K]) Get(probe I) (item I, found bool) { + it := s.m().a().Iterator() + if it.SeekExact(probe) { + return it.Key(), true + } + return item, false } // Len returns the number of items. @@ -297,10 +358,10 @@ func (s *Set[I, K]) String() string { return s.m().String() } func (s *Set[I, K]) Verify() error { return s.m().Verify() } // Iterator returns a new iterator positioned before the first item. -func (s *Set[I, K]) Iterator() Iterator[I, K, struct{}] { return s.m().Iterator() } +func (s *Set[I, K]) Iterator() SetIterator[I, K] { return SetIterator[I, K]{it: s.m().Iterator()} } // Cursor returns a new cursor positioned before the first item. -func (s *Set[I, K]) Cursor() Cursor[I, K, struct{}] { return s.m().Cursor() } +func (s *Set[I, K]) Cursor() SetCursor[I, K] { return SetCursor[I, K]{c: s.m().Cursor()} } // All returns an iterator over every item in order. func (s *Set[I, K]) All() iter.Seq[I] { return keys(s.m().All()) } @@ -335,3 +396,118 @@ func keys[T any](seq iter.Seq2[T, struct{}]) iter.Seq[T] { } } } + +// SetIterator iterates over a set. Every positioning method reports whether +// the iterator is now at an item. +type SetIterator[I, K any] struct { + it Iterator[I, K, struct{}] +} + +// Reset marks the iterator invalid, before the first item. +func (i *SetIterator[I, K]) Reset() { i.it.Reset() } + +// First positions the iterator at the smallest item. +func (i *SetIterator[I, K]) First() bool { return i.it.First() } + +// Last positions the iterator at the largest item. +func (i *SetIterator[I, K]) Last() bool { return i.it.Last() } + +// Next positions the iterator at the following item. +func (i *SetIterator[I, K]) Next() bool { return i.it.Next() } + +// Prev positions the iterator at the preceding item. +func (i *SetIterator[I, K]) Prev() bool { return i.it.Prev() } + +// SeekGE positions the iterator at the first item >= item. +func (i *SetIterator[I, K]) SeekGE(item I) bool { return i.it.SeekGE(item) } + +// SeekGT positions the iterator at the first item > item. +func (i *SetIterator[I, K]) SeekGT(item I) bool { return i.it.SeekGT(item) } + +// SeekLE positions the iterator at the last item <= item. +func (i *SetIterator[I, K]) SeekLE(item I) bool { return i.it.SeekLE(item) } + +// SeekLT positions the iterator at the last item < item. +func (i *SetIterator[I, K]) SeekLT(item I) bool { return i.it.SeekLT(item) } + +// SeekExact positions the iterator at item and reports whether it is in +// the set; if not, the iterator is at the first greater item. +func (i *SetIterator[I, K]) SeekExact(item I) bool { return i.it.SeekExact(item) } + +// Valid reports whether the iterator is at an item. +func (i *SetIterator[I, K]) Valid() bool { return i.it.Valid() } + +// Item returns the item at the iterator's position, which must be valid. +func (i *SetIterator[I, K]) Item() I { return i.it.Key() } + +// Compare compares two items with the set's comparison function. +func (i *SetIterator[I, K]) Compare(a, b I) int { return i.it.Compare(a, b) } + +// SetCursor is a SetIterator that can also mutate the set at its position; +// see aug.Cursor for the rules. +type SetCursor[I, K any] struct { + c Cursor[I, K, struct{}] +} + +// Reset marks the cursor invalid, before the first item. +func (c *SetCursor[I, K]) Reset() { c.c.Reset() } + +// First positions the cursor at the smallest item. +func (c *SetCursor[I, K]) First() bool { return c.c.First() } + +// Last positions the cursor at the largest item. +func (c *SetCursor[I, K]) Last() bool { return c.c.Last() } + +// Next positions the cursor at the following item. +func (c *SetCursor[I, K]) Next() bool { return c.c.Next() } + +// Prev positions the cursor at the preceding item. +func (c *SetCursor[I, K]) Prev() bool { return c.c.Prev() } + +// SeekGE positions the cursor at the first item >= item. +func (c *SetCursor[I, K]) SeekGE(item I) bool { return c.c.SeekGE(item) } + +// SeekGT positions the cursor at the first item > item. +func (c *SetCursor[I, K]) SeekGT(item I) bool { return c.c.SeekGT(item) } + +// SeekLE positions the cursor at the last item <= item. +func (c *SetCursor[I, K]) SeekLE(item I) bool { return c.c.SeekLE(item) } + +// SeekLT positions the cursor at the last item < item. +func (c *SetCursor[I, K]) SeekLT(item I) bool { return c.c.SeekLT(item) } + +// SeekExact positions the cursor at item and reports whether it is in the +// set; if not, the cursor is at the first greater item. +func (c *SetCursor[I, K]) SeekExact(item I) bool { return c.c.SeekExact(item) } + +// Valid reports whether the cursor is at an item. +func (c *SetCursor[I, K]) Valid() bool { return c.c.Valid() } + +// Item returns the item at the cursor's position, which must be valid. +func (c *SetCursor[I, K]) Item() I { return c.c.Key() } + +// Compare compares two items with the set's comparison function. +func (c *SetCursor[I, K]) Compare(a, b I) int { return c.c.Compare(a, b) } + +// Delete removes the current item and returns it, leaving the cursor on +// the following item or past the end; see aug.Cursor.Delete. +func (c *SetCursor[I, K]) Delete() I { + item, _ := c.c.Delete() + return item +} + +// Rekey replaces the current item with item, which may sort elsewhere, +// and returns any other item equal to it that was displaced. Afterwards +// the cursor is on item. +func (c *SetCursor[I, K]) Rekey(item I) (displaced I, ok bool) { + displaced, _, ok = c.c.Rekey(item) + return displaced, ok +} + +// Upsert inserts item, or replaces the equal item already present and +// returns it, using the cursor's position as a hint. Afterwards the cursor +// is on item. +func (c *SetCursor[I, K]) Upsert(item I) (replaced I, ok bool) { + replaced, _, ok = c.c.Upsert(item, struct{}{}) + return replaced, ok +} diff --git a/interval/interval_span_test.go b/interval/interval_span_test.go index b134817..e95d5e7 100644 --- a/interval/interval_span_test.go +++ b/interval/interval_span_test.go @@ -729,12 +729,13 @@ func BenchmarkBTreeIterFirstOverlap(b *testing.B) { } rng := newRNG(b) + it := tr.Overlaps(query(spans[0])) b.ResetTimer() for i := 0; i < b.N; i++ { j := rng.IntN(len(spans)) s := spans[j] - it := tr.Overlaps(query(s)) + it.Seek(query(s)) if testing.Verbose() { if !it.Valid() { b.Fatal("expected to find key") @@ -818,9 +819,10 @@ func BenchmarkBTreeIterOverlapScan(b *testing.B) { tr.Upsert(newLatch(spanWithEnd(i, i+size+1)), struct{}{}) } + ov := tr.Overlaps(query(randomSpan(rng, count))) b.ResetTimer() for i := 0; i < b.N; i++ { - for ov := tr.Overlaps(query(randomSpan(rng, count))); ov.Valid(); ov.Next() { + for ok := ov.Seek(query(randomSpan(rng, count))); ok; ok = ov.Next() { } } } diff --git a/interval/iterator.go b/interval/iterator.go index 8f15ec9..3051be3 100644 --- a/interval/iterator.go +++ b/interval/iterator.go @@ -123,20 +123,30 @@ type overlapScan[I, K, V any] struct { } func newOverlapIterator[I, K, V any](m *aug.Map[I, V, subtreeBound[K]], span Span[K]) OverlapIterator[I, K, V] { - i := OverlapIterator[I, K, V]{it: m.Iterator(), o: overlapScan[I, K, V]{span: span}} + i := OverlapIterator[I, K, V]{it: m.Iterator()} + i.Seek(span) + return i +} + +// Seek restarts the iterator at the first entry overlapping span and +// reports whether there is one. Reusing an iterator this way avoids the +// cost of constructing one per query. +func (i *OverlapIterator[I, K, V]) Seek(span Span[K]) bool { + i.it.Reset() + i.o = overlapScan[I, K, V]{span: span} ll := aug.LowLevel(&i.it) cfg := i.cfg() if span.empty(cfg.cmp) { - return i + return false } ll.IncrementPos() if !i.it.Valid() { - return i + return false } i.constrainMinSearchBounds() i.constrainMaxSearchBounds() i.findNextOverlap() - return i + return i.it.Valid() } func (i *OverlapIterator[I, K, V]) cfg() *updater[I, K, V] { diff --git a/interval/property_test.go b/interval/property_test.go index 92421ca..b71dd0d 100644 --- a/interval/property_test.go +++ b/interval/property_test.go @@ -152,7 +152,7 @@ func testOverlapProperty(t *testing.T, degree int) { break } } - removed := m.Delete(s) + _, removed := m.Delete(s) _, had := ref[s] if removed != had { t.Fatalf("step %d: Delete(%v) = %v, want %v", step, s, removed, had) @@ -190,29 +190,33 @@ func TestOverlapIteratorIsIndependent(t *testing.T) { m.Upsert(s) } ov := m.Overlaps(interval.HalfOpen(10, 16)) - if !ov.Valid() || ov.Key() != (span{10, 11}) { - t.Fatalf("Overlaps at %v", ov.Key()) + if !ov.Valid() || ov.Item() != (span{10, 11}) { + t.Fatalf("Overlaps at %v", ov.Item()) } it := m.Iterator() it.SeekGE(span{5, 6}) it.Prev() it.Next() - if !ov.Next() || ov.Key() != (span{15, 16}) { - t.Fatalf("Next after plain iteration at valid=%v %v", ov.Valid(), ov.Key()) + if !ov.Next() || ov.Item() != (span{15, 16}) { + t.Fatalf("Next after plain iteration at valid=%v %v", ov.Valid(), ov.Item()) } if ov.Next() || ov.Valid() { - t.Fatalf("Next past the last overlap is valid at %v", ov.Key()) + t.Fatalf("Next past the last overlap is valid at %v", ov.Item()) } - if !it.Valid() || it.Key() != (span{5, 6}) { - t.Fatalf("plain iterator disturbed: valid=%v %v", it.Valid(), it.Key()) + if !it.Valid() || it.Item() != (span{5, 6}) { + t.Fatalf("plain iterator disturbed: valid=%v %v", it.Valid(), it.Item()) } for _, q := range []interval.Span[int]{interval.HalfOpen(5, 5), interval.HalfOpen(6, 5)} { if e := m.Overlaps(q); e.Valid() { - t.Fatalf("empty span %v overlaps %v", q, e.Key()) + t.Fatalf("empty span %v overlaps %v", q, e.Item()) } } - if e := m.Overlaps(interval.Point(5)); !e.Valid() || e.Key() != (span{5, 6}) { - t.Fatalf("Point(5) = valid %v %v", e.Valid(), e.Key()) + if e := m.Overlaps(interval.Point(5)); !e.Valid() || e.Item() != (span{5, 6}) { + t.Fatalf("Point(5) = valid %v %v", e.Valid(), e.Item()) + } + // Seek reuses the iterator for another query. + if !ov.Seek(interval.HalfOpen(0, 6)) || ov.Item() != (span{0, 1}) || !ov.Next() || ov.Item() != (span{5, 6}) || ov.Next() { + t.Fatal("Seek did not restart the scan") } } @@ -322,3 +326,38 @@ func TestTieBreakCannotReorderStarts(t *testing.T) { t.Fatalf("got %v, want %v", got, want) } } + +// TestDefaultEqualityNormalizesPoints checks that the default tie-break +// treats every representation of a point at the same key as the same item. +func TestDefaultEqualityNormalizesPoints(t *testing.T) { + m := interval.NewSet(interval.BoundsOf[span](cmp.Compare[int])) + m.Upsert(span{5, 2}) // reversed: a point at 5 + if _, replaced := m.Upsert(span{5, 5}); !replaced { + t.Fatal("{5,5} did not replace the reversed point {5,2}") + } + if m.Len() != 1 { + t.Fatalf("len %d", m.Len()) + } + if got, ok := m.Get(span{5, 0}); !ok || got != (span{5, 5}) { + t.Fatalf("Get = %v %v", got, ok) + } + m.Upsert(span{5, 9}) // a range at the same start is a different item + if m.Len() != 2 { + t.Fatalf("len %d after adding a range", m.Len()) + } + // The set API: iterators and cursors expose items only. + it := m.Iterator() + if !it.First() || it.Item() != (span{5, 5}) || !it.Next() || it.Item() != (span{5, 9}) { + t.Fatal("set iterator order") + } + c := m.Cursor() + if !c.SeekExact(span{5, 5}) { + t.Fatal("SeekExact") + } + if removed := c.Delete(); removed != (span{5, 5}) || !c.Valid() || c.Item() != (span{5, 9}) { + t.Fatalf("cursor Delete = %v, now at valid=%v %v", removed, c.Valid(), c.Item()) + } + if removed, ok := m.Delete(span{5, 9}); !ok || removed != (span{5, 9}) { + t.Fatalf("Delete = %v %v", removed, ok) + } +} From 8923a2e3d8024aeef78bed0feca0c995c47f2050 Mon Sep 17 00:00:00 2001 From: Andrew Werner Date: Tue, 29 Sep 2026 09:12:43 -0400 Subject: [PATCH 6/7] interval: single-call bounds, validation at insert, cached updater; Next fast path - The point-or-range decision is HasEnd's alone: the default HasEnd normalises (an end not after the start makes a point) and a custom HasEnd is trusted, so computing a bound costs no comparison; an interval a custom HasEnd claims an end for that does not end after it starts panics at insertion or replacement. - OverlapIterator caches the updater instead of copying the Config and type-asserting on every step. - Iterator.Next returns without settling when the next entry is in the same leaf. Against v0.2.0 the interval and iteration benchmarks are now within noise or ahead (geomean -1.3%); the rest of the suite was already at parity. Co-Authored-By: Claude Fable 5.1 --- CHANGELOG.md | 37 ++++++++++++++++++------------ README.md | 2 +- aug/iterator.go | 3 +++ docs/design.md | 14 +++++++----- interval/aug.go | 24 +++++++++++++++----- interval/interval.go | 24 +++++++++++--------- interval/iterator.go | 6 ++--- interval/property_test.go | 47 +++++++++++++++++++++------------------ 8 files changed, 95 insertions(+), 62 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 913ecc0..00d8b13 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -18,10 +18,15 @@ `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 `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. + 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 @@ -33,11 +38,14 @@ ### 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. +- 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 @@ -57,11 +65,12 @@ ### 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. +- `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 diff --git a/README.md b/README.md index 771a79e..c5d125e 100644 --- a/README.md +++ b/README.md @@ -41,7 +41,7 @@ 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. 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; `Overlapping(span)` is the range-over-func form. +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 diff --git a/aug/iterator.go b/aug/iterator.go index d3a53de..0d658f2 100644 --- a/aug/iterator.go +++ b/aug/iterator.go @@ -170,6 +170,9 @@ func (i *Iterator[K, V, A]) Next() bool { if i.node.IsLeaf() { if i.pos < i.node.Count() { i.pos++ + if i.pos < i.node.Count() { + return true + } } i.settle() return i.pos < i.node.Count() diff --git a/docs/design.md b/docs/design.md index cf52902..4f0cf3f 100644 --- a/docs/design.md +++ b/docs/design.md @@ -90,14 +90,18 @@ spilled most of them to the GC. ## Augmentation Each node carries an `A`. An `Updater` is called with the node and an -`UpdateInfo` after every structural change: `Insertion` and `Removal` of an -entry below the node (with the moved subtree's aggregate on rebalances), +`UpdateInfo` after every change to the entries below it: `Insertion` and +`Removal` of an entry (with the moved subtree's aggregate on rebalances), `Split` (the node is the left half; the right half's aggregate is given), `Replacement` of an entry at its position (with an equal key, or, for an in-place `Cursor.Rekey`, a key that still sorts between the same -neighbours), and `Default` (recompute from scratch). The return value says whether the node's aggregate changed and so -whether ancestors need updating; the write paths stop propagating as soon -as it is false. +neighbours), and `Default` (recompute from scratch). The return value says +whether the node's aggregate changed and so whether ancestors need updating; +the write paths stop propagating as soon as it is false. A split, merge or +rebalance leaves the entries below a node unchanged, so an augmentation +that is a function of those entries needs nothing; one that depends on +the shape of the subtree implements `Restructurer` and is called on the +node whose children changed. Nodes start from the monoid's `Identity`. Most augmentations are commutative monoids (`CommutativeMonoid`: an `Identity`, `Of` and `Combine`): an aggregate is `Combine` folded over `Of` diff --git a/interval/aug.go b/interval/aug.go index 96a8915..7878447 100644 --- a/interval/aug.go +++ b/interval/aug.go @@ -35,6 +35,9 @@ func (u *updater[I, K, V]) Update( a := n.Aug() switch md.Action { case aug.Insertion: + if n.IsLeaf() && md.ModifiedOther == nil { + u.validate(md.RelevantKey) + } up := u.upperBound(md.RelevantKey) if child := md.ModifiedOther; child != nil { if up.compare(u.cmp, child.keyBound) < 0 { @@ -66,7 +69,10 @@ func (u *updater[I, K, V]) Update( return false } fallthrough - case aug.Default, aug.Replacement: + case aug.Replacement: + u.validate(md.RelevantKey) + fallthrough + case aug.Default: prev, prevSet := a.keyBound, a.set a.keyBound, a.set = u.findUpperBound(n) n.SetAug(a) @@ -85,13 +91,19 @@ type keyBound[K any] struct { // or its start, inclusive, for a point. An interval whose end is not after // its start is a point, so that leaf matching and subtree pruning agree. func (up *updater[I, K, V]) upperBound(interval I) keyBound[K] { - k := up.key(interval) if up.hasEnd(interval) { - if end := up.end(interval); up.cmp(end, k) > 0 { - return keyBound[K]{k: end} - } + return keyBound[K]{k: up.end(interval)} + } + return keyBound[K]{k: up.key(interval), inclusive: true} +} + +// validate panics if a custom HasEnd claims an end that is not after the +// start; the default HasEnd cannot. It is called once per inserted or +// replaced interval, at the node holding it. +func (up *updater[I, K, V]) validate(interval I) { + if up.hasEnd(interval) && up.cmp(up.end(interval), up.key(interval)) <= 0 { + panic("interval: an interval with an end must end after it starts; see Bounds.HasEnd") } - return keyBound[K]{k: k, inclusive: true} } // findUpperBound recomputes the bound of n's subtree; ok is false for an diff --git a/interval/interval.go b/interval/interval.go index 296e386..34c12e1 100644 --- a/interval/interval.go +++ b/interval/interval.go @@ -21,9 +21,9 @@ // which is O(k) when matches are adjacent and up to O(k log(n/k)) when // they are scattered. // -// A stored interval whose end is not after its start (or which has no -// end, see Bounds.HasEnd) is a point containing only its start key. -// Queries are Spans: HalfOpen(start, end) or Point(key). +// A stored interval is a range [Key, End) or a point at Key; Bounds.HasEnd +// tells them apart, and by default an interval whose End is not after its +// Key is a point. Queries are Spans: HalfOpen(start, end) or Point(key). package interval import ( @@ -53,10 +53,13 @@ type Bounds[I, K any] struct { // end (see HasEnd) is a point containing only its Key. End func(I) K - // HasEnd reports whether an interval has an end. It is optional: by - // default every interval has one, and an interval whose End is not - // after its Key is a point containing only its Key. A type that marks - // points some other way (a nil end, a flag) provides HasEnd. + // HasEnd reports whether an interval has an end; one without an end is + // a point containing only its Key. It is optional: by default an + // interval has an end when its End is after its Key, so an empty or + // reversed range is a point. A type that marks points some other way + // (a nil end, a flag) provides HasEnd, which then saves a comparison + // per bound; an interval for which it returns true must end after it + // starts, and inserting one that does not panics. HasEnd func(I) bool // TieBreak orders intervals with equal start keys and so defines which @@ -97,7 +100,7 @@ func (b Bounds[I, K]) withDefaults() Bounds[I, K] { panic("interval: Bounds.Compare, Key and End are required") } if b.HasEnd == nil { - b.HasEnd = func(I) bool { return true } + b.HasEnd = func(i I) bool { return b.Compare(b.End(i), b.Key(i)) > 0 } } if b.TieBreak == nil { b.TieBreak = func(x, y I) int { @@ -125,10 +128,9 @@ func (b Bounds[I, K]) withDefaults() Bounds[I, K] { return b } -// isPoint reports whether an interval covers only its start key: it has no -// end, or its end is not after its start. +// isPoint reports whether an interval covers only its start key. func (b Bounds[I, K]) isPoint(i I) bool { - return !b.HasEnd(i) || b.Compare(b.End(i), b.Key(i)) <= 0 + return !b.HasEnd(i) } // New constructs a Map over intervals described by b. See aug.WithDegree diff --git a/interval/iterator.go b/interval/iterator.go index 3051be3..3ae24e9 100644 --- a/interval/iterator.go +++ b/interval/iterator.go @@ -62,6 +62,7 @@ type Iterator[I, K, V any] = aug.Iterator[I, V, subtreeBound[K]] // an Iterator does. type OverlapIterator[I, K, V any] struct { it aug.Iterator[I, V, subtreeBound[K]] + u *updater[I, K, V] o overlapScan[I, K, V] } @@ -124,6 +125,7 @@ type overlapScan[I, K, V any] struct { func newOverlapIterator[I, K, V any](m *aug.Map[I, V, subtreeBound[K]], span Span[K]) OverlapIterator[I, K, V] { i := OverlapIterator[I, K, V]{it: m.Iterator()} + i.u = aug.LowLevel(&i.it).Config().Updater.(*updater[I, K, V]) i.Seek(span) return i } @@ -149,9 +151,7 @@ func (i *OverlapIterator[I, K, V]) Seek(span Span[K]) bool { return i.it.Valid() } -func (i *OverlapIterator[I, K, V]) cfg() *updater[I, K, V] { - return aug.LowLevel(&i.it).Config().Updater.(*updater[I, K, V]) -} +func (i *OverlapIterator[I, K, V]) cfg() *updater[I, K, V] { return i.u } // Valid reports whether the iterator is at an overlapping entry. func (i *OverlapIterator[I, K, V]) Valid() bool { return i.it.Valid() } diff --git a/interval/property_test.go b/interval/property_test.go index b71dd0d..78d90c1 100644 --- a/interval/property_test.go +++ b/interval/property_test.go @@ -262,51 +262,54 @@ func TestPointerEndpoints(t *testing.T) { } } -// TestEmptyIntervalsArePoints checks that an interval whose end is not -// after its start behaves as a point, wherever it sits in the tree. +// TestEmptyIntervalsArePoints checks that with the default HasEnd an +// interval whose end is not after its start is a point wherever it sits in +// the tree, and that a custom HasEnd claiming an end for such an interval +// is rejected at insertion. func TestEmptyIntervalsArePoints(t *testing.T) { - b := interval.BoundsOf[span](cmp.Compare[int]) - b.HasEnd = func(span) bool { return true } // [3,3) claims an end - m := interval.NewSet(b, aug.WithDegree(2)) + m := interval.NewSet(interval.BoundsOf[span](cmp.Compare[int]), aug.WithDegree(2)) m.Upsert(span{3, 3}) - collect := func(q span) (got []span) { - for s := range m.Overlapping(toSpan(q)) { + collect := func(q interval.Span[int]) (got []span) { + for s := range m.Overlapping(q) { got = append(got, s) } return got } - alone := collect(span{3, 4}) - if !slices.Equal(alone, []span{{3, 3}}) { - t.Fatalf("alone in a leaf: %v", alone) + if got := collect(interval.HalfOpen(3, 4)); !slices.Equal(got, []span{{3, 3}}) { + t.Fatalf("alone in a leaf: %v", got) } - if got := collect(span{2, 3}); len(got) != 0 { + if got := collect(interval.HalfOpen(2, 3)); len(got) != 0 { t.Fatalf("[2,3) should not cover the point 3: %v", got) } // Split the leaf many times over; the answer must not change. for i := range 200 { m.Upsert(span{i * 10, i*10 + 1}) } - if got := collect(span{3, 4}); !slices.Equal(got, []span{{0, 1}, {3, 3}}) && !slices.Equal(got, []span{{3, 3}}) { - t.Fatalf("after splits: %v", got) - } - found := false - for _, s := range collect(span{3, 4}) { - found = found || s == span{3, 3} - } - if !found { - t.Fatal("[3,3) lost after the leaf split") + if got := collect(interval.HalfOpen(3, 4)); !slices.Contains(got, span{3, 3}) { + t.Fatalf("[3,3) lost after the leaf split: %v", got) } // A reversed interval is also a point at its start. m.Upsert(span{50, 40}) - if got := collect(span{50, 51}); !slices.Contains(got, span{50, 40}) { + if got := collect(interval.Point(50)); !slices.Contains(got, span{50, 40}) { t.Fatalf("reversed interval not found as a point: %v", got) } - if got := collect(span{41, 49}); slices.Contains(got, span{50, 40}) { + if got := collect(interval.HalfOpen(41, 49)); slices.Contains(got, span{50, 40}) { t.Fatalf("reversed interval matched inside its reversed range: %v", got) } if err := m.Verify(); err != nil { t.Fatal(err) } + // A custom HasEnd that claims an end for an empty range is a bug the + // tree reports at insertion. + b := interval.BoundsOf[span](cmp.Compare[int]) + b.HasEnd = func(span) bool { return true } + bad := interval.NewSet(b) + defer func() { + if recover() == nil { + t.Fatal("inserting [3,3) with HasEnd true did not panic") + } + }() + bad.Upsert(span{3, 3}) } // TestTieBreakCannotReorderStarts checks that a tie-breaker preferring From b5208f49d4fd0b0eea89721320bd8a8e98b0576b Mon Sep 17 00:00:00 2001 From: Andrew Werner Date: Tue, 29 Sep 2026 09:15:55 -0400 Subject: [PATCH 7/7] Map.Lookup returns the stored entry; Set.Get and Contains use direct descents Set.Get had built an iterator and sought; it now uses Map.Lookup, a Get that also returns the stored key, and Contains uses Get directly. Co-Authored-By: Claude Fable 5.1 --- aug/map.go | 11 +++++++++-- aug/monoid_map.go | 3 +++ btree.go | 14 +++++++------- interval/interval.go | 14 +++++++------- orderstat/order_stat.go | 14 +++++++------- 5 files changed, 33 insertions(+), 23 deletions(-) diff --git a/aug/map.go b/aug/map.go index 0ccd4a4..9e10645 100644 --- a/aug/map.go +++ b/aug/map.go @@ -174,18 +174,25 @@ func (t *Map[K, V, A]) Len() int { // Get returns the value associated with the requested key, if it exists. func (t *Map[K, V, A]) Get(k K) (v V, ok bool) { + _, v, ok = t.Lookup(k) + return v, ok +} + +// Lookup returns the stored entry whose key compares equal to k, if any. +// The stored key may differ from k in fields the comparison ignores. +func (t *Map[K, V, A]) Lookup(k K) (key K, v V, ok bool) { n := t.root for n != nil { i, found := n.find(&t.cfg, k) if found { - return n.values[i], true + return n.keys[i], n.values[i], true } if n.IsLeaf() { break } n = n.children[i] } - return v, false + return key, v, false } // Compare compares two keys using the Map's comparison function. diff --git a/aug/monoid_map.go b/aug/monoid_map.go index c1cdf00..ae26cc8 100644 --- a/aug/monoid_map.go +++ b/aug/monoid_map.go @@ -70,6 +70,9 @@ func (t *MonoidMap[K, V, A]) Upsert(k K, v V) (replacedK K, replacedV V, replace // Get returns the value for k, if any. func (t *MonoidMap[K, V, A]) Get(k K) (v V, ok bool) { return t.a().Get(k) } +// Lookup returns the stored entry whose key compares equal to k, if any. +func (t *MonoidMap[K, V, A]) Lookup(k K) (key K, v V, ok bool) { return t.a().Lookup(k) } + // Len returns the number of entries. func (t *MonoidMap[K, V, A]) Len() int { return t.a().Len() } diff --git a/btree.go b/btree.go index 7ad83b7..4f98ced 100644 --- a/btree.go +++ b/btree.go @@ -76,6 +76,9 @@ func (m *Map[K, V]) Upsert(k K, v V) (replacedK K, replacedV V, replaced bool) { // Get returns the value for k, if any. func (m *Map[K, V]) Get(k K) (v V, ok bool) { return m.a().Get(k) } +// Lookup returns the stored entry whose key compares equal to k, if any. +func (m *Map[K, V]) Lookup(k K) (key K, v V, ok bool) { return m.a().Lookup(k) } + // Len returns the number of entries. func (m *Map[K, V]) Len() int { return m.a().Len() } @@ -156,18 +159,15 @@ func (s *Set[T]) Delete(item T) (removed T, ok bool) { // Contains reports whether an item equal to item is in the set. func (s *Set[T]) Contains(item T) bool { - _, found := s.Get(item) - return found + _, ok := s.m().a().Get(item) + return ok } // Get returns the item in the set equal to probe, if any. Items that // compare equal may differ in fields the comparison ignores. func (s *Set[T]) Get(probe T) (item T, found bool) { - it := s.m().a().Iterator() - if it.SeekExact(probe) { - return it.Key(), true - } - return item, false + item, _, found = s.m().a().Lookup(probe) + return item, found } // Len returns the number of items. diff --git a/interval/interval.go b/interval/interval.go index 34c12e1..257bd3d 100644 --- a/interval/interval.go +++ b/interval/interval.go @@ -190,6 +190,9 @@ func (m *Map[I, K, V]) Upsert(k I, v V) (replacedK I, replacedV V, replaced bool // Get returns the value for k, if any. func (m *Map[I, K, V]) Get(k I) (v V, ok bool) { return m.a().Get(k) } +// Lookup returns the stored entry whose key compares equal to k, if any. +func (m *Map[I, K, V]) Lookup(k I) (key I, v V, ok bool) { return m.a().Lookup(k) } + // Len returns the number of entries. func (m *Map[I, K, V]) Len() int { return m.a().Len() } @@ -327,18 +330,15 @@ func (s *Set[I, K]) Delete(item I) (removed I, ok bool) { // Contains reports whether an item equal to item is in the set. func (s *Set[I, K]) Contains(item I) bool { - _, found := s.Get(item) - return found + _, ok := s.m().a().Get(item) + return ok } // Get returns the item in the set equal to probe, if any. Items that // compare equal may differ in fields the comparison ignores. func (s *Set[I, K]) Get(probe I) (item I, found bool) { - it := s.m().a().Iterator() - if it.SeekExact(probe) { - return it.Key(), true - } - return item, false + item, _, found = s.m().a().Lookup(probe) + return item, found } // Len returns the number of items. diff --git a/orderstat/order_stat.go b/orderstat/order_stat.go index eb0c257..3560365 100644 --- a/orderstat/order_stat.go +++ b/orderstat/order_stat.go @@ -80,6 +80,9 @@ func (m *Map[K, V]) Upsert(k K, v V) (replacedK K, replacedV V, replaced bool) { // Get returns the value for k, if any. func (m *Map[K, V]) Get(k K) (v V, ok bool) { return m.a().Get(k) } +// Lookup returns the stored entry whose key compares equal to k, if any. +func (m *Map[K, V]) Lookup(k K) (key K, v V, ok bool) { return m.a().Lookup(k) } + // Len returns the number of entries. func (m *Map[K, V]) Len() int { return m.a().Len() } @@ -176,18 +179,15 @@ func (s *Set[T]) Delete(item T) (removed T, ok bool) { // Contains reports whether an item equal to item is in the set. func (s *Set[T]) Contains(item T) bool { - _, found := s.Get(item) - return found + _, ok := s.m().a().Get(item) + return ok } // Get returns the item in the set equal to probe, if any. Items that // compare equal may differ in fields the comparison ignores. func (s *Set[T]) Get(probe T) (item T, found bool) { - it := s.m().a().Iterator() - if it.SeekExact(probe) { - return it.Key(), true - } - return item, false + item, _, found = s.m().a().Lookup(probe) + return item, found } // Len returns the number of items.