Nelze vybrat více než 25 témat Téma musí začínat písmenem nebo číslem, může obsahovat pomlčky („-“) a může být dlouhé až 35 znaků.
 
 
 

1014 řádky
30 KiB

  1. // Copyright 2014 Google LLC
  2. // Modified 2018 by Jonathan Amsterdam (jbamsterdam@gmail.com)
  3. //
  4. // Licensed under the Apache License, Version 2.0 (the "License");
  5. // you may not use this file except in compliance with the License.
  6. // You may obtain a copy of the License at
  7. //
  8. // http://www.apache.org/licenses/LICENSE-2.0
  9. //
  10. // Unless required by applicable law or agreed to in writing, software
  11. // distributed under the License is distributed on an "AS IS" BASIS,
  12. // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13. // See the License for the specific language governing permissions and
  14. // limitations under the License.
  15. // Package btree implements in-memory B-Trees of arbitrary degree.
  16. //
  17. // This implementation is based on google/btree (http://github.com/google/btree), and
  18. // much of the code is taken from there. But the API has been changed significantly,
  19. // particularly around iteration, and support for indexing by position has been
  20. // added.
  21. //
  22. // btree implements an in-memory B-Tree for use as an ordered data structure.
  23. // It is not meant for persistent storage solutions.
  24. //
  25. // It has a flatter structure than an equivalent red-black or other binary tree,
  26. // which in some cases yields better memory usage and/or performance.
  27. // See some discussion on the matter here:
  28. // http://google-opensource.blogspot.com/2013/01/c-containers-that-save-memory-and-time.html
  29. // Note, though, that this project is in no way related to the C++ B-Tree
  30. // implementation written about there.
  31. //
  32. // Within this tree, each node contains a slice of items and a (possibly nil)
  33. // slice of children. For basic numeric values or raw structs, this can cause
  34. // efficiency differences when compared to equivalent C++ template code that
  35. // stores values in arrays within the node:
  36. // * Due to the overhead of storing values as interfaces (each
  37. // value needs to be stored as the value itself, then 2 words for the
  38. // interface pointing to that value and its type), resulting in higher
  39. // memory use.
  40. // * Since interfaces can point to values anywhere in memory, values are
  41. // most likely not stored in contiguous blocks, resulting in a higher
  42. // number of cache misses.
  43. // These issues don't tend to matter, though, when working with strings or other
  44. // heap-allocated structures, since C++-equivalent structures also must store
  45. // pointers and also distribute their values across the heap.
  46. package btree
  47. import (
  48. "sort"
  49. "sync"
  50. )
  51. // Key represents a key into the tree.
  52. type Key interface{}
  53. // Value represents a value in the tree.
  54. type Value interface{}
  55. // item is a key-value pair.
  56. type item struct {
  57. key Key
  58. value Value
  59. }
  60. type lessFunc func(interface{}, interface{}) bool
  61. // New creates a new B-Tree with the given degree and comparison function.
  62. //
  63. // New(2, less), for example, will create a 2-3-4 tree (each node contains 1-3 items
  64. // and 2-4 children).
  65. //
  66. // The less function tests whether the current item is less than the given argument.
  67. // It must provide a strict weak ordering.
  68. // If !less(a, b) && !less(b, a), we treat this to mean a == b (i.e. the tree
  69. // can hold only one of a or b).
  70. func New(degree int, less func(interface{}, interface{}) bool) *BTree {
  71. if degree <= 1 {
  72. panic("bad degree")
  73. }
  74. return &BTree{
  75. degree: degree,
  76. less: less,
  77. cow: &copyOnWriteContext{},
  78. }
  79. }
  80. // items stores items in a node.
  81. type items []item
  82. // insertAt inserts a value into the given index, pushing all subsequent values
  83. // forward.
  84. func (s *items) insertAt(index int, m item) {
  85. *s = append(*s, item{})
  86. if index < len(*s) {
  87. copy((*s)[index+1:], (*s)[index:])
  88. }
  89. (*s)[index] = m
  90. }
  91. // removeAt removes a value at a given index, pulling all subsequent values
  92. // back.
  93. func (s *items) removeAt(index int) item {
  94. m := (*s)[index]
  95. copy((*s)[index:], (*s)[index+1:])
  96. (*s)[len(*s)-1] = item{}
  97. *s = (*s)[:len(*s)-1]
  98. return m
  99. }
  100. // pop removes and returns the last element in the list.
  101. func (s *items) pop() item {
  102. index := len(*s) - 1
  103. out := (*s)[index]
  104. (*s)[index] = item{}
  105. *s = (*s)[:index]
  106. return out
  107. }
  108. var nilItems = make(items, 16)
  109. // truncate truncates this instance at index so that it contains only the
  110. // first index items. index must be less than or equal to length.
  111. func (s *items) truncate(index int) {
  112. var toClear items
  113. *s, toClear = (*s)[:index], (*s)[index:]
  114. for len(toClear) > 0 {
  115. toClear = toClear[copy(toClear, nilItems):]
  116. }
  117. }
  118. // find returns the index where an item with key should be inserted into this
  119. // list. 'found' is true if the item already exists in the list at the given
  120. // index.
  121. func (s items) find(k Key, less lessFunc) (index int, found bool) {
  122. i := sort.Search(len(s), func(i int) bool { return less(k, s[i].key) })
  123. // i is the smallest index of s for which k.Less(s[i].Key), or len(s).
  124. if i > 0 && !less(s[i-1].key, k) {
  125. return i - 1, true
  126. }
  127. return i, false
  128. }
  129. // children stores child nodes in a node.
  130. type children []*node
  131. // insertAt inserts a value into the given index, pushing all subsequent values
  132. // forward.
  133. func (s *children) insertAt(index int, n *node) {
  134. *s = append(*s, nil)
  135. if index < len(*s) {
  136. copy((*s)[index+1:], (*s)[index:])
  137. }
  138. (*s)[index] = n
  139. }
  140. // removeAt removes a value at a given index, pulling all subsequent values
  141. // back.
  142. func (s *children) removeAt(index int) *node {
  143. n := (*s)[index]
  144. copy((*s)[index:], (*s)[index+1:])
  145. (*s)[len(*s)-1] = nil
  146. *s = (*s)[:len(*s)-1]
  147. return n
  148. }
  149. // pop removes and returns the last element in the list.
  150. func (s *children) pop() (out *node) {
  151. index := len(*s) - 1
  152. out = (*s)[index]
  153. (*s)[index] = nil
  154. *s = (*s)[:index]
  155. return
  156. }
  157. var nilChildren = make(children, 16)
  158. // truncate truncates this instance at index so that it contains only the
  159. // first index children. index must be less than or equal to length.
  160. func (s *children) truncate(index int) {
  161. var toClear children
  162. *s, toClear = (*s)[:index], (*s)[index:]
  163. for len(toClear) > 0 {
  164. toClear = toClear[copy(toClear, nilChildren):]
  165. }
  166. }
  167. // node is an internal node in a tree.
  168. //
  169. // It must at all times maintain the invariant that either
  170. // * len(children) == 0, len(items) unconstrained
  171. // * len(children) == len(items) + 1
  172. type node struct {
  173. items items
  174. children children
  175. size int // number of items in the subtree: len(items) + sum over i of children[i].size
  176. cow *copyOnWriteContext
  177. }
  178. func (n *node) computeSize() int {
  179. sz := len(n.items)
  180. for _, c := range n.children {
  181. sz += c.size
  182. }
  183. return sz
  184. }
  185. func (n *node) mutableFor(cow *copyOnWriteContext) *node {
  186. if n.cow == cow {
  187. return n
  188. }
  189. out := cow.newNode()
  190. if cap(out.items) >= len(n.items) {
  191. out.items = out.items[:len(n.items)]
  192. } else {
  193. out.items = make(items, len(n.items), cap(n.items))
  194. }
  195. copy(out.items, n.items)
  196. // Copy children
  197. if cap(out.children) >= len(n.children) {
  198. out.children = out.children[:len(n.children)]
  199. } else {
  200. out.children = make(children, len(n.children), cap(n.children))
  201. }
  202. copy(out.children, n.children)
  203. out.size = n.size
  204. return out
  205. }
  206. func (n *node) mutableChild(i int) *node {
  207. c := n.children[i].mutableFor(n.cow)
  208. n.children[i] = c
  209. return c
  210. }
  211. // split splits the given node at the given index. The current node shrinks,
  212. // and this function returns the item that existed at that index and a new node
  213. // containing all items/children after it.
  214. func (n *node) split(i int) (item, *node) {
  215. item := n.items[i]
  216. next := n.cow.newNode()
  217. next.items = append(next.items, n.items[i+1:]...)
  218. n.items.truncate(i)
  219. if len(n.children) > 0 {
  220. next.children = append(next.children, n.children[i+1:]...)
  221. n.children.truncate(i + 1)
  222. }
  223. n.size = n.computeSize()
  224. next.size = next.computeSize()
  225. return item, next
  226. }
  227. // maybeSplitChild checks if a child should be split, and if so splits it.
  228. // Returns whether or not a split occurred.
  229. func (n *node) maybeSplitChild(i, maxItems int) bool {
  230. if len(n.children[i].items) < maxItems {
  231. return false
  232. }
  233. first := n.mutableChild(i)
  234. item, second := first.split(maxItems / 2)
  235. n.items.insertAt(i, item)
  236. n.children.insertAt(i+1, second)
  237. // The size of n doesn't change.
  238. return true
  239. }
  240. // insert inserts an item into the subtree rooted at this node, making sure
  241. // no nodes in the subtree exceed maxItems items. Should an equivalent item be
  242. // be found/replaced by insert, its value will be returned.
  243. //
  244. // If computeIndex is true, the third return value is the index of the value with respect to n.
  245. func (n *node) insert(m item, maxItems int, less lessFunc, computeIndex bool) (old Value, present bool, idx int) {
  246. i, found := n.items.find(m.key, less)
  247. if found {
  248. out := n.items[i]
  249. n.items[i] = m
  250. if computeIndex {
  251. idx = n.itemIndex(i)
  252. }
  253. return out.value, true, idx
  254. }
  255. if len(n.children) == 0 {
  256. n.items.insertAt(i, m)
  257. n.size++
  258. return old, false, i
  259. }
  260. if n.maybeSplitChild(i, maxItems) {
  261. inTree := n.items[i]
  262. switch {
  263. case less(m.key, inTree.key):
  264. // no change, we want first split node
  265. case less(inTree.key, m.key):
  266. i++ // we want second split node
  267. default:
  268. out := n.items[i]
  269. n.items[i] = m
  270. if computeIndex {
  271. idx = n.itemIndex(i)
  272. }
  273. return out.value, true, idx
  274. }
  275. }
  276. old, present, idx = n.mutableChild(i).insert(m, maxItems, less, computeIndex)
  277. if !present {
  278. n.size++
  279. }
  280. if computeIndex {
  281. idx += n.partialSize(i)
  282. }
  283. return old, present, idx
  284. }
  285. // get finds the given key in the subtree and returns the corresponding item, along with a boolean reporting
  286. // whether it was found.
  287. // If computeIndex is true, it also returns the index of the key relative to the node's subtree.
  288. func (n *node) get(k Key, computeIndex bool, less lessFunc) (item, bool, int) {
  289. i, found := n.items.find(k, less)
  290. if found {
  291. return n.items[i], true, n.itemIndex(i)
  292. }
  293. if len(n.children) > 0 {
  294. m, found, idx := n.children[i].get(k, computeIndex, less)
  295. if computeIndex && found {
  296. idx += n.partialSize(i)
  297. }
  298. return m, found, idx
  299. }
  300. return item{}, false, -1
  301. }
  302. // itemIndex returns the index w.r.t. n of the ith item in n.
  303. func (n *node) itemIndex(i int) int {
  304. if len(n.children) == 0 {
  305. return i
  306. }
  307. // Get the size of the node up to but not including the child to the right of
  308. // item i. Subtract 1 because the index is 0-based.
  309. return n.partialSize(i+1) - 1
  310. }
  311. // Returns the size of the non-leaf node up to but not including child i.
  312. func (n *node) partialSize(i int) int {
  313. var sz int
  314. for j, c := range n.children {
  315. if j == i {
  316. break
  317. }
  318. sz += c.size + 1
  319. }
  320. return sz
  321. }
  322. // cursorStackForKey returns a stack of cursors for the key, along with whether the key was found and the index.
  323. func (n *node) cursorStackForKey(k Key, cs cursorStack, less lessFunc) (cursorStack, bool, int) {
  324. i, found := n.items.find(k, less)
  325. cs.push(cursor{n, i})
  326. idx := i
  327. if found {
  328. if len(n.children) > 0 {
  329. idx = n.partialSize(i+1) - 1
  330. }
  331. return cs, true, idx
  332. }
  333. if len(n.children) > 0 {
  334. cs, found, idx := n.children[i].cursorStackForKey(k, cs, less)
  335. return cs, found, idx + n.partialSize(i)
  336. }
  337. return cs, false, idx
  338. }
  339. // at returns the item at the i'th position in the subtree rooted at n.
  340. // It assumes i is in range.
  341. func (n *node) at(i int) item {
  342. if len(n.children) == 0 {
  343. return n.items[i]
  344. }
  345. for j, c := range n.children {
  346. if i < c.size {
  347. return c.at(i)
  348. }
  349. i -= c.size
  350. if i == 0 {
  351. return n.items[j]
  352. }
  353. i--
  354. }
  355. panic("impossible")
  356. }
  357. // cursorStackForIndex returns a stack of cursors for the index.
  358. // It assumes i is in range.
  359. func (n *node) cursorStackForIndex(i int, cs cursorStack) cursorStack {
  360. if len(n.children) == 0 {
  361. return cs.push(cursor{n, i})
  362. }
  363. for j, c := range n.children {
  364. if i < c.size {
  365. return c.cursorStackForIndex(i, cs.push(cursor{n, j}))
  366. }
  367. i -= c.size
  368. if i == 0 {
  369. return cs.push(cursor{n, j})
  370. }
  371. i--
  372. }
  373. panic("impossible")
  374. }
  375. // toRemove details what item to remove in a node.remove call.
  376. type toRemove int
  377. const (
  378. removeItem toRemove = iota // removes the given item
  379. removeMin // removes smallest item in the subtree
  380. removeMax // removes largest item in the subtree
  381. )
  382. // remove removes an item from the subtree rooted at this node.
  383. func (n *node) remove(key Key, minItems int, typ toRemove, less lessFunc) (item, bool) {
  384. var i int
  385. var found bool
  386. switch typ {
  387. case removeMax:
  388. if len(n.children) == 0 {
  389. n.size--
  390. return n.items.pop(), true
  391. }
  392. i = len(n.items)
  393. case removeMin:
  394. if len(n.children) == 0 {
  395. n.size--
  396. return n.items.removeAt(0), true
  397. }
  398. i = 0
  399. case removeItem:
  400. i, found = n.items.find(key, less)
  401. if len(n.children) == 0 {
  402. if found {
  403. n.size--
  404. return n.items.removeAt(i), true
  405. }
  406. return item{}, false
  407. }
  408. default:
  409. panic("invalid type")
  410. }
  411. // If we get to here, we have children.
  412. if len(n.children[i].items) <= minItems {
  413. return n.growChildAndRemove(i, key, minItems, typ, less)
  414. }
  415. child := n.mutableChild(i)
  416. // Either we had enough items to begin with, or we've done some
  417. // merging/stealing, because we've got enough now and we're ready to return
  418. // stuff.
  419. if found {
  420. // The item exists at index 'i', and the child we've selected can give us a
  421. // predecessor, since if we've gotten here it's got > minItems items in it.
  422. out := n.items[i]
  423. // We use our special-case 'remove' call with typ=maxItem to pull the
  424. // predecessor of item i (the rightmost leaf of our immediate left child)
  425. // and set it into where we pulled the item from.
  426. n.items[i], _ = child.remove(nil, minItems, removeMax, less)
  427. n.size--
  428. return out, true
  429. }
  430. // Final recursive call. Once we're here, we know that the item isn't in this
  431. // node and that the child is big enough to remove from.
  432. m, removed := child.remove(key, minItems, typ, less)
  433. if removed {
  434. n.size--
  435. }
  436. return m, removed
  437. }
  438. // growChildAndRemove grows child 'i' to make sure it's possible to remove an
  439. // item from it while keeping it at minItems, then calls remove to actually
  440. // remove it.
  441. //
  442. // Most documentation says we have to do two sets of special casing:
  443. // 1) item is in this node
  444. // 2) item is in child
  445. // In both cases, we need to handle the two subcases:
  446. // A) node has enough values that it can spare one
  447. // B) node doesn't have enough values
  448. // For the latter, we have to check:
  449. // a) left sibling has node to spare
  450. // b) right sibling has node to spare
  451. // c) we must merge
  452. // To simplify our code here, we handle cases #1 and #2 the same:
  453. // If a node doesn't have enough items, we make sure it does (using a,b,c).
  454. // We then simply redo our remove call, and the second time (regardless of
  455. // whether we're in case 1 or 2), we'll have enough items and can guarantee
  456. // that we hit case A.
  457. func (n *node) growChildAndRemove(i int, key Key, minItems int, typ toRemove, less lessFunc) (item, bool) {
  458. if i > 0 && len(n.children[i-1].items) > minItems {
  459. // Steal from left child
  460. child := n.mutableChild(i)
  461. stealFrom := n.mutableChild(i - 1)
  462. stolenItem := stealFrom.items.pop()
  463. stealFrom.size--
  464. child.items.insertAt(0, n.items[i-1])
  465. child.size++
  466. n.items[i-1] = stolenItem
  467. if len(stealFrom.children) > 0 {
  468. c := stealFrom.children.pop()
  469. stealFrom.size -= c.size
  470. child.children.insertAt(0, c)
  471. child.size += c.size
  472. }
  473. } else if i < len(n.items) && len(n.children[i+1].items) > minItems {
  474. // steal from right child
  475. child := n.mutableChild(i)
  476. stealFrom := n.mutableChild(i + 1)
  477. stolenItem := stealFrom.items.removeAt(0)
  478. stealFrom.size--
  479. child.items = append(child.items, n.items[i])
  480. child.size++
  481. n.items[i] = stolenItem
  482. if len(stealFrom.children) > 0 {
  483. c := stealFrom.children.removeAt(0)
  484. stealFrom.size -= c.size
  485. child.children = append(child.children, c)
  486. child.size += c.size
  487. }
  488. } else {
  489. if i >= len(n.items) {
  490. i--
  491. }
  492. child := n.mutableChild(i)
  493. // merge with right child
  494. mergeItem := n.items.removeAt(i)
  495. mergeChild := n.children.removeAt(i + 1)
  496. child.items = append(child.items, mergeItem)
  497. child.items = append(child.items, mergeChild.items...)
  498. child.children = append(child.children, mergeChild.children...)
  499. child.size = child.computeSize()
  500. n.cow.freeNode(mergeChild)
  501. }
  502. return n.remove(key, minItems, typ, less)
  503. }
  504. // BTree is an implementation of a B-Tree.
  505. //
  506. // BTree stores item instances in an ordered structure, allowing easy insertion,
  507. // removal, and iteration.
  508. //
  509. // Write operations are not safe for concurrent mutation by multiple
  510. // goroutines, but Read operations are.
  511. type BTree struct {
  512. degree int
  513. less lessFunc
  514. root *node
  515. cow *copyOnWriteContext
  516. }
  517. // copyOnWriteContext pointers determine node ownership. A tree with a cow
  518. // context equivalent to a node's cow context is allowed to modify that node.
  519. // A tree whose write context does not match a node's is not allowed to modify
  520. // it, and must create a new, writable copy (IE: it's a Clone).
  521. //
  522. // When doing any write operation, we maintain the invariant that the current
  523. // node's context is equal to the context of the tree that requested the write.
  524. // We do this by, before we descend into any node, creating a copy with the
  525. // correct context if the contexts don't match.
  526. //
  527. // Since the node we're currently visiting on any write has the requesting
  528. // tree's context, that node is modifiable in place. Children of that node may
  529. // not share context, but before we descend into them, we'll make a mutable
  530. // copy.
  531. type copyOnWriteContext struct{ byte } // non-empty, because empty structs may have same addr
  532. // Clone clones the btree, lazily. Clone should not be called concurrently,
  533. // but the original tree (t) and the new tree (t2) can be used concurrently
  534. // once the Clone call completes.
  535. //
  536. // The internal tree structure of b is marked read-only and shared between t and
  537. // t2. Writes to both t and t2 use copy-on-write logic, creating new nodes
  538. // whenever one of b's original nodes would have been modified. Read operations
  539. // should have no performance degredation. Write operations for both t and t2
  540. // will initially experience minor slow-downs caused by additional allocs and
  541. // copies due to the aforementioned copy-on-write logic, but should converge to
  542. // the original performance characteristics of the original tree.
  543. func (t *BTree) Clone() *BTree {
  544. // Create two entirely new copy-on-write contexts.
  545. // This operation effectively creates three trees:
  546. // the original, shared nodes (old b.cow)
  547. // the new b.cow nodes
  548. // the new out.cow nodes
  549. cow1, cow2 := *t.cow, *t.cow
  550. out := *t
  551. t.cow = &cow1
  552. out.cow = &cow2
  553. return &out
  554. }
  555. // maxItems returns the max number of items to allow per node.
  556. func (t *BTree) maxItems() int {
  557. return t.degree*2 - 1
  558. }
  559. // minItems returns the min number of items to allow per node (ignored for the
  560. // root node).
  561. func (t *BTree) minItems() int {
  562. return t.degree - 1
  563. }
  564. var nodePool = sync.Pool{New: func() interface{} { return new(node) }}
  565. func (c *copyOnWriteContext) newNode() *node {
  566. n := nodePool.Get().(*node)
  567. n.cow = c
  568. return n
  569. }
  570. func (c *copyOnWriteContext) freeNode(n *node) {
  571. if n.cow == c {
  572. // clear to allow GC
  573. n.items.truncate(0)
  574. n.children.truncate(0)
  575. n.cow = nil
  576. nodePool.Put(n)
  577. }
  578. }
  579. // Set sets the given key to the given value in the tree. If the key is present in
  580. // the tree, its value is changed and the old value is returned along with a second
  581. // return value of true. If the key is not in the tree, it is added, and the second
  582. // return value is false.
  583. func (t *BTree) Set(k Key, v Value) (old Value, present bool) {
  584. old, present, _ = t.set(k, v, false)
  585. return old, present
  586. }
  587. // SetWithIndex sets the given key to the given value in the tree, and returns the
  588. // index at which it was inserted.
  589. func (t *BTree) SetWithIndex(k Key, v Value) (old Value, present bool, index int) {
  590. return t.set(k, v, true)
  591. }
  592. func (t *BTree) set(k Key, v Value, computeIndex bool) (old Value, present bool, idx int) {
  593. if t.root == nil {
  594. t.root = t.cow.newNode()
  595. t.root.items = append(t.root.items, item{k, v})
  596. t.root.size = 1
  597. return old, false, 0
  598. }
  599. t.root = t.root.mutableFor(t.cow)
  600. if len(t.root.items) >= t.maxItems() {
  601. sz := t.root.size
  602. item2, second := t.root.split(t.maxItems() / 2)
  603. oldroot := t.root
  604. t.root = t.cow.newNode()
  605. t.root.items = append(t.root.items, item2)
  606. t.root.children = append(t.root.children, oldroot, second)
  607. t.root.size = sz
  608. }
  609. return t.root.insert(item{k, v}, t.maxItems(), t.less, computeIndex)
  610. }
  611. // Delete removes the item with the given key, returning its value. The second return value
  612. // reports whether the key was found.
  613. func (t *BTree) Delete(k Key) (Value, bool) {
  614. m, removed := t.deleteItem(k, removeItem)
  615. return m.value, removed
  616. }
  617. // DeleteMin removes the smallest item in the tree and returns its key and value.
  618. // If the tree is empty, it returns zero values.
  619. func (t *BTree) DeleteMin() (Key, Value) {
  620. item, _ := t.deleteItem(nil, removeMin)
  621. return item.key, item.value
  622. }
  623. // DeleteMax removes the largest item in the tree and returns its key and value.
  624. // If the tree is empty, it returns zero values.
  625. func (t *BTree) DeleteMax() (Key, Value) {
  626. item, _ := t.deleteItem(nil, removeMax)
  627. return item.key, item.value
  628. }
  629. func (t *BTree) deleteItem(key Key, typ toRemove) (item, bool) {
  630. if t.root == nil || len(t.root.items) == 0 {
  631. return item{}, false
  632. }
  633. t.root = t.root.mutableFor(t.cow)
  634. out, removed := t.root.remove(key, t.minItems(), typ, t.less)
  635. if len(t.root.items) == 0 && len(t.root.children) > 0 {
  636. oldroot := t.root
  637. t.root = t.root.children[0]
  638. t.cow.freeNode(oldroot)
  639. }
  640. return out, removed
  641. }
  642. // Get returns the value for the given key in the tree, or the zero value if the
  643. // key is not in the tree.
  644. //
  645. // To distinguish a zero value from a key that is not present, use GetWithIndex.
  646. func (t *BTree) Get(k Key) Value {
  647. var z Value
  648. if t.root == nil {
  649. return z
  650. }
  651. item, ok, _ := t.root.get(k, false, t.less)
  652. if !ok {
  653. return z
  654. }
  655. return item.value
  656. }
  657. // GetWithIndex returns the value and index for the given key in the tree, or the
  658. // zero value and -1 if the key is not in the tree.
  659. func (t *BTree) GetWithIndex(k Key) (Value, int) {
  660. var z Value
  661. if t.root == nil {
  662. return z, -1
  663. }
  664. item, _, index := t.root.get(k, true, t.less)
  665. return item.value, index
  666. }
  667. // At returns the key and value at index i. The minimum item has index 0.
  668. // If i is outside the range [0, t.Len()), At panics.
  669. func (t *BTree) At(i int) (Key, Value) {
  670. if i < 0 || i >= t.Len() {
  671. panic("btree: index out of range")
  672. }
  673. item := t.root.at(i)
  674. return item.key, item.value
  675. }
  676. // Has reports whether the given key is in the tree.
  677. func (t *BTree) Has(k Key) bool {
  678. if t.root == nil {
  679. return false
  680. }
  681. _, ok, _ := t.root.get(k, false, t.less)
  682. return ok
  683. }
  684. // Min returns the smallest key in the tree and its value. If the tree is empty, it
  685. // returns zero values.
  686. func (t *BTree) Min() (Key, Value) {
  687. var k Key
  688. var v Value
  689. if t.root == nil {
  690. return k, v
  691. }
  692. n := t.root
  693. for len(n.children) > 0 {
  694. n = n.children[0]
  695. }
  696. if len(n.items) == 0 {
  697. return k, v
  698. }
  699. return n.items[0].key, n.items[0].value
  700. }
  701. // Max returns the largest key in the tree and its value. If the tree is empty, both
  702. // return values are zero values.
  703. func (t *BTree) Max() (Key, Value) {
  704. var k Key
  705. var v Value
  706. if t.root == nil {
  707. return k, v
  708. }
  709. n := t.root
  710. for len(n.children) > 0 {
  711. n = n.children[len(n.children)-1]
  712. }
  713. if len(n.items) == 0 {
  714. return k, v
  715. }
  716. m := n.items[len(n.items)-1]
  717. return m.key, m.value
  718. }
  719. // Len returns the number of items currently in the tree.
  720. func (t *BTree) Len() int {
  721. if t.root == nil {
  722. return 0
  723. }
  724. return t.root.size
  725. }
  726. // Before returns an iterator positioned just before k. After the first call to Next,
  727. // the Iterator will be at k, or at the key just greater than k if k is not in the tree.
  728. // Subsequent calls to Next will traverse the tree's items in ascending order.
  729. func (t *BTree) Before(k Key) *Iterator {
  730. if t.root == nil {
  731. return &Iterator{}
  732. }
  733. var cs cursorStack
  734. cs, found, idx := t.root.cursorStackForKey(k, cs, t.less)
  735. // If we found the key, the cursor stack is pointing to it. Since that is
  736. // the first element we want, don't advance the iterator on the initial call to Next.
  737. // If we haven't found the key, then the top of the cursor stack is either pointing at the
  738. // item just after k, in which case we do not want to move the iterator; or the index
  739. // is past the end of the items slice, in which case we do.
  740. var stay bool
  741. top := cs[len(cs)-1]
  742. if found {
  743. stay = true
  744. } else if top.index < len(top.node.items) {
  745. stay = true
  746. } else {
  747. idx--
  748. }
  749. return &Iterator{
  750. cursors: cs,
  751. stay: stay,
  752. descending: false,
  753. Index: idx,
  754. }
  755. }
  756. // After returns an iterator positioned just after k. After the first call to Next,
  757. // the Iterator will be at k, or at the key just less than k if k is not in the tree.
  758. // Subsequent calls to Next will traverse the tree's items in descending order.
  759. func (t *BTree) After(k Key) *Iterator {
  760. if t.root == nil {
  761. return &Iterator{}
  762. }
  763. var cs cursorStack
  764. cs, found, idx := t.root.cursorStackForKey(k, cs, t.less)
  765. // If we found the key, the cursor stack is pointing to it. Since that is
  766. // the first element we want, don't advance the iterator on the initial call to Next.
  767. // If we haven't found the key, the cursor stack is pointing just after the first item,
  768. // so we do want to advance.
  769. return &Iterator{
  770. cursors: cs,
  771. stay: found,
  772. descending: true,
  773. Index: idx,
  774. }
  775. }
  776. // BeforeIndex returns an iterator positioned just before the item with the given index.
  777. // The iterator will traverse the tree's items in ascending order.
  778. // If i is not in the range [0, tr.Len()], BeforeIndex panics.
  779. // Note that it is not an error to provide an index of tr.Len().
  780. func (t *BTree) BeforeIndex(i int) *Iterator {
  781. return t.indexIterator(i, false)
  782. }
  783. // AfterIndex returns an iterator positioned just after the item with the given index.
  784. // The iterator will traverse the tree's items in descending order.
  785. // If i is not in the range [0, tr.Len()], AfterIndex panics.
  786. // Note that it is not an error to provide an index of tr.Len().
  787. func (t *BTree) AfterIndex(i int) *Iterator {
  788. return t.indexIterator(i, true)
  789. }
  790. func (t *BTree) indexIterator(i int, descending bool) *Iterator {
  791. if i < 0 || i > t.Len() {
  792. panic("btree: index out of range")
  793. }
  794. if i == t.Len() {
  795. return &Iterator{}
  796. }
  797. var cs cursorStack
  798. return &Iterator{
  799. cursors: t.root.cursorStackForIndex(i, cs),
  800. stay: true,
  801. descending: descending,
  802. Index: i,
  803. }
  804. }
  805. // An Iterator supports traversing the items in the tree.
  806. type Iterator struct {
  807. Key Key
  808. Value Value
  809. // Index is the position of the item in the tree viewed as a sequence.
  810. // The minimum item has index zero.
  811. Index int
  812. cursors cursorStack // stack of nodes with indices; last element is the top
  813. stay bool // don't do anything on the first call to Next.
  814. descending bool // traverse the items in descending order
  815. }
  816. // Next advances the Iterator to the next item in the tree. If Next returns true,
  817. // the Iterator's Key, Value and Index fields refer to the next item. If Next returns
  818. // false, there are no more items and the values of Key, Value and Index are undefined.
  819. //
  820. // If the tree is modified during iteration, the behavior is undefined.
  821. func (it *Iterator) Next() bool {
  822. var more bool
  823. switch {
  824. case len(it.cursors) == 0:
  825. more = false
  826. case it.stay:
  827. it.stay = false
  828. more = true
  829. case it.descending:
  830. more = it.dec()
  831. default:
  832. more = it.inc()
  833. }
  834. if !more {
  835. return false
  836. }
  837. top := it.cursors[len(it.cursors)-1]
  838. item := top.node.items[top.index]
  839. it.Key = item.key
  840. it.Value = item.value
  841. return true
  842. }
  843. // When inc returns true, the top cursor on the stack refers to the new current item.
  844. func (it *Iterator) inc() bool {
  845. // Useful invariants for understanding this function:
  846. // - Leaf nodes have zero children, and zero or more items.
  847. // - Nonleaf nodes have one more child than item, and children[i] < items[i] < children[i+1].
  848. // - The current item in the iterator is top.node.items[top.index].
  849. it.Index++
  850. // If we are at a non-leaf node, the current item is items[i], so
  851. // now we want to continue with children[i+1], which must exist
  852. // by the node invariant. We want the minimum item in that child's subtree.
  853. top := it.cursors.incTop(1)
  854. for len(top.node.children) > 0 {
  855. top = cursor{top.node.children[top.index], 0}
  856. it.cursors.push(top)
  857. }
  858. // Here, we are at a leaf node. top.index points to
  859. // the new current item, if it's within the items slice.
  860. for top.index >= len(top.node.items) {
  861. // We've gone through everything in this node. Pop it off the stack.
  862. it.cursors.pop()
  863. // If the stack is now empty,we're past the last item in the tree.
  864. if it.cursors.empty() {
  865. return false
  866. }
  867. top = it.cursors.top()
  868. // The new top's index points to a child, which we've just finished
  869. // exploring. The next item is the one at the same index in the items slice.
  870. }
  871. // Here, the top cursor on the stack points to the new current item.
  872. return true
  873. }
  874. func (it *Iterator) dec() bool {
  875. // See the invariants for inc, above.
  876. it.Index--
  877. top := it.cursors.top()
  878. // If we are at a non-leaf node, the current item is items[i], so
  879. // now we want to continue with children[i]. We want the maximum item in that child's subtree.
  880. for len(top.node.children) > 0 {
  881. c := top.node.children[top.index]
  882. top = cursor{c, len(c.items)}
  883. it.cursors.push(top)
  884. }
  885. top = it.cursors.incTop(-1)
  886. // Here, we are at a leaf node. top.index points to
  887. // the new current item, if it's within the items slice.
  888. for top.index < 0 {
  889. // We've gone through everything in this node. Pop it off the stack.
  890. it.cursors.pop()
  891. // If the stack is now empty,we're past the last item in the tree.
  892. if it.cursors.empty() {
  893. return false
  894. }
  895. // The new top's index points to a child, which we've just finished
  896. // exploring. That child is to the right of the item we want to advance to,
  897. // so decrement the index.
  898. top = it.cursors.incTop(-1)
  899. }
  900. return true
  901. }
  902. // A cursor is effectively a pointer into a node. A stack of cursors identifies an item in the tree,
  903. // and makes it possible to move to the next or previous item efficiently.
  904. //
  905. // If the cursor is on the top of the stack, its index points into the node's items slice, selecting
  906. // the current item. Otherwise, the index points into the children slice and identifies the child
  907. // that is next in the stack.
  908. type cursor struct {
  909. node *node
  910. index int
  911. }
  912. // A cursorStack is a stack of cursors, representing a path of nodes from the root of the tree.
  913. type cursorStack []cursor
  914. func (s *cursorStack) push(c cursor) cursorStack {
  915. *s = append(*s, c)
  916. return *s
  917. }
  918. func (s *cursorStack) pop() cursor {
  919. last := len(*s) - 1
  920. t := (*s)[last]
  921. *s = (*s)[:last]
  922. return t
  923. }
  924. func (s *cursorStack) top() cursor {
  925. return (*s)[len(*s)-1]
  926. }
  927. func (s *cursorStack) empty() bool {
  928. return len(*s) == 0
  929. }
  930. // incTop increments top's index by n and returns it.
  931. func (s *cursorStack) incTop(n int) cursor {
  932. (*s)[len(*s)-1].index += n // Don't call top: modify the original, not a copy.
  933. return s.top()
  934. }