btree.go 26 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890
  1. // Copyright 2014 Google Inc.
  2. //
  3. // Licensed under the Apache License, Version 2.0 (the "License");
  4. // you may not use this file except in compliance with the License.
  5. // You may obtain a copy of the License at
  6. //
  7. // http://www.apache.org/licenses/LICENSE-2.0
  8. //
  9. // Unless required by applicable law or agreed to in writing, software
  10. // distributed under the License is distributed on an "AS IS" BASIS,
  11. // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. // See the License for the specific language governing permissions and
  13. // limitations under the License.
  14. // Package btree implements in-memory B-Trees of arbitrary degree.
  15. //
  16. // btree implements an in-memory B-Tree for use as an ordered data structure.
  17. // It is not meant for persistent storage solutions.
  18. //
  19. // It has a flatter structure than an equivalent red-black or other binary tree,
  20. // which in some cases yields better memory usage and/or performance.
  21. // See some discussion on the matter here:
  22. // http://google-opensource.blogspot.com/2013/01/c-containers-that-save-memory-and-time.html
  23. // Note, though, that this project is in no way related to the C++ B-Tree
  24. // implementation written about there.
  25. //
  26. // Within this tree, each node contains a slice of items and a (possibly nil)
  27. // slice of children. For basic numeric values or raw structs, this can cause
  28. // efficiency differences when compared to equivalent C++ template code that
  29. // stores values in arrays within the node:
  30. // * Due to the overhead of storing values as interfaces (each
  31. // value needs to be stored as the value itself, then 2 words for the
  32. // interface pointing to that value and its type), resulting in higher
  33. // memory use.
  34. // * Since interfaces can point to values anywhere in memory, values are
  35. // most likely not stored in contiguous blocks, resulting in a higher
  36. // number of cache misses.
  37. // These issues don't tend to matter, though, when working with strings or other
  38. // heap-allocated structures, since C++-equivalent structures also must store
  39. // pointers and also distribute their values across the heap.
  40. //
  41. // This implementation is designed to be a drop-in replacement to gollrb.LLRB
  42. // trees, (http://github.com/petar/gollrb), an excellent and probably the most
  43. // widely used ordered tree implementation in the Go ecosystem currently.
  44. // Its functions, therefore, exactly mirror those of
  45. // llrb.LLRB where possible. Unlike gollrb, though, we currently don't
  46. // support storing multiple equivalent values.
  47. package btree
  48. import (
  49. "fmt"
  50. "io"
  51. "sort"
  52. "strings"
  53. "sync"
  54. )
  55. // Item represents a single object in the tree.
  56. type Item interface {
  57. // Less tests whether the current item is less than the given argument.
  58. //
  59. // This must provide a strict weak ordering.
  60. // If !a.Less(b) && !b.Less(a), we treat this to mean a == b (i.e. we can only
  61. // hold one of either a or b in the tree).
  62. Less(than Item) bool
  63. }
  64. const (
  65. DefaultFreeListSize = 32
  66. )
  67. var (
  68. nilItems = make(items, 16)
  69. nilChildren = make(children, 16)
  70. )
  71. // FreeList represents a free list of btree nodes. By default each
  72. // BTree has its own FreeList, but multiple BTrees can share the same
  73. // FreeList.
  74. // Two Btrees using the same freelist are safe for concurrent write access.
  75. type FreeList struct {
  76. mu sync.Mutex
  77. freelist []*node
  78. }
  79. // NewFreeList creates a new free list.
  80. // size is the maximum size of the returned free list.
  81. func NewFreeList(size int) *FreeList {
  82. return &FreeList{freelist: make([]*node, 0, size)}
  83. }
  84. func (f *FreeList) newNode() (n *node) {
  85. f.mu.Lock()
  86. index := len(f.freelist) - 1
  87. if index < 0 {
  88. f.mu.Unlock()
  89. return new(node)
  90. }
  91. n = f.freelist[index]
  92. f.freelist[index] = nil
  93. f.freelist = f.freelist[:index]
  94. f.mu.Unlock()
  95. return
  96. }
  97. // freeNode adds the given node to the list, returning true if it was added
  98. // and false if it was discarded.
  99. func (f *FreeList) freeNode(n *node) (out bool) {
  100. f.mu.Lock()
  101. if len(f.freelist) < cap(f.freelist) {
  102. f.freelist = append(f.freelist, n)
  103. out = true
  104. }
  105. f.mu.Unlock()
  106. return
  107. }
  108. // ItemIterator allows callers of Ascend* to iterate in-order over portions of
  109. // the tree. When this function returns false, iteration will stop and the
  110. // associated Ascend* function will immediately return.
  111. type ItemIterator func(i Item) bool
  112. // New creates a new B-Tree with the given degree.
  113. //
  114. // New(2), for example, will create a 2-3-4 tree (each node contains 1-3 items
  115. // and 2-4 children).
  116. func New(degree int) *BTree {
  117. return NewWithFreeList(degree, NewFreeList(DefaultFreeListSize))
  118. }
  119. // NewWithFreeList creates a new B-Tree that uses the given node free list.
  120. func NewWithFreeList(degree int, f *FreeList) *BTree {
  121. if degree <= 1 {
  122. panic("bad degree")
  123. }
  124. return &BTree{
  125. degree: degree,
  126. cow: &copyOnWriteContext{freelist: f},
  127. }
  128. }
  129. // items stores items in a node.
  130. type items []Item
  131. // insertAt inserts a value into the given index, pushing all subsequent values
  132. // forward.
  133. func (s *items) insertAt(index int, item Item) {
  134. *s = append(*s, nil)
  135. if index < len(*s) {
  136. copy((*s)[index+1:], (*s)[index:])
  137. }
  138. (*s)[index] = item
  139. }
  140. // removeAt removes a value at a given index, pulling all subsequent values
  141. // back.
  142. func (s *items) removeAt(index int) Item {
  143. item := (*s)[index]
  144. copy((*s)[index:], (*s)[index+1:])
  145. (*s)[len(*s)-1] = nil
  146. *s = (*s)[:len(*s)-1]
  147. return item
  148. }
  149. // pop removes and returns the last element in the list.
  150. func (s *items) pop() (out Item) {
  151. index := len(*s) - 1
  152. out = (*s)[index]
  153. (*s)[index] = nil
  154. *s = (*s)[:index]
  155. return
  156. }
  157. // truncate truncates this instance at index so that it contains only the
  158. // first index items. index must be less than or equal to length.
  159. func (s *items) truncate(index int) {
  160. var toClear items
  161. *s, toClear = (*s)[:index], (*s)[index:]
  162. for len(toClear) > 0 {
  163. toClear = toClear[copy(toClear, nilItems):]
  164. }
  165. }
  166. // find returns the index where the given item should be inserted into this
  167. // list. 'found' is true if the item already exists in the list at the given
  168. // index.
  169. func (s items) find(item Item) (index int, found bool) {
  170. i := sort.Search(len(s), func(i int) bool {
  171. return item.Less(s[i])
  172. })
  173. if i > 0 && !s[i-1].Less(item) {
  174. return i - 1, true
  175. }
  176. return i, false
  177. }
  178. // children stores child nodes in a node.
  179. type children []*node
  180. // insertAt inserts a value into the given index, pushing all subsequent values
  181. // forward.
  182. func (s *children) insertAt(index int, n *node) {
  183. *s = append(*s, nil)
  184. if index < len(*s) {
  185. copy((*s)[index+1:], (*s)[index:])
  186. }
  187. (*s)[index] = n
  188. }
  189. // removeAt removes a value at a given index, pulling all subsequent values
  190. // back.
  191. func (s *children) removeAt(index int) *node {
  192. n := (*s)[index]
  193. copy((*s)[index:], (*s)[index+1:])
  194. (*s)[len(*s)-1] = nil
  195. *s = (*s)[:len(*s)-1]
  196. return n
  197. }
  198. // pop removes and returns the last element in the list.
  199. func (s *children) pop() (out *node) {
  200. index := len(*s) - 1
  201. out = (*s)[index]
  202. (*s)[index] = nil
  203. *s = (*s)[:index]
  204. return
  205. }
  206. // truncate truncates this instance at index so that it contains only the
  207. // first index children. index must be less than or equal to length.
  208. func (s *children) truncate(index int) {
  209. var toClear children
  210. *s, toClear = (*s)[:index], (*s)[index:]
  211. for len(toClear) > 0 {
  212. toClear = toClear[copy(toClear, nilChildren):]
  213. }
  214. }
  215. // node is an internal node in a tree.
  216. //
  217. // It must at all times maintain the invariant that either
  218. // * len(children) == 0, len(items) unconstrained
  219. // * len(children) == len(items) + 1
  220. type node struct {
  221. items items
  222. children children
  223. cow *copyOnWriteContext
  224. }
  225. func (n *node) mutableFor(cow *copyOnWriteContext) *node {
  226. if n.cow == cow {
  227. return n
  228. }
  229. out := cow.newNode()
  230. if cap(out.items) >= len(n.items) {
  231. out.items = out.items[:len(n.items)]
  232. } else {
  233. out.items = make(items, len(n.items), cap(n.items))
  234. }
  235. copy(out.items, n.items)
  236. // Copy children
  237. if cap(out.children) >= len(n.children) {
  238. out.children = out.children[:len(n.children)]
  239. } else {
  240. out.children = make(children, len(n.children), cap(n.children))
  241. }
  242. copy(out.children, n.children)
  243. return out
  244. }
  245. func (n *node) mutableChild(i int) *node {
  246. c := n.children[i].mutableFor(n.cow)
  247. n.children[i] = c
  248. return c
  249. }
  250. // split splits the given node at the given index. The current node shrinks,
  251. // and this function returns the item that existed at that index and a new node
  252. // containing all items/children after it.
  253. func (n *node) split(i int) (Item, *node) {
  254. item := n.items[i]
  255. next := n.cow.newNode()
  256. next.items = append(next.items, n.items[i+1:]...)
  257. n.items.truncate(i)
  258. if len(n.children) > 0 {
  259. next.children = append(next.children, n.children[i+1:]...)
  260. n.children.truncate(i + 1)
  261. }
  262. return item, next
  263. }
  264. // maybeSplitChild checks if a child should be split, and if so splits it.
  265. // Returns whether or not a split occurred.
  266. func (n *node) maybeSplitChild(i, maxItems int) bool {
  267. if len(n.children[i].items) < maxItems {
  268. return false
  269. }
  270. first := n.mutableChild(i)
  271. item, second := first.split(maxItems / 2)
  272. n.items.insertAt(i, item)
  273. n.children.insertAt(i+1, second)
  274. return true
  275. }
  276. // insert inserts an item into the subtree rooted at this node, making sure
  277. // no nodes in the subtree exceed maxItems items. Should an equivalent item be
  278. // be found/replaced by insert, it will be returned.
  279. func (n *node) insert(item Item, maxItems int) Item {
  280. i, found := n.items.find(item)
  281. if found {
  282. out := n.items[i]
  283. n.items[i] = item
  284. return out
  285. }
  286. if len(n.children) == 0 {
  287. n.items.insertAt(i, item)
  288. return nil
  289. }
  290. if n.maybeSplitChild(i, maxItems) {
  291. inTree := n.items[i]
  292. switch {
  293. case item.Less(inTree):
  294. // no change, we want first split node
  295. case inTree.Less(item):
  296. i++ // we want second split node
  297. default:
  298. out := n.items[i]
  299. n.items[i] = item
  300. return out
  301. }
  302. }
  303. return n.mutableChild(i).insert(item, maxItems)
  304. }
  305. // get finds the given key in the subtree and returns it.
  306. func (n *node) get(key Item) Item {
  307. i, found := n.items.find(key)
  308. if found {
  309. return n.items[i]
  310. } else if len(n.children) > 0 {
  311. return n.children[i].get(key)
  312. }
  313. return nil
  314. }
  315. // min returns the first item in the subtree.
  316. func min(n *node) Item {
  317. if n == nil {
  318. return nil
  319. }
  320. for len(n.children) > 0 {
  321. n = n.children[0]
  322. }
  323. if len(n.items) == 0 {
  324. return nil
  325. }
  326. return n.items[0]
  327. }
  328. // max returns the last item in the subtree.
  329. func max(n *node) Item {
  330. if n == nil {
  331. return nil
  332. }
  333. for len(n.children) > 0 {
  334. n = n.children[len(n.children)-1]
  335. }
  336. if len(n.items) == 0 {
  337. return nil
  338. }
  339. return n.items[len(n.items)-1]
  340. }
  341. // toRemove details what item to remove in a node.remove call.
  342. type toRemove int
  343. const (
  344. removeItem toRemove = iota // removes the given item
  345. removeMin // removes smallest item in the subtree
  346. removeMax // removes largest item in the subtree
  347. )
  348. // remove removes an item from the subtree rooted at this node.
  349. func (n *node) remove(item Item, minItems int, typ toRemove) Item {
  350. var i int
  351. var found bool
  352. switch typ {
  353. case removeMax:
  354. if len(n.children) == 0 {
  355. return n.items.pop()
  356. }
  357. i = len(n.items)
  358. case removeMin:
  359. if len(n.children) == 0 {
  360. return n.items.removeAt(0)
  361. }
  362. i = 0
  363. case removeItem:
  364. i, found = n.items.find(item)
  365. if len(n.children) == 0 {
  366. if found {
  367. return n.items.removeAt(i)
  368. }
  369. return nil
  370. }
  371. default:
  372. panic("invalid type")
  373. }
  374. // If we get to here, we have children.
  375. if len(n.children[i].items) <= minItems {
  376. return n.growChildAndRemove(i, item, minItems, typ)
  377. }
  378. child := n.mutableChild(i)
  379. // Either we had enough items to begin with, or we've done some
  380. // merging/stealing, because we've got enough now and we're ready to return
  381. // stuff.
  382. if found {
  383. // The item exists at index 'i', and the child we've selected can give us a
  384. // predecessor, since if we've gotten here it's got > minItems items in it.
  385. out := n.items[i]
  386. // We use our special-case 'remove' call with typ=maxItem to pull the
  387. // predecessor of item i (the rightmost leaf of our immediate left child)
  388. // and set it into where we pulled the item from.
  389. n.items[i] = child.remove(nil, minItems, removeMax)
  390. return out
  391. }
  392. // Final recursive call. Once we're here, we know that the item isn't in this
  393. // node and that the child is big enough to remove from.
  394. return child.remove(item, minItems, typ)
  395. }
  396. // growChildAndRemove grows child 'i' to make sure it's possible to remove an
  397. // item from it while keeping it at minItems, then calls remove to actually
  398. // remove it.
  399. //
  400. // Most documentation says we have to do two sets of special casing:
  401. // 1) item is in this node
  402. // 2) item is in child
  403. // In both cases, we need to handle the two subcases:
  404. // A) node has enough values that it can spare one
  405. // B) node doesn't have enough values
  406. // For the latter, we have to check:
  407. // a) left sibling has node to spare
  408. // b) right sibling has node to spare
  409. // c) we must merge
  410. // To simplify our code here, we handle cases #1 and #2 the same:
  411. // If a node doesn't have enough items, we make sure it does (using a,b,c).
  412. // We then simply redo our remove call, and the second time (regardless of
  413. // whether we're in case 1 or 2), we'll have enough items and can guarantee
  414. // that we hit case A.
  415. func (n *node) growChildAndRemove(i int, item Item, minItems int, typ toRemove) Item {
  416. if i > 0 && len(n.children[i-1].items) > minItems {
  417. // Steal from left child
  418. child := n.mutableChild(i)
  419. stealFrom := n.mutableChild(i - 1)
  420. stolenItem := stealFrom.items.pop()
  421. child.items.insertAt(0, n.items[i-1])
  422. n.items[i-1] = stolenItem
  423. if len(stealFrom.children) > 0 {
  424. child.children.insertAt(0, stealFrom.children.pop())
  425. }
  426. } else if i < len(n.items) && len(n.children[i+1].items) > minItems {
  427. // steal from right child
  428. child := n.mutableChild(i)
  429. stealFrom := n.mutableChild(i + 1)
  430. stolenItem := stealFrom.items.removeAt(0)
  431. child.items = append(child.items, n.items[i])
  432. n.items[i] = stolenItem
  433. if len(stealFrom.children) > 0 {
  434. child.children = append(child.children, stealFrom.children.removeAt(0))
  435. }
  436. } else {
  437. if i >= len(n.items) {
  438. i--
  439. }
  440. child := n.mutableChild(i)
  441. // merge with right child
  442. mergeItem := n.items.removeAt(i)
  443. mergeChild := n.children.removeAt(i + 1)
  444. child.items = append(child.items, mergeItem)
  445. child.items = append(child.items, mergeChild.items...)
  446. child.children = append(child.children, mergeChild.children...)
  447. n.cow.freeNode(mergeChild)
  448. }
  449. return n.remove(item, minItems, typ)
  450. }
  451. type direction int
  452. const (
  453. descend = direction(-1)
  454. ascend = direction(+1)
  455. )
  456. // iterate provides a simple method for iterating over elements in the tree.
  457. //
  458. // When ascending, the 'start' should be less than 'stop' and when descending,
  459. // the 'start' should be greater than 'stop'. Setting 'includeStart' to true
  460. // will force the iterator to include the first item when it equals 'start',
  461. // thus creating a "greaterOrEqual" or "lessThanEqual" rather than just a
  462. // "greaterThan" or "lessThan" queries.
  463. func (n *node) iterate(dir direction, start, stop Item, includeStart bool, hit bool, iter ItemIterator) (bool, bool) {
  464. var ok, found bool
  465. var index int
  466. switch dir {
  467. case ascend:
  468. if start != nil {
  469. index, _ = n.items.find(start)
  470. }
  471. for i := index; i < len(n.items); i++ {
  472. if len(n.children) > 0 {
  473. if hit, ok = n.children[i].iterate(dir, start, stop, includeStart, hit, iter); !ok {
  474. return hit, false
  475. }
  476. }
  477. if !includeStart && !hit && start != nil && !start.Less(n.items[i]) {
  478. hit = true
  479. continue
  480. }
  481. hit = true
  482. if stop != nil && !n.items[i].Less(stop) {
  483. return hit, false
  484. }
  485. if !iter(n.items[i]) {
  486. return hit, false
  487. }
  488. }
  489. if len(n.children) > 0 {
  490. if hit, ok = n.children[len(n.children)-1].iterate(dir, start, stop, includeStart, hit, iter); !ok {
  491. return hit, false
  492. }
  493. }
  494. case descend:
  495. if start != nil {
  496. index, found = n.items.find(start)
  497. if !found {
  498. index = index - 1
  499. }
  500. } else {
  501. index = len(n.items) - 1
  502. }
  503. for i := index; i >= 0; i-- {
  504. if start != nil && !n.items[i].Less(start) {
  505. if !includeStart || hit || start.Less(n.items[i]) {
  506. continue
  507. }
  508. }
  509. if len(n.children) > 0 {
  510. if hit, ok = n.children[i+1].iterate(dir, start, stop, includeStart, hit, iter); !ok {
  511. return hit, false
  512. }
  513. }
  514. if stop != nil && !stop.Less(n.items[i]) {
  515. return hit, false // continue
  516. }
  517. hit = true
  518. if !iter(n.items[i]) {
  519. return hit, false
  520. }
  521. }
  522. if len(n.children) > 0 {
  523. if hit, ok = n.children[0].iterate(dir, start, stop, includeStart, hit, iter); !ok {
  524. return hit, false
  525. }
  526. }
  527. }
  528. return hit, true
  529. }
  530. // Used for testing/debugging purposes.
  531. func (n *node) print(w io.Writer, level int) {
  532. fmt.Fprintf(w, "%sNODE:%v\n", strings.Repeat(" ", level), n.items)
  533. for _, c := range n.children {
  534. c.print(w, level+1)
  535. }
  536. }
  537. // BTree is an implementation of a B-Tree.
  538. //
  539. // BTree stores Item instances in an ordered structure, allowing easy insertion,
  540. // removal, and iteration.
  541. //
  542. // Write operations are not safe for concurrent mutation by multiple
  543. // goroutines, but Read operations are.
  544. type BTree struct {
  545. degree int
  546. length int
  547. root *node
  548. cow *copyOnWriteContext
  549. }
  550. // copyOnWriteContext pointers determine node ownership... a tree with a write
  551. // context equivalent to a node's write context is allowed to modify that node.
  552. // A tree whose write context does not match a node's is not allowed to modify
  553. // it, and must create a new, writable copy (IE: it's a Clone).
  554. //
  555. // When doing any write operation, we maintain the invariant that the current
  556. // node's context is equal to the context of the tree that requested the write.
  557. // We do this by, before we descend into any node, creating a copy with the
  558. // correct context if the contexts don't match.
  559. //
  560. // Since the node we're currently visiting on any write has the requesting
  561. // tree's context, that node is modifiable in place. Children of that node may
  562. // not share context, but before we descend into them, we'll make a mutable
  563. // copy.
  564. type copyOnWriteContext struct {
  565. freelist *FreeList
  566. }
  567. // Clone clones the btree, lazily. Clone should not be called concurrently,
  568. // but the original tree (t) and the new tree (t2) can be used concurrently
  569. // once the Clone call completes.
  570. //
  571. // The internal tree structure of b is marked read-only and shared between t and
  572. // t2. Writes to both t and t2 use copy-on-write logic, creating new nodes
  573. // whenever one of b's original nodes would have been modified. Read operations
  574. // should have no performance degredation. Write operations for both t and t2
  575. // will initially experience minor slow-downs caused by additional allocs and
  576. // copies due to the aforementioned copy-on-write logic, but should converge to
  577. // the original performance characteristics of the original tree.
  578. func (t *BTree) Clone() (t2 *BTree) {
  579. // Create two entirely new copy-on-write contexts.
  580. // This operation effectively creates three trees:
  581. // the original, shared nodes (old b.cow)
  582. // the new b.cow nodes
  583. // the new out.cow nodes
  584. cow1, cow2 := *t.cow, *t.cow
  585. out := *t
  586. t.cow = &cow1
  587. out.cow = &cow2
  588. return &out
  589. }
  590. // maxItems returns the max number of items to allow per node.
  591. func (t *BTree) maxItems() int {
  592. return t.degree*2 - 1
  593. }
  594. // minItems returns the min number of items to allow per node (ignored for the
  595. // root node).
  596. func (t *BTree) minItems() int {
  597. return t.degree - 1
  598. }
  599. func (c *copyOnWriteContext) newNode() (n *node) {
  600. n = c.freelist.newNode()
  601. n.cow = c
  602. return
  603. }
  604. type freeType int
  605. const (
  606. ftFreelistFull freeType = iota // node was freed (available for GC, not stored in freelist)
  607. ftStored // node was stored in the freelist for later use
  608. ftNotOwned // node was ignored by COW, since it's owned by another one
  609. )
  610. // freeNode frees a node within a given COW context, if it's owned by that
  611. // context. It returns what happened to the node (see freeType const
  612. // documentation).
  613. func (c *copyOnWriteContext) freeNode(n *node) freeType {
  614. if n.cow == c {
  615. // clear to allow GC
  616. n.items.truncate(0)
  617. n.children.truncate(0)
  618. n.cow = nil
  619. if c.freelist.freeNode(n) {
  620. return ftStored
  621. } else {
  622. return ftFreelistFull
  623. }
  624. } else {
  625. return ftNotOwned
  626. }
  627. }
  628. // ReplaceOrInsert adds the given item to the tree. If an item in the tree
  629. // already equals the given one, it is removed from the tree and returned.
  630. // Otherwise, nil is returned.
  631. //
  632. // nil cannot be added to the tree (will panic).
  633. func (t *BTree) ReplaceOrInsert(item Item) Item {
  634. if item == nil {
  635. panic("nil item being added to BTree")
  636. }
  637. if t.root == nil {
  638. t.root = t.cow.newNode()
  639. t.root.items = append(t.root.items, item)
  640. t.length++
  641. return nil
  642. } else {
  643. t.root = t.root.mutableFor(t.cow)
  644. if len(t.root.items) >= t.maxItems() {
  645. item2, second := t.root.split(t.maxItems() / 2)
  646. oldroot := t.root
  647. t.root = t.cow.newNode()
  648. t.root.items = append(t.root.items, item2)
  649. t.root.children = append(t.root.children, oldroot, second)
  650. }
  651. }
  652. out := t.root.insert(item, t.maxItems())
  653. if out == nil {
  654. t.length++
  655. }
  656. return out
  657. }
  658. // Delete removes an item equal to the passed in item from the tree, returning
  659. // it. If no such item exists, returns nil.
  660. func (t *BTree) Delete(item Item) Item {
  661. return t.deleteItem(item, removeItem)
  662. }
  663. // DeleteMin removes the smallest item in the tree and returns it.
  664. // If no such item exists, returns nil.
  665. func (t *BTree) DeleteMin() Item {
  666. return t.deleteItem(nil, removeMin)
  667. }
  668. // DeleteMax removes the largest item in the tree and returns it.
  669. // If no such item exists, returns nil.
  670. func (t *BTree) DeleteMax() Item {
  671. return t.deleteItem(nil, removeMax)
  672. }
  673. func (t *BTree) deleteItem(item Item, typ toRemove) Item {
  674. if t.root == nil || len(t.root.items) == 0 {
  675. return nil
  676. }
  677. t.root = t.root.mutableFor(t.cow)
  678. out := t.root.remove(item, t.minItems(), typ)
  679. if len(t.root.items) == 0 && len(t.root.children) > 0 {
  680. oldroot := t.root
  681. t.root = t.root.children[0]
  682. t.cow.freeNode(oldroot)
  683. }
  684. if out != nil {
  685. t.length--
  686. }
  687. return out
  688. }
  689. // AscendRange calls the iterator for every value in the tree within the range
  690. // [greaterOrEqual, lessThan), until iterator returns false.
  691. func (t *BTree) AscendRange(greaterOrEqual, lessThan Item, iterator ItemIterator) {
  692. if t.root == nil {
  693. return
  694. }
  695. t.root.iterate(ascend, greaterOrEqual, lessThan, true, false, iterator)
  696. }
  697. // AscendLessThan calls the iterator for every value in the tree within the range
  698. // [first, pivot), until iterator returns false.
  699. func (t *BTree) AscendLessThan(pivot Item, iterator ItemIterator) {
  700. if t.root == nil {
  701. return
  702. }
  703. t.root.iterate(ascend, nil, pivot, false, false, iterator)
  704. }
  705. // AscendGreaterOrEqual calls the iterator for every value in the tree within
  706. // the range [pivot, last], until iterator returns false.
  707. func (t *BTree) AscendGreaterOrEqual(pivot Item, iterator ItemIterator) {
  708. if t.root == nil {
  709. return
  710. }
  711. t.root.iterate(ascend, pivot, nil, true, false, iterator)
  712. }
  713. // Ascend calls the iterator for every value in the tree within the range
  714. // [first, last], until iterator returns false.
  715. func (t *BTree) Ascend(iterator ItemIterator) {
  716. if t.root == nil {
  717. return
  718. }
  719. t.root.iterate(ascend, nil, nil, false, false, iterator)
  720. }
  721. // DescendRange calls the iterator for every value in the tree within the range
  722. // [lessOrEqual, greaterThan), until iterator returns false.
  723. func (t *BTree) DescendRange(lessOrEqual, greaterThan Item, iterator ItemIterator) {
  724. if t.root == nil {
  725. return
  726. }
  727. t.root.iterate(descend, lessOrEqual, greaterThan, true, false, iterator)
  728. }
  729. // DescendLessOrEqual calls the iterator for every value in the tree within the range
  730. // [pivot, first], until iterator returns false.
  731. func (t *BTree) DescendLessOrEqual(pivot Item, iterator ItemIterator) {
  732. if t.root == nil {
  733. return
  734. }
  735. t.root.iterate(descend, pivot, nil, true, false, iterator)
  736. }
  737. // DescendGreaterThan calls the iterator for every value in the tree within
  738. // the range [last, pivot), until iterator returns false.
  739. func (t *BTree) DescendGreaterThan(pivot Item, iterator ItemIterator) {
  740. if t.root == nil {
  741. return
  742. }
  743. t.root.iterate(descend, nil, pivot, false, false, iterator)
  744. }
  745. // Descend calls the iterator for every value in the tree within the range
  746. // [last, first], until iterator returns false.
  747. func (t *BTree) Descend(iterator ItemIterator) {
  748. if t.root == nil {
  749. return
  750. }
  751. t.root.iterate(descend, nil, nil, false, false, iterator)
  752. }
  753. // Get looks for the key item in the tree, returning it. It returns nil if
  754. // unable to find that item.
  755. func (t *BTree) Get(key Item) Item {
  756. if t.root == nil {
  757. return nil
  758. }
  759. return t.root.get(key)
  760. }
  761. // Min returns the smallest item in the tree, or nil if the tree is empty.
  762. func (t *BTree) Min() Item {
  763. return min(t.root)
  764. }
  765. // Max returns the largest item in the tree, or nil if the tree is empty.
  766. func (t *BTree) Max() Item {
  767. return max(t.root)
  768. }
  769. // Has returns true if the given key is in the tree.
  770. func (t *BTree) Has(key Item) bool {
  771. return t.Get(key) != nil
  772. }
  773. // Len returns the number of items currently in the tree.
  774. func (t *BTree) Len() int {
  775. return t.length
  776. }
  777. // Clear removes all items from the btree. If addNodesToFreelist is true,
  778. // t's nodes are added to its freelist as part of this call, until the freelist
  779. // is full. Otherwise, the root node is simply dereferenced and the subtree
  780. // left to Go's normal GC processes.
  781. //
  782. // This can be much faster
  783. // than calling Delete on all elements, because that requires finding/removing
  784. // each element in the tree and updating the tree accordingly. It also is
  785. // somewhat faster than creating a new tree to replace the old one, because
  786. // nodes from the old tree are reclaimed into the freelist for use by the new
  787. // one, instead of being lost to the garbage collector.
  788. //
  789. // This call takes:
  790. // O(1): when addNodesToFreelist is false, this is a single operation.
  791. // O(1): when the freelist is already full, it breaks out immediately
  792. // O(freelist size): when the freelist is empty and the nodes are all owned
  793. // by this tree, nodes are added to the freelist until full.
  794. // O(tree size): when all nodes are owned by another tree, all nodes are
  795. // iterated over looking for nodes to add to the freelist, and due to
  796. // ownership, none are.
  797. func (t *BTree) Clear(addNodesToFreelist bool) {
  798. if t.root != nil && addNodesToFreelist {
  799. t.root.reset(t.cow)
  800. }
  801. t.root, t.length = nil, 0
  802. }
  803. // reset returns a subtree to the freelist. It breaks out immediately if the
  804. // freelist is full, since the only benefit of iterating is to fill that
  805. // freelist up. Returns true if parent reset call should continue.
  806. func (n *node) reset(c *copyOnWriteContext) bool {
  807. for _, child := range n.children {
  808. if !child.reset(c) {
  809. return false
  810. }
  811. }
  812. return c.freeNode(n) != ftFreelistFull
  813. }
  814. // Int implements the Item interface for integers.
  815. type Int int
  816. // Less returns true if int(a) < int(b).
  817. func (a Int) Less(b Item) bool {
  818. return a < b.(Int)
  819. }