diff --git a/container/gtree/gtree_btree.go b/container/gtree/gtree_btree.go index fd6c06ce40d..f8283d6872c 100644 --- a/container/gtree/gtree_btree.go +++ b/container/gtree/gtree_btree.go @@ -7,32 +7,21 @@ package gtree import ( - "bytes" - "context" "fmt" - "strings" + "github.com/emirpasic/gods/trees/btree" "github.com/gogf/gf/v2/container/gvar" - "github.com/gogf/gf/v2/internal/intlog" - "github.com/gogf/gf/v2/internal/json" "github.com/gogf/gf/v2/internal/rwmutex" + "github.com/gogf/gf/v2/text/gstr" "github.com/gogf/gf/v2/util/gconv" ) // BTree holds elements of the B-tree. type BTree struct { mu rwmutex.RWMutex - root *BTreeNode comparator func(v1, v2 interface{}) int - size int // Total number of keys in the tree m int // order (maximum number of children) -} - -// BTreeNode is a single element within the tree. -type BTreeNode struct { - Parent *BTreeNode - Entries []*BTreeEntry // Contained keys in node - Children []*BTreeNode // Children nodes + tree *btree.Tree } // BTreeEntry represents the key-value pair contained within nodes. @@ -46,13 +35,11 @@ type BTreeEntry struct { // which is false in default. // Note that the `m` must be greater or equal than 3, or else it panics. func NewBTree(m int, comparator func(v1, v2 interface{}) int, safe ...bool) *BTree { - if m < 3 { - panic("Invalid order, should be at least 3") - } return &BTree{ - comparator: comparator, mu: rwmutex.Create(safe...), m: m, + comparator: comparator, + tree: btree.NewWith(m, comparator), } } @@ -81,21 +68,6 @@ func (tree *BTree) Set(key interface{}, value interface{}) { tree.doSet(key, value) } -// doSet inserts key-value pair node into the tree. -// If key already exists, then its value is updated with the new value. -func (tree *BTree) doSet(key interface{}, value interface{}) { - entry := &BTreeEntry{Key: key, Value: value} - if tree.root == nil { - tree.root = &BTreeNode{Entries: []*BTreeEntry{entry}, Children: []*BTreeNode{}} - tree.size++ - return - } - - if tree.insert(tree.root, entry) { - tree.size++ - } -} - // Sets batch sets key-values to the tree. func (tree *BTree) Sets(data map[interface{}]interface{}) { tree.mu.Lock() @@ -107,39 +79,19 @@ func (tree *BTree) Sets(data map[interface{}]interface{}) { // Get searches the node in the tree by `key` and returns its value or nil if key is not found in tree. func (tree *BTree) Get(key interface{}) (value interface{}) { - value, _ = tree.Search(key) - return -} - -// doSetWithLockCheck checks whether value of the key exists with mutex.Lock, -// if not exists, set value to the map with given `key`, -// or else just return the existing value. -// -// When setting value, if `value` is type of , -// it will be executed with mutex.Lock of the hash map, -// and its return value will be set to the map with `key`. -// -// It returns value with given `key`. -func (tree *BTree) doSetWithLockCheck(key interface{}, value interface{}) interface{} { tree.mu.Lock() defer tree.mu.Unlock() - if entry := tree.doSearch(key); entry != nil { - return entry.Value - } - if f, ok := value.(func() interface{}); ok { - value = f() - } - if value != nil { - tree.doSet(key, value) - } - return value + value, _ = tree.doGet(key) + return } // GetOrSet returns the value by key, // or sets value with given `value` if it does not exist and then returns this value. func (tree *BTree) GetOrSet(key interface{}, value interface{}) interface{} { - if v, ok := tree.Search(key); !ok { - return tree.doSetWithLockCheck(key, value) + tree.mu.Lock() + defer tree.mu.Unlock() + if v, ok := tree.doGet(key); !ok { + return tree.doSet(key, value) } else { return v } @@ -149,8 +101,10 @@ func (tree *BTree) GetOrSet(key interface{}, value interface{}) interface{} { // or sets value with returned value of callback function `f` if it does not exist // and then returns this value. func (tree *BTree) GetOrSetFunc(key interface{}, f func() interface{}) interface{} { - if v, ok := tree.Search(key); !ok { - return tree.doSetWithLockCheck(key, f()) + tree.mu.Lock() + defer tree.mu.Unlock() + if v, ok := tree.doGet(key); !ok { + return tree.doSet(key, f()) } else { return v } @@ -163,8 +117,10 @@ func (tree *BTree) GetOrSetFunc(key interface{}, f func() interface{}) interface // GetOrSetFuncLock differs with GetOrSetFunc function is that it executes function `f` // with mutex.Lock of the hash map. func (tree *BTree) GetOrSetFuncLock(key interface{}, f func() interface{}) interface{} { - if v, ok := tree.Search(key); !ok { - return tree.doSetWithLockCheck(key, f) + tree.mu.Lock() + defer tree.mu.Unlock() + if v, ok := tree.doGet(key); !ok { + return tree.doSet(key, f) } else { return v } @@ -197,8 +153,10 @@ func (tree *BTree) GetVarOrSetFuncLock(key interface{}, f func() interface{}) *g // SetIfNotExist sets `value` to the map if the `key` does not exist, and then returns true. // It returns false if `key` exists, and `value` would be ignored. func (tree *BTree) SetIfNotExist(key interface{}, value interface{}) bool { - if !tree.Contains(key) { - tree.doSetWithLockCheck(key, value) + tree.mu.Lock() + defer tree.mu.Unlock() + if _, ok := tree.doGet(key); !ok { + tree.doSet(key, value) return true } return false @@ -207,8 +165,10 @@ func (tree *BTree) SetIfNotExist(key interface{}, value interface{}) bool { // SetIfNotExistFunc sets value with return value of callback function `f`, and then returns true. // It returns false if `key` exists, and `value` would be ignored. func (tree *BTree) SetIfNotExistFunc(key interface{}, f func() interface{}) bool { - if !tree.Contains(key) { - tree.doSetWithLockCheck(key, f()) + tree.mu.Lock() + defer tree.mu.Unlock() + if _, ok := tree.doGet(key); !ok { + tree.doSet(key, f()) return true } return false @@ -220,31 +180,31 @@ func (tree *BTree) SetIfNotExistFunc(key interface{}, f func() interface{}) bool // SetIfNotExistFuncLock differs with SetIfNotExistFunc function is that // it executes function `f` with mutex.Lock of the hash map. func (tree *BTree) SetIfNotExistFuncLock(key interface{}, f func() interface{}) bool { - if !tree.Contains(key) { - tree.doSetWithLockCheck(key, f) + tree.mu.Lock() + defer tree.mu.Unlock() + if _, ok := tree.doGet(key); !ok { + tree.doSet(key, f) return true } return false } +// Search searches the tree with given `key`. +// Second return parameter `found` is true if key was found, otherwise false. +func (tree *BTree) Search(key interface{}) (value interface{}, found bool) { + tree.mu.RLock() + defer tree.mu.RUnlock() + return tree.tree.Get(key) +} + // Contains checks whether `key` exists in the tree. func (tree *BTree) Contains(key interface{}) bool { - _, ok := tree.Search(key) + tree.mu.RLock() + defer tree.mu.RUnlock() + _, ok := tree.doGet(key) return ok } -// doRemove removes the node from the tree by key. -// Key should adhere to the comparator's type assertion, otherwise method panics. -func (tree *BTree) doRemove(key interface{}) (value interface{}) { - node, index, found := tree.searchRecursively(tree.root, key) - if found { - value = node.Entries[index].Value - tree.delete(node, index) - tree.size-- - } - return -} - // Remove removes the node from the tree by `key`. func (tree *BTree) Remove(key interface{}) (value interface{}) { tree.mu.Lock() @@ -263,42 +223,36 @@ func (tree *BTree) Removes(keys []interface{}) { // IsEmpty returns true if tree does not contain any nodes func (tree *BTree) IsEmpty() bool { - return tree.Size() == 0 + tree.mu.RLock() + defer tree.mu.RUnlock() + return tree.tree.Size() == 0 } // Size returns number of nodes in the tree. func (tree *BTree) Size() int { tree.mu.RLock() defer tree.mu.RUnlock() - return tree.size + return tree.tree.Size() } // Keys returns all keys in asc order. func (tree *BTree) Keys() []interface{} { - keys := make([]interface{}, tree.Size()) - index := 0 - tree.IteratorAsc(func(key, value interface{}) bool { - keys[index] = key - index++ - return true - }) - return keys + tree.mu.RLock() + defer tree.mu.RUnlock() + return tree.tree.Keys() } // Values returns all values in asc order based on the key. func (tree *BTree) Values() []interface{} { - values := make([]interface{}, tree.Size()) - index := 0 - tree.IteratorAsc(func(key, value interface{}) bool { - values[index] = value - index++ - return true - }) - return values + tree.mu.RLock() + defer tree.mu.RUnlock() + return tree.tree.Values() } // Map returns all key-value items as map. func (tree *BTree) Map() map[interface{}]interface{} { + tree.mu.RLock() + defer tree.mu.RUnlock() m := make(map[interface{}]interface{}, tree.Size()) tree.IteratorAsc(func(key, value interface{}) bool { m[key] = value @@ -309,6 +263,8 @@ func (tree *BTree) Map() map[interface{}]interface{} { // MapStrAny returns all key-value items as map[string]interface{}. func (tree *BTree) MapStrAny() map[string]interface{} { + tree.mu.RLock() + defer tree.mu.RUnlock() m := make(map[string]interface{}, tree.Size()) tree.IteratorAsc(func(key, value interface{}) bool { m[gconv.String(key)] = value @@ -321,16 +277,14 @@ func (tree *BTree) MapStrAny() map[string]interface{} { func (tree *BTree) Clear() { tree.mu.Lock() defer tree.mu.Unlock() - tree.root = nil - tree.size = 0 + tree.tree.Clear() } // Replace the data of the tree with given `data`. func (tree *BTree) Replace(data map[interface{}]interface{}) { tree.mu.Lock() defer tree.mu.Unlock() - tree.root = nil - tree.size = 0 + tree.tree.Clear() for k, v := range data { tree.doSet(k, v) } @@ -340,65 +294,42 @@ func (tree *BTree) Replace(data map[interface{}]interface{}) { func (tree *BTree) Height() int { tree.mu.RLock() defer tree.mu.RUnlock() - return tree.root.height() + return tree.tree.Height() } // Left returns the left-most (min) entry or nil if tree is empty. func (tree *BTree) Left() *BTreeEntry { tree.mu.RLock() defer tree.mu.RUnlock() - node := tree.left(tree.root) - if node != nil { - return node.Entries[0] + node := tree.tree.Left() + if node == nil || node.Entries == nil || len(node.Entries) == 0 { + return nil + } + return &BTreeEntry{ + Key: node.Entries[0].Key, + Value: node.Entries[0].Value, } - return nil } // Right returns the right-most (max) entry or nil if tree is empty. func (tree *BTree) Right() *BTreeEntry { tree.mu.RLock() defer tree.mu.RUnlock() - node := tree.right(tree.root) - if node != nil { - return node.Entries[len(node.Entries)-1] + node := tree.tree.Right() + if node == nil || node.Entries == nil || len(node.Entries) == 0 { + return nil + } + return &BTreeEntry{ + Key: node.Entries[len(node.Entries)-1].Key, + Value: node.Entries[len(node.Entries)-1].Value, } - return nil } // String returns a string representation of container (for debugging purposes) func (tree *BTree) String() string { - if tree == nil { - return "" - } tree.mu.RLock() defer tree.mu.RUnlock() - var buffer bytes.Buffer - if tree.size != 0 { - tree.output(&buffer, tree.root, 0, true) - } - return buffer.String() -} - -// Search searches the tree with given `key`. -// Second return parameter `found` is true if key was found, otherwise false. -func (tree *BTree) Search(key interface{}) (value interface{}, found bool) { - tree.mu.RLock() - defer tree.mu.RUnlock() - node, index, found := tree.searchRecursively(tree.root, key) - if found { - return node.Entries[index].Value, true - } - return nil, false -} - -// Search searches the tree with given `key` without mutex. -// It returns the entry if found or otherwise nil. -func (tree *BTree) doSearch(key interface{}) *BTreeEntry { - node, index, found := tree.searchRecursively(tree.root, key) - if found { - return node.Entries[index] - } - return nil + return gstr.Replace(tree.tree.String(), "BTree\n", "") } // Print prints the tree to stdout. @@ -421,11 +352,13 @@ func (tree *BTree) IteratorFrom(key interface{}, match bool, f func(key, value i func (tree *BTree) IteratorAsc(f func(key, value interface{}) bool) { tree.mu.RLock() defer tree.mu.RUnlock() - node := tree.left(tree.root) - if node == nil { - return + it := tree.tree.Iterator() + for it.Begin(); it.Next(); { + index, value := it.Key(), it.Value() + if ok := f(index, value); !ok { + break + } } - tree.doIteratorAsc(node, node.Entries[0], 0, f) } // IteratorAscFrom iterates the tree readonly in ascending order with given callback function `f`. @@ -435,59 +368,13 @@ func (tree *BTree) IteratorAsc(f func(key, value interface{}) bool) { func (tree *BTree) IteratorAscFrom(key interface{}, match bool, f func(key, value interface{}) bool) { tree.mu.RLock() defer tree.mu.RUnlock() - node, index, found := tree.searchRecursively(tree.root, key) - if match { - if found { - tree.doIteratorAsc(node, node.Entries[index], index, f) - } - } else { - if index >= 0 && index < len(node.Entries) { - tree.doIteratorAsc(node, node.Entries[index], index, f) - } - } -} - -func (tree *BTree) doIteratorAsc(node *BTreeNode, entry *BTreeEntry, index int, f func(key, value interface{}) bool) { - first := true -loop: - if entry == nil { - return - } - if !f(entry.Key, entry.Value) { + var keys = tree.tree.Keys() + index, isIterator := tree.iteratorFromGetIndex(key, keys, match) + if !isIterator { return } - // Find current entry position in current node - if !first { - index, _ = tree.search(node, entry.Key) - } else { - first = false - } - // Try to go down to the child right of the current entry - if index+1 < len(node.Children) { - node = node.Children[index+1] - // Try to go down to the child left of the current node - for len(node.Children) > 0 { - node = node.Children[0] - } - // Return the left-most entry - entry = node.Entries[0] - goto loop - } - // Above assures that we have reached a leaf node, so return the next entry in current node (if any) - if index+1 < len(node.Entries) { - entry = node.Entries[index+1] - goto loop - } - // Reached leaf node and there are no entries to the right of the current entry, so go up to the parent - for node.Parent != nil { - node = node.Parent - // Find next entry position in current node (note: search returns the first equal or bigger than entry) - index, _ = tree.search(node, entry.Key) - // Check that there is a next entry position in current node - if index < len(node.Entries) { - entry = node.Entries[index] - goto loop - } + for ; index < len(keys); index++ { + f(keys[index], tree.Get(keys[index])) } } @@ -496,13 +383,13 @@ loop: func (tree *BTree) IteratorDesc(f func(key, value interface{}) bool) { tree.mu.RLock() defer tree.mu.RUnlock() - node := tree.right(tree.root) - if node == nil { - return + it := tree.tree.Iterator() + for it.End(); it.Prev(); { + index, value := it.Key(), it.Value() + if ok := f(index, value); !ok { + break + } } - index := len(node.Entries) - 1 - entry := node.Entries[index] - tree.doIteratorDesc(node, entry, index, f) } // IteratorDescFrom iterates the tree readonly in descending order with given callback function `f`. @@ -512,468 +399,70 @@ func (tree *BTree) IteratorDesc(f func(key, value interface{}) bool) { func (tree *BTree) IteratorDescFrom(key interface{}, match bool, f func(key, value interface{}) bool) { tree.mu.RLock() defer tree.mu.RUnlock() - node, index, found := tree.searchRecursively(tree.root, key) - if match { - if found { - tree.doIteratorDesc(node, node.Entries[index], index, f) - } - } else { - if index >= 0 && index < len(node.Entries) { - tree.doIteratorDesc(node, node.Entries[index], index, f) - } - } -} - -// IteratorDesc iterates the tree readonly in descending order with given callback function `f`. -// If `f` returns true, then it continues iterating; or false to stop. -func (tree *BTree) doIteratorDesc(node *BTreeNode, entry *BTreeEntry, index int, f func(key, value interface{}) bool) { - first := true -loop: - if entry == nil { - return - } - if !f(entry.Key, entry.Value) { - return - } - // Find current entry position in current node - if !first { - index, _ = tree.search(node, entry.Key) - } else { - first = false - } - // Try to go down to the child left of the current entry - if index < len(node.Children) { - node = node.Children[index] - // Try to go down to the child right of the current node - for len(node.Children) > 0 { - node = node.Children[len(node.Children)-1] - } - // Return the right-most entry - entry = node.Entries[len(node.Entries)-1] - goto loop - } - // Above assures that we have reached a leaf node, so return the previous entry in current node (if any) - if index-1 >= 0 { - entry = node.Entries[index-1] - goto loop - } - - // Reached leaf node and there are no entries to the left of the current entry, so go up to the parent - for node.Parent != nil { - node = node.Parent - // Find previous entry position in current node (note: search returns the first equal or bigger than entry) - index, _ = tree.search(node, entry.Key) - // Check that there is a previous entry position in current node - if index-1 >= 0 { - entry = node.Entries[index-1] - goto loop - } - } -} - -func (tree *BTree) output(buffer *bytes.Buffer, node *BTreeNode, level int, isTail bool) { - for e := 0; e < len(node.Entries)+1; e++ { - if e < len(node.Children) { - tree.output(buffer, node.Children[e], level+1, true) - } - if e < len(node.Entries) { - if _, err := buffer.WriteString(strings.Repeat(" ", level)); err != nil { - intlog.Errorf(context.TODO(), `%+v`, err) - } - if _, err := buffer.WriteString(fmt.Sprintf("%v", node.Entries[e].Key) + "\n"); err != nil { - intlog.Errorf(context.TODO(), `%+v`, err) - } - } - } -} - -func (node *BTreeNode) height() int { - h := 0 - n := node - for ; n != nil; n = n.Children[0] { - h++ - if len(n.Children) == 0 { - break - } - } - return h -} - -func (tree *BTree) isLeaf(node *BTreeNode) bool { - return len(node.Children) == 0 -} - -// func (tree *BTree) isFull(node *BTreeNode) bool { -// return len(node.Entries) == tree.maxEntries() -// } - -func (tree *BTree) shouldSplit(node *BTreeNode) bool { - return len(node.Entries) > tree.maxEntries() -} - -func (tree *BTree) maxChildren() int { - return tree.m -} - -func (tree *BTree) minChildren() int { - return (tree.m + 1) / 2 // ceil(m/2) -} - -func (tree *BTree) maxEntries() int { - return tree.maxChildren() - 1 -} - -func (tree *BTree) minEntries() int { - return tree.minChildren() - 1 -} - -func (tree *BTree) middle() int { - // "-1" to favor right nodes to have more keys when splitting - return (tree.m - 1) / 2 -} - -// search does search only within the single node among its entries -func (tree *BTree) search(node *BTreeNode, key interface{}) (index int, found bool) { - low, mid, high := 0, 0, len(node.Entries)-1 - for low <= high { - mid = low + (high-low)/2 - compare := tree.getComparator()(key, node.Entries[mid].Key) - switch { - case compare > 0: - low = mid + 1 - case compare < 0: - high = mid - 1 - case compare == 0: - return mid, true - } - } - return low, false -} - -// searchRecursively searches recursively down the tree starting at the startNode -func (tree *BTree) searchRecursively(startNode *BTreeNode, key interface{}) (node *BTreeNode, index int, found bool) { - if tree.size == 0 { - return nil, -1, false - } - node = startNode - for { - index, found = tree.search(node, key) - if found { - return node, index, true - } - if tree.isLeaf(node) { - return node, index, false - } - node = node.Children[index] - } -} - -func (tree *BTree) insert(node *BTreeNode, entry *BTreeEntry) (inserted bool) { - if tree.isLeaf(node) { - return tree.insertIntoLeaf(node, entry) - } - return tree.insertIntoInternal(node, entry) -} - -func (tree *BTree) insertIntoLeaf(node *BTreeNode, entry *BTreeEntry) (inserted bool) { - insertPosition, found := tree.search(node, entry.Key) - if found { - node.Entries[insertPosition] = entry - return false - } - // Insert entry's key in the middle of the node - node.Entries = append(node.Entries, nil) - copy(node.Entries[insertPosition+1:], node.Entries[insertPosition:]) - node.Entries[insertPosition] = entry - tree.split(node) - return true -} - -func (tree *BTree) insertIntoInternal(node *BTreeNode, entry *BTreeEntry) (inserted bool) { - insertPosition, found := tree.search(node, entry.Key) - if found { - node.Entries[insertPosition] = entry - return false - } - return tree.insert(node.Children[insertPosition], entry) -} - -func (tree *BTree) split(node *BTreeNode) { - if !tree.shouldSplit(node) { - return - } - - if node == tree.root { - tree.splitRoot() + var keys = tree.tree.Keys() + index, isIterator := tree.iteratorFromGetIndex(key, keys, match) + if !isIterator { return } - - tree.splitNonRoot(node) -} - -func (tree *BTree) splitNonRoot(node *BTreeNode) { - middle := tree.middle() - parent := node.Parent - - left := &BTreeNode{Entries: append([]*BTreeEntry(nil), node.Entries[:middle]...), Parent: parent} - right := &BTreeNode{Entries: append([]*BTreeEntry(nil), node.Entries[middle+1:]...), Parent: parent} - - // Move children from the node to be split into left and right nodes - if !tree.isLeaf(node) { - left.Children = append([]*BTreeNode(nil), node.Children[:middle+1]...) - right.Children = append([]*BTreeNode(nil), node.Children[middle+1:]...) - setParent(left.Children, left) - setParent(right.Children, right) - } - - insertPosition, _ := tree.search(parent, node.Entries[middle].Key) - - // Insert middle key into parent - parent.Entries = append(parent.Entries, nil) - copy(parent.Entries[insertPosition+1:], parent.Entries[insertPosition:]) - parent.Entries[insertPosition] = node.Entries[middle] - - // Set child left of inserted key in parent to the created left node - parent.Children[insertPosition] = left - - // Set child right of inserted key in parent to the created right node - parent.Children = append(parent.Children, nil) - copy(parent.Children[insertPosition+2:], parent.Children[insertPosition+1:]) - parent.Children[insertPosition+1] = right - - tree.split(parent) -} - -func (tree *BTree) splitRoot() { - middle := tree.middle() - left := &BTreeNode{Entries: append([]*BTreeEntry(nil), tree.root.Entries[:middle]...)} - right := &BTreeNode{Entries: append([]*BTreeEntry(nil), tree.root.Entries[middle+1:]...)} - - // Move children from the node to be split into left and right nodes - if !tree.isLeaf(tree.root) { - left.Children = append([]*BTreeNode(nil), tree.root.Children[:middle+1]...) - right.Children = append([]*BTreeNode(nil), tree.root.Children[middle+1:]...) - setParent(left.Children, left) - setParent(right.Children, right) - } - - // Root is a node with one entry and two children (left and right) - newRoot := &BTreeNode{ - Entries: []*BTreeEntry{tree.root.Entries[middle]}, - Children: []*BTreeNode{left, right}, - } - - left.Parent = newRoot - right.Parent = newRoot - tree.root = newRoot -} - -func setParent(nodes []*BTreeNode, parent *BTreeNode) { - for _, node := range nodes { - node.Parent = parent + for ; index >= 0; index-- { + f(keys[index], tree.Get(keys[index])) } } -func (tree *BTree) left(node *BTreeNode) *BTreeNode { - if tree.size == 0 { - return nil - } - current := node - for { - if tree.isLeaf(current) { - return current - } - current = current.Children[0] - } -} - -func (tree *BTree) right(node *BTreeNode) *BTreeNode { - if tree.size == 0 { - return nil - } - current := node - for { - if tree.isLeaf(current) { - return current - } - current = current.Children[len(current.Children)-1] - } -} - -// leftSibling returns the node's left sibling and child index (in parent) if it exists, otherwise (nil,-1) -// key is any of keys in node (could even be deleted). -func (tree *BTree) leftSibling(node *BTreeNode, key interface{}) (*BTreeNode, int) { - if node.Parent != nil { - index, _ := tree.search(node.Parent, key) - index-- - if index >= 0 && index < len(node.Parent.Children) { - return node.Parent.Children[index], index - } - } - return nil, -1 -} - -// rightSibling returns the node's right sibling and child index (in parent) if it exists, otherwise (nil,-1) -// key is any of keys in node (could even be deleted). -func (tree *BTree) rightSibling(node *BTreeNode, key interface{}) (*BTreeNode, int) { - if node.Parent != nil { - index, _ := tree.search(node.Parent, key) - index++ - if index < len(node.Parent.Children) { - return node.Parent.Children[index], index - } - } - return nil, -1 -} - -// delete deletes an entry in node at entries' index -// ref.: https://en.wikipedia.org/wiki/B-tree#Deletion -func (tree *BTree) delete(node *BTreeNode, index int) { - // deleting from a leaf node - if tree.isLeaf(node) { - deletedKey := node.Entries[index].Key - tree.deleteEntry(node, index) - tree.reBalance(node, deletedKey) - if len(tree.root.Entries) == 0 { - tree.root = nil - } - return - } - - // deleting from an internal node - leftLargestNode := tree.right(node.Children[index]) // largest node in the left sub-tree (assumed to exist) - leftLargestEntryIndex := len(leftLargestNode.Entries) - 1 - node.Entries[index] = leftLargestNode.Entries[leftLargestEntryIndex] - deletedKey := leftLargestNode.Entries[leftLargestEntryIndex].Key - tree.deleteEntry(leftLargestNode, leftLargestEntryIndex) - tree.reBalance(leftLargestNode, deletedKey) +// MarshalJSON implements the interface MarshalJSON for json.Marshal. +func (tree *BTree) MarshalJSON() (jsonBytes []byte, err error) { + tree.mu.RLock() + defer tree.mu.RUnlock() + return tree.tree.MarshalJSON() } -// reBalance reBalances the tree after deletion if necessary and returns true, otherwise false. -// Note that we first delete the entry and then call reBalance, thus the passed deleted key as reference. -func (tree *BTree) reBalance(node *BTreeNode, deletedKey interface{}) { - // check if re-balancing is needed - if node == nil || len(node.Entries) >= tree.minEntries() { - return - } - - // try to borrow from left sibling - leftSibling, leftSiblingIndex := tree.leftSibling(node, deletedKey) - if leftSibling != nil && len(leftSibling.Entries) > tree.minEntries() { - // rotate right - node.Entries = append([]*BTreeEntry{node.Parent.Entries[leftSiblingIndex]}, node.Entries...) // prepend parent's separator entry to node's entries - node.Parent.Entries[leftSiblingIndex] = leftSibling.Entries[len(leftSibling.Entries)-1] - tree.deleteEntry(leftSibling, len(leftSibling.Entries)-1) - if !tree.isLeaf(leftSibling) { - leftSiblingRightMostChild := leftSibling.Children[len(leftSibling.Children)-1] - leftSiblingRightMostChild.Parent = node - node.Children = append([]*BTreeNode{leftSiblingRightMostChild}, node.Children...) - tree.deleteChild(leftSibling, len(leftSibling.Children)-1) - } - return - } - - // try to borrow from right sibling - rightSibling, rightSiblingIndex := tree.rightSibling(node, deletedKey) - if rightSibling != nil && len(rightSibling.Entries) > tree.minEntries() { - // rotate left - node.Entries = append(node.Entries, node.Parent.Entries[rightSiblingIndex-1]) // append parent's separator entry to node's entries - node.Parent.Entries[rightSiblingIndex-1] = rightSibling.Entries[0] - tree.deleteEntry(rightSibling, 0) - if !tree.isLeaf(rightSibling) { - rightSiblingLeftMostChild := rightSibling.Children[0] - rightSiblingLeftMostChild.Parent = node - node.Children = append(node.Children, rightSiblingLeftMostChild) - tree.deleteChild(rightSibling, 0) - } - return - } - - // merge with siblings - if rightSibling != nil { - // merge with right sibling - node.Entries = append(node.Entries, node.Parent.Entries[rightSiblingIndex-1]) - node.Entries = append(node.Entries, rightSibling.Entries...) - deletedKey = node.Parent.Entries[rightSiblingIndex-1].Key - tree.deleteEntry(node.Parent, rightSiblingIndex-1) - tree.appendChildren(node.Parent.Children[rightSiblingIndex], node) - tree.deleteChild(node.Parent, rightSiblingIndex) - } else if leftSibling != nil { - // merge with left sibling - entries := append([]*BTreeEntry(nil), leftSibling.Entries...) - entries = append(entries, node.Parent.Entries[leftSiblingIndex]) - node.Entries = append(entries, node.Entries...) - deletedKey = node.Parent.Entries[leftSiblingIndex].Key - tree.deleteEntry(node.Parent, leftSiblingIndex) - tree.prependChildren(node.Parent.Children[leftSiblingIndex], node) - tree.deleteChild(node.Parent, leftSiblingIndex) +// doSet inserts key-value pair node into the tree. +// If key already exists, then its value is updated with the new value. +// If `value` is type of , +// it will be executed and its return value will be set to the map with `key`. +// +// It returns value with given `key`. +func (tree *BTree) doSet(key interface{}, value interface{}) interface{} { + if f, ok := value.(func() interface{}); ok { + value = f() } - - // make the merged node the root if its parent was the root and the root is empty - if node.Parent == tree.root && len(tree.root.Entries) == 0 { - tree.root = node - node.Parent = nil - return + if value == nil { + return value } - - // parent might be underflow, so try to reBalance if necessary - tree.reBalance(node.Parent, deletedKey) -} - -func (tree *BTree) prependChildren(fromNode *BTreeNode, toNode *BTreeNode) { - children := append([]*BTreeNode(nil), fromNode.Children...) - toNode.Children = append(children, toNode.Children...) - setParent(fromNode.Children, toNode) -} - -func (tree *BTree) appendChildren(fromNode *BTreeNode, toNode *BTreeNode) { - toNode.Children = append(toNode.Children, fromNode.Children...) - setParent(fromNode.Children, toNode) + tree.tree.Put(key, value) + return value } -func (tree *BTree) deleteEntry(node *BTreeNode, index int) { - copy(node.Entries[index:], node.Entries[index+1:]) - node.Entries[len(node.Entries)-1] = nil - node.Entries = node.Entries[:len(node.Entries)-1] +// doGet get the value from the tree by key. +func (tree *BTree) doGet(key interface{}) (value interface{}, ok bool) { + return tree.tree.Get(key) } -func (tree *BTree) deleteChild(node *BTreeNode, index int) { - if index >= len(node.Children) { - return - } - copy(node.Children[index:], node.Children[index+1:]) - node.Children[len(node.Children)-1] = nil - node.Children = node.Children[:len(node.Children)-1] +// doRemove removes the node from the tree by key. +// Key should adhere to the comparator's type assertion, otherwise method panics. +func (tree *BTree) doRemove(key interface{}) (value interface{}) { + value, _ = tree.tree.Get(key) + tree.tree.Remove(key) + return } -// MarshalJSON implements the interface MarshalJSON for json.Marshal. -func (tree BTree) MarshalJSON() (jsonBytes []byte, err error) { - if tree.root == nil { - return []byte("null"), nil - } - buffer := bytes.NewBuffer(nil) - buffer.WriteByte('{') - tree.Iterator(func(key, value interface{}) bool { - valueBytes, valueJsonErr := json.Marshal(value) - if valueJsonErr != nil { - err = valueJsonErr - return false +// iteratorFromGetIndex returns the index of the key in the keys slice. +// The parameter `match` specifies whether starting iterating if the `key` is fully matched, +// or else using index searching iterating. +// If `isIterator` is true, iterator is available; or else not. +func (tree *BTree) iteratorFromGetIndex(key interface{}, keys []interface{}, match bool) (index int, isIterator bool) { + if match { + for i, k := range keys { + if k == key { + isIterator = true + index = i + } } - if buffer.Len() > 1 { - buffer.WriteByte(',') + } else { + if i, ok := key.(int); ok { + isIterator = true + index = i } - buffer.WriteString(fmt.Sprintf(`"%v":%s`, key, valueBytes)) - return true - }) - buffer.WriteByte('}') - return buffer.Bytes(), nil -} - -// getComparator returns the comparator if it's previously set, -// or else it panics. -func (tree *BTree) getComparator() func(a, b interface{}) int { - if tree.comparator == nil { - panic("comparator is missing for tree") } - return tree.comparator + return } diff --git a/container/gtree/gtree_z_example_btree_test.go b/container/gtree/gtree_z_example_btree_test.go index fb5a551a2e9..4c06f058ae1 100644 --- a/container/gtree/gtree_z_example_btree_test.go +++ b/container/gtree/gtree_z_example_btree_test.go @@ -453,7 +453,7 @@ func ExampleBTree_String() { fmt.Println(tree.String()) // Output: - // key0 + // key0 // key1 // key2 // key3 @@ -484,7 +484,7 @@ func ExampleBTree_Print() { tree.Print() // Output: - // key0 + // key0 // key1 // key2 // key3 diff --git a/go.mod b/go.mod index e4602e51f7d..7a79d171108 100644 --- a/go.mod +++ b/go.mod @@ -20,6 +20,7 @@ require ( ) require ( + github.com/emirpasic/gods v1.18.1 // indirect github.com/go-logr/logr v1.2.3 // indirect github.com/go-logr/stdr v1.2.2 // indirect github.com/mattn/go-colorable v0.1.13 // indirect diff --git a/go.sum b/go.sum index 48f46182bb3..b0003fe6b3e 100644 --- a/go.sum +++ b/go.sum @@ -3,6 +3,8 @@ github.com/BurntSushi/toml v1.3.2/go.mod h1:CxXYINrC8qIiEnFrOxCa7Jy5BFHlXnUU2pbi github.com/clbanning/mxj/v2 v2.7.0 h1:WA/La7UGCanFe5NpHF0Q3DNtnCsVoxbPKuyBNHWRyME= github.com/clbanning/mxj/v2 v2.7.0/go.mod h1:hNiWqW14h+kc+MdF9C6/YoRfjEJoR3ou6tn/Qo+ve2s= github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= +github.com/emirpasic/gods v1.18.1 h1:FXtiHYKDGKCW2KzwZKx0iC0PQmdlorYgdFG9jPXJ1Bc= +github.com/emirpasic/gods v1.18.1/go.mod h1:8tpGGwCnJ5H4r6BWwaV6OrWmMoPhUl5jm/FMNAnJvWQ= github.com/fatih/color v1.16.0 h1:zmkK9Ngbjj+K0yRhTVONQh1p/HknKYSlNT+vZCzyokM= github.com/fatih/color v1.16.0/go.mod h1:fL2Sau1YI5c0pdGEVCbKQbLXB6edEj1ZgiY4NijnWvE= github.com/fsnotify/fsnotify v1.7.0 h1:8JEhPFa5W2WU7YfeZzPNqzMP6Lwt7L2715Ggo0nosvA=