majorly clean up the code

btree.go

@@ -18,13 +18,13 @@ import "github.com/ajwerner/btree/internal/abstract"
 
 // Map is a ordered map from K to V.
 type Map[K, V any] struct {
-	abstract.Map[K, V, struct{}, aug[K], *aug[K]]
+	abstract.Map[K, V, struct{}]
 }
 
 // MakeMap constructs a new Map with the provided comparison function.
 func MakeMap[K, V any](cmp func(K, K) int) Map[K, V] {
 	return Map[K, V]{
-		Map: abstract.MakeMap[K, V, struct{}, aug[K]](struct{}{}, cmp),
+		Map: abstract.MakeMap[K, V, struct{}](cmp, nil),
 	}
 }
 
@@ -59,13 +59,3 @@ func (t *Set[K]) Delete(item K) (removed bool) {
 	_, _, removed = t.Map.Delete(item)
 	return removed
 }
-
-type aug[K any] struct{}
-
-func (a *aug[T]) Update(
-	*abstract.Config[T, struct{}],
-	abstract.Node[T, *aug[T]],
-	abstract.UpdateMeta[T, aug[T]],
-) bool {
-	return false
-}

internal/abstract/aug.go

@@ -18,6 +18,7 @@ package abstract
 // Node represents an abstraction of a node exposed to the
 // augmentation and low-level iteration primitives.
 type Node[K, A any] interface {
+	GetA() *A
 
 	// IsLeaf returns whether this node is a leaf.
 	IsLeaf() bool
@@ -32,19 +33,7 @@ type Node[K, A any] interface {
 	// GetChild returns the augmentation of the child at the given position. It
 	// may be called on non-leaf nodes with values in [0, Count()] if Count()
 	// is greater than 0.
-	GetChild(i int16) A
-}
-
-// Aug is a data structure which augments a node of the tree. It is updated
-// when the structure or contents of the subtree rooted at the current node
-// changes.
-type Aug[K, Aux, A any] interface {
-	*A
-
-	// Update is used to update the state of the node augmentation in response
-	// to a mutation to the tree. See Action and UpdateMeta for the semantics.
-	// The method must return true if the augmentation's value changed.
-	Update(*Config[K, Aux], Node[K, *A], UpdateMeta[K, A]) (changed bool)
+	GetChild(i int16) *A
 }
 
 // Action is used to classify the type of Update in order to permit various

internal/abstract/aug_btree.go

@@ -39,24 +39,24 @@ const (
 //
 // Write operations are not safe for concurrent mutation by multiple
 // goroutines, but Read operations are.
-type Map[K, V, Aux, A any, AP Aug[K, Aux, A]] struct {
-	root   *node[K, V, Aux, A, AP]
+type Map[K, V, A any] struct {
+	root   *node[K, V, A]
 	length int
-	cfg    config[K, V, Aux, A, AP]
+	cfg    config[K, V, A]
 }
 
 // MakeMap constructs a new Map.
-func MakeMap[K, V, Aux, A any, AP Aug[K, Aux, A]](aux Aux, cmp func(K, K) int) Map[K, V, Aux, A, AP] {
-	return Map[K, V, Aux, A, AP]{
-		cfg: makeConfig[K, V, Aux, A, AP](aux, cmp),
+func MakeMap[K, V, A any](cmp func(K, K) int, up Updater[K, A]) Map[K, V, A] {
+	return Map[K, V, A]{
+		cfg: makeConfig[K, V](cmp, up),
 	}
 }
 
 // Reset removes all items from the AugBTree. In doing so, it allows memory
 // held by the AugBTree to be recycled. Failure to call this method before
 // letting a AugBTree be GCed is safe in that it won't cause a memory leak,
 // but it will prevent AugBTree nodes from being efficiently re-used.
-func (t *Map[K, V, Aux, A, AP]) Reset() {
+func (t *Map[K, V, A]) Reset() {
 	if t.root != nil {
 		t.root.decRef(t.cfg.np, true /* recursive */)
 		t.root = nil
@@ -65,7 +65,7 @@ func (t *Map[K, V, Aux, A, AP]) Reset() {
 }
 
 // Clone clones the AugBTree, lazily. It does so in constant time.
-func (t *Map[K, V, Aux, A, AP]) Clone() Map[K, V, Aux, A, AP] {
+func (t *Map[K, V, A]) Clone() Map[K, V, A] {
 	if t.root != nil {
 		// Incrementing the reference count on the root node is sufficient to
 		// ensure that no node in the cloned tree can be mutated by an actor
@@ -87,7 +87,7 @@ func (t *Map[K, V, Aux, A, AP]) Clone() Map[K, V, Aux, A, AP] {
 }
 
 // Delete removes an item equal to the passed in item from the tree.
-func (t *Map[K, V, Aux, A, AP]) Delete(k K) (removedK K, v V, found bool) {
+func (t *Map[K, V, A]) Delete(k K) (removedK K, v V, found bool) {
 	if t.root == nil || t.root.count == 0 {
 		return removedK, v, false
 	}
@@ -108,7 +108,7 @@ func (t *Map[K, V, Aux, A, AP]) Delete(k K) (removedK K, v V, found bool) {
 
 // Upsert adds the given item to the tree. If an item in the tree already equals
 // the given one, it is replaced with the new item.
-func (t *Map[K, V, Aux, A, AP]) Upsert(item K, value V) (replacedK K, replacedV V, replaced bool) {
+func (t *Map[K, V, A]) Upsert(item K, value V) (replacedK K, replacedV V, replaced bool) {
 	if t.root == nil {
 		t.root = t.cfg.np.getLeafNode()
 	} else if t.root.count >= MaxEntries {
@@ -134,14 +134,14 @@ func (t *Map[K, V, Aux, A, AP]) Upsert(item K, value V) (replacedK K, replacedV
 // MakeIter 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.
-func (t *Map[K, V, Aux, A, AP]) Iterator() Iterator[K, V, Aux, A, AP] {
-	it := Iterator[K, V, Aux, A, AP]{r: t}
+func (t *Map[K, V, A]) Iterator() Iterator[K, V, A] {
+	it := Iterator[K, V, A]{r: t}
 	it.Reset()
 	return it
 }
 
 // Height returns the height of the tree.
-func (t *Map[K, V, Aux, A, AP]) Height() int {
+func (t *Map[K, V, A]) Height() int {
 	if t.root == nil {
 		return 0
 	}
@@ -155,12 +155,12 @@ func (t *Map[K, V, Aux, A, AP]) Height() int {
 }
 
 // Len returns the number of items currently in the tree.
-func (t *Map[K, V, Aux, A, AP]) Len() int {
+func (t *Map[K, V, A]) Len() int {
 	return t.length
 }
 
 // Get returns the value associated with the requested key, if it exists.
-func (t *Map[K, V, Aux, A, AP]) Get(k K) (v V, ok bool) {
+func (t *Map[K, V, A]) Get(k K) (v V, ok bool) {
 	it := t.Iterator()
 	it.SeekGE(k)
 	if it.Valid() && it.Compare(it.Cur(), k) == 0 {
@@ -171,7 +171,7 @@ func (t *Map[K, V, Aux, A, AP]) Get(k K) (v V, ok bool) {
 
 // String returns a string description of the tree. The format is
 // similar to the https://en.wikipedia.org/wiki/Newick_format.
-func (t *Map[K, V, Aux, A, AP]) String() string {
+func (t *Map[K, V, A]) String() string {
 	if t.length == 0 {
 		return ";"
 	}

internal/abstract/config.go

@@ -17,28 +17,31 @@ package abstract
 // Config is used to configure the tree. It consists of a comparison function
 // for keys and any auxiliary data provided by the instantiator. It is provided
 // on the iterator and passed to the augmentation's Update method.
-type Config[K, Aux any] struct {
+type Config[K, A any] struct {
 
-	// Config is the configuration provided by the instantiator of the
-	// tree.
-	Aux Aux
+	// Updater is used to update the augmentations to the tree.
+	Updater Updater[K, A]
 
 	cmp func(K, K) int
 }
 
+type Updater[K, A any] interface {
+	Update(Node[K, A], UpdateMeta[K, A]) (changed bool)
+}
+
 // Compare compares two values using the same comparison function as the Map.
 func (c *Config[K, Aux]) Compare(a, b K) int { return c.cmp(a, b) }
 
-type config[K, V, Aux, A any, AP Aug[K, Aux, A]] struct {
-	Config[K, Aux]
-	np *nodePool[K, V, Aux, A, AP]
+type config[K, V, A any] struct {
+	Config[K, A]
+	np *nodePool[K, V, A]
 }
 
-func makeConfig[K, V, Aux, A any, AP Aug[K, Aux, A]](
-	aux Aux, cmp func(K, K) int,
-) (c config[K, V, Aux, A, AP]) {
-	c.Aux = aux
+func makeConfig[K, V, A any](
+	cmp func(K, K) int, up Updater[K, A],
+) (c config[K, V, A]) {
+	c.Updater = up
 	c.cmp = cmp
-	c.np = getNodePool[K, V, Aux, A, AP]()
+	c.np = getNodePool[K, V, A]()
 	return c
 }

internal/abstract/iter_stack.go

@@ -17,27 +17,27 @@ package abstract
 
 // iterStack represents a stack of (node, pos) tuples, which captures
 // iteration state as an Iterator descends a AugBTree.
-type iterStack[K, V, Aux, A any, AP Aug[K, Aux, A]] struct {
-	a    iterStackArr[K, V, Aux, A, AP]
+type iterStack[K, V, A any] struct {
+	a    iterStackArr[K, V, A]
 	aLen int16 // -1 when using s
-	s    []iterFrame[K, V, Aux, A, AP]
+	s    []iterFrame[K, V, A]
 }
 
 const iterStackDepth = 6
 
 // Used to avoid allocations for stacks below a certain size.
-type iterStackArr[K, V, Aux, A any, AP Aug[K, Aux, A]] [iterStackDepth]iterFrame[K, V, Aux, A, AP]
+type iterStackArr[K, V, A any] [iterStackDepth]iterFrame[K, V, A]
 
-type iterFrame[K, V, Aux, A any, AP Aug[K, Aux, A]] struct {
-	*node[K, V, Aux, A, AP]
+type iterFrame[K, V, A any] struct {
+	*node[K, V, A]
 	pos int16
 }
 
-func (is *iterStack[K, V, Aux, A, AP]) push(f iterFrame[K, V, Aux, A, AP]) {
+func (is *iterStack[K, V, A]) push(f iterFrame[K, V, A]) {
 	if is.aLen == -1 {
 		is.s = append(is.s, f)
 	} else if int(is.aLen) == len(is.a) {
-		is.s = make([](iterFrame[K, V, Aux, A, AP]), int(is.aLen)+1, 2*int(is.aLen))
+		is.s = make([](iterFrame[K, V, A]), int(is.aLen)+1, 2*int(is.aLen))
 		copy(is.s, is.a[:])
 		is.s[int(is.aLen)] = f
 		is.aLen = -1
@@ -47,7 +47,7 @@ func (is *iterStack[K, V, Aux, A, AP]) push(f iterFrame[K, V, Aux, A, AP]) {
 	}
 }
 
-func (is *iterStack[K, V, Aux, A, AP]) pop() iterFrame[K, V, Aux, A, AP] {
+func (is *iterStack[K, V, A]) pop() iterFrame[K, V, A] {
 	if is.aLen == -1 {
 		f := is.s[len(is.s)-1]
 		is.s = is.s[:len(is.s)-1]
@@ -57,14 +57,14 @@ func (is *iterStack[K, V, Aux, A, AP]) pop() iterFrame[K, V, Aux, A, AP] {
 	return is.a[is.aLen]
 }
 
-func (is *iterStack[K, V, Aux, A, AP]) len() int {
+func (is *iterStack[K, V, A]) len() int {
 	if is.aLen == -1 {
 		return len(is.s)
 	}
 	return int(is.aLen)
 }
 
-func (is *iterStack[K, V, Aux, A, AP]) reset() {
+func (is *iterStack[K, V, A]) reset() {
 	if is.aLen == -1 {
 		is.s = is.s[:0]
 	} else {

internal/abstract/iterator.go

@@ -16,31 +16,31 @@
 package abstract
 
 // Iterator is responsible for search and traversal within a AugBTree.
-type Iterator[K, V, Aux, A any, AP Aug[K, Aux, A]] struct {
-	r *Map[K, V, Aux, A, AP]
-	iterFrame[K, V, Aux, A, AP]
-	s iterStack[K, V, Aux, A, AP]
+type Iterator[K, V, A any] struct {
+	r *Map[K, V, A]
+	iterFrame[K, V, A]
+	s iterStack[K, V, A]
 }
 
-func (i *Iterator[K, V, Aux, A, AP]) lowLevel() *LowLevelIterator[K, V, Aux, A, AP] {
-	return (*LowLevelIterator[K, V, Aux, A, AP])(i)
+func (i *Iterator[K, V, A]) lowLevel() *LowLevelIterator[K, V, A] {
+	return (*LowLevelIterator[K, V, A])(i)
 }
 
 // Compare compares two keys using the same comparison function as the map.
-func (i *Iterator[K, V, Aux, A, AP]) Compare(a, b K) int {
+func (i *Iterator[K, V, A]) Compare(a, b K) int {
 	return i.r.cfg.cmp(a, b)
 }
 
 // Reset marks the iterator as invalid and clears any state.
-func (i *Iterator[K, V, Aux, A, AP]) Reset() {
+func (i *Iterator[K, V, A]) Reset() {
 	i.node = i.r.root
 	i.pos = -1
 	i.s.reset()
 }
 
 // SeekGE seeks to the first key greater-than or equal to the provided
 // key.
-func (i *Iterator[K, V, Aux, A, AP]) SeekGE(key K) {
+func (i *Iterator[K, V, A]) SeekGE(key K) {
 	i.Reset()
 	if i.node == nil {
 		return
@@ -63,7 +63,7 @@ func (i *Iterator[K, V, Aux, A, AP]) SeekGE(key K) {
 }
 
 // SeekLT seeks to the first key less-than the provided key.
-func (i *Iterator[K, V, Aux, A, AP]) SeekLT(key K) {
+func (i *Iterator[K, V, A]) SeekLT(key K) {
 	i.Reset()
 	if i.node == nil {
 		return
@@ -81,7 +81,7 @@ func (i *Iterator[K, V, Aux, A, AP]) SeekLT(key K) {
 }
 
 // First seeks to the first key in the AugBTree.
-func (i *Iterator[K, V, Aux, A, AP]) First() {
+func (i *Iterator[K, V, A]) First() {
 	i.Reset()
 	i.pos = 0
 	if i.node == nil {
@@ -95,7 +95,7 @@ func (i *Iterator[K, V, Aux, A, AP]) First() {
 }
 
 // Last seeks to the last key in the AugBTree.
-func (i *Iterator[K, V, Aux, A, AP]) Last() {
+func (i *Iterator[K, V, A]) Last() {
 	i.Reset()
 	if i.node == nil {
 		return
@@ -110,7 +110,7 @@ func (i *Iterator[K, V, Aux, A, AP]) Last() {
 
 // Next positions the Iterator to the key immediately following
 // its current position.
-func (i *Iterator[K, V, Aux, A, AP]) Next() {
+func (i *Iterator[K, V, A]) Next() {
 	if i.node == nil {
 		return
 	}
@@ -136,7 +136,7 @@ func (i *Iterator[K, V, Aux, A, AP]) Next() {
 
 // Prev positions the Iterator to the key immediately preceding
 // its current position.
-func (i *Iterator[K, V, Aux, A, AP]) Prev() {
+func (i *Iterator[K, V, A]) Prev() {
 	if i.node == nil {
 		return
 	}
@@ -162,18 +162,18 @@ func (i *Iterator[K, V, Aux, A, AP]) Prev() {
 }
 
 // Valid returns whether the Iterator is positioned at a valid position.
-func (i *Iterator[K, V, Aux, A, AP]) Valid() bool {
+func (i *Iterator[K, V, A]) Valid() bool {
 	return i.node != nil && i.pos >= 0 && i.pos < i.count
 }
 
 // Cur returns the key at the Iterator's current position. It is illegal
 // to call Key if the Iterator is not valid.
-func (i *Iterator[K, V, Aux, A, AP]) Cur() K {
+func (i *Iterator[K, V, A]) Cur() K {
 	return i.keys[i.pos]
 }
 
 // Value returns the value at the Iterator's current position. It is illegal
 // to call Value if the Iterator is not valid.
-func (i *Iterator[K, V, Aux, A, AP]) Value() V {
+func (i *Iterator[K, V, A]) Value() V {
 	return i.values[i.pos]
 }