gtree_btree.go 28 KB

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