diff --git a/design/44253-generic-array-sizes/bad_array_parameterization_btree.go2 b/design/44253-generic-array-sizes/bad_array_parameterization_btree.go2 new file mode 100644 index 00000000..7a0f2c3c --- /dev/null +++ b/design/44253-generic-array-sizes/bad_array_parameterization_btree.go2 @@ -0,0 +1,690 @@ +package main + +import "fmt" + +// These constants are the wart. +const ( + degree = 16 + maxItems = 2*degree - 1 + minItems = degree - 1 +) + +func main() { + bt := NewBTree[int, string](func(i, j int) bool { return i < j }) + bt.Insert(3, "world") + bt.Insert(1, "hello") + bt.Insert(2, ",") + bt.Iterate(func(it Iterator[int, string]) { + for it.First(); it.Valid(); it.Next() { + fmt.Println(it.Cur()) + } + }) + +} + +type LessFn[K any] func(a, b K) bool + +type OrderedMap[K, V any] interface { + Insert(K, V) bool + Find(K) (V, bool) + Remove(K) (removed bool) + Iterate(func(Iterator[K, V])) +} + +type Iterator[K, V any] interface { + First() + Last() + Next() + Prev() + SeekGE(K) + SeekLT(K) + Valid() bool + Cur() (K, V) +} + +type nodeI[K, V, N any] interface { + type *N + find(K, LessFn[K]) (idx int, found bool) + insert(K, V, LessFn[K]) (replaced bool) + remove(K, LessFn[K]) (K, V, bool) + len() int16 + at(idx int16) (K, V) + child(idx int16) *N + isLeaf() bool +} + +func NewBTree[K, V any](cmp LessFn[K]) OrderedMap[K, V] { + type N = node[K, V, [maxItems]K, [maxItems]V] + return &btree[K, V, N, *N]{ + cmp: cmp, + newNode: func(isLeaf bool) *N { + return &N{ + leaf: isLeaf, + } + }, + } +} + +type btree[K, V, N any, NP nodeI[K, V, N]] struct { + len int + cmp LessFn[K] + root NP + newNode func(isLeaf bool) NP +} + +func (bt *btree[K, V, N, NP]) Find(k K) (v V, found bool) { + bt.Iterate(func(it Iterator[K, V]) { + it.SeekGE(k) + if !it.Valid() { + return + } + fk, fv := it.Cur() + if !bt.cmp(fk, k) && !bt.cmp(k, fk) { + v, found = fv, true + } + }) + return v, found +} + +func (bt *btree[K, V, N, NP]) Iterate(f func(it Iterator[K, V])) { + it := iterator[K, V, N, NP]{r: bt} + f(&it) +} + +type Arr[T any] interface { + type [maxItems]T +} + +type node[K, V any, KA Arr[K], VA Arr[V]] struct { + count int16 + leaf bool + keys [maxItems]K + vals [maxItems]V + children [maxItems + 1]*node[K, V, KA, VA] +} + +func (n *node[K, V, KA, KV]) at(index int16) (K, V) { + return n.keys[index], n.vals[index] +} + +func (n *node[K, V, KA, KV]) child(index int16) *node[K, V, KA, KV] { + return n.children[index] +} + +func (n *node[K, V, KA, KV]) len() int16 { + return n.count +} + +func (n *node[K, V, KA, KV]) isLeaf() bool { + return n.leaf +} + +func (t *btree[K, V, N, NP]) Insert(k K, v V) (replaced bool) { + if t.root == nil { + t.root = t.newNode(true /* isLeaf */) + } + return t.root.insert(k, v, t.cmp) +} + +func (t *btree[K, V, N, NP]) Remove(k K) (removed bool) { + if t.root == nil { + return false + } + _, _, removed = t.root.remove(k, t.cmp) + return removed +} + +func (n *node[K, V, KA, KV]) insertAt(index int, k K, v V, nd *node[K, V, KA, KV]) { + if index < int(n.count) { + copy(n.keys[index+1:n.count+1], n.keys[index:n.count]) + copy(n.vals[index+1:n.count+1], n.vals[index:n.count]) + if !n.leaf { + copy(n.children[index+2:n.count+2], n.children[index+1:n.count+1]) + } + } + n.keys[index] = k + n.vals[index] = v + if !n.leaf { + n.children[index+1] = nd + } + n.count++ +} + +func (n *node[K, V, KA, KV]) pushBack(k K, v V, nd *node[K, V, KA, KV]) { + n.keys[n.count] = k + n.vals[n.count] = v + if !n.leaf { + n.children[n.count+1] = nd + } + n.count++ +} + +func (n *node[K, V, KA, KV]) pushFront(k K, v V, nd *node[K, V, KA, KV]) { + if !n.leaf { + copy(n.children[1:n.count+2], n.children[:n.count+1]) + n.children[0] = nd + } + copy(n.keys[1:n.count+1], n.keys[:n.count]) + copy(n.vals[1:n.count+1], n.vals[:n.count]) + n.keys[0] = k + n.vals[0] = v + n.count++ +} + +// removeAt removes a value at a given index, pulling all subsequent vals +// back. +func (n *node[K, V, KA, KV]) removeAt(index int) (rk K, rv V, child *node[K, V, KA, KV]) { + if !n.leaf { + child = n.children[index+1] + copy(n.children[index+1:n.count], n.children[index+2:n.count+1]) + n.children[n.count] = nil + } + n.count-- + rk, rv = n.keys[index], n.vals[index] + copy(n.keys[index:n.count], n.keys[index+1:n.count+1]) + copy(n.vals[index:n.count], n.vals[index+1:n.count+1]) + n.clear(n.count) + return rk, rv, child +} + +func (n *node[K, V, KA, KV]) clear(idx int16) { + var zk K + n.keys[n.count] = zk + var zv V + n.vals[n.count] = zv +} + +// popBack removes and returns the last element in the list. +func (n *node[K, V, KA, KV]) popBack() (rk K, rv V, child *node[K, V, KA, KV]) { + n.count-- + rk, rv = n.keys[n.count], n.vals[n.count] + n.clear(n.count) + if n.leaf { + return rk, rv, nil + } + child = n.children[n.count+1] + n.children[n.count+1] = nil + return rk, rv, child +} + +// popFront removes and returns the first element in the list. +func (n *node[K, V, KA, KV]) popFront() (rk K, rv V, child *node[K, V, KA, KV]) { + n.count-- + if !n.leaf { + child = n.children[0] + copy(n.children[:n.count+1], n.children[1:n.count+2]) + n.children[n.count+1] = nil + } + rk, rv = n.keys[0], n.vals[0] + copy(n.keys[:n.count], n.keys[1:n.count+1]) + copy(n.vals[:n.count], n.vals[1:n.count+1]) + n.clear(n.count) + return rk, rv, child +} + +// find returns the index where the given item should be inserted into this +// list. 'found' is true if the item already exists in the list at the given +// index. +func (n *node[K, V, KA, KV]) find(k K, cmp LessFn[K]) (index int, found bool) { + // Logic copied from sort.Search. Inlining this gave + // an 11% speedup on BenchmarkBTreeDeleteInsert. + i, j := 0, int(n.count) + for i < j { + h := int(uint(i+j) >> 1) // avoid overflow when computing h + // i ≤ h < j + if cmp(k, n.keys[h]) { + j = h + } else if cmp(n.keys[h], k) { + i = h + 1 + } else { + return h, true + } + } + return i, false +} + +// split splits the given node at the given index. The current node shrinks, +// and this function returns the k K, v Vhat existed at that index and a new +// node containing all items/children after it. +// +// Before: +// +// +-----------+ +// | x y z | +// +--/-/-\-\--+ +// +// After: +// +// +-----------+ +// | y | +// +----/-\----+ +// / \ +// v v +// +-----------+ +-----------+ +// | x | | z | +// +-----------+ +-----------+ +// +func (n *node[K, V, KA, KV]) split(i int) (k K, v V, next *node[K, V, KA, KV]) { + k, v = n.keys[i], n.vals[i] + next = &node[K, V, KA, KV]{ + leaf: n.leaf, + } + next.count = n.count - int16(i+1) + copy(next.keys[:], n.keys[i+1:n.count]) + copy(next.vals[:], n.vals[i+1:n.count]) + var zk K + var zv V + for j := int16(i); j < n.count; j++ { + n.keys[j] = zk + n.vals[j] = zv + } + if !n.leaf { + copy(next.children[:], n.children[i+1:n.count+1]) + for j := int16(i + 1); j <= n.count; j++ { + n.children[j] = nil + } + } + n.count = int16(i) + return k, v, next +} + +// insert inserts an item into the btree rooted at this node, making sure no +// nodes in the btree exceed maxItems items. Returns true if an existing item +// was replaced and false if an item was inserted. +func (n *node[K, V, KA, KV]) insert(k K, v V, cmp LessFn[K]) (replaced bool) { + i, found := n.find(k, cmp) + if found { + n.vals[i] = v + return true + } + if n.leaf { + n.insertAt(i, k, v, nil) + return false + } + if n.children[i].count >= maxItems { + splitK, splitV, splitNode := n.children[i].split(maxItems / 2) + n.insertAt(i, splitK, splitV, splitNode) + if cmp(k, n.keys[i]) { + // no change, we want first split node + } else if cmp(n.keys[i], k) { + i++ // we want second split node + } else { + n.keys[i] = k + n.vals[i] = v + return true + } + } + return n.children[i].insert(k, v, cmp) +} + +// removeMax removes and returns the maximum item from the suAugBTree rooted at +// this node. +func (n *node[K, V, KA, KV]) removeMax() (k K, v V) { + if n.leaf { + n.count-- + k, v = n.keys[n.count], n.vals[n.count] + n.clear(n.count) + return k, v + } + // Recurse into max child. + i := int(n.count) + if n.children[i].count <= minItems { + // Child not large enough to remove from. + n.rebalanceOrMerge(i) + return n.removeMax() // redo + } + child := n.children[i] + return child.removeMax() +} + +// rebalanceOrMerge grows child 'i' to ensure it has sufficient room to remove +// an item from it while keeping it at or above minItems. +func (n *node[K, V, KA, KV]) rebalanceOrMerge(i int) { + switch { + case i > 0 && n.children[i-1].count > minItems: + // Rebalance from left sibling. + // + // +-----------+ + // | y | + // +----/-\----+ + // / \ + // v v + // +-----------+ +-----------+ + // | x | | | + // +----------\+ +-----------+ + // \ + // v + // a + // + // After: + // + // +-----------+ + // | x | + // +----/-\----+ + // / \ + // v v + // +-----------+ +-----------+ + // | | | y | + // +-----------+ +/----------+ + // / + // v + // a + // + left := n.children[i-1] + child := n.children[i] + xLak, xLav, grandChild := left.popBack() + yLak, yLav := n.keys[i-1], n.vals[i-1] + child.pushFront(yLak, yLav, grandChild) + n.keys[i-1] = xLak + n.vals[i-1] = xLav + + case i < int(n.count) && n.children[i+1].count > minItems: + // Rebalance from right sibling. + // + // +-----------+ + // | y | + // +----/-\----+ + // / \ + // v v + // +-----------+ +-----------+ + // | | | x | + // +-----------+ +/----------+ + // / + // v + // a + // + // After: + // + // +-----------+ + // | x | + // +----/-\----+ + // / \ + // v v + // +-----------+ +-----------+ + // | y | | | + // +----------\+ +-----------+ + // \ + // v + // a + // + right := n.children[i+1] + child := n.children[i] + xLak, xLav, grandChild := right.popFront() + yLak, yLav := n.keys[i], n.vals[i] + child.pushBack(yLak, yLav, grandChild) + n.keys[i] = xLak + n.vals[i] = xLav + + default: + // Merge with either the left or right sibling. + // + // +-----------+ + // | u y v | + // +----/-\----+ + // / \ + // v v + // +-----------+ +-----------+ + // | x | | z | + // +-----------+ +-----------+ + // + // After: + // + // +-----------+ + // | u v | + // +-----|-----+ + // | + // v + // +-----------+ + // | x y z | + // +-----------+ + // + if i >= int(n.count) { + i = int(n.count - 1) + } + child := n.children[i] + mergeK, mergeV, mergeChild := n.removeAt(i) + child.keys[child.count] = mergeK + child.vals[child.count] = mergeV + copy(child.keys[child.count+1:], mergeChild.keys[:mergeChild.count]) + copy(child.vals[child.count+1:], mergeChild.vals[:mergeChild.count]) + if !child.leaf { + copy(child.children[child.count+1:], mergeChild.children[:mergeChild.count+1]) + } + child.count += mergeChild.count + 1 + } +} + +// remove removes an item from the suAugBTree rooted at this node. Returns the +// k K, v Vhat was removed or nil if no matching item was found. +func (n *node[K, V, KA, KV]) remove(k K, cmp LessFn[K]) (rk K, rv V, found bool) { + i, found := n.find(k, cmp) + if n.leaf { + if found { + rk, rv, _ = n.removeAt(i) + return rk, rv, true + } + return rk, rv, false + } + if n.children[i].count <= minItems { + // Child not large enough to remove from. + n.rebalanceOrMerge(i) + return n.remove(k, cmp) // redo + } + child := n.children[i] + if found { + // Replace the item being removed with the max item in our left child. + rk, rv = n.keys[i], n.vals[i] + n.keys[i], n.vals[i] = child.removeMax() + return rk, rv, true + } + return child.remove(k, cmp) +} + +// Iterator is responsible for search and traversal within a BTree. +type iterator[K, V, N any, NP nodeI[K, V, N]] struct { + r *btree[K, V, N, NP] + n NP + pos int16 + s iterStack[K, V, N, NP] +} + +// Cur returns the item at the Iterator's current position. It is illegal +// to call Cur if the Iterator is not valid. +func (i *iterator[K, V, N, NP]) Cur() (K, V) { + return i.n.at(i.pos) +} + +// SeekGE seeks to the first item greater-than or equal to the provided +// item. +func (i *iterator[K, V, N, NP]) SeekGE(k K) { + i.reset() + if i.n == nil { + return + } + for { + pos, found := i.n.find(k, i.r.cmp) + i.pos = int16(pos) + if found { + return + } + if i.n.isLeaf() { + if i.pos == i.n.len() { + i.Next() + } + return + } + i.descend(i.n, i.pos) + } +} + +// SeekLT seeks to the first item less-than the provided item. +func (i *iterator[K, V, N, NP]) SeekLT(k K) { + i.reset() + if i.n == nil { + return + } + for { + pos, found := i.n.find(k, i.r.cmp) + i.pos = int16(pos) + if found || i.n.isLeaf() { + i.Prev() + return + } + i.descend(i.n, i.pos) + } +} + +func (i *iterator[K, V, N, NP]) reset() { + i.n = i.r.root + i.pos = -1 + i.s.reset() +} + +func (i *iterator[K, V, N, NP]) descend(n NP, pos int16) { + i.s.push(iterFrame[K, V, N, NP]{n: n, pos: pos}) + i.n = n.child(pos) + i.pos = 0 +} + +// ascend ascends up to the current node's parent and resets the position +// to the one previously set for this parent node. +func (i *iterator[K, V, N, NP]) ascend() { + f := i.s.pop() + i.n = f.n + i.pos = f.pos +} + +// First seeks to the first item in the AugBTree. +func (i *iterator[K, V, N, NP]) First() { + i.reset() + if i.n == nil { + return + } + for !i.n.isLeaf() { + i.descend(i.n, 0) + } + i.pos = 0 +} + +// Last seeks to the last item in the AugBTree. +func (i *iterator[K, V, N, NP]) Last() { + i.reset() + if i.n == nil { + return + } + for !i.n.isLeaf() { + i.descend(i.n, i.n.len()) + } + i.pos = i.n.len() - 1 +} + +// Next positions the Iterator to the item immediately following +// its current position. +func (i *iterator[K, V, N, NP]) Next() { + if i.n == nil { + return + } + + if i.n.isLeaf() { + i.pos++ + if i.pos < i.n.len() { + return + } + for i.s.len() > 0 && i.pos >= i.n.len() { + i.ascend() + } + return + } + + i.descend(i.n, i.pos+1) + for !i.n.isLeaf() { + i.descend(i.n, 0) + } + i.pos = 0 +} + +// Prev positions the Iterator to the item immediately preceding +// its current position. +func (i *iterator[K, V, N, NP]) Prev() { + if i.n == nil { + return + } + + if i.n.isLeaf() { + i.pos-- + if i.pos >= 0 { + return + } + for i.s.len() > 0 && i.pos < 0 { + i.ascend() + i.pos-- + } + return + } + + i.descend(i.n, i.pos) + for !i.n.isLeaf() { + i.descend(i.n, i.n.len()) + } + i.pos = i.n.len() - 1 +} + +// Valid returns whether the Iterator is positioned at a valid position. +func (i *iterator[K, V, N, NP]) Valid() bool { + return i.n != nil && i.pos >= 0 && i.pos < i.n.len() +} + +// iterStack represents a stack of (node, pos) tuples, which captures +// iteration state as an Iterator descends a AugBTree. +type iterStack[K, V, N any, NP nodeI[K, V, N]] struct { + a iterStackArr[K, V, N, NP] + aLen int16 // -1 when using s + s []iterFrame[K, V, N, NP] +} + +// Used to avoid allocations for stacks below a certain size. +type iterStackArr[K, V, N any, NP nodeI[K, V, N]] [3]iterFrame[K, V, N, NP] + +type iterFrame[K, V, N any, NP nodeI[K, V, N]] struct { + n NP + pos int16 +} + +func (is *iterStack[K, V, N, NP]) push(f iterFrame[K, V, N, NP]) { + if is.aLen == -1 { + is.s = append(is.s, f) + } else if int(is.aLen) == len(is.a) { + is.s = make([](iterFrame[K, V, N, NP]), int(is.aLen)+1, 2*int(is.aLen)) + copy(is.s, is.a[:]) + is.s[int(is.aLen)] = f + is.aLen = -1 + } else { + is.a[is.aLen] = f + is.aLen++ + } +} + +func (is *iterStack[K, V, N, NP]) pop() iterFrame[K, V, N, NP] { + if is.aLen == -1 { + f := is.s[len(is.s)-1] + is.s = is.s[:len(is.s)-1] + return f + } + is.aLen-- + return is.a[is.aLen] +} + +func (is *iterStack[K, V, N, NP]) len() int { + if is.aLen == -1 { + return len(is.s) + } + return int(is.aLen) +} + +func (is *iterStack[K, V, N, NP]) reset() { + if is.aLen == -1 { + is.s = is.s[:0] + } else { + is.aLen = 0 + } +} diff --git a/design/44253-generic-array-sizes/parameterized_node_btree.go2 b/design/44253-generic-array-sizes/parameterized_node_btree.go2 new file mode 100644 index 00000000..f4e96aa0 --- /dev/null +++ b/design/44253-generic-array-sizes/parameterized_node_btree.go2 @@ -0,0 +1,685 @@ +package main + +import "fmt" + +// These constants are the wart. +const ( + degree = 16 + maxItems = 2*degree - 1 + minItems = degree - 1 +) + +func main() { + bt := NewBTree[int, string](func(i, j int) bool { return i < j }) + bt.Insert(3, "world") + bt.Insert(1, "hello") + bt.Insert(2, ",") + bt.Iterate(func(it Iterator[int, string]) { + for it.First(); it.Valid(); it.Next() { + fmt.Println(it.Cur()) + } + }) + +} + +type LessFn[K any] func(a, b K) bool + +type OrderedMap[K, V any] interface { + Insert(K, V) bool + Find(K) (V, bool) + Remove(K) (removed bool) + Iterate(func(Iterator[K, V])) +} + +type Iterator[K, V any] interface { + First() + Last() + Next() + Prev() + SeekGE(K) + SeekLT(K) + Valid() bool + Cur() (K, V) +} + +type nodeI[K, V, N any] interface { + type *N + find(K, LessFn[K]) (idx int, found bool) + insert(K, V, LessFn[K]) (replaced bool) + remove(K, LessFn[K]) (K, V, bool) + len() int16 + at(idx int16) (K, V) + child(idx int16) *N + isLeaf() bool +} + +func NewBTree[K, V any](cmp LessFn[K]) OrderedMap[K, V] { + return &btree[K, V, node[K, V], *node[K, V]]{ + cmp: cmp, + newNode: func(isLeaf bool) *node[K, V] { + return &node[K, V]{ + leaf: isLeaf, + } + }, + } +} + +type btree[K, V, N any, NP nodeI[K, V, N]] struct { + len int + cmp LessFn[K] + root NP + newNode func(isLeaf bool) NP +} + +func (bt *btree[K, V, N, NP]) Find(k K) (v V, found bool) { + bt.Iterate(func(it Iterator[K, V]) { + it.SeekGE(k) + if !it.Valid() { + return + } + fk, fv := it.Cur() + if !bt.cmp(fk, k) && !bt.cmp(k, fk) { + v, found = fv, true + } + }) + return v, found +} + +func (bt *btree[K, V, N, NP]) Iterate(f func(it Iterator[K, V])) { + it := iterator[K, V, N, NP]{r: bt} + f(&it) +} + +func (t *btree[K, V, N, NP]) Insert(k K, v V) (replaced bool) { + if t.root == nil { + t.root = t.newNode(true /* isLeaf */) + } + return t.root.insert(k, v, t.cmp) +} + +type node[K, V any] struct { + count int16 + leaf bool + keys [maxItems]K + vals [maxItems]V + children [maxItems + 1]*node[K, V] +} + +func (n *node[K, V]) at(index int16) (K, V) { + return n.keys[index], n.vals[index] +} + +func (n *node[K, V]) child(index int16) *node[K, V] { + return n.children[index] +} + +func (n *node[K, V]) len() int16 { + return n.count +} + +func (n *node[K, V]) isLeaf() bool { + return n.leaf +} + +func (t *btree[K, V, N, NP]) Remove(k K) (removed bool) { + if t.root == nil { + return false + } + _, _, removed = t.root.remove(k, t.cmp) + return removed +} + +func (n *node[K, V]) insertAt(index int, k K, v V, nd *node[K, V]) { + if index < int(n.count) { + copy(n.keys[index+1:n.count+1], n.keys[index:n.count]) + copy(n.vals[index+1:n.count+1], n.vals[index:n.count]) + if !n.leaf { + copy(n.children[index+2:n.count+2], n.children[index+1:n.count+1]) + } + } + n.keys[index] = k + n.vals[index] = v + if !n.leaf { + n.children[index+1] = nd + } + n.count++ +} + +func (n *node[K, V]) pushBack(k K, v V, nd *node[K, V]) { + n.keys[n.count] = k + n.vals[n.count] = v + if !n.leaf { + n.children[n.count+1] = nd + } + n.count++ +} + +func (n *node[K, V]) pushFront(k K, v V, nd *node[K, V]) { + if !n.leaf { + copy(n.children[1:n.count+2], n.children[:n.count+1]) + n.children[0] = nd + } + copy(n.keys[1:n.count+1], n.keys[:n.count]) + copy(n.vals[1:n.count+1], n.vals[:n.count]) + n.keys[0] = k + n.vals[0] = v + n.count++ +} + +// removeAt removes a value at a given index, pulling all subsequent vals +// back. +func (n *node[K, V]) removeAt(index int) (rk K, rv V, child *node[K, V]) { + if !n.leaf { + child = n.children[index+1] + copy(n.children[index+1:n.count], n.children[index+2:n.count+1]) + n.children[n.count] = nil + } + n.count-- + rk, rv = n.keys[index], n.vals[index] + copy(n.keys[index:n.count], n.keys[index+1:n.count+1]) + copy(n.vals[index:n.count], n.vals[index+1:n.count+1]) + n.clear(n.count) + return rk, rv, child +} + +func (n *node[K, V]) clear(idx int16) { + var zk K + n.keys[n.count] = zk + var zv V + n.vals[n.count] = zv +} + +// popBack removes and returns the last element in the list. +func (n *node[K, V]) popBack() (rk K, rv V, child *node[K, V]) { + n.count-- + rk, rv = n.keys[n.count], n.vals[n.count] + n.clear(n.count) + if n.leaf { + return rk, rv, nil + } + child = n.children[n.count+1] + n.children[n.count+1] = nil + return rk, rv, child +} + +// popFront removes and returns the first element in the list. +func (n *node[K, V]) popFront() (rk K, rv V, child *node[K, V]) { + n.count-- + if !n.leaf { + child = n.children[0] + copy(n.children[:n.count+1], n.children[1:n.count+2]) + n.children[n.count+1] = nil + } + rk, rv = n.keys[0], n.vals[0] + copy(n.keys[:n.count], n.keys[1:n.count+1]) + copy(n.vals[:n.count], n.vals[1:n.count+1]) + n.clear(n.count) + return rk, rv, child +} + +// find returns the index where the given item should be inserted into this +// list. 'found' is true if the item already exists in the list at the given +// index. +func (n *node[K, V]) find(k K, cmp LessFn[K]) (index int, found bool) { + // Logic copied from sort.Search. Inlining this gave + // an 11% speedup on BenchmarkBTreeDeleteInsert. + i, j := 0, int(n.count) + for i < j { + h := int(uint(i+j) >> 1) // avoid overflow when computing h + // i ≤ h < j + if cmp(k, n.keys[h]) { + j = h + } else if cmp(n.keys[h], k) { + i = h + 1 + } else { + return h, true + } + } + return i, false +} + +// split splits the given node at the given index. The current node shrinks, +// and this function returns the k K, v Vhat existed at that index and a new +// node containing all items/children after it. +// +// Before: +// +// +-----------+ +// | x y z | +// +--/-/-\-\--+ +// +// After: +// +// +-----------+ +// | y | +// +----/-\----+ +// / \ +// v v +// +-----------+ +-----------+ +// | x | | z | +// +-----------+ +-----------+ +// +func (n *node[K, V]) split(i int) (k K, v V, next *node[K, V]) { + k, v = n.keys[i], n.vals[i] + next = &node[K, V]{ + leaf: n.leaf, + } + next.count = n.count - int16(i+1) + copy(next.keys[:], n.keys[i+1:n.count]) + copy(next.vals[:], n.vals[i+1:n.count]) + var zk K + var zv V + for j := int16(i); j < n.count; j++ { + n.keys[j] = zk + n.vals[j] = zv + } + if !n.leaf { + copy(next.children[:], n.children[i+1:n.count+1]) + for j := int16(i + 1); j <= n.count; j++ { + n.children[j] = nil + } + } + n.count = int16(i) + return k, v, next +} + +// insert inserts an item into the btree rooted at this node, making sure no +// nodes in the btree exceed maxItems items. Returns true if an existing item +// was replaced and false if an item was inserted. +func (n *node[K, V]) insert(k K, v V, cmp LessFn[K]) (replaced bool) { + i, found := n.find(k, cmp) + if found { + n.vals[i] = v + return true + } + if n.leaf { + n.insertAt(i, k, v, nil) + return false + } + if n.children[i].count >= maxItems { + splitK, splitV, splitNode := n.children[i].split(maxItems / 2) + n.insertAt(i, splitK, splitV, splitNode) + if cmp(k, n.keys[i]) { + // no change, we want first split node + } else if cmp(n.keys[i], k) { + i++ // we want second split node + } else { + n.keys[i] = k + n.vals[i] = v + return true + } + } + return n.children[i].insert(k, v, cmp) +} + +// removeMax removes and returns the maximum item from the btree rooted at +// this node. +func (n *node[K, V]) removeMax() (k K, v V) { + if n.leaf { + n.count-- + k, v = n.keys[n.count], n.vals[n.count] + n.clear(n.count) + return k, v + } + // Recurse into max child. + i := int(n.count) + if n.children[i].count <= minItems { + // Child not large enough to remove from. + n.rebalanceOrMerge(i) + return n.removeMax() // redo + } + child := n.children[i] + return child.removeMax() +} + +// rebalanceOrMerge grows child 'i' to ensure it has sufficient room to remove +// an item from it while keeping it at or above minItems. +func (n *node[K, V]) rebalanceOrMerge(i int) { + switch { + case i > 0 && n.children[i-1].count > minItems: + // Rebalance from left sibling. + // + // +-----------+ + // | y | + // +----/-\----+ + // / \ + // v v + // +-----------+ +-----------+ + // | x | | | + // +----------\+ +-----------+ + // \ + // v + // a + // + // After: + // + // +-----------+ + // | x | + // +----/-\----+ + // / \ + // v v + // +-----------+ +-----------+ + // | | | y | + // +-----------+ +/----------+ + // / + // v + // a + // + left := n.children[i-1] + child := n.children[i] + xLak, xLav, grandChild := left.popBack() + yLak, yLav := n.keys[i-1], n.vals[i-1] + child.pushFront(yLak, yLav, grandChild) + n.keys[i-1] = xLak + n.vals[i-1] = xLav + + case i < int(n.count) && n.children[i+1].count > minItems: + // Rebalance from right sibling. + // + // +-----------+ + // | y | + // +----/-\----+ + // / \ + // v v + // +-----------+ +-----------+ + // | | | x | + // +-----------+ +/----------+ + // / + // v + // a + // + // After: + // + // +-----------+ + // | x | + // +----/-\----+ + // / \ + // v v + // +-----------+ +-----------+ + // | y | | | + // +----------\+ +-----------+ + // \ + // v + // a + // + right := n.children[i+1] + child := n.children[i] + xLak, xLav, grandChild := right.popFront() + yLak, yLav := n.keys[i], n.vals[i] + child.pushBack(yLak, yLav, grandChild) + n.keys[i] = xLak + n.vals[i] = xLav + + default: + // Merge with either the left or right sibling. + // + // +-----------+ + // | u y v | + // +----/-\----+ + // / \ + // v v + // +-----------+ +-----------+ + // | x | | z | + // +-----------+ +-----------+ + // + // After: + // + // +-----------+ + // | u v | + // +-----|-----+ + // | + // v + // +-----------+ + // | x y z | + // +-----------+ + // + if i >= int(n.count) { + i = int(n.count - 1) + } + child := n.children[i] + mergeK, mergeV, mergeChild := n.removeAt(i) + child.keys[child.count] = mergeK + child.vals[child.count] = mergeV + copy(child.keys[child.count+1:], mergeChild.keys[:mergeChild.count]) + copy(child.vals[child.count+1:], mergeChild.vals[:mergeChild.count]) + if !child.leaf { + copy(child.children[child.count+1:], mergeChild.children[:mergeChild.count+1]) + } + child.count += mergeChild.count + 1 + } +} + +// remove removes an item from the btree rooted at this node. Returns the +// k K, v Vhat was removed or nil if no matching item was found. +func (n *node[K, V]) remove(k K, cmp LessFn[K]) (rk K, rv V, found bool) { + i, found := n.find(k, cmp) + if n.leaf { + if found { + rk, rv, _ = n.removeAt(i) + return rk, rv, true + } + return rk, rv, false + } + if n.children[i].count <= minItems { + // Child not large enough to remove from. + n.rebalanceOrMerge(i) + return n.remove(k, cmp) // redo + } + child := n.children[i] + if found { + // Replace the item being removed with the max item in our left child. + rk, rv = n.keys[i], n.vals[i] + n.keys[i], n.vals[i] = child.removeMax() + return rk, rv, true + } + return child.remove(k, cmp) +} + +// Iterator is responsible for search and traversal within a BTree. +type iterator[K, V, N any, NP nodeI[K, V, N]] struct { + r *btree[K, V, N, NP] + n NP + pos int16 + s iterStack[K, V, N, NP] +} + +// Cur returns the item at the Iterator's current position. It is illegal +// to call Cur if the Iterator is not valid. +func (i *iterator[K, V, N, NP]) Cur() (K, V) { + return i.n.at(i.pos) +} + +// SeekGE seeks to the first item greater-than or equal to the provided +// item. +func (i *iterator[K, V, N, NP]) SeekGE(k K) { + i.reset() + if i.n == nil { + return + } + for { + pos, found := i.n.find(k, i.r.cmp) + i.pos = int16(pos) + if found { + return + } + if i.n.isLeaf() { + if i.pos == i.n.len() { + i.Next() + } + return + } + i.descend(i.n, i.pos) + } +} + +// SeekLT seeks to the first item less-than the provided item. +func (i *iterator[K, V, N, NP]) SeekLT(k K) { + i.reset() + if i.n == nil { + return + } + for { + pos, found := i.n.find(k, i.r.cmp) + i.pos = int16(pos) + if found || i.n.isLeaf() { + i.Prev() + return + } + i.descend(i.n, i.pos) + } +} + +func (i *iterator[K, V, N, NP]) reset() { + i.n = i.r.root + i.pos = -1 + i.s.reset() +} + +func (i *iterator[K, V, N, NP]) descend(n NP, pos int16) { + i.s.push(iterFrame[K, V, N, NP]{n: n, pos: pos}) + i.n = n.child(pos) + i.pos = 0 +} + +// ascend ascends up to the current node's parent and resets the position +// to the one previously set for this parent node. +func (i *iterator[K, V, N, NP]) ascend() { + f := i.s.pop() + i.n = f.n + i.pos = f.pos +} + +// First seeks to the first item in the btree. +func (i *iterator[K, V, N, NP]) First() { + i.reset() + if i.n == nil { + return + } + for !i.n.isLeaf() { + i.descend(i.n, 0) + } + i.pos = 0 +} + +// Last seeks to the last item in the btree. +func (i *iterator[K, V, N, NP]) Last() { + i.reset() + if i.n == nil { + return + } + for !i.n.isLeaf() { + i.descend(i.n, i.n.len()) + } + i.pos = i.n.len() - 1 +} + +// Next positions the Iterator to the item immediately following +// its current position. +func (i *iterator[K, V, N, NP]) Next() { + if i.n == nil { + return + } + + if i.n.isLeaf() { + i.pos++ + if i.pos < i.n.len() { + return + } + for i.s.len() > 0 && i.pos >= i.n.len() { + i.ascend() + } + return + } + + i.descend(i.n, i.pos+1) + for !i.n.isLeaf() { + i.descend(i.n, 0) + } + i.pos = 0 +} + +// Prev positions the Iterator to the item immediately preceding +// its current position. +func (i *iterator[K, V, N, NP]) Prev() { + if i.n == nil { + return + } + + if i.n.isLeaf() { + i.pos-- + if i.pos >= 0 { + return + } + for i.s.len() > 0 && i.pos < 0 { + i.ascend() + i.pos-- + } + return + } + + i.descend(i.n, i.pos) + for !i.n.isLeaf() { + i.descend(i.n, i.n.len()) + } + i.pos = i.n.len() - 1 +} + +// Valid returns whether the Iterator is positioned at a valid position. +func (i *iterator[K, V, N, NP]) Valid() bool { + return i.n != nil && i.pos >= 0 && i.pos < i.n.len() +} + +// iterStack represents a stack of (node, pos) tuples, which captures +// iteration state as an Iterator descends a btree. +type iterStack[K, V, N any, NP nodeI[K, V, N]] struct { + a iterStackArr[K, V, N, NP] + aLen int16 // -1 when using s + s []iterFrame[K, V, N, NP] +} + +// Used to avoid allocations for stacks below a certain size. +type iterStackArr[K, V, N any, NP nodeI[K, V, N]] [3]iterFrame[K, V, N, NP] + +type iterFrame[K, V, N any, NP nodeI[K, V, N]] struct { + n NP + pos int16 +} + +func (is *iterStack[K, V, N, NP]) push(f iterFrame[K, V, N, NP]) { + if is.aLen == -1 { + is.s = append(is.s, f) + } else if int(is.aLen) == len(is.a) { + is.s = make([](iterFrame[K, V, N, NP]), int(is.aLen)+1, 2*int(is.aLen)) + copy(is.s, is.a[:]) + is.s[int(is.aLen)] = f + is.aLen = -1 + } else { + is.a[is.aLen] = f + is.aLen++ + } +} + +func (is *iterStack[K, V, N, NP]) pop() iterFrame[K, V, N, NP] { + if is.aLen == -1 { + f := is.s[len(is.s)-1] + is.s = is.s[:len(is.s)-1] + return f + } + is.aLen-- + return is.a[is.aLen] +} + +func (is *iterStack[K, V, N, NP]) len() int { + if is.aLen == -1 { + return len(is.s) + } + return int(is.aLen) +} + +func (is *iterStack[K, V, N, NP]) reset() { + if is.aLen == -1 { + is.s = is.s[:0] + } else { + is.aLen = 0 + } +} diff --git a/design/44253-generic-array-sizes/proposal.md b/design/44253-generic-array-sizes/proposal.md new file mode 100644 index 00000000..6de19f75 --- /dev/null +++ b/design/44253-generic-array-sizes/proposal.md @@ -0,0 +1,250 @@ +# Proposal: Generic parameterization of array sizes + +Author(s): Andrew Werner + +Last updated: 2/13/2021 + +## Abstract + +With the type parameters generics proposal has been accepted, though not yet +fully specified or implemented. That proposal lists the following omission: + +> No parameterization on non-type values such as constants. This arises most +obviously for arrays, where it might sometimes be convenient to write type +`Matrix[n int] [n][n]float64`. It might also sometimes be useful to specify +significant values for a container type, such as a default value for elements. + +This proposal seeks to resolve this limitation by (a) specifying when `len` can +be used as a compile-time constant and (b) adding syntax to specify constraints +for all arrays of a given type in type lists. + +## Background + +An important property of the generics proposal is that it enables the creation +of libraries of specialized container data structures. The existence of such +libraries will help developers write more efficient code as these data +structures will be able to allocate fewer object and provide greater access +locality. [This Google blog post][Block Based Data Structures] about block-based +C++ data drives home the point. + +The justification is laid out in the omission of the type parameter proposal. +The motivation that I've stumbled upon is in trying to implement a B-Tree +and allowing the client to dictate the degree of the node. + +One initial idea would be to allow the client to provide the actual array +which will be backing the data inside the node as a type parameter. This might +actually be okay in some data structure user cases but in a B-Tree it's bad +because we still would like to instantiate an array for the pointers and that +array needs to have a size that is a function of the data array. + +The proposal here seeks to make it possible for clients to provide default +values for array sizes of generic data structures in a way that is minimally +invasive to the concepts which go already has. The shorthand comment stated +in the Omission of the Type Parameter Proposal waves its hand at what feels +like a number of new and complex concepts for the language. + +## Proposal + +This proposals attempts to side-step questions of how one might provide a +scalar value in a type constraint by not ever providing a scalar directly. +This proposal recognizes that constants can be used to specify array lengths. +It also notes that the value of `len()` can be computed as a compile-time +constant in some cases. Lastly, it observes that type lists could be extended +minimally to indicate a constraint that a type is an array of a given type +without constraining the length of the array. + +### The vanilla generic B-Tree + +Let's explore the example of a generic B-Tree with a fixed-size buffer. Find +such an example [here](https://go2goplay.golang.org/p/A5auAIdW2ZR). + +```go +// These constants are the wart. +const ( + degree = 16 + maxItems = 2*degree - 1 + minItems = degree - 1 +) + +func NewBTree[K, V any](cmp LessFn[K]) OrderedMap[K, V] { + return &btree[K, V]{cmp: cmp} +} + +type btree[K, V any] struct { + cmp LessFn[K] + root *node[K, V] +} + +// ... + +type node[K, V any] struct { + count int16 + leaf bool + keys [maxItems]K + vals [maxItems]V + children [maxItems + 1]*node[K, V] +} +``` + +### Parameterized nodes + +Then we allow parameterization on the node type within the btree implementation +so that different node concrete types with different memory layouts may be +used. Find an example of this generalization +[here](https://go2goplay.golang.org/p/TFn9BujIlc3). + +```go +type nodeI[K, V, N any] interface { + type *N + find(K, LessFn[K]) (idx int, found bool) + insert(K, V, LessFn[K]) (replaced bool) + remove(K, LessFn[K]) (K, V, bool) + len() int16 + at(idx int16) (K, V) + child(idx int16) *N + isLeaf() bool +} + +func NewBTree[K, V any](cmp LessFn[K]) OrderedMap[K, V] { + type N = node[K, V] + return &btree[K, V, N, *N]{ + cmp: cmp, + newNode: func(isLeaf bool) *N { + return &N{leaf: isLeaf} + }, + } +} + +type btree[K, V, N any, NP nodeI[K, V, N]] struct { + len int + cmp LessFn[K] + root NP + newNode func(isLeaf bool) NP +} + +type node[K, V any] struct { + count int16 + leaf bool + keys [maxItems]K + vals [maxItems]V + children [maxItems + 1]*node[K, V] +} +``` + +This still ends up using constants and there's no really easy +way around that. You might want to parameterize the arrays into the node like +in [this example](https://go2goplay.golang.org/p/JGgyabtu_9F). This still +doesn't tell a story about how to relate the children array to the items. + +### The proposal to parameterize the arrays + +Instead, we'd like to find a way to express the idea that there's a size +constant which can be used in the type definitions. The proposal would +result in an implementation that looked like +[this](https://go2goplay.golang.org/p/4o36RLxF73C) + +```go + +// StructArr is a constraint that says that a type is an array of empty +// structs of any length. +type StructArr interface { + type [...]struct{} +} + +type btree[K, V, N any, NP nodeI[K, V, N]] struct { + len int + cmp LessFn[K] + root NP + newNode func(isLeaf bool) NP +} + +// NewBTree constructs a generic BTree-backed map with degree 16. +func NewBTree[K, V any](cmp LessFn[K]) OrderedMap[K, V] { + const defaultDegree = 16 + return NewBTreeWithDegree[K, V, [defaultDegree]struct{}](cmp) +} + +// NewBTreeWithDegree constructs a generic BTree-backed map with degree equal +// to the length of the array used as type parameter A. +func NewBTreeWithDegree[K, V any, A StructArr](cmp LessFn[K]) OrderedMap[K, V] { + type N = node[K, V, A] + return &btree[K, V, N, *N]{ + cmp: cmp, + newNode: func(isLeaf bool) *N { + return &N{leaf: isLeaf} + }, + } +} + +type node[K, V any, A StructArr] struct { + count int16 + leaf bool + keys [2*len(A) - 1]K + values [2*len(A) - 1]V + children [2 * len(A)]*node[K, V, A] +} +``` +### The Matrix example + +The example of the omission in type parameter proposal could be achieved in +the following way: + +```go +type Dim interface { + type [...]struct{} +} + +type SquareFloatMatrix2D[D Dim] [len(D)][len(D)]float64 +``` + +### Summary + +1) Support type list constraints to express that a type is an array + + +```go +// Array expresses a constraint that a type is an array of T of any +// size. +type Array[T any] interface { + type [...]T +} +``` + +2) Support a compile-time constant expression for `len([...]T)` + +This handy syntax would permit parameterization of arrays relative to other +array types. Note that the constant expression `len` function on array types +could actually be implemented today using `unsafe.Sizeof` by a parameterization +over an array whose members have non-zero size. For example `len` could be +written as `unsafe.Sizeof([...]T)/unsafe.Sizeof(T)` so long as +`unsafe.Sizeof(T) > 0`. + +## Rationale + +This approach is simpler than generally providing a constant scalar expression +parameterization of generic types. Of the two elements of the proposal, neither +feels particularly out of line with the design of the language or its concepts. +The `[...]T` syntax exists in the language to imply length inference for array +literals and is not a hard to imagine concept. It is the deeper requirement to +make this proposal work. + +One potential downside of this proposal is that we're not really using the +array for anything other than its size which may feel awkward. For that reason +I've opted to use a constraint which forces the array to use `struct{}` values +to indicate that the structure of the elements isn't relevant. This awkwardness +feels justified to side-step introduces scalars to type parameters. + +## Compatibility + +This proposal is fully backwards compatible with all of the language and also +the now accepted type parameters proposal. + +## Implementation + +Neither of the two features of this proposal feel particularly onerous to +implement. My guess is that the `[...]T` type list constraint would be extremely +straightforward given an implementation of type parameters. The `len` +implementation is also likely to be straightforward given the existence of +both compile-time evaluation of `len` expressions on array types which exist +in the language and the constant nature of `unsafe.Sizeof`. Maybe there'd be +some pain in deferring the expression evaluation until after type checking. diff --git a/design/44253-generic-array-sizes/proposal_btree.go2 b/design/44253-generic-array-sizes/proposal_btree.go2 new file mode 100644 index 00000000..e69de29b diff --git a/design/44253-generic-array-sizes/vanilla_btree.go2 b/design/44253-generic-array-sizes/vanilla_btree.go2 new file mode 100644 index 00000000..27b894ca --- /dev/null +++ b/design/44253-generic-array-sizes/vanilla_btree.go2 @@ -0,0 +1,654 @@ +package main + +import "fmt" + +// These constants are the wart. +const ( + degree = 16 + maxItems = 2*degree - 1 + minItems = degree - 1 +) + +func main() { + bt := NewBTree[int, string](func(i, j int) bool { return i < j }) + bt.Insert(2, "world") + bt.Insert(1, "hello") + bt.Iterate(func(it Iterator[int, string]) { + for it.First(); it.Valid(); it.Next() { + fmt.Println(it.CurV()) + } + }) + +} + +type LessFn[K any] func(a, b K) bool + +type OrderedMap[K, V any] interface { + Insert(K, V) bool + Find(K) (V, bool) + Remove(K) (removed bool) + Iterate(func(Iterator[K, V])) +} + +type Iterator[K, V any] interface { + First() + Last() + Next() + Prev() + SeekGE(K) + SeekLT(K) + Valid() bool + CurK() K + CurV() V +} + +func NewBTree[K, V any](cmp LessFn[K]) OrderedMap[K, V] { + return &btree[K, V]{cmp: cmp} +} + +type btree[K, V any] struct { + len int + cmp LessFn[K] + root *node[K, V] +} + +func (bt *btree[K, V]) Find(k K) (v V, found bool) { + bt.Iterate(func(it Iterator[K, V]) { + it.SeekGE(k) + if !it.Valid() { + return + } + fk := it.CurK() + if !bt.cmp(fk, k) && !bt.cmp(k, fk) { + v, found = it.CurV(), true + } + }) + return v, found +} + +func (bt *btree[K, V]) Iterate(f func(it Iterator[K, V])) { + it := iterator[K, V]{r: bt} + f(&it) +} + +func (t *btree[K, V]) Insert(k K, v V) (replaced bool) { + if t.root == nil { + t.root = &node[K, V]{leaf: true} + } + return t.root.insert(k, v, t.cmp) +} + +func (t *btree[K, V]) Remove(k K) (removed bool) { + if t.root == nil { + return false + } + _, _, removed = t.root.remove(k, t.cmp) + return removed +} + +type node[K, V any] struct { + count int16 + leaf bool + keys [maxItems]K + values [maxItems]V + children [maxItems + 1]*node[K, V] +} + +func (n *node[K, V]) insertAt(index int, k K, v V, nd *node[K, V]) { + if index < int(n.count) { + copy(n.keys[index+1:n.count+1], n.keys[index:n.count]) + copy(n.values[index+1:n.count+1], n.values[index:n.count]) + if !n.leaf { + copy(n.children[index+2:n.count+2], n.children[index+1:n.count+1]) + } + } + n.keys[index] = k + n.values[index] = v + if !n.leaf { + n.children[index+1] = nd + } + n.count++ +} + +func (n *node[K, V]) pushBack(k K, v V, nd *node[K, V]) { + n.keys[n.count] = k + n.values[n.count] = v + if !n.leaf { + n.children[n.count+1] = nd + } + n.count++ +} + +func (n *node[K, V]) pushFront(k K, v V, nd *node[K, V]) { + if !n.leaf { + copy(n.children[1:n.count+2], n.children[:n.count+1]) + n.children[0] = nd + } + copy(n.keys[1:n.count+1], n.keys[:n.count]) + copy(n.values[1:n.count+1], n.values[:n.count]) + n.keys[0] = k + n.values[0] = v + n.count++ +} + +// removeAt removes a value at a given index, pulling all subsequent values +// back. +func (n *node[K, V]) removeAt(index int) (rk K, rv V, child *node[K, V]) { + if !n.leaf { + child = n.children[index+1] + copy(n.children[index+1:n.count], n.children[index+2:n.count+1]) + n.children[n.count] = nil + } + n.count-- + rk, rv = n.keys[index], n.values[index] + copy(n.keys[index:n.count], n.keys[index+1:n.count+1]) + copy(n.values[index:n.count], n.values[index+1:n.count+1]) + n.clear(n.count) + return rk, rv, child +} + +func (n *node[K, V]) clear(idx int16) { + var zk K + n.keys[n.count] = zk + var zv V + n.values[n.count] = zv +} + +// popBack removes and returns the last element in the list. +func (n *node[K, V]) popBack() (rk K, rv V, child *node[K, V]) { + n.count-- + rk, rv = n.keys[n.count], n.values[n.count] + n.clear(n.count) + if n.leaf { + return rk, rv, nil + } + child = n.children[n.count+1] + n.children[n.count+1] = nil + return rk, rv, child +} + +// popFront removes and returns the first element in the list. +func (n *node[K, V]) popFront() (rk K, rv V, child *node[K, V]) { + n.count-- + if !n.leaf { + child = n.children[0] + copy(n.children[:n.count+1], n.children[1:n.count+2]) + n.children[n.count+1] = nil + } + rk, rv = n.keys[0], n.values[0] + copy(n.keys[:n.count], n.keys[1:n.count+1]) + copy(n.values[:n.count], n.values[1:n.count+1]) + n.clear(n.count) + return rk, rv, child +} + +// find returns the index where the given item should be inserted into this +// list. 'found' is true if the item already exists in the list at the given +// index. +func (n *node[K, V]) find(k K, cmp LessFn[K]) (index int, found bool) { + // Logic copied from sort.Search. Inlining this gave + // an 11% speedup on BenchmarkBTreeDeleteInsert. + i, j := 0, int(n.count) + for i < j { + h := int(uint(i+j) >> 1) // avoid overflow when computing h + // i ≤ h < j + if cmp(k, n.keys[h]) { + j = h + } else if cmp(n.keys[h], k) { + i = h + 1 + } else { + return h, true + } + } + return i, false +} + +// split splits the given node at the given index. The current node shrinks, +// and this function returns the k K, v Vhat existed at that index and a new +// node containing all items/children after it. +// +// Before: +// +// +-----------+ +// | x y z | +// +--/-/-\-\--+ +// +// After: +// +// +-----------+ +// | y | +// +----/-\----+ +// / \ +// v v +// +-----------+ +-----------+ +// | x | | z | +// +-----------+ +-----------+ +// +func (n *node[K, V]) split(i int) (k K, v V, next *node[K, V]) { + k, v = n.keys[i], n.values[i] + next = &node[K, V]{ + leaf: n.leaf, + } + next.count = n.count - int16(i+1) + copy(next.keys[:], n.keys[i+1:n.count]) + copy(next.values[:], n.values[i+1:n.count]) + var zk K + var zv V + for j := int16(i); j < n.count; j++ { + n.keys[j] = zk + n.values[j] = zv + } + if !n.leaf { + copy(next.children[:], n.children[i+1:n.count+1]) + for j := int16(i + 1); j <= n.count; j++ { + n.children[j] = nil + } + } + n.count = int16(i) + return k, v, next +} + +// insert inserts an item into the btree rooted at this node, making sure no +// nodes in the btree exceed maxItems items. Returns true if an existing item +// was replaced and false if an item was inserted. +func (n *node[K, V]) insert(k K, v V, cmp LessFn[K]) (replaced bool) { + i, found := n.find(k, cmp) + if found { + n.values[i] = v + return true + } + if n.leaf { + n.insertAt(i, k, v, nil) + return false + } + if n.children[i].count >= maxItems { + splitK, splitV, splitNode := n.children[i].split(maxItems / 2) + n.insertAt(i, splitK, splitV, splitNode) + if cmp(k, n.keys[i]) { + // no change, we want first split node + } else if cmp(n.keys[i], k) { + i++ // we want second split node + } else { + n.keys[i] = k + n.values[i] = v + return true + } + } + return n.children[i].insert(k, v, cmp) +} + +// removeMax removes and returns the maximum item from the tree rooted at +// this node. +func (n *node[K, V]) removeMax() (k K, v V) { + if n.leaf { + n.count-- + k, v = n.keys[n.count], n.values[n.count] + n.clear(n.count) + return k, v + } + // Recurse into max child. + i := int(n.count) + if n.children[i].count <= minItems { + // Child not large enough to remove from. + n.rebalanceOrMerge(i) + return n.removeMax() // redo + } + child := n.children[i] + return child.removeMax() +} + +// rebalanceOrMerge grows child 'i' to ensure it has sufficient room to remove +// an item from it while keeping it at or above minItems. +func (n *node[K, V]) rebalanceOrMerge(i int) { + switch { + case i > 0 && n.children[i-1].count > minItems: + // Rebalance from left sibling. + // + // +-----------+ + // | y | + // +----/-\----+ + // / \ + // v v + // +-----------+ +-----------+ + // | x | | | + // +----------\+ +-----------+ + // \ + // v + // a + // + // After: + // + // +-----------+ + // | x | + // +----/-\----+ + // / \ + // v v + // +-----------+ +-----------+ + // | | | y | + // +-----------+ +/----------+ + // / + // v + // a + // + left := n.children[i-1] + child := n.children[i] + xLak, xLav, grandChild := left.popBack() + yLak, yLav := n.keys[i-1], n.values[i-1] + child.pushFront(yLak, yLav, grandChild) + n.keys[i-1] = xLak + n.values[i-1] = xLav + + case i < int(n.count) && n.children[i+1].count > minItems: + // Rebalance from right sibling. + // + // +-----------+ + // | y | + // +----/-\----+ + // / \ + // v v + // +-----------+ +-----------+ + // | | | x | + // +-----------+ +/----------+ + // / + // v + // a + // + // After: + // + // +-----------+ + // | x | + // +----/-\----+ + // / \ + // v v + // +-----------+ +-----------+ + // | y | | | + // +----------\+ +-----------+ + // \ + // v + // a + // + right := n.children[i+1] + child := n.children[i] + xLak, xLav, grandChild := right.popFront() + yLak, yLav := n.keys[i], n.values[i] + child.pushBack(yLak, yLav, grandChild) + n.keys[i] = xLak + n.values[i] = xLav + + default: + // Merge with either the left or right sibling. + // + // +-----------+ + // | u y v | + // +----/-\----+ + // / \ + // v v + // +-----------+ +-----------+ + // | x | | z | + // +-----------+ +-----------+ + // + // After: + // + // +-----------+ + // | u v | + // +-----|-----+ + // | + // v + // +-----------+ + // | x y z | + // +-----------+ + // + if i >= int(n.count) { + i = int(n.count - 1) + } + child := n.children[i] + mergeK, mergeV, mergeChild := n.removeAt(i) + child.keys[child.count] = mergeK + child.values[child.count] = mergeV + copy(child.keys[child.count+1:], mergeChild.keys[:mergeChild.count]) + copy(child.values[child.count+1:], mergeChild.values[:mergeChild.count]) + if !child.leaf { + copy(child.children[child.count+1:], mergeChild.children[:mergeChild.count+1]) + } + child.count += mergeChild.count + 1 + } +} + +// remove removes an item from the tree rooted at this node. Returns the +// k K, v Vhat was removed or nil if no matching item was found. +func (n *node[K, V]) remove(k K, cmp LessFn[K]) (rk K, rv V, found bool) { + i, found := n.find(k, cmp) + if n.leaf { + if found { + rk, rv, _ = n.removeAt(i) + return rk, rv, true + } + return rk, rv, false + } + if n.children[i].count <= minItems { + // Child not large enough to remove from. + n.rebalanceOrMerge(i) + return n.remove(k, cmp) // redo + } + child := n.children[i] + if found { + // Replace the item being removed with the max item in our left child. + rk, rv = n.keys[i], n.values[i] + n.keys[i], n.values[i] = child.removeMax() + return rk, rv, true + } + return child.remove(k, cmp) +} + +// iterStack represents a stack of (node, pos) tuples, which captures +// iteration state as an Iterator descends a btree. +type iterStack[K, V any] struct { + a iterStackArr[K, V] + aLen int16 // -1 when using s + s []iterFrame[K, V] +} + +// Used to avoid allocations for stacks below a certain size. +type iterStackArr[K, V any] [3]iterFrame[K, V] + +type iterFrame[K, V any] struct { + n *node[K, V] + pos int16 +} + +func (is *iterStack[K, V]) push(f iterFrame[K, V]) { + if is.aLen == -1 { + is.s = append(is.s, f) + } else if int(is.aLen) == len(is.a) { + is.s = make([](iterFrame[K, V]), int(is.aLen)+1, 2*int(is.aLen)) + copy(is.s, is.a[:]) + is.s[int(is.aLen)] = f + is.aLen = -1 + } else { + is.a[is.aLen] = f + is.aLen++ + } +} + +func (is *iterStack[K, V]) pop() iterFrame[K, V] { + if is.aLen == -1 { + f := is.s[len(is.s)-1] + is.s = is.s[:len(is.s)-1] + return f + } + is.aLen-- + return is.a[is.aLen] +} + +func (is *iterStack[K, V]) len() int { + if is.aLen == -1 { + return len(is.s) + } + return int(is.aLen) +} + +func (is *iterStack[K, V]) reset() { + if is.aLen == -1 { + is.s = is.s[:0] + } else { + is.aLen = 0 + } +} + +// Iterator is responsible for search and traversal within a BTree. +type iterator[K, V any] struct { + r *btree[K, V] + n *node[K, V] + pos int16 + s iterStack[K, V] +} + +func (i *iterator[K, V]) reset() { + i.n = i.r.root + i.pos = -1 + i.s.reset() +} + +func (i *iterator[K, V]) descend(n *node[K, V], pos int16) { + i.s.push(iterFrame[K, V]{n: n, pos: pos}) + i.n = n.children[pos] + i.pos = 0 +} + +// ascend ascends up to the current node's parent and resets the position +// to the one previously set for this parent node. +func (i *iterator[K, V]) ascend() { + f := i.s.pop() + i.n = f.n + i.pos = f.pos +} + +// SeekGE seeks to the first item greater-than or equal to the provided +// item. +func (i *iterator[K, V]) SeekGE(k K) { + i.reset() + if i.n == nil { + return + } + for { + pos, found := i.n.find(k, i.r.cmp) + i.pos = int16(pos) + if found { + return + } + if i.n.leaf { + if i.pos == i.n.count { + i.Next() + } + return + } + i.descend(i.n, i.pos) + } +} + +// SeekLT seeks to the first item less-than the provided item. +func (i *iterator[K, V]) SeekLT(k K) { + i.reset() + if i.n == nil { + return + } + for { + pos, found := i.n.find(k, i.r.cmp) + i.pos = int16(pos) + if found || i.n.leaf { + i.Prev() + return + } + i.descend(i.n, i.pos) + } +} + +// First seeks to the first item in the btree. +func (i *iterator[K, V]) First() { + i.reset() + if i.n == nil { + return + } + for !i.n.leaf { + i.descend(i.n, 0) + } + i.pos = 0 +} + +// Last seeks to the last item in the btree. +func (i *iterator[K, V]) Last() { + i.reset() + if i.n == nil { + return + } + for !i.n.leaf { + i.descend(i.n, i.n.count) + } + i.pos = i.n.count - 1 +} + +// Next positions the Iterator to the item immediately following +// its current position. +func (i *iterator[K, V]) Next() { + if i.n == nil { + return + } + + if i.n.leaf { + i.pos++ + if i.pos < i.n.count { + return + } + for i.s.len() > 0 && i.pos >= i.n.count { + i.ascend() + } + return + } + + i.descend(i.n, i.pos+1) + for !i.n.leaf { + i.descend(i.n, 0) + } + i.pos = 0 +} + +// Prev positions the Iterator to the item immediately preceding +// its current position. +func (i *iterator[K, V]) Prev() { + if i.n == nil { + return + } + + if i.n.leaf { + i.pos-- + if i.pos >= 0 { + return + } + for i.s.len() > 0 && i.pos < 0 { + i.ascend() + i.pos-- + } + return + } + + i.descend(i.n, i.pos) + for !i.n.leaf { + i.descend(i.n, i.n.count) + } + i.pos = i.n.count - 1 +} + +// Valid returns whether the Iterator is positioned at a valid position. +func (i *iterator[K, V]) Valid() bool { + return i.n != nil && i.pos >= 0 && i.pos < i.n.count +} + +// Cur returns the item at the Iterator's current position. It is illegal +// to call Cur if the Iterator is not valid. +func (i *iterator[K, V]) CurK() K { + return i.n.keys[i.pos] +} + +func (i *iterator[K, V]) CurV() V { + return i.n.values[i.pos] +}