gtree_btree.go 29 KB

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