Jun 18, 2022

golang泛型在btree的应用与优化

zzjin

golanggenericsbtree

14817 Words

2022-06-18 00:30 +0800

近几个月更新的 golang1.18 支持了泛型,正好前两天看到了 google 官方的 golang-btree 更新了 v2,使用泛型generics代替了interface{},号称提升了 40% 的性能

Show me the code

https://github.com/google/btree

重新整理后的函数级别的完整的diff (原始MR点这里):

展开函数层面的完整diff

btree.go → btree_generic.go RENAMED

@@ -1,4 +1,4 @@
-// Copyright 2014 Google Inc.
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
@@ -12,8 +12,13 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
-//go:build !go1.18
-// +build !go1.18
 // Package btree implements in-memory B-Trees of arbitrary degree.
 //
@@ -48,6 +53,11 @@
 // Its functions, therefore, exactly mirror those of
 // llrb.LLRB where possible. Unlike gollrb, though, we currently don't
 // support storing multiple equivalent values.
 package btree
 import (
@@ -72,32 +82,27 @@
 DefaultFreeListSize = 32
 )
-var (
-nilItems = make(items, 16)
-nilChildren = make(children, 16)
-)
-// FreeList represents a free list of btree nodes. By default each
 // BTree has its own FreeList, but multiple BTrees can share the same
-// FreeList.
 // Two Btrees using the same freelist are safe for concurrent write access.
-type FreeList struct {
 mu sync.Mutex
-freelist []*node
 }
-// NewFreeList creates a new free list.
 // size is the maximum size of the returned free list.
-func NewFreeList(size int) *FreeList {
-return &FreeList{freelist: make([]*node, 0, size)}
 }
-func (f *FreeList) newNode() (n *node) {
 f.mu.Lock()
 index := len(f.freelist) - 1
 if index < 0 {
 f.mu.Unlock()
-return new(node)
 }
 n = f.freelist[index]
 f.freelist[index] = nil
@@ -106,9 +111,7 @@
 return
 }
-// freeNode adds the given node to the list, returning true if it was added
-// and false if it was discarded.
-func (f *FreeList) freeNode(n *node) (out bool) {
 f.mu.Lock()
 if len(f.freelist) < cap(f.freelist) {
 f.freelist = append(f.freelist, n)
@@ -118,37 +121,55 @@
 return
 }
-// ItemIterator allows callers of Ascend* to iterate in-order over portions of
 // the tree. When this function returns false, iteration will stop and the
 // associated Ascend* function will immediately return.
-type ItemIterator func(i Item) bool
-// New creates a new B-Tree with the given degree.
 //
-// New(2), for example, will create a 2-3-4 tree (each node contains 1-3 items
 // and 2-4 children).
-func New(degree int) *BTree {
-return NewWithFreeList(degree, NewFreeList(DefaultFreeListSize))
 }
-// NewWithFreeList creates a new B-Tree that uses the given node free list.
-func NewWithFreeList(degree int, f *FreeList) *BTree {
 if degree <= 1 {
 panic("bad degree")
 }
-return &BTree{
 degree: degree,
-cow: &copyOnWriteContext{freelist: f},
 }
 }
 // items stores items in a node.
-type items []Item
 // insertAt inserts a value into the given index, pushing all subsequent values
 // forward.
-func (s *items) insertAt(index int, item Item) {
-*s = append(*s, nil)
 if index < len(*s) {
 copy((*s)[index+1:], (*s)[index:])
 }
@@ -157,100 +178,61 @@
 // removeAt removes a value at a given index, pulling all subsequent values
 // back.
-func (s *items) removeAt(index int) Item {
 item := (*s)[index]
 copy((*s)[index:], (*s)[index+1:])
-(*s)[len(*s)-1] = nil
 *s = (*s)[:len(*s)-1]
 return item
 }
 // pop removes and returns the last element in the list.
-func (s *items) pop() (out Item) {
 index := len(*s) - 1
 out = (*s)[index]
-(*s)[index] = nil
 *s = (*s)[:index]
 return
 }
 // truncate truncates this instance at index so that it contains only the
 // first index items. index must be less than or equal to length.
-func (s *items) truncate(index int) {
-var toClear items
 *s, toClear = (*s)[:index], (*s)[index:]
-for len(toClear) > 0 {
-toClear = toClear[copy(toClear, nilItems):]
 }
 }
 // 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 (s items) find(item Item) (index int, found bool) {
 i := sort.Search(len(s), func(i int) bool {
-return item.Less(s[i])
 })
-if i > 0 && !s[i-1].Less(item) {
 return i - 1, true
 }
 return i, false
 }
-// children stores child nodes in a node.
-type children []*node
-// insertAt inserts a value into the given index, pushing all subsequent values
-// forward.
-func (s *children) insertAt(index int, n *node) {
-*s = append(*s, nil)
-if index < len(*s) {
-copy((*s)[index+1:], (*s)[index:])
-}
-(*s)[index] = n
-}
-// removeAt removes a value at a given index, pulling all subsequent values
-// back.
-func (s *children) removeAt(index int) *node {
-n := (*s)[index]
-copy((*s)[index:], (*s)[index+1:])
-(*s)[len(*s)-1] = nil
-*s = (*s)[:len(*s)-1]
-return n
-}
-// pop removes and returns the last element in the list.
-func (s *children) pop() (out *node) {
-index := len(*s) - 1
-out = (*s)[index]
-(*s)[index] = nil
-*s = (*s)[:index]
-return
-}
-// truncate truncates this instance at index so that it contains only the
-// first index children. index must be less than or equal to length.
-func (s *children) truncate(index int) {
-var toClear children
-*s, toClear = (*s)[:index], (*s)[index:]
-for len(toClear) > 0 {
-toClear = toClear[copy(toClear, nilChildren):]
-}
-}
 // node is an internal node in a tree.
 //
 // It must at all times maintain the invariant that either
 // * len(children) == 0, len(items) unconstrained
 // * len(children) == len(items) + 1
-type node struct {
-items items
-children children
-cow *copyOnWriteContext
 }
-func (n *node) mutableFor(cow *copyOnWriteContext) *node {
 if n.cow == cow {
 return n
 }
@@ -258,20 +240,20 @@
 if cap(out.items) >= len(n.items) {
 out.items = out.items[:len(n.items)]
 } else {
-out.items = make(items, len(n.items), cap(n.items))
 }
 copy(out.items, n.items)
 // Copy children
 if cap(out.children) >= len(n.children) {
 out.children = out.children[:len(n.children)]
 } else {
-out.children = make(children, len(n.children), cap(n.children))
 }
 copy(out.children, n.children)
 return out
 }
-func (n *node) mutableChild(i int) *node {
 c := n.children[i].mutableFor(n.cow)
 n.children[i] = c
 return c
@@ -280,7 +262,7 @@
 // split splits the given node at the given index. The current node shrinks,
 // and this function returns the item that existed at that index and a new node
 // containing all items/children after it.
-func (n *node) split(i int) (Item, *node) {
 item := n.items[i]
 next := n.cow.newNode()
 next.items = append(next.items, n.items[i+1:]...)
@@ -294,7 +276,7 @@
 // maybeSplitChild checks if a child should be split, and if so splits it.
 // Returns whether or not a split occurred.
-func (n *node) maybeSplitChild(i, maxItems int) bool {
 if len(n.children[i].items) < maxItems {
 return false
 }
@@ -308,70 +290,70 @@
 // insert inserts an item into the subtree rooted at this node, making sure
 // no nodes in the subtree exceed maxItems items. Should an equivalent item be
 // be found/replaced by insert, it will be returned.
-func (n *node) insert(item Item, maxItems int) Item {
-i, found := n.items.find(item)
 if found {
 out := n.items[i]
 n.items[i] = item
-return out
 }
 if len(n.children) == 0 {
 n.items.insertAt(i, item)
-return nil
 }
 if n.maybeSplitChild(i, maxItems) {
 inTree := n.items[i]
 switch {
-case item.Less(inTree):
 // no change, we want first split node
-case inTree.Less(item):
 i++ // we want second split node
 default:
 out := n.items[i]
 n.items[i] = item
-return out
 }
 }
 return n.mutableChild(i).insert(item, maxItems)
 }
 // get finds the given key in the subtree and returns it.
-func (n *node) get(key Item) Item {
-i, found := n.items.find(key)
 if found {
-return n.items[i]
 } else if len(n.children) > 0 {
 return n.children[i].get(key)
 }
-return nil
 }
 // min returns the first item in the subtree.
-func min(n *node) Item {
 if n == nil {
-return nil
 }
 for len(n.children) > 0 {
 n = n.children[0]
 }
 if len(n.items) == 0 {
-return nil
 }
-return n.items[0]
 }
 // max returns the last item in the subtree.
-func max(n *node) Item {
 if n == nil {
-return nil
 }
 for len(n.children) > 0 {
 n = n.children[len(n.children)-1]
 }
 if len(n.items) == 0 {
-return nil
 }
-return n.items[len(n.items)-1]
 }
 // toRemove details what item to remove in a node.remove call.
@@ -384,27 +366,27 @@
 )
 // remove removes an item from the subtree rooted at this node.
-func (n *node) remove(item Item, minItems int, typ toRemove) Item {
 var i int
 var found bool
 switch typ {
 case removeMax:
 if len(n.children) == 0 {
-return n.items.pop()
 }
 i = len(n.items)
 case removeMin:
 if len(n.children) == 0 {
-return n.items.removeAt(0)
 }
 i = 0
 case removeItem:
-i, found = n.items.find(item)
 if len(n.children) == 0 {
 if found {
-return n.items.removeAt(i)
 }
-return nil
 }
 default:
 panic("invalid type")
@@ -424,8 +406,9 @@
 // We use our special-case 'remove' call with typ=maxItem to pull the
 // predecessor of item i (the rightmost leaf of our immediate left child)
 // and set it into where we pulled the item from.
-n.items[i] = child.remove(nil, minItems, removeMax)
-return out
 }
 // Final recursive call. Once we're here, we know that the item isn't in this
 // node and that the child is big enough to remove from.
@@ -451,7 +434,7 @@
 // We then simply redo our remove call, and the second time (regardless of
 // whether we're in case 1 or 2), we'll have enough items and can guarantee
 // that we hit case A.
-func (n *node) growChildAndRemove(i int, item Item, minItems int, typ toRemove) Item {
 if i > 0 && len(n.children[i-1].items) > minItems {
 // Steal from left child
 child := n.mutableChild(i)
@@ -495,6 +478,18 @@
 ascend = direction(+1)
 )
 // iterate provides a simple method for iterating over elements in the tree.
 //
 // When ascending, the 'start' should be less than 'stop' and when descending,
@@ -502,13 +497,13 @@
 // will force the iterator to include the first item when it equals 'start',
 // thus creating a "greaterOrEqual" or "lessThanEqual" rather than just a
 // "greaterThan" or "lessThan" queries.
-func (n *node) iterate(dir direction, start, stop Item, includeStart bool, hit bool, iter ItemIterator) (bool, bool) {
 var ok, found bool
 var index int
 switch dir {
 case ascend:
-if start != nil {
-index, _ = n.items.find(start)
 }
 for i := index; i < len(n.items); i++ {
 if len(n.children) > 0 {
@@ -516,12 +511,12 @@
 return hit, false
 }
 }
-if !includeStart && !hit && start != nil && !start.Less(n.items[i]) {
 hit = true
 continue
 }
 hit = true
-if stop != nil && !n.items[i].Less(stop) {
 return hit, false
 }
 if !iter(n.items[i]) {
@@ -534,8 +529,8 @@
 }
 }
 case descend:
-if start != nil {
-index, found = n.items.find(start)
 if !found {
 index = index - 1
 }
@@ -543,8 +538,8 @@
 index = len(n.items) - 1
 }
 for i := index; i >= 0; i-- {
-if start != nil && !n.items[i].Less(start) {
-if !includeStart || hit || start.Less(n.items[i]) {
 continue
 }
 }
@@ -553,7 +548,7 @@
 return hit, false
 }
 }
-if stop != nil && !stop.Less(n.items[i]) {
 return hit, false // continue
 }
 hit = true
@@ -570,28 +565,32 @@
 return hit, true
 }
-// Used for testing/debugging purposes.
-func (n *node) print(w io.Writer, level int) {
 fmt.Fprintf(w, "%sNODE:%v\n", strings.Repeat(" ", level), n.items)
 for _, c := range n.children {
 c.print(w, level+1)
 }
 }
-// BTree is an implementation of a B-Tree.
 //
-// BTree stores Item instances in an ordered structure, allowing easy insertion,
 // removal, and iteration.
 //
 // Write operations are not safe for concurrent mutation by multiple
 // goroutines, but Read operations are.
-type BTree struct {
 degree int
 length int
-root *node
-cow *copyOnWriteContext
 }
 // copyOnWriteContext pointers determine node ownership... a tree with a write
 // context equivalent to a node's write context is allowed to modify that node.
 // A tree whose write context does not match a node's is not allowed to modify
@@ -606,8 +605,9 @@
 // tree's context, that node is modifiable in place. Children of that node may
 // not share context, but before we descend into them, we'll make a mutable
 // copy.
-type copyOnWriteContext struct {
-freelist *FreeList
 }
 // Clone clones the btree, lazily. Clone should not be called concurrently,
@@ -621,7 +621,7 @@
 // will initially experience minor slow-downs caused by additional allocs and
 // copies due to the aforementioned copy-on-write logic, but should converge to
 // the original performance characteristics of the original tree.
-func (t *BTree) Clone() (t2 *BTree) {
 // Create two entirely new copy-on-write contexts.
 // This operation effectively creates three trees:
 // the original, shared nodes (old b.cow)
@@ -635,17 +635,17 @@
 }
 // maxItems returns the max number of items to allow per node.
-func (t *BTree) maxItems() int {
 return t.degree*2 - 1
 }
 // minItems returns the min number of items to allow per node (ignored for the
 // root node).
-func (t *BTree) minItems() int {
 return t.degree - 1
 }
-func (c *copyOnWriteContext) newNode() (n *node) {
 n = c.freelist.newNode()
 n.cow = c
 return
@@ -662,7 +662,7 @@
 // freeNode frees a node within a given COW context, if it's owned by that
 // context. It returns what happened to the node (see freeType const
 // documentation).
-func (c *copyOnWriteContext) freeNode(n *node) freeType {
 if n.cow == c {
 // clear to allow GC
 n.items.truncate(0)
@@ -679,19 +679,16 @@
 }
 // ReplaceOrInsert adds the given item to the tree. If an item in the tree
-// already equals the given one, it is removed from the tree and returned.
-// Otherwise, nil is returned.
 //
 // nil cannot be added to the tree (will panic).
-func (t *BTree) ReplaceOrInsert(item Item) Item {
-if item == nil {
-panic("nil item being added to BTree")
-}
 if t.root == nil {
 t.root = t.cow.newNode()
 t.root.items = append(t.root.items, item)
 t.length++
-return nil
 } else {
 t.root = t.root.mutableFor(t.cow)
 if len(t.root.items) >= t.maxItems() {
@@ -702,146 +699,149 @@
 t.root.children = append(t.root.children, oldroot, second)
 }
 }
-out := t.root.insert(item, t.maxItems())
-if out == nil {
 t.length++
 }
-return out
 }
 // Delete removes an item equal to the passed in item from the tree, returning
-// it. If no such item exists, returns nil.
-func (t *BTree) Delete(item Item) Item {
 return t.deleteItem(item, removeItem)
 }
 // DeleteMin removes the smallest item in the tree and returns it.
-// If no such item exists, returns nil.
-func (t *BTree) DeleteMin() Item {
-return t.deleteItem(nil, removeMin)
 }
 // DeleteMax removes the largest item in the tree and returns it.
-// If no such item exists, returns nil.
-func (t *BTree) DeleteMax() Item {
-return t.deleteItem(nil, removeMax)
 }
-func (t *BTree) deleteItem(item Item, typ toRemove) Item {
 if t.root == nil || len(t.root.items) == 0 {
-return nil
 }
 t.root = t.root.mutableFor(t.cow)
-out := t.root.remove(item, t.minItems(), typ)
 if len(t.root.items) == 0 && len(t.root.children) > 0 {
 oldroot := t.root
 t.root = t.root.children[0]
 t.cow.freeNode(oldroot)
 }
-if out != nil {
 t.length--
 }
-return out
 }
 // AscendRange calls the iterator for every value in the tree within the range
 // [greaterOrEqual, lessThan), until iterator returns false.
-func (t *BTree) AscendRange(greaterOrEqual, lessThan Item, iterator ItemIterator) {
 if t.root == nil {
 return
 }
-t.root.iterate(ascend, greaterOrEqual, lessThan, true, false, iterator)
 }
 // AscendLessThan calls the iterator for every value in the tree within the range
 // [first, pivot), until iterator returns false.
-func (t *BTree) AscendLessThan(pivot Item, iterator ItemIterator) {
 if t.root == nil {
 return
 }
-t.root.iterate(ascend, nil, pivot, false, false, iterator)
 }
 // AscendGreaterOrEqual calls the iterator for every value in the tree within
 // the range [pivot, last], until iterator returns false.
-func (t *BTree) AscendGreaterOrEqual(pivot Item, iterator ItemIterator) {
 if t.root == nil {
 return
 }
-t.root.iterate(ascend, pivot, nil, true, false, iterator)
 }
 // Ascend calls the iterator for every value in the tree within the range
 // [first, last], until iterator returns false.
-func (t *BTree) Ascend(iterator ItemIterator) {
 if t.root == nil {
 return
 }
-t.root.iterate(ascend, nil, nil, false, false, iterator)
 }
 // DescendRange calls the iterator for every value in the tree within the range
 // [lessOrEqual, greaterThan), until iterator returns false.
-func (t *BTree) DescendRange(lessOrEqual, greaterThan Item, iterator ItemIterator) {
 if t.root == nil {
 return
 }
-t.root.iterate(descend, lessOrEqual, greaterThan, true, false, iterator)
 }
 // DescendLessOrEqual calls the iterator for every value in the tree within the range
 // [pivot, first], until iterator returns false.
-func (t *BTree) DescendLessOrEqual(pivot Item, iterator ItemIterator) {
 if t.root == nil {
 return
 }
-t.root.iterate(descend, pivot, nil, true, false, iterator)
 }
 // DescendGreaterThan calls the iterator for every value in the tree within
 // the range [last, pivot), until iterator returns false.
-func (t *BTree) DescendGreaterThan(pivot Item, iterator ItemIterator) {
 if t.root == nil {
 return
 }
-t.root.iterate(descend, nil, pivot, false, false, iterator)
 }
 // Descend calls the iterator for every value in the tree within the range
 // [last, first], until iterator returns false.
-func (t *BTree) Descend(iterator ItemIterator) {
 if t.root == nil {
 return
 }
-t.root.iterate(descend, nil, nil, false, false, iterator)
 }
-// Get looks for the key item in the tree, returning it. It returns nil if
-// unable to find that item.
-func (t *BTree) Get(key Item) Item {
 if t.root == nil {
-return nil
 }
 return t.root.get(key)
 }
-// Min returns the smallest item in the tree, or nil if the tree is empty.
-func (t *BTree) Min() Item {
 return min(t.root)
 }
-// Max returns the largest item in the tree, or nil if the tree is empty.
-func (t *BTree) Max() Item {
 return max(t.root)
 }
 // Has returns true if the given key is in the tree.
-func (t *BTree) Has(key Item) bool {
-return t.Get(key) != nil
 }
 // Len returns the number of items currently in the tree.
-func (t *BTree) Len() int {
 return t.length
 }
@@ -865,7 +865,7 @@
 // O(tree size): when all nodes are owned by another tree, all nodes are
 // iterated over looking for nodes to add to the freelist, and due to
 // ownership, none are.
-func (t *BTree) Clear(addNodesToFreelist bool) {
 if t.root != nil && addNodesToFreelist {
 t.root.reset(t.cow)
 }
@@ -875,7 +875,7 @@
 // reset returns a subtree to the freelist. It breaks out immediately if the
 // freelist is full, since the only benefit of iterating is to fill that
 // freelist up. Returns true if parent reset call should continue.
-func (n *node) reset(c *copyOnWriteContext) bool {
 for _, child := range n.children {
 if !child.reset(c) {
 return false
@@ -891,3 +891,193 @@
 func (a Int) Less(b Item) bool {
 return a < b.(Int)
 }

+// Copyright 2014-2022 Google Inc.
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
 // See the License for the specific language governing permissions and
 // limitations under the License.
+//go:build go1.18
+// +build go1.18
+// In Go 1.18 and beyond, a BTreeG generic is created, and BTree is a specific
+// instantiation of that generic for the Item interface, with a backwards-
+// compatible API. Before go1.18, generics are not supported,
+// and BTree is just an implementation based around the Item interface.
 // Package btree implements in-memory B-Trees of arbitrary degree.
 //
 // Its functions, therefore, exactly mirror those of
 // llrb.LLRB where possible. Unlike gollrb, though, we currently don't
 // support storing multiple equivalent values.
+//
+// There are two implementations; those suffixed with 'G' are generics, usable
+// for any type, and require a passed-in "less" function to define their ordering.
+// Those without this prefix are specific to the 'Item' interface, and use
+// its 'Less' function for ordering.
 package btree
 import (
 DefaultFreeListSize = 32
 )
+// FreeListG represents a free list of btree nodes. By default each
 // BTree has its own FreeList, but multiple BTrees can share the same
+// FreeList, in particular when they're created with Clone.
 // Two Btrees using the same freelist are safe for concurrent write access.
+type FreeListG[T any] struct {
 mu sync.Mutex
+freelist []*node[T]
 }
+// NewFreeListG creates a new free list.
 // size is the maximum size of the returned free list.
+func NewFreeListG[T any](size int) *FreeListG[T] {
+return &FreeListG[T]{freelist: make([]*node[T], 0, size)}
 }
+func (f *FreeListG[T]) newNode() (n *node[T]) {
 f.mu.Lock()
 index := len(f.freelist) - 1
 if index < 0 {
 f.mu.Unlock()
+return new(node[T])
 }
 n = f.freelist[index]
 f.freelist[index] = nil
 return
 }
+func (f *FreeListG[T]) freeNode(n *node[T]) (out bool) {
 f.mu.Lock()
 if len(f.freelist) < cap(f.freelist) {
 f.freelist = append(f.freelist, n)
 return
 }
+// ItemIteratorG allows callers of {A/De}scend* to iterate in-order over portions of
 // the tree. When this function returns false, iteration will stop and the
 // associated Ascend* function will immediately return.
+type ItemIteratorG[T any] func(item T) bool
+// Ordered represents the set of types for which the '<' operator work.
+type Ordered interface {
+~int | ~int8 | ~int16 | ~int32 | ~int64 | ~uint | ~uint8 | ~uint16 | ~uint32 | ~uint64 | ~float32 | ~float64 | ~string
+}
+// Less[T] returns a default LessFunc that uses the '<' operator for types that support it.
+func Less[T Ordered]() LessFunc[T] {
+return func(a, b T) bool { return a < b }
+}
+// NewOrderedG creates a new B-Tree for ordered types.
+func NewOrderedG[T Ordered](degree int) *BTreeG[T] {
+return NewG[T](degree, Less[T]())
+}
+// NewG creates a new B-Tree with the given degree.
 //
+// NewG(2), for example, will create a 2-3-4 tree (each node contains 1-3 items
 // and 2-4 children).
+//
+// The passed-in LessFunc determines how objects of type T are ordered.
+func NewG[T any](degree int, less LessFunc[T]) *BTreeG[T] {
+return NewWithFreeListG(degree, less, NewFreeListG[T](DefaultFreeListSize))
 }
+// NewWithFreeListG creates a new B-Tree that uses the given node free list.
+func NewWithFreeListG[T any](degree int, less LessFunc[T], f *FreeListG[T]) *BTreeG[T] {
 if degree <= 1 {
 panic("bad degree")
 }
+return &BTreeG[T]{
 degree: degree,
+cow: &copyOnWriteContext[T]{freelist: f, less: less},
 }
 }
 // items stores items in a node.
+type items[T any] []T
 // insertAt inserts a value into the given index, pushing all subsequent values
 // forward.
+func (s *items[T]) insertAt(index int, item T) {
+var zero T
+*s = append(*s, zero)
 if index < len(*s) {
 copy((*s)[index+1:], (*s)[index:])
 }
 // removeAt removes a value at a given index, pulling all subsequent values
 // back.
+func (s *items[T]) removeAt(index int) T {
 item := (*s)[index]
 copy((*s)[index:], (*s)[index+1:])
+var zero T
+(*s)[len(*s)-1] = zero
 *s = (*s)[:len(*s)-1]
 return item
 }
 // pop removes and returns the last element in the list.
+func (s *items[T]) pop() (out T) {
 index := len(*s) - 1
 out = (*s)[index]
+var zero T
+(*s)[index] = zero
 *s = (*s)[:index]
 return
 }
 // truncate truncates this instance at index so that it contains only the
 // first index items. index must be less than or equal to length.
+func (s *items[T]) truncate(index int) {
+var toClear items[T]
 *s, toClear = (*s)[:index], (*s)[index:]
+var zero T
+for i := 0; i < len(toClear); i++ {
+toClear[i] = zero
 }
 }
 // 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 (s items[T]) find(item T, less func(T, T) bool) (index int, found bool) {
 i := sort.Search(len(s), func(i int) bool {
+return less(item, s[i])
 })
+if i > 0 && !less(s[i-1], item) {
 return i - 1, true
 }
 return i, false
 }
 // node is an internal node in a tree.
 //
 // It must at all times maintain the invariant that either
 // * len(children) == 0, len(items) unconstrained
 // * len(children) == len(items) + 1
+type node[T any] struct {
+items items[T]
+children items[*node[T]]
+cow *copyOnWriteContext[T]
 }
+func (n *node[T]) mutableFor(cow *copyOnWriteContext[T]) *node[T] {
 if n.cow == cow {
 return n
 }
 if cap(out.items) >= len(n.items) {
 out.items = out.items[:len(n.items)]
 } else {
+out.items = make(items[T], len(n.items), cap(n.items))
 }
 copy(out.items, n.items)
 // Copy children
 if cap(out.children) >= len(n.children) {
 out.children = out.children[:len(n.children)]
 } else {
+out.children = make(items[*node[T]], len(n.children), cap(n.children))
 }
 copy(out.children, n.children)
 return out
 }
+func (n *node[T]) mutableChild(i int) *node[T] {
 c := n.children[i].mutableFor(n.cow)
 n.children[i] = c
 return c
 // split splits the given node at the given index. The current node shrinks,
 // and this function returns the item that existed at that index and a new node
 // containing all items/children after it.
+func (n *node[T]) split(i int) (T, *node[T]) {
 item := n.items[i]
 next := n.cow.newNode()
 next.items = append(next.items, n.items[i+1:]...)
 // maybeSplitChild checks if a child should be split, and if so splits it.
 // Returns whether or not a split occurred.
+func (n *node[T]) maybeSplitChild(i, maxItems int) bool {
 if len(n.children[i].items) < maxItems {
 return false
 }
 // insert inserts an item into the subtree rooted at this node, making sure
 // no nodes in the subtree exceed maxItems items. Should an equivalent item be
 // be found/replaced by insert, it will be returned.
+func (n *node[T]) insert(item T, maxItems int) (_ T, _ bool) {
+i, found := n.items.find(item, n.cow.less)
 if found {
 out := n.items[i]
 n.items[i] = item
+return out, true
 }
 if len(n.children) == 0 {
 n.items.insertAt(i, item)
+return
 }
 if n.maybeSplitChild(i, maxItems) {
 inTree := n.items[i]
 switch {
+case n.cow.less(item, inTree):
 // no change, we want first split node
+case n.cow.less(inTree, item):
 i++ // we want second split node
 default:
 out := n.items[i]
 n.items[i] = item
+return out, true
 }
 }
 return n.mutableChild(i).insert(item, maxItems)
 }
 // get finds the given key in the subtree and returns it.
+func (n *node[T]) get(key T) (_ T, _ bool) {
+i, found := n.items.find(key, n.cow.less)
 if found {
+return n.items[i], true
 } else if len(n.children) > 0 {
 return n.children[i].get(key)
 }
+return
 }
 // min returns the first item in the subtree.
+func min[T any](n *node[T]) (_ T, found bool) {
 if n == nil {
+return
 }
 for len(n.children) > 0 {
 n = n.children[0]
 }
 if len(n.items) == 0 {
+return
 }
+return n.items[0], true
 }
 // max returns the last item in the subtree.
+func max[T any](n *node[T]) (_ T, found bool) {
 if n == nil {
+return
 }
 for len(n.children) > 0 {
 n = n.children[len(n.children)-1]
 }
 if len(n.items) == 0 {
+return
 }
+return n.items[len(n.items)-1], true
 }
 // toRemove details what item to remove in a node.remove call.
 )
 // remove removes an item from the subtree rooted at this node.
+func (n *node[T]) remove(item T, minItems int, typ toRemove) (_ T, _ bool) {
 var i int
 var found bool
 switch typ {
 case removeMax:
 if len(n.children) == 0 {
+return n.items.pop(), true
 }
 i = len(n.items)
 case removeMin:
 if len(n.children) == 0 {
+return n.items.removeAt(0), true
 }
 i = 0
 case removeItem:
+i, found = n.items.find(item, n.cow.less)
 if len(n.children) == 0 {
 if found {
+return n.items.removeAt(i), true
 }
+return
 }
 default:
 panic("invalid type")
 // We use our special-case 'remove' call with typ=maxItem to pull the
 // predecessor of item i (the rightmost leaf of our immediate left child)
 // and set it into where we pulled the item from.
+var zero T
+n.items[i], _ = child.remove(zero, minItems, removeMax)
+return out, true
 }
 // Final recursive call. Once we're here, we know that the item isn't in this
 // node and that the child is big enough to remove from.
 // We then simply redo our remove call, and the second time (regardless of
 // whether we're in case 1 or 2), we'll have enough items and can guarantee
 // that we hit case A.
+func (n *node[T]) growChildAndRemove(i int, item T, minItems int, typ toRemove) (T, bool) {
 if i > 0 && len(n.children[i-1].items) > minItems {
 // Steal from left child
 child := n.mutableChild(i)
 ascend = direction(+1)
 )
+type optionalItem[T any] struct {
+item T
+valid bool
+}
+func optional[T any](item T) optionalItem[T] {
+return optionalItem[T]{item: item, valid: true}
+}
+func empty[T any]() optionalItem[T] {
+return optionalItem[T]{}
+}
 // iterate provides a simple method for iterating over elements in the tree.
 //
 // When ascending, the 'start' should be less than 'stop' and when descending,
 // will force the iterator to include the first item when it equals 'start',
 // thus creating a "greaterOrEqual" or "lessThanEqual" rather than just a
 // "greaterThan" or "lessThan" queries.
+func (n *node[T]) iterate(dir direction, start, stop optionalItem[T], includeStart bool, hit bool, iter ItemIteratorG[T]) (bool, bool) {
 var ok, found bool
 var index int
 switch dir {
 case ascend:
+if start.valid {
+index, _ = n.items.find(start.item, n.cow.less)
 }
 for i := index; i < len(n.items); i++ {
 if len(n.children) > 0 {
 return hit, false
 }
 }
+if !includeStart && !hit && start.valid && !n.cow.less(start.item, n.items[i]) {
 hit = true
 continue
 }
 hit = true
+if stop.valid && !n.cow.less(n.items[i], stop.item) {
 return hit, false
 }
 if !iter(n.items[i]) {
 }
 }
 case descend:
+if start.valid {
+index, found = n.items.find(start.item, n.cow.less)
 if !found {
 index = index - 1
 }
 index = len(n.items) - 1
 }
 for i := index; i >= 0; i-- {
+if start.valid && !n.cow.less(n.items[i], start.item) {
+if !includeStart || hit || n.cow.less(start.item, n.items[i]) {
 continue
 }
 }
 return hit, false
 }
 }
+if stop.valid && !n.cow.less(stop.item, n.items[i]) {
 return hit, false // continue
 }
 hit = true
 return hit, true
 }
+// print is used for testing/debugging purposes.
+func (n *node[T]) print(w io.Writer, level int) {
 fmt.Fprintf(w, "%sNODE:%v\n", strings.Repeat(" ", level), n.items)
 for _, c := range n.children {
 c.print(w, level+1)
 }
 }
+// BTreeG is a generic implementation of a B-Tree.
 //
+// BTreeG stores items of type T in an ordered structure, allowing easy insertion,
 // removal, and iteration.
 //
 // Write operations are not safe for concurrent mutation by multiple
 // goroutines, but Read operations are.
+type BTreeG[T any] struct {
 degree int
 length int
+root *node[T]
+cow *copyOnWriteContext[T]
 }
+// LessFunc[T] determines how to order a type 'T'. It should implement a strict
+// ordering, and should return true if within that ordering, 'a' < 'b'.
+type LessFunc[T any] func(a, b T) bool
 // copyOnWriteContext pointers determine node ownership... a tree with a write
 // context equivalent to a node's write context is allowed to modify that node.
 // A tree whose write context does not match a node's is not allowed to modify
 // tree's context, that node is modifiable in place. Children of that node may
 // not share context, but before we descend into them, we'll make a mutable
 // copy.
+type copyOnWriteContext[T any] struct {
+freelist *FreeListG[T]
+less LessFunc[T]
 }
 // Clone clones the btree, lazily. Clone should not be called concurrently,
 // will initially experience minor slow-downs caused by additional allocs and
 // copies due to the aforementioned copy-on-write logic, but should converge to
 // the original performance characteristics of the original tree.
+func (t *BTreeG[T]) Clone() (t2 *BTreeG[T]) {
 // Create two entirely new copy-on-write contexts.
 // This operation effectively creates three trees:
 // the original, shared nodes (old b.cow)
 }
 // maxItems returns the max number of items to allow per node.
+func (t *BTreeG[T]) maxItems() int {
 return t.degree*2 - 1
 }
 // minItems returns the min number of items to allow per node (ignored for the
 // root node).
+func (t *BTreeG[T]) minItems() int {
 return t.degree - 1
 }
+func (c *copyOnWriteContext[T]) newNode() (n *node[T]) {
 n = c.freelist.newNode()
 n.cow = c
 return
 // freeNode frees a node within a given COW context, if it's owned by that
 // context. It returns what happened to the node (see freeType const
 // documentation).
+func (c *copyOnWriteContext[T]) freeNode(n *node[T]) freeType {
 if n.cow == c {
 // clear to allow GC
 n.items.truncate(0)
 }
 // ReplaceOrInsert adds the given item to the tree. If an item in the tree
+// already equals the given one, it is removed from the tree and returned,
+// and the second return value is true. Otherwise, (zeroValue, false)
 //
 // nil cannot be added to the tree (will panic).
+func (t *BTreeG[T]) ReplaceOrInsert(item T) (_ T, _ bool) {
 if t.root == nil {
 t.root = t.cow.newNode()
 t.root.items = append(t.root.items, item)
 t.length++
+return
 } else {
 t.root = t.root.mutableFor(t.cow)
 if len(t.root.items) >= t.maxItems() {
 t.root.children = append(t.root.children, oldroot, second)
 }
 }
+out, outb := t.root.insert(item, t.maxItems())
+if !outb {
 t.length++
 }
+return out, outb
 }
 // Delete removes an item equal to the passed in item from the tree, returning
+// it. If no such item exists, returns (zeroValue, false).
+func (t *BTreeG[T]) Delete(item T) (T, bool) {
 return t.deleteItem(item, removeItem)
 }
 // DeleteMin removes the smallest item in the tree and returns it.
+// If no such item exists, returns (zeroValue, false).
+func (t *BTreeG[T]) DeleteMin() (T, bool) {
+var zero T
+return t.deleteItem(zero, removeMin)
 }
 // DeleteMax removes the largest item in the tree and returns it.
+// If no such item exists, returns (zeroValue, false).
+func (t *BTreeG[T]) DeleteMax() (T, bool) {
+var zero T
+return t.deleteItem(zero, removeMax)
 }
+func (t *BTreeG[T]) deleteItem(item T, typ toRemove) (_ T, _ bool) {
 if t.root == nil || len(t.root.items) == 0 {
+return
 }
 t.root = t.root.mutableFor(t.cow)
+out, outb := t.root.remove(item, t.minItems(), typ)
 if len(t.root.items) == 0 && len(t.root.children) > 0 {
 oldroot := t.root
 t.root = t.root.children[0]
 t.cow.freeNode(oldroot)
 }
+if outb {
 t.length--
 }
+return out, outb
 }
 // AscendRange calls the iterator for every value in the tree within the range
 // [greaterOrEqual, lessThan), until iterator returns false.
+func (t *BTreeG[T]) AscendRange(greaterOrEqual, lessThan T, iterator ItemIteratorG[T]) {
 if t.root == nil {
 return
 }
+t.root.iterate(ascend, optional[T](greaterOrEqual), optional[T](lessThan), true, false, iterator)
 }
 // AscendLessThan calls the iterator for every value in the tree within the range
 // [first, pivot), until iterator returns false.
+func (t *BTreeG[T]) AscendLessThan(pivot T, iterator ItemIteratorG[T]) {
 if t.root == nil {
 return
 }
+t.root.iterate(ascend, empty[T](), optional(pivot), false, false, iterator)
 }
 // AscendGreaterOrEqual calls the iterator for every value in the tree within
 // the range [pivot, last], until iterator returns false.
+func (t *BTreeG[T]) AscendGreaterOrEqual(pivot T, iterator ItemIteratorG[T]) {
 if t.root == nil {
 return
 }
+t.root.iterate(ascend, optional[T](pivot), empty[T](), true, false, iterator)
 }
 // Ascend calls the iterator for every value in the tree within the range
 // [first, last], until iterator returns false.
+func (t *BTreeG[T]) Ascend(iterator ItemIteratorG[T]) {
 if t.root == nil {
 return
 }
+t.root.iterate(ascend, empty[T](), empty[T](), false, false, iterator)
 }
 // DescendRange calls the iterator for every value in the tree within the range
 // [lessOrEqual, greaterThan), until iterator returns false.
+func (t *BTreeG[T]) DescendRange(lessOrEqual, greaterThan T, iterator ItemIteratorG[T]) {
 if t.root == nil {
 return
 }
+t.root.iterate(descend, optional[T](lessOrEqual), optional[T](greaterThan), true, false, iterator)
 }
 // DescendLessOrEqual calls the iterator for every value in the tree within the range
 // [pivot, first], until iterator returns false.
+func (t *BTreeG[T]) DescendLessOrEqual(pivot T, iterator ItemIteratorG[T]) {
 if t.root == nil {
 return
 }
+t.root.iterate(descend, optional[T](pivot), empty[T](), true, false, iterator)
 }
 // DescendGreaterThan calls the iterator for every value in the tree within
 // the range [last, pivot), until iterator returns false.
+func (t *BTreeG[T]) DescendGreaterThan(pivot T, iterator ItemIteratorG[T]) {
 if t.root == nil {
 return
 }
+t.root.iterate(descend, empty[T](), optional[T](pivot), false, false, iterator)
 }
 // Descend calls the iterator for every value in the tree within the range
 // [last, first], until iterator returns false.
+func (t *BTreeG[T]) Descend(iterator ItemIteratorG[T]) {
 if t.root == nil {
 return
 }
+t.root.iterate(descend, empty[T](), empty[T](), false, false, iterator)
 }
+// Get looks for the key item in the tree, returning it. It returns
+// (zeroValue, false) if unable to find that item.
+func (t *BTreeG[T]) Get(key T) (_ T, _ bool) {
 if t.root == nil {
+return
 }
 return t.root.get(key)
 }
+// Min returns the smallest item in the tree, or (zeroValue, false) if the tree is empty.
+func (t *BTreeG[T]) Min() (_ T, _ bool) {
 return min(t.root)
 }
+// Max returns the largest item in the tree, or (zeroValue, false) if the tree is empty.
+func (t *BTreeG[T]) Max() (_ T, _ bool) {
 return max(t.root)
 }
 // Has returns true if the given key is in the tree.
+func (t *BTreeG[T]) Has(key T) bool {
+_, ok := t.Get(key)
+return ok
 }
 // Len returns the number of items currently in the tree.
+func (t *BTreeG[T]) Len() int {
 return t.length
 }
 // O(tree size): when all nodes are owned by another tree, all nodes are
 // iterated over looking for nodes to add to the freelist, and due to
 // ownership, none are.
+func (t *BTreeG[T]) Clear(addNodesToFreelist bool) {
 if t.root != nil && addNodesToFreelist {
 t.root.reset(t.cow)
 }
 // reset returns a subtree to the freelist. It breaks out immediately if the
 // freelist is full, since the only benefit of iterating is to fill that
 // freelist up. Returns true if parent reset call should continue.
+func (n *node[T]) reset(c *copyOnWriteContext[T]) bool {
 for _, child := range n.children {
 if !child.reset(c) {
 return false
 func (a Int) Less(b Item) bool {
 return a < b.(Int)
 }
+// BTree is an implementation of a B-Tree.
+//
+// BTree stores Item instances in an ordered structure, allowing easy insertion,
+// removal, and iteration.
+//
+// Write operations are not safe for concurrent mutation by multiple
+// goroutines, but Read operations are.
+type BTree BTreeG[Item]
+var itemLess LessFunc[Item] = func(a, b Item) bool {
+return a.Less(b)
+}
+// New creates a new B-Tree with the given degree.
+//
+// New(2), for example, will create a 2-3-4 tree (each node contains 1-3 items
+// and 2-4 children).
+func New(degree int) *BTree {
+return (*BTree)(NewG[Item](degree, itemLess))
+}
+// FreeList represents a free list of btree nodes. By default each
+// BTree has its own FreeList, but multiple BTrees can share the same
+// FreeList.
+// Two Btrees using the same freelist are safe for concurrent write access.
+type FreeList FreeListG[Item]
+// NewFreeList creates a new free list.
+// size is the maximum size of the returned free list.
+func NewFreeList(size int) *FreeList {
+return (*FreeList)(NewFreeListG[Item](size))
+}
+// NewWithFreeList creates a new B-Tree that uses the given node free list.
+func NewWithFreeList(degree int, f *FreeList) *BTree {
+return (*BTree)(NewWithFreeListG[Item](degree, itemLess, (*FreeListG[Item])(f)))
+}
+// ItemIterator allows callers of Ascend* to iterate in-order over portions of
+// the tree. When this function returns false, iteration will stop and the
+// associated Ascend* function will immediately return.
+type ItemIterator ItemIteratorG[Item]
+// Clone clones the btree, lazily. Clone should not be called concurrently,
+// but the original tree (t) and the new tree (t2) can be used concurrently
+// once the Clone call completes.
+//
+// The internal tree structure of b is marked read-only and shared between t and
+// t2. Writes to both t and t2 use copy-on-write logic, creating new nodes
+// whenever one of b's original nodes would have been modified. Read operations
+// should have no performance degredation. Write operations for both t and t2
+// will initially experience minor slow-downs caused by additional allocs and
+// copies due to the aforementioned copy-on-write logic, but should converge to
+// the original performance characteristics of the original tree.
+func (t *BTree) Clone() (t2 *BTree) {
+return (*BTree)((*BTreeG[Item])(t).Clone())
+}
+// Delete removes an item equal to the passed in item from the tree, returning
+// it. If no such item exists, returns nil.
+func (t *BTree) Delete(item Item) Item {
+i, _ := (*BTreeG[Item])(t).Delete(item)
+return i
+}
+// DeleteMax removes the largest item in the tree and returns it.
+// If no such item exists, returns nil.
+func (t *BTree) DeleteMax() Item {
+i, _ := (*BTreeG[Item])(t).DeleteMax()
+return i
+}
+// DeleteMin removes the smallest item in the tree and returns it.
+// If no such item exists, returns nil.
+func (t *BTree) DeleteMin() Item {
+i, _ := (*BTreeG[Item])(t).DeleteMin()
+return i
+}
+// Get looks for the key item in the tree, returning it. It returns nil if
+// unable to find that item.
+func (t *BTree) Get(key Item) Item {
+i, _ := (*BTreeG[Item])(t).Get(key)
+return i
+}
+// Max returns the largest item in the tree, or nil if the tree is empty.
+func (t *BTree) Max() Item {
+i, _ := (*BTreeG[Item])(t).Max()
+return i
+}
+// Min returns the smallest item in the tree, or nil if the tree is empty.
+func (t *BTree) Min() Item {
+i, _ := (*BTreeG[Item])(t).Min()
+return i
+}
+// Has returns true if the given key is in the tree.
+func (t *BTree) Has(key Item) bool {
+return (*BTreeG[Item])(t).Has(key)
+}
+// ReplaceOrInsert adds the given item to the tree. If an item in the tree
+// already equals the given one, it is removed from the tree and returned.
+// Otherwise, nil is returned.
+//
+// nil cannot be added to the tree (will panic).
+func (t *BTree) ReplaceOrInsert(item Item) Item {
+i, _ := (*BTreeG[Item])(t).ReplaceOrInsert(item)
+return i
+}
+// AscendRange calls the iterator for every value in the tree within the range
+// [greaterOrEqual, lessThan), until iterator returns false.
+func (t *BTree) AscendRange(greaterOrEqual, lessThan Item, iterator ItemIterator) {
+(*BTreeG[Item])(t).AscendRange(greaterOrEqual, lessThan, (ItemIteratorG[Item])(iterator))
+}
+// AscendLessThan calls the iterator for every value in the tree within the range
+// [first, pivot), until iterator returns false.
+func (t *BTree) AscendLessThan(pivot Item, iterator ItemIterator) {
+(*BTreeG[Item])(t).AscendLessThan(pivot, (ItemIteratorG[Item])(iterator))
+}
+// AscendGreaterOrEqual calls the iterator for every value in the tree within
+// the range [pivot, last], until iterator returns false.
+func (t *BTree) AscendGreaterOrEqual(pivot Item, iterator ItemIterator) {
+(*BTreeG[Item])(t).AscendGreaterOrEqual(pivot, (ItemIteratorG[Item])(iterator))
+}
+// Ascend calls the iterator for every value in the tree within the range
+// [first, last], until iterator returns false.
+func (t *BTree) Ascend(iterator ItemIterator) {
+(*BTreeG[Item])(t).Ascend((ItemIteratorG[Item])(iterator))
+}
+// DescendRange calls the iterator for every value in the tree within the range
+// [lessOrEqual, greaterThan), until iterator returns false.
+func (t *BTree) DescendRange(lessOrEqual, greaterThan Item, iterator ItemIterator) {
+(*BTreeG[Item])(t).DescendRange(lessOrEqual, greaterThan, (ItemIteratorG[Item])(iterator))
+}
+// DescendLessOrEqual calls the iterator for every value in the tree within the range
+// [pivot, first], until iterator returns false.
+func (t *BTree) DescendLessOrEqual(pivot Item, iterator ItemIterator) {
+(*BTreeG[Item])(t).DescendLessOrEqual(pivot, (ItemIteratorG[Item])(iterator))
+}
+// DescendGreaterThan calls the iterator for every value in the tree within
+// the range [last, pivot), until iterator returns false.
+func (t *BTree) DescendGreaterThan(pivot Item, iterator ItemIterator) {
+(*BTreeG[Item])(t).DescendGreaterThan(pivot, (ItemIteratorG[Item])(iterator))
+}
+// Descend calls the iterator for every value in the tree within the range
+// [last, first], until iterator returns false.
+func (t *BTree) Descend(iterator ItemIterator) {
+(*BTreeG[Item])(t).Descend((ItemIteratorG[Item])(iterator))
+}
+// Len returns the number of items currently in the tree.
+func (t *BTree) Len() int {
+return (*BTreeG[Item])(t).Len()
+}
+// Clear removes all items from the btree. If addNodesToFreelist is true,
+// t's nodes are added to its freelist as part of this call, until the freelist
+// is full. Otherwise, the root node is simply dereferenced and the subtree
+// left to Go's normal GC processes.
+//
+// This can be much faster
+// than calling Delete on all elements, because that requires finding/removing
+// each element in the tree and updating the tree accordingly. It also is
+// somewhat faster than creating a new tree to replace the old one, because
+// nodes from the old tree are reclaimed into the freelist for use by the new
+// one, instead of being lost to the garbage collector.
+//
+// This call takes:
+// O(1): when addNodesToFreelist is false, this is a single operation.
+// O(1): when the freelist is already full, it breaks out immediately
+// O(freelist size): when the freelist is empty and the nodes are all owned
+// by this tree, nodes are added to the freelist until full.
+// O(tree size): when all nodes are owned by another tree, all nodes are
+// iterated over looking for nodes to add to the freelist, and due to
+// ownership, none are.
+func (t *BTree) Clear(addNodesToFreelist bool) {
+(*BTreeG[Item])(t).Clear(addNodesToFreelist)
+}

简述:

google实现的btree包,本身其实已经做了很多的优化,其中每个外部可以调用的函数都有对应的Benchmark,这对我们分析泛型带来的优势提供了很大的帮助.

本文不讨论btree本身的优劣,单纯的从一个可以进行数据操作的集合函数库来着重讨论泛型带来的golang性能优化.

如果对golang的1.18新带来的泛型(generics)不太了解,可以先阅读一下 https://go.dev/doc/tutorial/generics (官方) 或者 https://segmentfault.com/a/1190000041634906 (中文讲解)

实现解析:

如上文diff代码能看到的,泛型化的btree代码,主要就是将原来的 item/node 替换成 item[T]/node[T]

@@ -72,32 +82,27 @@
-type FreeList struct {
+type FreeListG[T any] struct {
 	mu       sync.Mutex
-	freelist []*node
+	freelist []*node[T]
 }
 
 // size is the maximum size of the returned free list.
-func NewFreeList(size int) *FreeList {
-	return &FreeList{freelist: make([]*node, 0, size)}
+func NewFreeListG[T any](size int) *FreeListG[T] {
+	return &FreeListG[T]{freelist: make([]*node[T], 0, size)}
 }
 
-func (f *FreeList) newNode() (n *node) {
+func (f *FreeListG[T]) newNode() (n *node[T]) {
 	f.mu.Lock()
 	index := len(f.freelist) - 1
 	if index < 0 {
 		f.mu.Unlock()
-		return new(node)
+		return new(node[T])
 	}
 	n = f.freelist[index]
 	f.freelist[index] = nil

这一块就是标准的泛型更新,不需要特别的说明.

全新定义了泛型里面可以被排序的类型集合与他们的generics的比较函数

// LessFunc[T] 函数是用来定义泛型的类型形参(Type parameter) 'T'是如何比较大小的.
// 返回的顺序必须是严密的,且当'a' < 'b'是返回true.
type LessFunc[T any] func(a, b T) bool

// Ordered 定义了golang原生可以被比较大小的内置类型
// @Note: 这里主要是做前向兼容,如果有自定义的类型,实现了`Less`方法也可以用
type Ordered interface {
	~int | ~int8 | ~int16 | ~int32 | ~int64 | ~uint | ~uint8 | ~uint16 | ~uint32 | ~uint64 | ~float32 | ~float64 | ~string
}

// Less[T] 就是内置的Less函数,直接返回支持的泛型类型形参的比较结果
func Less[T Ordered]() LessFunc[T] {
	return func(a, b T) bool { return a < b }
}

相对的,老版本直接是定义了一个Item的interface{}来实现同样的Less函数

// Item 表示树上的一个独立对象.
type Item interface {
	// Less 函数必须要实现判断当前值是否小于传递进来的值.
	// 相等的值只会保存一个
	Less(than Item) bool
}

修改内部函数的返回值

@@ -308,70 +290,70 @@
-func (n *node) insert(item Item, maxItems int) Item {
-	i, found := n.items.find(item)
+func (n *node[T]) insert(item T, maxItems int) (_ T, _ bool) {
+	i, found := n.items.find(item, n.cow.less)
 	if found {
 		out := n.items[i]
 		n.items[i] = item
-		return out
+		return out, true
 	}
 	if len(n.children) == 0 {
 		n.items.insertAt(i, item)
-		return nil
+		return
 	}

@@ -424,8 +406,9 @@
 		// 在这里使用的了一个内部的'remove'调用并且传递了一个'maxItem'参数,
 		// 拉取i的前一项(最右边的子节点)并设置到拉取的当前节点上
-		n.items[i] = child.remove(nil, minItems, removeMax)
-		return out
+		var zero T
+		n.items[i], _ = child.remove(zero, minItems, removeMax)
+		return out, true
 	}

之前的函数实现返回Item,Item是一个interface{}(1.18之后是一个基本接口),之前可以通过直接返回一个nil表示没有找到数据或者操作失败

但是泛型之后,不同的类型形参不能表示这种状态,所以在返回泛型的形参T之外,还增加返回了一个bool类型来表示操作结果

这里的Item(实际是一个interface{}接口)是性能瓶颈与优化的核心,后续我们回详细展开聊到

最后再增加亿点点兼容代码

// 前向兼容的方式,定义泛型下的'BTree'是一个'BTreeG'的Item约束
type BTree BTreeG[Item]

// 定义默认的函数,就是上文#2定义好的比较函数的实例
var itemLess LessFunc[Item] = func(a, b Item) bool {
	return a.Less(b)
}

// New 创建一个新的'BTree',内部转换为创建一个'BTreeG[item]',同时传递less函数
func New(degree int) *BTree {
	return (*BTree)(NewG[Item](degree, itemLess))
}
...更多的函数

性能测试结果

MR里面提到使用泛型改造之后,对于官方的bench函数(int存储),性能能优化40%+!

这里写一个简单的Bench来对比两种实现方式的性能

Code

package main

import (
	"math/rand"
	"testing"
	"time"

	"github.com/google/btree"
)

const treeSize = 50000

func init() {
	rand.Seed(time.Now().Unix())
}

func BenchmarkBTree(b *testing.B) {
	b.StopTimer()
	insertP := rand.Perm(treeSize)
	b.StartTimer()
	i := 0
	for i < b.N {
		tr := btree.New(32)
		for _, item := range insertP {
			tr.ReplaceOrInsert(btree.Int(item))
			i++
			if i >= b.N {
				return
			}
		}
	}
}

func BenchmarkBTreeG(b *testing.B) {
	b.StopTimer()
	insertP := rand.Perm(treeSize)
	b.StartTimer()
	i := 0
	for i < b.N {
		tr := btree.NewOrderedG[int](32)
		for _, item := range insertP {
			tr.ReplaceOrInsert(item)
			i++
			if i >= b.N {
				return
			}
		}
	}
}

Bench执行之后的结果对比(使用github-codespace机器)

Running tool: /usr/local/go/bin/go test -benchmem -run=^$ -bench ^(BenchmarkBTree|BenchmarkBTreeG)$ gh/btree -count=1 -failfast

goos: linux
goarch: amd64
pkg: gh/btree
cpu: Intel(R) Xeon(R) Platinum 8272CL CPU @ 2.60GHz
BenchmarkBTree-4    	 3254396	       373.8 ns/op	      44 B/op	       1 allocs/op
BenchmarkBTreeG-4   	 5487584	       215.7 ns/op	      19 B/op	       0 allocs/op
PASS
ok  	gh/btree	5.106s

完整的官方全部bench对比结果:(再次感叹下原来版本的btree已经非常优化了,alloc(s)已经很小)

go test -bench -count=5 -failfast > [old|new].txt
benchstst old.txt new.txt

name	old time/op	new time/op	delta
Insert-20	175ns ± 1%	131ns ± 1%	-25.23%	(p=0.008 n=5+5)
Seek-20	141ns ± 2%	82ns ± 6%	-42.10%	(p=0.008 n=5+5)
DeleteInsert-20	357ns ± 1%	272ns ± 1%	-24.01%	(p=0.008 n=5+5)
DeleteInsertCloneOnce-20	354ns ± 1%	269ns ± 1%	-24.01%	(p=0.008 n=5+5)
DeleteInsertCloneEachTime-20	1.53µs ± 1%	0.97µs ± 7%	-36.70%	(p=0.008 n=5+5)
Delete-20	196ns ± 1%	145ns ± 2%	-25.91%	(p=0.008 n=5+5)
Get-20	155ns ± 2%	108ns ± 0%	-30.03%	(p=0.008 n=5+5)
GetCloneEachTime-20	239ns ± 1%	180ns ± 2%	-24.78%	(p=0.008 n=5+5)
Ascend-20	47.7µs ± 1%	35.0µs ± 1%	-26.63%	(p=0.008 n=5+5)
Descend-20	47.3µs ± 1%	34.1µs ± 1%	-27.88%	(p=0.008 n=5+5)
AscendRange-20	91.7µs ± 2%	54.7µs ± 1%	-40.36%	(p=0.008 n=5+5)
DescendRange-20	130µs ± 5%	72µs ± 1%	-44.78%	(p=0.008 n=5+5)
AscendGreaterOrEqual-20	56.9µs ± 1%	42.0µs ± 1%	-26.19%	(p=0.008 n=5+5)
DescendLessOrEqual-20	94.0µs ± 2%	57.4µs ± 1%	-38.89%	(p=0.008 n=5+5)
DeleteAndRestore/CopyBigFreeList-20	5.08ms ± 1%	3.60ms ± 0%	-29.02%	(p=0.016 n=5+4)
DeleteAndRestore/Copy-20	5.41ms ± 3%	3.66ms ± 0%	-32.37%	(p=0.008 n=5+5)
DeleteAndRestore/ClearBigFreelist-20	2.92ms ± 1%	2.16ms ± 1%	-26.09%	(p=0.008 n=5+5)
DeleteAndRestore/Clear-20	3.14ms ± 1%	2.24ms ± 1%	-28.91%	(p=0.008 n=5+5)

name	old alloc/op	new alloc/op	delta
Insert-20	36.0B ± 3%	18.4B ± 3%	-48.89%	(p=0.008 n=5+5)
Seek-20	7.00B ± 0%	0.00B	-100.00%	(p=0.008 n=5+5)
DeleteInsert-20	0.00B	0.00B	~	(all equal)
DeleteInsertCloneOnce-20	0.00B	0.00B	~	(all equal)
DeleteInsertCloneEachTime-20	2.97kB ± 1%	1.95kB ± 4%	-34.30%	(p=0.008 n=5+5)
Delete-20	0.00B	0.00B	~	(all equal)
Get-20	0.00B	0.00B	~	(all equal)
GetCloneEachTime-20	64.0B ± 0%	64.0B ± 0%	~	(all equal)
Ascend-20	0.00B	0.00B	~	(all equal)
Descend-20	0.00B	0.00B	~	(all equal)
AscendRange-20	0.00B	0.00B	~	(all equal)
DescendRange-20	0.00B	0.00B	~	(all equal)
AscendGreaterOrEqual-20	0.00B	0.00B	~	(all equal)
DescendLessOrEqual-20	0.00B	0.00B	~	(all equal)
DeleteAndRestore/CopyBigFreeList-20	274kB ± 0%	142kB ± 0%	-48.20%	(p=0.008 n=5+5)
DeleteAndRestore/Copy-20	876kB ± 0%	443kB ± 0%	-49.36%	(p=0.008 n=5+5)
DeleteAndRestore/ClearBigFreelist-20	631B ± 1%	510B ± 5%	-19.06%	(p=0.008 n=5+5)
DeleteAndRestore/Clear-20	553kB ± 0%	278kB ± 0%	-49.77%	(p=0.000 n=5+4)

name	old allocs/op	new allocs/op	delta
Insert-20	0.00	0.00	~	(all equal)
Seek-20	0.00	0.00	~	(all equal)
DeleteInsert-20	0.00	0.00	~	(all equal)
DeleteInsertCloneOnce-20	0.00	0.00	~	(all equal)
DeleteInsertCloneEachTime-20	11.0 ± 0%	11.0 ± 0%	~	(all equal)
Delete-20	0.00	0.00	~	(all equal)
Get-20	0.00	0.00	~	(all equal)
GetCloneEachTime-20	3.00 ± 0%	3.00 ± 0%	~	(all equal)
Ascend-20	0.00	0.00	~	(all equal)
Descend-20	0.00	0.00	~	(all equal)
AscendRange-20	0.00	0.00	~	(all equal)
DescendRange-20	0.00	0.00	~	(all equal)
AscendGreaterOrEqual-20	0.00	0.00	~	(all equal)
DescendLessOrEqual-20	0.00	0.00	~	(all equal)
DeleteAndRestore/CopyBigFreeList-20	12.0 ± 0%	14.0 ± 0%	+16.67%	(p=0.008 n=5+5)
DeleteAndRestore/Copy-20	1.18k ± 0%	1.13k ± 0%	-4.15%	(p=0.008 n=5+5)
DeleteAndRestore/ClearBigFreelist-20	1.00 ± 0%	1.00 ± 0%	~	(all equal)
DeleteAndRestore/Clear-20	1.07k ± 0%	1.02k ± 0%	-4.49%	(p=0.008 n=5+5)

原理分析与说明

细心的人已经发现了端倪,Seek的bench结果显示泛型之后的代码的alloc/op直接将为了0

google/btree的bench使用的是一个int类型,但是传递的都是一个Item接口.在bench的开头,创建了一个填充了随机数的int数组(同时转换为btree.Int,一个Item的实现).

type node struct {
	items    items
	children children
	cow      *copyOnWriteContext
}
type items []Item
type Item interface {
	Less(than Item) bool
}
type children []*node

当我们调用一个函数的时候,比如Seek实际内部调用AscendGreaterOrEqual实际内部去循环t.root.iterate()

bench代码模拟了一个大树的随机数据创建,const benchmarkTreeSize = 10000个int已经在函数的堆(heap)上了,正常来说,我们调用函数传参都是简单的复制了指向这个int地址的一个指针而已

但是当btree传递数据参数变成Item之后,情况就发生了变化,这个时候,传递的int地址的指针,变成了一个Item的intertface{}的指针,这个时候参数就会逃逸到栈上(escape to heap)

这种情况其实很常见,比如最常用到的printf系列函数,他接收的永远是interface{}作为参数,golang的编译器就会认为这类的函数,任意传递的v都会被转换成interface{}之后复制到栈上,再进行处理

这其实很好理解,v(aka Item in btree)可以是任何类型,传递形参为了保证安全,会复制到一个叫做interface{}的万能圣杯里,再去传递

// fmt.Printf 接收任意个interface{}变量,所以在go编译阶段,会逃逸到栈heap上
func Printf(format string, a ...interface{}) (n int, err error) {}

// btree.AscendGreaterOrEqual 接收一个Item变量,实际还是一个实现了Less函数的interface{},同样会被编译器转换成interface{}的指针,这个时候就会逃逸到栈上
func (t *BTree) AscendGreaterOrEqual(pivot Item, iterator ItemIterator) {}

// 泛型实现,在编译阶段就会解析成实际的int类型,此时的pivot就是int,不会发生逃逸
func (t *BTreeG[T]) AscendGreaterOrEqual(pivot T, iterator ItemIteratorG[T]) {}

想要检测是否发生逃逸也很简单,直接使用golang自带的go build -gcflags="-m" .

golang这些年来做了很多编译的智能推断与优化,实际情况没有这么简单,举一个编译器可以优化的例子:
如果函数虽然接收的是一个interface{},但是实际内部只用了一种类型断言的时候,go编译器是不会发生逃逸到栈的处理的
func IsUIntStr(v interface{}) string {
  if vi, ok := v.(int); ok && vi > 0 {
     return "true"
  }
  return "false"
}
在这种情况下,是不会发生额外的复制行为的

回到正题,btree接收的类型是Item(interface{}),当有大量的Item对象调用的时候,需要复制到栈(heap)的数据也会很大.

这种情况在bench里面更加明显,Item实际内部只是一个int的占用空间,但是发生函数调用时,需要生成一个指向这个int的完整的interface{},每调用一个参数,增加的指针比实际int的内存占用空间还多的多

一个非空的interface{} 至少包含
一个itab指针,指向一个对象的类型信息
一个unsafe.Pointer,指向实际的data地址
@see: https://github.com/golang/go/blob/master/src/reflect/value.go#L200

使用了泛型之后,上述的函数在编译阶段就会被解析成:

func (t *BTreeG[int]) AscendGreaterOrEqual(pivot int, iterator ItemIteratorG[int]) {}

这样永远传递的都是指向int实际地址的一个确定指针,指挥复制一次指针,在函数的堆内就能完成操作

大量的对象逃逸到栈(heap)的时候,由于heap是随着函数调用随机不连续的调用与复制(copy+alloc),进一步加大gc了的扫描与处理时间,导致最终的性能下降

关于为什么逃逸到堆(heap)之后的go性能会下降,和对应的cpu与alloc相关的详细说明,可以参考下面这些,这里就不再赘述了:

引申

从上面的分析可以看到,泛型在处理interface{}的问题上提供了一个新的优化思路,在我目前使用的小工具合集的utils里面也有一定的借鉴思路

可以假设有这样一个非常常用的需求: 把一组数组去重,需要同时支持[]string,[]int64,[]float64等等

以前的go代码只有有两种思路:

使用 // go generate 使用模板自动生成N个函数但是这样会带来一个问题,不同函数的签名方法名肯定是不一样的,比如UniqStringSlice或者是UniqInt64Slice,对于调用方来说维护或者查询需要的函数也不简单
使用 switch .(type) 断言,全部代码写到一个函数里面去但是老的版本,只能使用一个interface{}去接收slice,好处是一个函数走天下,但是需要通过runtime+reflect的方式去判断具体传入的是怎么样 interface{}=>[]interface{}{}=> reflect.Slice(i).Interface()多种转换,性能必然下降

有了泛型(generics)之后,这类需求就非常容易了:(这也是在我的项目升级之后的utils库里)

func UniqSlice[T constraints.Ordered](list []T) []T {
	target := make([]T, 0)
	listMap := make(map[T]bool)
	for _, v := range list {
		if _, ok := listMap[v]; !ok {
			listMap[v] = true
			target = append(target, v)
		}
	}
	return target
}

总结

看起来泛型的引入没有对google/btree的代码结构发生大的变化,但是实际测试结果却让性能好了不少,可以预见到很多使用slice+interface{}的场景,更改为泛型(generics)之后都能带来代码量的减少和一定的性能的提升

后续希望可以看到更多的go代码使用泛型来进行优化,或者给出更多的解决实际问题的思路.