github.com/wangyougui/gf/v2@v2.6.5/container/gtree/gtree_btree.go (about)

     1  // Copyright GoFrame Author(https://goframe.org). All Rights Reserved.
     2  //
     3  // This Source Code Form is subject to the terms of the MIT License.
     4  // If a copy of the MIT was not distributed with this file,
     5  // You can obtain one at https://github.com/wangyougui/gf.
     6  
     7  package gtree
     8  
     9  import (
    10  	"bytes"
    11  	"context"
    12  	"fmt"
    13  	"strings"
    14  
    15  	"github.com/wangyougui/gf/v2/container/gvar"
    16  	"github.com/wangyougui/gf/v2/internal/intlog"
    17  	"github.com/wangyougui/gf/v2/internal/json"
    18  	"github.com/wangyougui/gf/v2/internal/rwmutex"
    19  	"github.com/wangyougui/gf/v2/util/gconv"
    20  )
    21  
    22  // BTree holds elements of the B-tree.
    23  type BTree struct {
    24  	mu         rwmutex.RWMutex
    25  	root       *BTreeNode
    26  	comparator func(v1, v2 interface{}) int
    27  	size       int // Total number of keys in the tree
    28  	m          int // order (maximum number of children)
    29  }
    30  
    31  // BTreeNode is a single element within the tree.
    32  type BTreeNode struct {
    33  	Parent   *BTreeNode
    34  	Entries  []*BTreeEntry // Contained keys in node
    35  	Children []*BTreeNode  // Children nodes
    36  }
    37  
    38  // BTreeEntry represents the key-value pair contained within nodes.
    39  type BTreeEntry struct {
    40  	Key   interface{}
    41  	Value interface{}
    42  }
    43  
    44  // NewBTree instantiates a B-tree with `m` (maximum number of children) and a custom key comparator.
    45  // The parameter `safe` is used to specify whether using tree in concurrent-safety,
    46  // which is false in default.
    47  // Note that the `m` must be greater or equal than 3, or else it panics.
    48  func NewBTree(m int, comparator func(v1, v2 interface{}) int, safe ...bool) *BTree {
    49  	if m < 3 {
    50  		panic("Invalid order, should be at least 3")
    51  	}
    52  	return &BTree{
    53  		comparator: comparator,
    54  		mu:         rwmutex.Create(safe...),
    55  		m:          m,
    56  	}
    57  }
    58  
    59  // NewBTreeFrom instantiates a B-tree with `m` (maximum number of children), a custom key comparator and data map.
    60  // The parameter `safe` is used to specify whether using tree in concurrent-safety,
    61  // which is false in default.
    62  func NewBTreeFrom(m int, comparator func(v1, v2 interface{}) int, data map[interface{}]interface{}, safe ...bool) *BTree {
    63  	tree := NewBTree(m, comparator, safe...)
    64  	for k, v := range data {
    65  		tree.doSet(k, v)
    66  	}
    67  	return tree
    68  }
    69  
    70  // Clone returns a new tree with a copy of current tree.
    71  func (tree *BTree) Clone() *BTree {
    72  	newTree := NewBTree(tree.m, tree.comparator, tree.mu.IsSafe())
    73  	newTree.Sets(tree.Map())
    74  	return newTree
    75  }
    76  
    77  // Set inserts key-value item into the tree.
    78  func (tree *BTree) Set(key interface{}, value interface{}) {
    79  	tree.mu.Lock()
    80  	defer tree.mu.Unlock()
    81  	tree.doSet(key, value)
    82  }
    83  
    84  // doSet inserts key-value pair node into the tree.
    85  // If key already exists, then its value is updated with the new value.
    86  func (tree *BTree) doSet(key interface{}, value interface{}) {
    87  	entry := &BTreeEntry{Key: key, Value: value}
    88  	if tree.root == nil {
    89  		tree.root = &BTreeNode{Entries: []*BTreeEntry{entry}, Children: []*BTreeNode{}}
    90  		tree.size++
    91  		return
    92  	}
    93  
    94  	if tree.insert(tree.root, entry) {
    95  		tree.size++
    96  	}
    97  }
    98  
    99  // Sets batch sets key-values to the tree.
   100  func (tree *BTree) Sets(data map[interface{}]interface{}) {
   101  	tree.mu.Lock()
   102  	defer tree.mu.Unlock()
   103  	for k, v := range data {
   104  		tree.doSet(k, v)
   105  	}
   106  }
   107  
   108  // Get searches the node in the tree by `key` and returns its value or nil if key is not found in tree.
   109  func (tree *BTree) Get(key interface{}) (value interface{}) {
   110  	value, _ = tree.Search(key)
   111  	return
   112  }
   113  
   114  // doSetWithLockCheck checks whether value of the key exists with mutex.Lock,
   115  // if not exists, set value to the map with given `key`,
   116  // or else just return the existing value.
   117  //
   118  // When setting value, if `value` is type of <func() interface {}>,
   119  // it will be executed with mutex.Lock of the hash map,
   120  // and its return value will be set to the map with `key`.
   121  //
   122  // It returns value with given `key`.
   123  func (tree *BTree) doSetWithLockCheck(key interface{}, value interface{}) interface{} {
   124  	tree.mu.Lock()
   125  	defer tree.mu.Unlock()
   126  	if entry := tree.doSearch(key); entry != nil {
   127  		return entry.Value
   128  	}
   129  	if f, ok := value.(func() interface{}); ok {
   130  		value = f()
   131  	}
   132  	if value != nil {
   133  		tree.doSet(key, value)
   134  	}
   135  	return value
   136  }
   137  
   138  // GetOrSet returns the value by key,
   139  // or sets value with given `value` if it does not exist and then returns this value.
   140  func (tree *BTree) GetOrSet(key interface{}, value interface{}) interface{} {
   141  	if v, ok := tree.Search(key); !ok {
   142  		return tree.doSetWithLockCheck(key, value)
   143  	} else {
   144  		return v
   145  	}
   146  }
   147  
   148  // GetOrSetFunc returns the value by key,
   149  // or sets value with returned value of callback function `f` if it does not exist
   150  // and then returns this value.
   151  func (tree *BTree) GetOrSetFunc(key interface{}, f func() interface{}) interface{} {
   152  	if v, ok := tree.Search(key); !ok {
   153  		return tree.doSetWithLockCheck(key, f())
   154  	} else {
   155  		return v
   156  	}
   157  }
   158  
   159  // GetOrSetFuncLock returns the value by key,
   160  // or sets value with returned value of callback function `f` if it does not exist
   161  // and then returns this value.
   162  //
   163  // GetOrSetFuncLock differs with GetOrSetFunc function is that it executes function `f`
   164  // with mutex.Lock of the hash map.
   165  func (tree *BTree) GetOrSetFuncLock(key interface{}, f func() interface{}) interface{} {
   166  	if v, ok := tree.Search(key); !ok {
   167  		return tree.doSetWithLockCheck(key, f)
   168  	} else {
   169  		return v
   170  	}
   171  }
   172  
   173  // GetVar returns a gvar.Var with the value by given `key`.
   174  // The returned gvar.Var is un-concurrent safe.
   175  func (tree *BTree) GetVar(key interface{}) *gvar.Var {
   176  	return gvar.New(tree.Get(key))
   177  }
   178  
   179  // GetVarOrSet returns a gvar.Var with result from GetVarOrSet.
   180  // The returned gvar.Var is un-concurrent safe.
   181  func (tree *BTree) GetVarOrSet(key interface{}, value interface{}) *gvar.Var {
   182  	return gvar.New(tree.GetOrSet(key, value))
   183  }
   184  
   185  // GetVarOrSetFunc returns a gvar.Var with result from GetOrSetFunc.
   186  // The returned gvar.Var is un-concurrent safe.
   187  func (tree *BTree) GetVarOrSetFunc(key interface{}, f func() interface{}) *gvar.Var {
   188  	return gvar.New(tree.GetOrSetFunc(key, f))
   189  }
   190  
   191  // GetVarOrSetFuncLock returns a gvar.Var with result from GetOrSetFuncLock.
   192  // The returned gvar.Var is un-concurrent safe.
   193  func (tree *BTree) GetVarOrSetFuncLock(key interface{}, f func() interface{}) *gvar.Var {
   194  	return gvar.New(tree.GetOrSetFuncLock(key, f))
   195  }
   196  
   197  // SetIfNotExist sets `value` to the map if the `key` does not exist, and then returns true.
   198  // It returns false if `key` exists, and `value` would be ignored.
   199  func (tree *BTree) SetIfNotExist(key interface{}, value interface{}) bool {
   200  	if !tree.Contains(key) {
   201  		tree.doSetWithLockCheck(key, value)
   202  		return true
   203  	}
   204  	return false
   205  }
   206  
   207  // SetIfNotExistFunc sets value with return value of callback function `f`, and then returns true.
   208  // It returns false if `key` exists, and `value` would be ignored.
   209  func (tree *BTree) SetIfNotExistFunc(key interface{}, f func() interface{}) bool {
   210  	if !tree.Contains(key) {
   211  		tree.doSetWithLockCheck(key, f())
   212  		return true
   213  	}
   214  	return false
   215  }
   216  
   217  // SetIfNotExistFuncLock sets value with return value of callback function `f`, and then returns true.
   218  // It returns false if `key` exists, and `value` would be ignored.
   219  //
   220  // SetIfNotExistFuncLock differs with SetIfNotExistFunc function is that
   221  // it executes function `f` with mutex.Lock of the hash map.
   222  func (tree *BTree) SetIfNotExistFuncLock(key interface{}, f func() interface{}) bool {
   223  	if !tree.Contains(key) {
   224  		tree.doSetWithLockCheck(key, f)
   225  		return true
   226  	}
   227  	return false
   228  }
   229  
   230  // Contains checks whether `key` exists in the tree.
   231  func (tree *BTree) Contains(key interface{}) bool {
   232  	_, ok := tree.Search(key)
   233  	return ok
   234  }
   235  
   236  // doRemove removes the node from the tree by key.
   237  // Key should adhere to the comparator's type assertion, otherwise method panics.
   238  func (tree *BTree) doRemove(key interface{}) (value interface{}) {
   239  	node, index, found := tree.searchRecursively(tree.root, key)
   240  	if found {
   241  		value = node.Entries[index].Value
   242  		tree.delete(node, index)
   243  		tree.size--
   244  	}
   245  	return
   246  }
   247  
   248  // Remove removes the node from the tree by `key`.
   249  func (tree *BTree) Remove(key interface{}) (value interface{}) {
   250  	tree.mu.Lock()
   251  	defer tree.mu.Unlock()
   252  	return tree.doRemove(key)
   253  }
   254  
   255  // Removes batch deletes values of the tree by `keys`.
   256  func (tree *BTree) Removes(keys []interface{}) {
   257  	tree.mu.Lock()
   258  	defer tree.mu.Unlock()
   259  	for _, key := range keys {
   260  		tree.doRemove(key)
   261  	}
   262  }
   263  
   264  // IsEmpty returns true if tree does not contain any nodes
   265  func (tree *BTree) IsEmpty() bool {
   266  	return tree.Size() == 0
   267  }
   268  
   269  // Size returns number of nodes in the tree.
   270  func (tree *BTree) Size() int {
   271  	tree.mu.RLock()
   272  	defer tree.mu.RUnlock()
   273  	return tree.size
   274  }
   275  
   276  // Keys returns all keys in asc order.
   277  func (tree *BTree) Keys() []interface{} {
   278  	keys := make([]interface{}, tree.Size())
   279  	index := 0
   280  	tree.IteratorAsc(func(key, value interface{}) bool {
   281  		keys[index] = key
   282  		index++
   283  		return true
   284  	})
   285  	return keys
   286  }
   287  
   288  // Values returns all values in asc order based on the key.
   289  func (tree *BTree) Values() []interface{} {
   290  	values := make([]interface{}, tree.Size())
   291  	index := 0
   292  	tree.IteratorAsc(func(key, value interface{}) bool {
   293  		values[index] = value
   294  		index++
   295  		return true
   296  	})
   297  	return values
   298  }
   299  
   300  // Map returns all key-value items as map.
   301  func (tree *BTree) Map() map[interface{}]interface{} {
   302  	m := make(map[interface{}]interface{}, tree.Size())
   303  	tree.IteratorAsc(func(key, value interface{}) bool {
   304  		m[key] = value
   305  		return true
   306  	})
   307  	return m
   308  }
   309  
   310  // MapStrAny returns all key-value items as map[string]interface{}.
   311  func (tree *BTree) MapStrAny() map[string]interface{} {
   312  	m := make(map[string]interface{}, tree.Size())
   313  	tree.IteratorAsc(func(key, value interface{}) bool {
   314  		m[gconv.String(key)] = value
   315  		return true
   316  	})
   317  	return m
   318  }
   319  
   320  // Clear removes all nodes from the tree.
   321  func (tree *BTree) Clear() {
   322  	tree.mu.Lock()
   323  	defer tree.mu.Unlock()
   324  	tree.root = nil
   325  	tree.size = 0
   326  }
   327  
   328  // Replace the data of the tree with given `data`.
   329  func (tree *BTree) Replace(data map[interface{}]interface{}) {
   330  	tree.mu.Lock()
   331  	defer tree.mu.Unlock()
   332  	tree.root = nil
   333  	tree.size = 0
   334  	for k, v := range data {
   335  		tree.doSet(k, v)
   336  	}
   337  }
   338  
   339  // Height returns the height of the tree.
   340  func (tree *BTree) Height() int {
   341  	tree.mu.RLock()
   342  	defer tree.mu.RUnlock()
   343  	return tree.root.height()
   344  }
   345  
   346  // Left returns the left-most (min) entry or nil if tree is empty.
   347  func (tree *BTree) Left() *BTreeEntry {
   348  	tree.mu.RLock()
   349  	defer tree.mu.RUnlock()
   350  	node := tree.left(tree.root)
   351  	if node != nil {
   352  		return node.Entries[0]
   353  	}
   354  	return nil
   355  }
   356  
   357  // Right returns the right-most (max) entry or nil if tree is empty.
   358  func (tree *BTree) Right() *BTreeEntry {
   359  	tree.mu.RLock()
   360  	defer tree.mu.RUnlock()
   361  	node := tree.right(tree.root)
   362  	if node != nil {
   363  		return node.Entries[len(node.Entries)-1]
   364  	}
   365  	return nil
   366  }
   367  
   368  // String returns a string representation of container (for debugging purposes)
   369  func (tree *BTree) String() string {
   370  	if tree == nil {
   371  		return ""
   372  	}
   373  	tree.mu.RLock()
   374  	defer tree.mu.RUnlock()
   375  	var buffer bytes.Buffer
   376  	if tree.size != 0 {
   377  		tree.output(&buffer, tree.root, 0, true)
   378  	}
   379  	return buffer.String()
   380  }
   381  
   382  // Search searches the tree with given `key`.
   383  // Second return parameter `found` is true if key was found, otherwise false.
   384  func (tree *BTree) Search(key interface{}) (value interface{}, found bool) {
   385  	tree.mu.RLock()
   386  	defer tree.mu.RUnlock()
   387  	node, index, found := tree.searchRecursively(tree.root, key)
   388  	if found {
   389  		return node.Entries[index].Value, true
   390  	}
   391  	return nil, false
   392  }
   393  
   394  // Search searches the tree with given `key` without mutex.
   395  // It returns the entry if found or otherwise nil.
   396  func (tree *BTree) doSearch(key interface{}) *BTreeEntry {
   397  	node, index, found := tree.searchRecursively(tree.root, key)
   398  	if found {
   399  		return node.Entries[index]
   400  	}
   401  	return nil
   402  }
   403  
   404  // Print prints the tree to stdout.
   405  func (tree *BTree) Print() {
   406  	fmt.Println(tree.String())
   407  }
   408  
   409  // Iterator is alias of IteratorAsc.
   410  func (tree *BTree) Iterator(f func(key, value interface{}) bool) {
   411  	tree.IteratorAsc(f)
   412  }
   413  
   414  // IteratorFrom is alias of IteratorAscFrom.
   415  func (tree *BTree) IteratorFrom(key interface{}, match bool, f func(key, value interface{}) bool) {
   416  	tree.IteratorAscFrom(key, match, f)
   417  }
   418  
   419  // IteratorAsc iterates the tree readonly in ascending order with given callback function `f`.
   420  // If `f` returns true, then it continues iterating; or false to stop.
   421  func (tree *BTree) IteratorAsc(f func(key, value interface{}) bool) {
   422  	tree.mu.RLock()
   423  	defer tree.mu.RUnlock()
   424  	node := tree.left(tree.root)
   425  	if node == nil {
   426  		return
   427  	}
   428  	tree.doIteratorAsc(node, node.Entries[0], 0, f)
   429  }
   430  
   431  // IteratorAscFrom iterates the tree readonly in ascending order with given callback function `f`.
   432  // The parameter `key` specifies the start entry for iterating. The `match` specifies whether
   433  // starting iterating if the `key` is fully matched, or else using index searching iterating.
   434  // If `f` returns true, then it continues iterating; or false to stop.
   435  func (tree *BTree) IteratorAscFrom(key interface{}, match bool, f func(key, value interface{}) bool) {
   436  	tree.mu.RLock()
   437  	defer tree.mu.RUnlock()
   438  	node, index, found := tree.searchRecursively(tree.root, key)
   439  	if match {
   440  		if found {
   441  			tree.doIteratorAsc(node, node.Entries[index], index, f)
   442  		}
   443  	} else {
   444  		if index >= 0 && index < len(node.Entries) {
   445  			tree.doIteratorAsc(node, node.Entries[index], index, f)
   446  		}
   447  	}
   448  }
   449  
   450  func (tree *BTree) doIteratorAsc(node *BTreeNode, entry *BTreeEntry, index int, f func(key, value interface{}) bool) {
   451  	first := true
   452  loop:
   453  	if entry == nil {
   454  		return
   455  	}
   456  	if !f(entry.Key, entry.Value) {
   457  		return
   458  	}
   459  	// Find current entry position in current node
   460  	if !first {
   461  		index, _ = tree.search(node, entry.Key)
   462  	} else {
   463  		first = false
   464  	}
   465  	// Try to go down to the child right of the current entry
   466  	if index+1 < len(node.Children) {
   467  		node = node.Children[index+1]
   468  		// Try to go down to the child left of the current node
   469  		for len(node.Children) > 0 {
   470  			node = node.Children[0]
   471  		}
   472  		// Return the left-most entry
   473  		entry = node.Entries[0]
   474  		goto loop
   475  	}
   476  	// Above assures that we have reached a leaf node, so return the next entry in current node (if any)
   477  	if index+1 < len(node.Entries) {
   478  		entry = node.Entries[index+1]
   479  		goto loop
   480  	}
   481  	// Reached leaf node and there are no entries to the right of the current entry, so go up to the parent
   482  	for node.Parent != nil {
   483  		node = node.Parent
   484  		// Find next entry position in current node (note: search returns the first equal or bigger than entry)
   485  		index, _ = tree.search(node, entry.Key)
   486  		// Check that there is a next entry position in current node
   487  		if index < len(node.Entries) {
   488  			entry = node.Entries[index]
   489  			goto loop
   490  		}
   491  	}
   492  }
   493  
   494  // IteratorDesc iterates the tree readonly in descending order with given callback function `f`.
   495  // If `f` returns true, then it continues iterating; or false to stop.
   496  func (tree *BTree) IteratorDesc(f func(key, value interface{}) bool) {
   497  	tree.mu.RLock()
   498  	defer tree.mu.RUnlock()
   499  	node := tree.right(tree.root)
   500  	if node == nil {
   501  		return
   502  	}
   503  	index := len(node.Entries) - 1
   504  	entry := node.Entries[index]
   505  	tree.doIteratorDesc(node, entry, index, f)
   506  }
   507  
   508  // IteratorDescFrom iterates the tree readonly in descending order with given callback function `f`.
   509  // The parameter `key` specifies the start entry for iterating. The `match` specifies whether
   510  // starting iterating if the `key` is fully matched, or else using index searching iterating.
   511  // If `f` returns true, then it continues iterating; or false to stop.
   512  func (tree *BTree) IteratorDescFrom(key interface{}, match bool, f func(key, value interface{}) bool) {
   513  	tree.mu.RLock()
   514  	defer tree.mu.RUnlock()
   515  	node, index, found := tree.searchRecursively(tree.root, key)
   516  	if match {
   517  		if found {
   518  			tree.doIteratorDesc(node, node.Entries[index], index, f)
   519  		}
   520  	} else {
   521  		if index >= 0 && index < len(node.Entries) {
   522  			tree.doIteratorDesc(node, node.Entries[index], index, f)
   523  		}
   524  	}
   525  }
   526  
   527  // IteratorDesc iterates the tree readonly in descending order with given callback function `f`.
   528  // If `f` returns true, then it continues iterating; or false to stop.
   529  func (tree *BTree) doIteratorDesc(node *BTreeNode, entry *BTreeEntry, index int, f func(key, value interface{}) bool) {
   530  	first := true
   531  loop:
   532  	if entry == nil {
   533  		return
   534  	}
   535  	if !f(entry.Key, entry.Value) {
   536  		return
   537  	}
   538  	// Find current entry position in current node
   539  	if !first {
   540  		index, _ = tree.search(node, entry.Key)
   541  	} else {
   542  		first = false
   543  	}
   544  	// Try to go down to the child left of the current entry
   545  	if index < len(node.Children) {
   546  		node = node.Children[index]
   547  		// Try to go down to the child right of the current node
   548  		for len(node.Children) > 0 {
   549  			node = node.Children[len(node.Children)-1]
   550  		}
   551  		// Return the right-most entry
   552  		entry = node.Entries[len(node.Entries)-1]
   553  		goto loop
   554  	}
   555  	// Above assures that we have reached a leaf node, so return the previous entry in current node (if any)
   556  	if index-1 >= 0 {
   557  		entry = node.Entries[index-1]
   558  		goto loop
   559  	}
   560  
   561  	// Reached leaf node and there are no entries to the left of the current entry, so go up to the parent
   562  	for node.Parent != nil {
   563  		node = node.Parent
   564  		// Find previous entry position in current node (note: search returns the first equal or bigger than entry)
   565  		index, _ = tree.search(node, entry.Key)
   566  		// Check that there is a previous entry position in current node
   567  		if index-1 >= 0 {
   568  			entry = node.Entries[index-1]
   569  			goto loop
   570  		}
   571  	}
   572  }
   573  
   574  func (tree *BTree) output(buffer *bytes.Buffer, node *BTreeNode, level int, isTail bool) {
   575  	for e := 0; e < len(node.Entries)+1; e++ {
   576  		if e < len(node.Children) {
   577  			tree.output(buffer, node.Children[e], level+1, true)
   578  		}
   579  		if e < len(node.Entries) {
   580  			if _, err := buffer.WriteString(strings.Repeat("    ", level)); err != nil {
   581  				intlog.Errorf(context.TODO(), `%+v`, err)
   582  			}
   583  			if _, err := buffer.WriteString(fmt.Sprintf("%v", node.Entries[e].Key) + "\n"); err != nil {
   584  				intlog.Errorf(context.TODO(), `%+v`, err)
   585  			}
   586  		}
   587  	}
   588  }
   589  
   590  func (node *BTreeNode) height() int {
   591  	h := 0
   592  	n := node
   593  	for ; n != nil; n = n.Children[0] {
   594  		h++
   595  		if len(n.Children) == 0 {
   596  			break
   597  		}
   598  	}
   599  	return h
   600  }
   601  
   602  func (tree *BTree) isLeaf(node *BTreeNode) bool {
   603  	return len(node.Children) == 0
   604  }
   605  
   606  // func (tree *BTree) isFull(node *BTreeNode) bool {
   607  //	return len(node.Entries) == tree.maxEntries()
   608  // }
   609  
   610  func (tree *BTree) shouldSplit(node *BTreeNode) bool {
   611  	return len(node.Entries) > tree.maxEntries()
   612  }
   613  
   614  func (tree *BTree) maxChildren() int {
   615  	return tree.m
   616  }
   617  
   618  func (tree *BTree) minChildren() int {
   619  	return (tree.m + 1) / 2 // ceil(m/2)
   620  }
   621  
   622  func (tree *BTree) maxEntries() int {
   623  	return tree.maxChildren() - 1
   624  }
   625  
   626  func (tree *BTree) minEntries() int {
   627  	return tree.minChildren() - 1
   628  }
   629  
   630  func (tree *BTree) middle() int {
   631  	// "-1" to favor right nodes to have more keys when splitting
   632  	return (tree.m - 1) / 2
   633  }
   634  
   635  // search does search only within the single node among its entries
   636  func (tree *BTree) search(node *BTreeNode, key interface{}) (index int, found bool) {
   637  	low, mid, high := 0, 0, len(node.Entries)-1
   638  	for low <= high {
   639  		mid = low + (high-low)/2
   640  		compare := tree.getComparator()(key, node.Entries[mid].Key)
   641  		switch {
   642  		case compare > 0:
   643  			low = mid + 1
   644  		case compare < 0:
   645  			high = mid - 1
   646  		case compare == 0:
   647  			return mid, true
   648  		}
   649  	}
   650  	return low, false
   651  }
   652  
   653  // searchRecursively searches recursively down the tree starting at the startNode
   654  func (tree *BTree) searchRecursively(startNode *BTreeNode, key interface{}) (node *BTreeNode, index int, found bool) {
   655  	if tree.size == 0 {
   656  		return nil, -1, false
   657  	}
   658  	node = startNode
   659  	for {
   660  		index, found = tree.search(node, key)
   661  		if found {
   662  			return node, index, true
   663  		}
   664  		if tree.isLeaf(node) {
   665  			return node, index, false
   666  		}
   667  		node = node.Children[index]
   668  	}
   669  }
   670  
   671  func (tree *BTree) insert(node *BTreeNode, entry *BTreeEntry) (inserted bool) {
   672  	if tree.isLeaf(node) {
   673  		return tree.insertIntoLeaf(node, entry)
   674  	}
   675  	return tree.insertIntoInternal(node, entry)
   676  }
   677  
   678  func (tree *BTree) insertIntoLeaf(node *BTreeNode, entry *BTreeEntry) (inserted bool) {
   679  	insertPosition, found := tree.search(node, entry.Key)
   680  	if found {
   681  		node.Entries[insertPosition] = entry
   682  		return false
   683  	}
   684  	// Insert entry's key in the middle of the node
   685  	node.Entries = append(node.Entries, nil)
   686  	copy(node.Entries[insertPosition+1:], node.Entries[insertPosition:])
   687  	node.Entries[insertPosition] = entry
   688  	tree.split(node)
   689  	return true
   690  }
   691  
   692  func (tree *BTree) insertIntoInternal(node *BTreeNode, entry *BTreeEntry) (inserted bool) {
   693  	insertPosition, found := tree.search(node, entry.Key)
   694  	if found {
   695  		node.Entries[insertPosition] = entry
   696  		return false
   697  	}
   698  	return tree.insert(node.Children[insertPosition], entry)
   699  }
   700  
   701  func (tree *BTree) split(node *BTreeNode) {
   702  	if !tree.shouldSplit(node) {
   703  		return
   704  	}
   705  
   706  	if node == tree.root {
   707  		tree.splitRoot()
   708  		return
   709  	}
   710  
   711  	tree.splitNonRoot(node)
   712  }
   713  
   714  func (tree *BTree) splitNonRoot(node *BTreeNode) {
   715  	middle := tree.middle()
   716  	parent := node.Parent
   717  
   718  	left := &BTreeNode{Entries: append([]*BTreeEntry(nil), node.Entries[:middle]...), Parent: parent}
   719  	right := &BTreeNode{Entries: append([]*BTreeEntry(nil), node.Entries[middle+1:]...), Parent: parent}
   720  
   721  	// Move children from the node to be split into left and right nodes
   722  	if !tree.isLeaf(node) {
   723  		left.Children = append([]*BTreeNode(nil), node.Children[:middle+1]...)
   724  		right.Children = append([]*BTreeNode(nil), node.Children[middle+1:]...)
   725  		setParent(left.Children, left)
   726  		setParent(right.Children, right)
   727  	}
   728  
   729  	insertPosition, _ := tree.search(parent, node.Entries[middle].Key)
   730  
   731  	// Insert middle key into parent
   732  	parent.Entries = append(parent.Entries, nil)
   733  	copy(parent.Entries[insertPosition+1:], parent.Entries[insertPosition:])
   734  	parent.Entries[insertPosition] = node.Entries[middle]
   735  
   736  	// Set child left of inserted key in parent to the created left node
   737  	parent.Children[insertPosition] = left
   738  
   739  	// Set child right of inserted key in parent to the created right node
   740  	parent.Children = append(parent.Children, nil)
   741  	copy(parent.Children[insertPosition+2:], parent.Children[insertPosition+1:])
   742  	parent.Children[insertPosition+1] = right
   743  
   744  	tree.split(parent)
   745  }
   746  
   747  func (tree *BTree) splitRoot() {
   748  	middle := tree.middle()
   749  	left := &BTreeNode{Entries: append([]*BTreeEntry(nil), tree.root.Entries[:middle]...)}
   750  	right := &BTreeNode{Entries: append([]*BTreeEntry(nil), tree.root.Entries[middle+1:]...)}
   751  
   752  	// Move children from the node to be split into left and right nodes
   753  	if !tree.isLeaf(tree.root) {
   754  		left.Children = append([]*BTreeNode(nil), tree.root.Children[:middle+1]...)
   755  		right.Children = append([]*BTreeNode(nil), tree.root.Children[middle+1:]...)
   756  		setParent(left.Children, left)
   757  		setParent(right.Children, right)
   758  	}
   759  
   760  	// Root is a node with one entry and two children (left and right)
   761  	newRoot := &BTreeNode{
   762  		Entries:  []*BTreeEntry{tree.root.Entries[middle]},
   763  		Children: []*BTreeNode{left, right},
   764  	}
   765  
   766  	left.Parent = newRoot
   767  	right.Parent = newRoot
   768  	tree.root = newRoot
   769  }
   770  
   771  func setParent(nodes []*BTreeNode, parent *BTreeNode) {
   772  	for _, node := range nodes {
   773  		node.Parent = parent
   774  	}
   775  }
   776  
   777  func (tree *BTree) left(node *BTreeNode) *BTreeNode {
   778  	if tree.size == 0 {
   779  		return nil
   780  	}
   781  	current := node
   782  	for {
   783  		if tree.isLeaf(current) {
   784  			return current
   785  		}
   786  		current = current.Children[0]
   787  	}
   788  }
   789  
   790  func (tree *BTree) right(node *BTreeNode) *BTreeNode {
   791  	if tree.size == 0 {
   792  		return nil
   793  	}
   794  	current := node
   795  	for {
   796  		if tree.isLeaf(current) {
   797  			return current
   798  		}
   799  		current = current.Children[len(current.Children)-1]
   800  	}
   801  }
   802  
   803  // leftSibling returns the node's left sibling and child index (in parent) if it exists, otherwise (nil,-1)
   804  // key is any of keys in node (could even be deleted).
   805  func (tree *BTree) leftSibling(node *BTreeNode, key interface{}) (*BTreeNode, int) {
   806  	if node.Parent != nil {
   807  		index, _ := tree.search(node.Parent, key)
   808  		index--
   809  		if index >= 0 && index < len(node.Parent.Children) {
   810  			return node.Parent.Children[index], index
   811  		}
   812  	}
   813  	return nil, -1
   814  }
   815  
   816  // rightSibling returns the node's right sibling and child index (in parent) if it exists, otherwise (nil,-1)
   817  // key is any of keys in node (could even be deleted).
   818  func (tree *BTree) rightSibling(node *BTreeNode, key interface{}) (*BTreeNode, int) {
   819  	if node.Parent != nil {
   820  		index, _ := tree.search(node.Parent, key)
   821  		index++
   822  		if index < len(node.Parent.Children) {
   823  			return node.Parent.Children[index], index
   824  		}
   825  	}
   826  	return nil, -1
   827  }
   828  
   829  // delete deletes an entry in node at entries' index
   830  // ref.: https://en.wikipedia.org/wiki/B-tree#Deletion
   831  func (tree *BTree) delete(node *BTreeNode, index int) {
   832  	// deleting from a leaf node
   833  	if tree.isLeaf(node) {
   834  		deletedKey := node.Entries[index].Key
   835  		tree.deleteEntry(node, index)
   836  		tree.reBalance(node, deletedKey)
   837  		if len(tree.root.Entries) == 0 {
   838  			tree.root = nil
   839  		}
   840  		return
   841  	}
   842  
   843  	// deleting from an internal node
   844  	leftLargestNode := tree.right(node.Children[index]) // largest node in the left sub-tree (assumed to exist)
   845  	leftLargestEntryIndex := len(leftLargestNode.Entries) - 1
   846  	node.Entries[index] = leftLargestNode.Entries[leftLargestEntryIndex]
   847  	deletedKey := leftLargestNode.Entries[leftLargestEntryIndex].Key
   848  	tree.deleteEntry(leftLargestNode, leftLargestEntryIndex)
   849  	tree.reBalance(leftLargestNode, deletedKey)
   850  }
   851  
   852  // reBalance reBalances the tree after deletion if necessary and returns true, otherwise false.
   853  // Note that we first delete the entry and then call reBalance, thus the passed deleted key as reference.
   854  func (tree *BTree) reBalance(node *BTreeNode, deletedKey interface{}) {
   855  	// check if re-balancing is needed
   856  	if node == nil || len(node.Entries) >= tree.minEntries() {
   857  		return
   858  	}
   859  
   860  	// try to borrow from left sibling
   861  	leftSibling, leftSiblingIndex := tree.leftSibling(node, deletedKey)
   862  	if leftSibling != nil && len(leftSibling.Entries) > tree.minEntries() {
   863  		// rotate right
   864  		node.Entries = append([]*BTreeEntry{node.Parent.Entries[leftSiblingIndex]}, node.Entries...) // prepend parent's separator entry to node's entries
   865  		node.Parent.Entries[leftSiblingIndex] = leftSibling.Entries[len(leftSibling.Entries)-1]
   866  		tree.deleteEntry(leftSibling, len(leftSibling.Entries)-1)
   867  		if !tree.isLeaf(leftSibling) {
   868  			leftSiblingRightMostChild := leftSibling.Children[len(leftSibling.Children)-1]
   869  			leftSiblingRightMostChild.Parent = node
   870  			node.Children = append([]*BTreeNode{leftSiblingRightMostChild}, node.Children...)
   871  			tree.deleteChild(leftSibling, len(leftSibling.Children)-1)
   872  		}
   873  		return
   874  	}
   875  
   876  	// try to borrow from right sibling
   877  	rightSibling, rightSiblingIndex := tree.rightSibling(node, deletedKey)
   878  	if rightSibling != nil && len(rightSibling.Entries) > tree.minEntries() {
   879  		// rotate left
   880  		node.Entries = append(node.Entries, node.Parent.Entries[rightSiblingIndex-1]) // append parent's separator entry to node's entries
   881  		node.Parent.Entries[rightSiblingIndex-1] = rightSibling.Entries[0]
   882  		tree.deleteEntry(rightSibling, 0)
   883  		if !tree.isLeaf(rightSibling) {
   884  			rightSiblingLeftMostChild := rightSibling.Children[0]
   885  			rightSiblingLeftMostChild.Parent = node
   886  			node.Children = append(node.Children, rightSiblingLeftMostChild)
   887  			tree.deleteChild(rightSibling, 0)
   888  		}
   889  		return
   890  	}
   891  
   892  	// merge with siblings
   893  	if rightSibling != nil {
   894  		// merge with right sibling
   895  		node.Entries = append(node.Entries, node.Parent.Entries[rightSiblingIndex-1])
   896  		node.Entries = append(node.Entries, rightSibling.Entries...)
   897  		deletedKey = node.Parent.Entries[rightSiblingIndex-1].Key
   898  		tree.deleteEntry(node.Parent, rightSiblingIndex-1)
   899  		tree.appendChildren(node.Parent.Children[rightSiblingIndex], node)
   900  		tree.deleteChild(node.Parent, rightSiblingIndex)
   901  	} else if leftSibling != nil {
   902  		// merge with left sibling
   903  		entries := append([]*BTreeEntry(nil), leftSibling.Entries...)
   904  		entries = append(entries, node.Parent.Entries[leftSiblingIndex])
   905  		node.Entries = append(entries, node.Entries...)
   906  		deletedKey = node.Parent.Entries[leftSiblingIndex].Key
   907  		tree.deleteEntry(node.Parent, leftSiblingIndex)
   908  		tree.prependChildren(node.Parent.Children[leftSiblingIndex], node)
   909  		tree.deleteChild(node.Parent, leftSiblingIndex)
   910  	}
   911  
   912  	// make the merged node the root if its parent was the root and the root is empty
   913  	if node.Parent == tree.root && len(tree.root.Entries) == 0 {
   914  		tree.root = node
   915  		node.Parent = nil
   916  		return
   917  	}
   918  
   919  	// parent might be underflow, so try to reBalance if necessary
   920  	tree.reBalance(node.Parent, deletedKey)
   921  }
   922  
   923  func (tree *BTree) prependChildren(fromNode *BTreeNode, toNode *BTreeNode) {
   924  	children := append([]*BTreeNode(nil), fromNode.Children...)
   925  	toNode.Children = append(children, toNode.Children...)
   926  	setParent(fromNode.Children, toNode)
   927  }
   928  
   929  func (tree *BTree) appendChildren(fromNode *BTreeNode, toNode *BTreeNode) {
   930  	toNode.Children = append(toNode.Children, fromNode.Children...)
   931  	setParent(fromNode.Children, toNode)
   932  }
   933  
   934  func (tree *BTree) deleteEntry(node *BTreeNode, index int) {
   935  	copy(node.Entries[index:], node.Entries[index+1:])
   936  	node.Entries[len(node.Entries)-1] = nil
   937  	node.Entries = node.Entries[:len(node.Entries)-1]
   938  }
   939  
   940  func (tree *BTree) deleteChild(node *BTreeNode, index int) {
   941  	if index >= len(node.Children) {
   942  		return
   943  	}
   944  	copy(node.Children[index:], node.Children[index+1:])
   945  	node.Children[len(node.Children)-1] = nil
   946  	node.Children = node.Children[:len(node.Children)-1]
   947  }
   948  
   949  // MarshalJSON implements the interface MarshalJSON for json.Marshal.
   950  func (tree BTree) MarshalJSON() (jsonBytes []byte, err error) {
   951  	if tree.root == nil {
   952  		return []byte("null"), nil
   953  	}
   954  	buffer := bytes.NewBuffer(nil)
   955  	buffer.WriteByte('{')
   956  	tree.Iterator(func(key, value interface{}) bool {
   957  		valueBytes, valueJsonErr := json.Marshal(value)
   958  		if valueJsonErr != nil {
   959  			err = valueJsonErr
   960  			return false
   961  		}
   962  		if buffer.Len() > 1 {
   963  			buffer.WriteByte(',')
   964  		}
   965  		buffer.WriteString(fmt.Sprintf(`"%v":%s`, key, valueBytes))
   966  		return true
   967  	})
   968  	buffer.WriteByte('}')
   969  	return buffer.Bytes(), nil
   970  }
   971  
   972  // getComparator returns the comparator if it's previously set,
   973  // or else it panics.
   974  func (tree *BTree) getComparator() func(a, b interface{}) int {
   975  	if tree.comparator == nil {
   976  		panic("comparator is missing for tree")
   977  	}
   978  	return tree.comparator
   979  }