Merge pull request #9574 from gyuho/btree

vendor: upgrade "google/btree"

Gopkg.lock

@@ -96,7 +96,7 @@
 [[projects]]
   name = "github.com/google/btree"
   packages = ["."]
-  revision = "925471ac9e2131377a91e1595defec898166fe49"
+  revision = "e89373fe6b4a7413d7acd6da1725b83ef713e6e4"
 
 [[projects]]
   name = "github.com/gorilla/websocket"

vendor/github.com/google/btree/btree.go

@@ -52,6 +52,7 @@ import (
 	"io"
 	"sort"
 	"strings"
+	"sync"
 )
 
 // Item represents a single object in the tree.
@@ -76,8 +77,9 @@ var (
 // 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 not safe for concurrent write access.
+// Two Btrees using the same freelist are safe for concurrent write access.
 type FreeList struct {
+	mu       sync.Mutex
 	freelist []*node
 }
 
@@ -88,20 +90,29 @@ func NewFreeList(size int) *FreeList {
 }
 
 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
 	f.freelist = f.freelist[:index]
+	f.mu.Unlock()
 	return
 }
 
-func (f *FreeList) freeNode(n *node) {
+// 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)
+		out = true
 	}
+	f.mu.Unlock()
+	return
 }
 
 // ItemIterator allows callers of Ascend* to iterate in-order over portions of
@@ -123,8 +134,8 @@ func NewWithFreeList(degree int, f *FreeList) *BTree {
 		panic("bad degree")
 	}
 	return &BTree{
-		degree:   degree,
-		freelist: f,
+		degree: degree,
+		cow:    &copyOnWriteContext{freelist: f},
 	}
 }
 
@@ -233,7 +244,34 @@ func (s *children) truncate(index int) {
 type node struct {
 	items    items
 	children children
-	t        *BTree
+	cow      *copyOnWriteContext
+}
+
+func (n *node) mutableFor(cow *copyOnWriteContext) *node {
+	if n.cow == cow {
+		return n
+	}
+	out := cow.newNode()
+	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
 }
 
 // split splits the given node at the given index.  The current node shrinks,
@@ -241,7 +279,7 @@ type node struct {
 // containing all items/children after it.
 func (n *node) split(i int) (Item, *node) {
 	item := n.items[i]
-	next := n.t.newNode()
+	next := n.cow.newNode()
 	next.items = append(next.items, n.items[i+1:]...)
 	n.items.truncate(i)
 	if len(n.children) > 0 {
@@ -257,7 +295,7 @@ func (n *node) maybeSplitChild(i, maxItems int) bool {
 	if len(n.children[i].items) < maxItems {
 		return false
 	}
-	first := n.children[i]
+	first := n.mutableChild(i)
 	item, second := first.split(maxItems / 2)
 	n.items.insertAt(i, item)
 	n.children.insertAt(i+1, second)
@@ -291,7 +329,7 @@ func (n *node) insert(item Item, maxItems int) Item {
 			return out
 		}
 	}
-	return n.children[i].insert(item, maxItems)
+	return n.mutableChild(i).insert(item, maxItems)
 }
 
 // get finds the given key in the subtree and returns it.
@@ -369,10 +407,10 @@ func (n *node) remove(item Item, minItems int, typ toRemove) Item {
 		panic("invalid type")
 	}
 	// If we get to here, we have children.
-	child := n.children[i]
-	if len(child.items) <= minItems {
+	if len(n.children[i].items) <= minItems {
 		return n.growChildAndRemove(i, item, minItems, typ)
 	}
+	child := n.mutableChild(i)
 	// Either we had enough items to begin with, or we've done some
 	// merging/stealing, because we've got enough now and we're ready to return
 	// stuff.
@@ -411,10 +449,10 @@ func (n *node) remove(item Item, minItems int, typ toRemove) Item {
 // 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 {
-	child := n.children[i]
 	if i > 0 && len(n.children[i-1].items) > minItems {
 		// Steal from left child
-		stealFrom := n.children[i-1]
+		child := n.mutableChild(i)
+		stealFrom := n.mutableChild(i - 1)
 		stolenItem := stealFrom.items.pop()
 		child.items.insertAt(0, n.items[i-1])
 		n.items[i-1] = stolenItem
@@ -423,7 +461,8 @@ func (n *node) growChildAndRemove(i int, item Item, minItems int, typ toRemove)
 		}
 	} else if i < len(n.items) && len(n.children[i+1].items) > minItems {
 		// steal from right child
-		stealFrom := n.children[i+1]
+		child := n.mutableChild(i)
+		stealFrom := n.mutableChild(i + 1)
 		stolenItem := stealFrom.items.removeAt(0)
 		child.items = append(child.items, n.items[i])
 		n.items[i] = stolenItem
@@ -433,15 +472,15 @@ func (n *node) growChildAndRemove(i int, item Item, minItems int, typ toRemove)
 	} else {
 		if i >= len(n.items) {
 			i--
-			child = n.children[i]
 		}
+		child := n.mutableChild(i)
 		// merge with right child
 		mergeItem := n.items.removeAt(i)
 		mergeChild := n.children.removeAt(i + 1)
 		child.items = append(child.items, mergeItem)
 		child.items = append(child.items, mergeChild.items...)
 		child.children = append(child.children, mergeChild.children...)
-		n.t.freeNode(mergeChild)
+		n.cow.freeNode(mergeChild)
 	}
 	return n.remove(item, minItems, typ)
 }
@@ -535,12 +574,54 @@ func (n *node) print(w io.Writer, level int) {
 // Write operations are not safe for concurrent mutation by multiple
 // goroutines, but Read operations are.
 type BTree struct {
-	degree   int
-	length   int
-	root     *node
+	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
+// it, and must create a new, writable copy (IE: it's a Clone).
+//
+// When doing any write operation, we maintain the invariant that the current
+// node's context is equal to the context of the tree that requested the write.
+// We do this by, before we descend into any node, creating a copy with the
+// correct context if the contexts don't match.
+//
+// Since the node we're currently visiting on any write has the requesting
+// 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,
+// 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) {
+	// Create two entirely new copy-on-write contexts.
+	// This operation effectively creates three trees:
+	//   the original, shared nodes (old b.cow)
+	//   the new b.cow nodes
+	//   the new out.cow nodes
+	cow1, cow2 := *t.cow, *t.cow
+	out := *t
+	t.cow = &cow1
+	out.cow = &cow2
+	return &out
+}
+
 // maxItems returns the max number of items to allow per node.
 func (t *BTree) maxItems() int {
 	return t.degree*2 - 1
@@ -552,18 +633,37 @@ func (t *BTree) minItems() int {
 	return t.degree - 1
 }
 
-func (t *BTree) newNode() (n *node) {
-	n = t.freelist.newNode()
-	n.t = t
+func (c *copyOnWriteContext) newNode() (n *node) {
+	n = c.freelist.newNode()
+	n.cow = c
 	return
 }
 
-func (t *BTree) freeNode(n *node) {
-	// clear to allow GC
-	n.items.truncate(0)
-	n.children.truncate(0)
-	n.t = nil // clear to allow GC
-	t.freelist.freeNode(n)
+type freeType int
+
+const (
+	ftFreelistFull freeType = iota // node was freed (available for GC, not stored in freelist)
+	ftStored                       // node was stored in the freelist for later use
+	ftNotOwned                     // node was ignored by COW, since it's owned by another one
+)
+
+// 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)
+		n.children.truncate(0)
+		n.cow = nil
+		if c.freelist.freeNode(n) {
+			return ftStored
+		} else {
+			return ftFreelistFull
+		}
+	} else {
+		return ftNotOwned
+	}
 }
 
 // ReplaceOrInsert adds the given item to the tree.  If an item in the tree
@@ -576,16 +676,19 @@ func (t *BTree) ReplaceOrInsert(item Item) Item {
 		panic("nil item being added to BTree")
 	}
 	if t.root == nil {
-		t.root = t.newNode()
+		t.root = t.cow.newNode()
 		t.root.items = append(t.root.items, item)
 		t.length++
 		return nil
-	} else if len(t.root.items) >= t.maxItems() {
-		item2, second := t.root.split(t.maxItems() / 2)
-		oldroot := t.root
-		t.root = t.newNode()
-		t.root.items = append(t.root.items, item2)
-		t.root.children = append(t.root.children, oldroot, second)
+	} else {
+		t.root = t.root.mutableFor(t.cow)
+		if len(t.root.items) >= t.maxItems() {
+			item2, second := t.root.split(t.maxItems() / 2)
+			oldroot := t.root
+			t.root = t.cow.newNode()
+			t.root.items = append(t.root.items, item2)
+			t.root.children = append(t.root.children, oldroot, second)
+		}
 	}
 	out := t.root.insert(item, t.maxItems())
 	if out == nil {
@@ -616,11 +719,12 @@ 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.freeNode(oldroot)
+		t.cow.freeNode(oldroot)
 	}
 	if out != nil {
 		t.length--
@@ -729,6 +833,45 @@ func (t *BTree) Len() int {
 	return t.length
 }
 
+// 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) {
+	if t.root != nil && addNodesToFreelist {
+		t.root.reset(t.cow)
+	}
+	t.root, t.length = nil, 0
+}
+
+// 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
+		}
+	}
+	return c.freeNode(n) != ftFreelistFull
+}
+
 // Int implements the Item interface for integers.
 type Int int