btree.go 26 KB

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