You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

btree.go 28KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122112311241125112611271128112911301131113211331134113511361137113811391140114111421143114411451146114711481149115011511152115311541155115611571158115911601161116211631164116511661167116811691170117111721173117411751176117711781179118011811182118311841185118611871188118911901191119211931194119511961197119811991200120112021203120412051206120712081209121012111212121312141215121612171218121912201221122212231224122512261227122812291230123112321233123412351236123712381239124012411242124312441245124612471248124912501251125212531254125512561257125812591260126112621263126412651266126712681269127012711272127312741275
  1. // Copyright 2020 Joshua J Baker. All rights reserved.
  2. // Use of this source code is governed by an MIT-style license that can be
  3. // found in the LICENSE file at https://github.com/tidwall/btree/LICENSE
  4. ///////////////////////////////////////////////////////////////////////////////
  5. // BEGIN PARAMS
  6. ///////////////////////////////////////////////////////////////////////////////
  7. package btree
  8. import "sync"
  9. // degree is the B-Tree degree, which is equal to maximum number of children
  10. // pre node times two.
  11. // The default is 128, which means each node can have 255 items and 256 child
  12. // nodes.
  13. const degree = 128
  14. // kind is the item type.
  15. // It's important to use the equal symbol, which tells Go to create an alias of
  16. // the type, rather than creating an entirely new type.
  17. type kind = interface{}
  18. // contextKind is the kind of context that can be passed to NewOptions and the
  19. // less function
  20. type contextKind = interface{}
  21. // less returns true if A is less than B.
  22. // The value of context will be whatever was passed to NewOptions through the
  23. // Options.Context field, otherwise nil if the field was not set.
  24. func less(a, b kind, context contextKind) bool {
  25. return context.(func(a, b contextKind) bool)(a, b)
  26. }
  27. // BTree aliases
  28. // These are aliases to the local bTree types and functions, which are exported
  29. // to allow for public use at a package level.
  30. // Rename them if desired, or comment them out to make the library private.
  31. type BTree = bTree
  32. type Options = bOptions
  33. type PathHint = bPathHint
  34. type Iter = bIter
  35. func New(less func(a, b kind) bool) *bTree { return bNew() }
  36. func NewOptions(opts bOptions) *bTree { return bNewOptions(opts) }
  37. // The functions below, which begin with "test*", are required by the
  38. // btree_test.go file. If you choose not use include the btree_test.go file in
  39. // your project then these functions may be omitted.
  40. // testCustomSeed can be used to generate a custom random seed for testing.
  41. // Returning false will use time.Now().UnixNano()
  42. func testCustomSeed() (seed int64, ok bool) {
  43. return 0, false
  44. }
  45. // testMakeItem must return a valid item for testing.
  46. // It's required that the returned item maintains equal order as the
  47. // provided int, such that:
  48. // testMakeItem(0) < testMakeItem(1) < testMakeItem(2) < testMakeItem(10)
  49. func testMakeItem(x int) (item kind) {
  50. return x
  51. }
  52. // testNewBTree must return an operational btree for testing.
  53. func testNewBTree() *bTree {
  54. return bNewOptions(bOptions{
  55. Context: func(a, b contextKind) bool {
  56. if a == nil {
  57. return b != nil
  58. } else if b == nil {
  59. return false
  60. }
  61. return a.(int) < b.(int)
  62. },
  63. })
  64. }
  65. ///////////////////////////////////////////////////////////////////////////////
  66. // END PARAMS
  67. ///////////////////////////////////////////////////////////////////////////////
  68. // Do not edit code below this line.
  69. const maxItems = degree*2 - 1 // max items per node. max children is +1
  70. const minItems = maxItems / 2
  71. type bTree struct {
  72. mu *sync.RWMutex
  73. cow *cow
  74. root *node
  75. count int
  76. ctx contextKind
  77. locks bool
  78. empty kind
  79. }
  80. type node struct {
  81. cow *cow
  82. count int
  83. items []kind
  84. children *[]*node
  85. }
  86. type cow struct {
  87. _ int // cannot be an empty struct
  88. }
  89. func (tr *bTree) newNode(leaf bool) *node {
  90. n := &node{cow: tr.cow}
  91. if !leaf {
  92. n.children = new([]*node)
  93. }
  94. return n
  95. }
  96. // leaf returns true if the node is a leaf.
  97. func (n *node) leaf() bool {
  98. return n.children == nil
  99. }
  100. // PathHint is a utility type used with the *Hint() functions. Hints provide
  101. // faster operations for clustered keys.
  102. type bPathHint struct {
  103. used [8]bool
  104. path [8]uint8
  105. }
  106. type bOptions struct {
  107. NoLocks bool
  108. Context contextKind
  109. }
  110. // New returns a new BTree
  111. func bNew() *bTree {
  112. return bNewOptions(bOptions{})
  113. }
  114. func bNewOptions(opts bOptions) *bTree {
  115. tr := new(bTree)
  116. tr.cow = new(cow)
  117. tr.mu = new(sync.RWMutex)
  118. tr.ctx = opts.Context
  119. tr.locks = !opts.NoLocks
  120. return tr
  121. }
  122. // Less is a convenience function that performs a comparison of two items
  123. // using the same "less" function provided to New.
  124. func (tr *bTree) Less(a, b kind) bool {
  125. return less(a, b, tr.ctx)
  126. }
  127. func (tr *bTree) find(n *node, key kind,
  128. hint *bPathHint, depth int,
  129. ) (index int, found bool) {
  130. if hint == nil {
  131. // fast path for no hinting
  132. low := 0
  133. high := len(n.items)
  134. for low < high {
  135. mid := (low + high) / 2
  136. if !tr.Less(key, n.items[mid]) {
  137. low = mid + 1
  138. } else {
  139. high = mid
  140. }
  141. }
  142. if low > 0 && !tr.Less(n.items[low-1], key) {
  143. return low - 1, true
  144. }
  145. return low, false
  146. }
  147. // Try using hint.
  148. // Best case finds the exact match, updates the hint and returns.
  149. // Worst case, updates the low and high bounds to binary search between.
  150. low := 0
  151. high := len(n.items) - 1
  152. if depth < 8 && hint.used[depth] {
  153. index = int(hint.path[depth])
  154. if index >= len(n.items) {
  155. // tail item
  156. if tr.Less(n.items[len(n.items)-1], key) {
  157. index = len(n.items)
  158. goto path_match
  159. }
  160. index = len(n.items) - 1
  161. }
  162. if tr.Less(key, n.items[index]) {
  163. if index == 0 || tr.Less(n.items[index-1], key) {
  164. goto path_match
  165. }
  166. high = index - 1
  167. } else if tr.Less(n.items[index], key) {
  168. low = index + 1
  169. } else {
  170. found = true
  171. goto path_match
  172. }
  173. }
  174. // Do a binary search between low and high
  175. // keep on going until low > high, where the guarantee on low is that
  176. // key >= items[low - 1]
  177. for low <= high {
  178. mid := low + ((high+1)-low)/2
  179. // if key >= n.items[mid], low = mid + 1
  180. // which implies that key >= everything below low
  181. if !tr.Less(key, n.items[mid]) {
  182. low = mid + 1
  183. } else {
  184. high = mid - 1
  185. }
  186. }
  187. // if low > 0, n.items[low - 1] >= key,
  188. // we have from before that key >= n.items[low - 1]
  189. // therefore key = n.items[low - 1],
  190. // and we have found the entry for key.
  191. // Otherwise we must keep searching for the key in index `low`.
  192. if low > 0 && !tr.Less(n.items[low-1], key) {
  193. index = low - 1
  194. found = true
  195. } else {
  196. index = low
  197. found = false
  198. }
  199. path_match:
  200. if depth < 8 {
  201. hint.used[depth] = true
  202. var pathIndex uint8
  203. if n.leaf() && found {
  204. pathIndex = uint8(index + 1)
  205. } else {
  206. pathIndex = uint8(index)
  207. }
  208. if pathIndex != hint.path[depth] {
  209. hint.path[depth] = pathIndex
  210. for i := depth + 1; i < 8; i++ {
  211. hint.used[i] = false
  212. }
  213. }
  214. }
  215. return index, found
  216. }
  217. // SetHint sets or replace a value for a key using a path hint
  218. func (tr *bTree) SetHint(item kind, hint *bPathHint) (prev kind, replaced bool) {
  219. if tr.lock() {
  220. defer tr.unlock()
  221. }
  222. return tr.setHint(item, hint)
  223. }
  224. func (tr *bTree) setHint(item kind, hint *bPathHint) (prev kind, replaced bool) {
  225. if tr.root == nil {
  226. tr.root = tr.newNode(true)
  227. tr.root.items = append([]kind{}, item)
  228. tr.root.count = 1
  229. tr.count = 1
  230. return tr.empty, false
  231. }
  232. prev, replaced, split := tr.nodeSet(&tr.root, item, hint, 0)
  233. if split {
  234. left := tr.cowLoad(&tr.root)
  235. right, median := tr.nodeSplit(left)
  236. tr.root = tr.newNode(false)
  237. *tr.root.children = make([]*node, 0, maxItems+1)
  238. *tr.root.children = append([]*node{}, left, right)
  239. tr.root.items = append([]kind{}, median)
  240. tr.root.updateCount()
  241. return tr.setHint(item, hint)
  242. }
  243. if replaced {
  244. return prev, true
  245. }
  246. tr.count++
  247. return tr.empty, false
  248. }
  249. // Set or replace a value for a key
  250. func (tr *bTree) Set(item kind) (kind, bool) {
  251. return tr.SetHint(item, nil)
  252. }
  253. func (tr *bTree) nodeSplit(n *node) (right *node, median kind) {
  254. i := maxItems / 2
  255. median = n.items[i]
  256. // left node
  257. left := tr.newNode(n.leaf())
  258. left.items = make([]kind, len(n.items[:i]), maxItems/2)
  259. copy(left.items, n.items[:i])
  260. if !n.leaf() {
  261. *left.children = make([]*node, len((*n.children)[:i+1]), maxItems+1)
  262. copy(*left.children, (*n.children)[:i+1])
  263. }
  264. left.updateCount()
  265. // right node
  266. right = tr.newNode(n.leaf())
  267. right.items = make([]kind, len(n.items[i+1:]), maxItems/2)
  268. copy(right.items, n.items[i+1:])
  269. if !n.leaf() {
  270. *right.children = make([]*node, len((*n.children)[i+1:]), maxItems+1)
  271. copy(*right.children, (*n.children)[i+1:])
  272. }
  273. right.updateCount()
  274. *n = *left
  275. return right, median
  276. }
  277. func (n *node) updateCount() {
  278. n.count = len(n.items)
  279. if !n.leaf() {
  280. for i := 0; i < len(*n.children); i++ {
  281. n.count += (*n.children)[i].count
  282. }
  283. }
  284. }
  285. // This operation should not be inlined because it's expensive and rarely
  286. // called outside of heavy copy-on-write situations. Marking it "noinline"
  287. // allows for the parent cowLoad to be inlined.
  288. // go:noinline
  289. func (tr *bTree) copy(n *node) *node {
  290. n2 := new(node)
  291. n2.cow = tr.cow
  292. n2.count = n.count
  293. n2.items = make([]kind, len(n.items), cap(n.items))
  294. copy(n2.items, n.items)
  295. if !n.leaf() {
  296. n2.children = new([]*node)
  297. *n2.children = make([]*node, len(*n.children), maxItems+1)
  298. copy(*n2.children, *n.children)
  299. }
  300. return n2
  301. }
  302. // cowLoad loads the provided node and, if needed, performs a copy-on-write.
  303. func (tr *bTree) cowLoad(cn **node) *node {
  304. if (*cn).cow != tr.cow {
  305. *cn = tr.copy(*cn)
  306. }
  307. return *cn
  308. }
  309. func (tr *bTree) nodeSet(cn **node, item kind,
  310. hint *bPathHint, depth int,
  311. ) (prev kind, replaced bool, split bool) {
  312. n := tr.cowLoad(cn)
  313. i, found := tr.find(n, item, hint, depth)
  314. if found {
  315. prev = n.items[i]
  316. n.items[i] = item
  317. return prev, true, false
  318. }
  319. if n.leaf() {
  320. if len(n.items) == maxItems {
  321. return tr.empty, false, true
  322. }
  323. n.items = append(n.items, tr.empty)
  324. copy(n.items[i+1:], n.items[i:])
  325. n.items[i] = item
  326. n.count++
  327. return tr.empty, false, false
  328. }
  329. prev, replaced, split = tr.nodeSet(&(*n.children)[i], item, hint, depth+1)
  330. if split {
  331. if len(n.items) == maxItems {
  332. return tr.empty, false, true
  333. }
  334. right, median := tr.nodeSplit((*n.children)[i])
  335. *n.children = append(*n.children, nil)
  336. copy((*n.children)[i+1:], (*n.children)[i:])
  337. (*n.children)[i+1] = right
  338. n.items = append(n.items, tr.empty)
  339. copy(n.items[i+1:], n.items[i:])
  340. n.items[i] = median
  341. return tr.nodeSet(&n, item, hint, depth)
  342. }
  343. if !replaced {
  344. n.count++
  345. }
  346. return prev, replaced, false
  347. }
  348. func (tr *bTree) Scan(iter func(item kind) bool) {
  349. if tr.rlock() {
  350. defer tr.runlock()
  351. }
  352. if tr.root == nil {
  353. return
  354. }
  355. tr.root.scan(iter)
  356. }
  357. func (n *node) scan(iter func(item kind) bool) bool {
  358. if n.leaf() {
  359. for i := 0; i < len(n.items); i++ {
  360. if !iter(n.items[i]) {
  361. return false
  362. }
  363. }
  364. return true
  365. }
  366. for i := 0; i < len(n.items); i++ {
  367. if !(*n.children)[i].scan(iter) {
  368. return false
  369. }
  370. if !iter(n.items[i]) {
  371. return false
  372. }
  373. }
  374. return (*n.children)[len(*n.children)-1].scan(iter)
  375. }
  376. // Get a value for key
  377. func (tr *bTree) Get(key kind) (kind, bool) {
  378. return tr.GetHint(key, nil)
  379. }
  380. // GetHint gets a value for key using a path hint
  381. func (tr *bTree) GetHint(key kind, hint *bPathHint) (kind, bool) {
  382. if tr.rlock() {
  383. defer tr.runlock()
  384. }
  385. if tr.root == nil {
  386. return tr.empty, false
  387. }
  388. n := tr.root
  389. depth := 0
  390. for {
  391. i, found := tr.find(n, key, hint, depth)
  392. if found {
  393. return n.items[i], true
  394. }
  395. if n.children == nil {
  396. return tr.empty, false
  397. }
  398. n = (*n.children)[i]
  399. depth++
  400. }
  401. }
  402. // Len returns the number of items in the tree
  403. func (tr *bTree) Len() int {
  404. return tr.count
  405. }
  406. // Delete a value for a key
  407. func (tr *bTree) Delete(key kind) (kind, bool) {
  408. return tr.DeleteHint(key, nil)
  409. }
  410. // DeleteHint deletes a value for a key using a path hint
  411. func (tr *bTree) DeleteHint(key kind, hint *bPathHint) (kind, bool) {
  412. if tr.lock() {
  413. defer tr.unlock()
  414. }
  415. return tr.deleteHint(key, hint)
  416. }
  417. func (tr *bTree) deleteHint(key kind, hint *bPathHint) (kind, bool) {
  418. if tr.root == nil {
  419. return tr.empty, false
  420. }
  421. prev, deleted := tr.delete(&tr.root, false, key, hint, 0)
  422. if !deleted {
  423. return tr.empty, false
  424. }
  425. if len(tr.root.items) == 0 && !tr.root.leaf() {
  426. tr.root = (*tr.root.children)[0]
  427. }
  428. tr.count--
  429. if tr.count == 0 {
  430. tr.root = nil
  431. }
  432. return prev, true
  433. }
  434. func (tr *bTree) delete(cn **node, max bool, key kind,
  435. hint *bPathHint, depth int,
  436. ) (kind, bool) {
  437. n := tr.cowLoad(cn)
  438. var i int
  439. var found bool
  440. if max {
  441. i, found = len(n.items)-1, true
  442. } else {
  443. i, found = tr.find(n, key, hint, depth)
  444. }
  445. if n.leaf() {
  446. if found {
  447. // found the items at the leaf, remove it and return.
  448. prev := n.items[i]
  449. copy(n.items[i:], n.items[i+1:])
  450. n.items[len(n.items)-1] = tr.empty
  451. n.items = n.items[:len(n.items)-1]
  452. n.count--
  453. return prev, true
  454. }
  455. return tr.empty, false
  456. }
  457. var prev kind
  458. var deleted bool
  459. if found {
  460. if max {
  461. i++
  462. prev, deleted = tr.delete(&(*n.children)[i], true, tr.empty, nil, 0)
  463. } else {
  464. prev = n.items[i]
  465. maxItem, _ := tr.delete(&(*n.children)[i], true, tr.empty, nil, 0)
  466. deleted = true
  467. n.items[i] = maxItem
  468. }
  469. } else {
  470. prev, deleted = tr.delete(&(*n.children)[i], max, key, hint, depth+1)
  471. }
  472. if !deleted {
  473. return tr.empty, false
  474. }
  475. n.count--
  476. if len((*n.children)[i].items) < minItems {
  477. tr.nodeRebalance(n, i)
  478. }
  479. return prev, true
  480. }
  481. // nodeRebalance rebalances the child nodes following a delete operation.
  482. // Provide the index of the child node with the number of items that fell
  483. // below minItems.
  484. func (tr *bTree) nodeRebalance(n *node, i int) {
  485. if i == len(n.items) {
  486. i--
  487. }
  488. // ensure copy-on-write
  489. left := tr.cowLoad(&(*n.children)[i])
  490. right := tr.cowLoad(&(*n.children)[i+1])
  491. if len(left.items)+len(right.items) < maxItems {
  492. // Merges the left and right children nodes together as a single node
  493. // that includes (left,item,right), and places the contents into the
  494. // existing left node. Delete the right node altogether and move the
  495. // following items and child nodes to the left by one slot.
  496. // merge (left,item,right)
  497. left.items = append(left.items, n.items[i])
  498. left.items = append(left.items, right.items...)
  499. if !left.leaf() {
  500. *left.children = append(*left.children, *right.children...)
  501. }
  502. left.count += right.count + 1
  503. // move the items over one slot
  504. copy(n.items[i:], n.items[i+1:])
  505. n.items[len(n.items)-1] = tr.empty
  506. n.items = n.items[:len(n.items)-1]
  507. // move the children over one slot
  508. copy((*n.children)[i+1:], (*n.children)[i+2:])
  509. (*n.children)[len(*n.children)-1] = nil
  510. (*n.children) = (*n.children)[:len(*n.children)-1]
  511. } else if len(left.items) > len(right.items) {
  512. // move left -> right over one slot
  513. // Move the item of the parent node at index into the right-node first
  514. // slot, and move the left-node last item into the previously moved
  515. // parent item slot.
  516. right.items = append(right.items, tr.empty)
  517. copy(right.items[1:], right.items)
  518. right.items[0] = n.items[i]
  519. right.count++
  520. n.items[i] = left.items[len(left.items)-1]
  521. left.items[len(left.items)-1] = tr.empty
  522. left.items = left.items[:len(left.items)-1]
  523. left.count--
  524. if !left.leaf() {
  525. // move the left-node last child into the right-node first slot
  526. *right.children = append(*right.children, nil)
  527. copy((*right.children)[1:], *right.children)
  528. (*right.children)[0] = (*left.children)[len(*left.children)-1]
  529. (*left.children)[len(*left.children)-1] = nil
  530. (*left.children) = (*left.children)[:len(*left.children)-1]
  531. left.count -= (*right.children)[0].count
  532. right.count += (*right.children)[0].count
  533. }
  534. } else {
  535. // move left <- right over one slot
  536. // Same as above but the other direction
  537. left.items = append(left.items, n.items[i])
  538. left.count++
  539. n.items[i] = right.items[0]
  540. copy(right.items, right.items[1:])
  541. right.items[len(right.items)-1] = tr.empty
  542. right.items = right.items[:len(right.items)-1]
  543. right.count--
  544. if !left.leaf() {
  545. *left.children = append(*left.children, (*right.children)[0])
  546. copy(*right.children, (*right.children)[1:])
  547. (*right.children)[len(*right.children)-1] = nil
  548. *right.children = (*right.children)[:len(*right.children)-1]
  549. left.count += (*left.children)[len(*left.children)-1].count
  550. right.count -= (*left.children)[len(*left.children)-1].count
  551. }
  552. }
  553. }
  554. // Ascend the tree within the range [pivot, last]
  555. // Pass nil for pivot to scan all item in ascending order
  556. // Return false to stop iterating
  557. func (tr *bTree) Ascend(pivot kind, iter func(item kind) bool) {
  558. if tr.rlock() {
  559. defer tr.runlock()
  560. }
  561. if tr.root == nil {
  562. return
  563. }
  564. tr.ascend(tr.root, pivot, nil, 0, iter)
  565. }
  566. // The return value of this function determines whether we should keep iterating
  567. // upon this functions return.
  568. func (tr *bTree) ascend(n *node, pivot kind,
  569. hint *bPathHint, depth int, iter func(item kind) bool,
  570. ) bool {
  571. i, found := tr.find(n, pivot, hint, depth)
  572. if !found {
  573. if !n.leaf() {
  574. if !tr.ascend((*n.children)[i], pivot, hint, depth+1, iter) {
  575. return false
  576. }
  577. }
  578. }
  579. // We are either in the case that
  580. // - node is found, we should iterate through it starting at `i`,
  581. // the index it was located at.
  582. // - node is not found, and TODO: fill in.
  583. for ; i < len(n.items); i++ {
  584. if !iter(n.items[i]) {
  585. return false
  586. }
  587. if !n.leaf() {
  588. if !(*n.children)[i+1].scan(iter) {
  589. return false
  590. }
  591. }
  592. }
  593. return true
  594. }
  595. func (tr *bTree) Reverse(iter func(item kind) bool) {
  596. if tr.rlock() {
  597. defer tr.runlock()
  598. }
  599. if tr.root == nil {
  600. return
  601. }
  602. tr.root.reverse(iter)
  603. }
  604. func (n *node) reverse(iter func(item kind) bool) bool {
  605. if n.leaf() {
  606. for i := len(n.items) - 1; i >= 0; i-- {
  607. if !iter(n.items[i]) {
  608. return false
  609. }
  610. }
  611. return true
  612. }
  613. if !(*n.children)[len(*n.children)-1].reverse(iter) {
  614. return false
  615. }
  616. for i := len(n.items) - 1; i >= 0; i-- {
  617. if !iter(n.items[i]) {
  618. return false
  619. }
  620. if !(*n.children)[i].reverse(iter) {
  621. return false
  622. }
  623. }
  624. return true
  625. }
  626. // Descend the tree within the range [pivot, first]
  627. // Pass nil for pivot to scan all item in descending order
  628. // Return false to stop iterating
  629. func (tr *bTree) Descend(pivot kind, iter func(item kind) bool) {
  630. if tr.rlock() {
  631. defer tr.runlock()
  632. }
  633. if tr.root == nil {
  634. return
  635. }
  636. tr.descend(tr.root, pivot, nil, 0, iter)
  637. }
  638. func (tr *bTree) descend(n *node, pivot kind,
  639. hint *bPathHint, depth int, iter func(item kind) bool,
  640. ) bool {
  641. i, found := tr.find(n, pivot, hint, depth)
  642. if !found {
  643. if !n.leaf() {
  644. if !tr.descend((*n.children)[i], pivot, hint, depth+1, iter) {
  645. return false
  646. }
  647. }
  648. i--
  649. }
  650. for ; i >= 0; i-- {
  651. if !iter(n.items[i]) {
  652. return false
  653. }
  654. if !n.leaf() {
  655. if !(*n.children)[i].reverse(iter) {
  656. return false
  657. }
  658. }
  659. }
  660. return true
  661. }
  662. // Load is for bulk loading pre-sorted items
  663. func (tr *bTree) Load(item kind) (kind, bool) {
  664. if tr.lock() {
  665. defer tr.unlock()
  666. }
  667. if tr.root == nil {
  668. return tr.setHint(item, nil)
  669. }
  670. n := tr.cowLoad(&tr.root)
  671. for {
  672. n.count++ // optimistically update counts
  673. if n.leaf() {
  674. if len(n.items) < maxItems {
  675. if tr.Less(n.items[len(n.items)-1], item) {
  676. n.items = append(n.items, item)
  677. tr.count++
  678. return tr.empty, false
  679. }
  680. }
  681. break
  682. }
  683. n = tr.cowLoad(&(*n.children)[len(*n.children)-1])
  684. }
  685. // revert the counts
  686. n = tr.root
  687. for {
  688. n.count--
  689. if n.leaf() {
  690. break
  691. }
  692. n = (*n.children)[len(*n.children)-1]
  693. }
  694. return tr.setHint(item, nil)
  695. }
  696. // Min returns the minimum item in tree.
  697. // Returns nil if the tree has no items.
  698. func (tr *bTree) Min() (kind, bool) {
  699. if tr.rlock() {
  700. defer tr.runlock()
  701. }
  702. if tr.root == nil {
  703. return tr.empty, false
  704. }
  705. n := tr.root
  706. for {
  707. if n.leaf() {
  708. return n.items[0], true
  709. }
  710. n = (*n.children)[0]
  711. }
  712. }
  713. // Max returns the maximum item in tree.
  714. // Returns nil if the tree has no items.
  715. func (tr *bTree) Max() (kind, bool) {
  716. if tr.rlock() {
  717. defer tr.runlock()
  718. }
  719. if tr.root == nil {
  720. return tr.empty, false
  721. }
  722. n := tr.root
  723. for {
  724. if n.leaf() {
  725. return n.items[len(n.items)-1], true
  726. }
  727. n = (*n.children)[len(*n.children)-1]
  728. }
  729. }
  730. // PopMin removes the minimum item in tree and returns it.
  731. // Returns nil if the tree has no items.
  732. func (tr *bTree) PopMin() (kind, bool) {
  733. if tr.lock() {
  734. defer tr.unlock()
  735. }
  736. if tr.root == nil {
  737. return tr.empty, false
  738. }
  739. n := tr.cowLoad(&tr.root)
  740. var item kind
  741. for {
  742. n.count-- // optimistically update counts
  743. if n.leaf() {
  744. item = n.items[0]
  745. if len(n.items) == minItems {
  746. break
  747. }
  748. copy(n.items[:], n.items[1:])
  749. n.items[len(n.items)-1] = tr.empty
  750. n.items = n.items[:len(n.items)-1]
  751. tr.count--
  752. if tr.count == 0 {
  753. tr.root = nil
  754. }
  755. return item, true
  756. }
  757. n = tr.cowLoad(&(*n.children)[0])
  758. }
  759. // revert the counts
  760. n = tr.root
  761. for {
  762. n.count++
  763. if n.leaf() {
  764. break
  765. }
  766. n = (*n.children)[0]
  767. }
  768. return tr.deleteHint(item, nil)
  769. }
  770. // PopMax removes the minimum item in tree and returns it.
  771. // Returns nil if the tree has no items.
  772. func (tr *bTree) PopMax() (kind, bool) {
  773. if tr.lock() {
  774. defer tr.unlock()
  775. }
  776. if tr.root == nil {
  777. return tr.empty, false
  778. }
  779. n := tr.cowLoad(&tr.root)
  780. var item kind
  781. for {
  782. n.count-- // optimistically update counts
  783. if n.leaf() {
  784. item = n.items[len(n.items)-1]
  785. if len(n.items) == minItems {
  786. break
  787. }
  788. n.items[len(n.items)-1] = tr.empty
  789. n.items = n.items[:len(n.items)-1]
  790. tr.count--
  791. if tr.count == 0 {
  792. tr.root = nil
  793. }
  794. return item, true
  795. }
  796. n = tr.cowLoad(&(*n.children)[len(*n.children)-1])
  797. }
  798. // revert the counts
  799. n = tr.root
  800. for {
  801. n.count++
  802. if n.leaf() {
  803. break
  804. }
  805. n = (*n.children)[len(*n.children)-1]
  806. }
  807. return tr.deleteHint(item, nil)
  808. }
  809. // GetAt returns the value at index.
  810. // Return nil if the tree is empty or the index is out of bounds.
  811. func (tr *bTree) GetAt(index int) (kind, bool) {
  812. if tr.rlock() {
  813. defer tr.runlock()
  814. }
  815. if tr.root == nil || index < 0 || index >= tr.count {
  816. return tr.empty, false
  817. }
  818. n := tr.root
  819. for {
  820. if n.leaf() {
  821. return n.items[index], true
  822. }
  823. i := 0
  824. for ; i < len(n.items); i++ {
  825. if index < (*n.children)[i].count {
  826. break
  827. } else if index == (*n.children)[i].count {
  828. return n.items[i], true
  829. }
  830. index -= (*n.children)[i].count + 1
  831. }
  832. n = (*n.children)[i]
  833. }
  834. }
  835. // DeleteAt deletes the item at index.
  836. // Return nil if the tree is empty or the index is out of bounds.
  837. func (tr *bTree) DeleteAt(index int) (kind, bool) {
  838. if tr.lock() {
  839. defer tr.unlock()
  840. }
  841. if tr.root == nil || index < 0 || index >= tr.count {
  842. return tr.empty, false
  843. }
  844. var pathbuf [8]uint8 // track the path
  845. path := pathbuf[:0]
  846. var item kind
  847. n := tr.cowLoad(&tr.root)
  848. outer:
  849. for {
  850. n.count-- // optimistically update counts
  851. if n.leaf() {
  852. // the index is the item position
  853. item = n.items[index]
  854. if len(n.items) == minItems {
  855. path = append(path, uint8(index))
  856. break outer
  857. }
  858. copy(n.items[index:], n.items[index+1:])
  859. n.items[len(n.items)-1] = tr.empty
  860. n.items = n.items[:len(n.items)-1]
  861. tr.count--
  862. if tr.count == 0 {
  863. tr.root = nil
  864. }
  865. return item, true
  866. }
  867. i := 0
  868. for ; i < len(n.items); i++ {
  869. if index < (*n.children)[i].count {
  870. break
  871. } else if index == (*n.children)[i].count {
  872. item = n.items[i]
  873. path = append(path, uint8(i))
  874. break outer
  875. }
  876. index -= (*n.children)[i].count + 1
  877. }
  878. path = append(path, uint8(i))
  879. n = tr.cowLoad(&(*n.children)[i])
  880. }
  881. // revert the counts
  882. var hint bPathHint
  883. n = tr.root
  884. for i := 0; i < len(path); i++ {
  885. if i < len(hint.path) {
  886. hint.path[i] = uint8(path[i])
  887. hint.used[i] = true
  888. }
  889. n.count++
  890. if !n.leaf() {
  891. n = (*n.children)[uint8(path[i])]
  892. }
  893. }
  894. return tr.deleteHint(item, &hint)
  895. }
  896. // Height returns the height of the tree.
  897. // Returns zero if tree has no items.
  898. func (tr *bTree) Height() int {
  899. if tr.rlock() {
  900. defer tr.runlock()
  901. }
  902. var height int
  903. if tr.root != nil {
  904. n := tr.root
  905. for {
  906. height++
  907. if n.leaf() {
  908. break
  909. }
  910. n = (*n.children)[0]
  911. }
  912. }
  913. return height
  914. }
  915. // Walk iterates over all items in tree, in order.
  916. // The items param will contain one or more items.
  917. func (tr *bTree) Walk(iter func(item []kind) bool) {
  918. if tr.rlock() {
  919. defer tr.runlock()
  920. }
  921. if tr.root != nil {
  922. tr.root.walk(iter)
  923. }
  924. }
  925. func (n *node) walk(iter func(item []kind) bool) bool {
  926. if n.leaf() {
  927. if !iter(n.items) {
  928. return false
  929. }
  930. } else {
  931. for i := 0; i < len(n.items); i++ {
  932. (*n.children)[i].walk(iter)
  933. if !iter(n.items[i : i+1]) {
  934. return false
  935. }
  936. }
  937. (*n.children)[len(n.items)].walk(iter)
  938. }
  939. return true
  940. }
  941. // Copy the tree. This is a copy-on-write operation and is very fast because
  942. // it only performs a shadowed copy.
  943. func (tr *bTree) Copy() *bTree {
  944. if tr.lock() {
  945. defer tr.unlock()
  946. }
  947. tr.cow = new(cow)
  948. tr2 := new(bTree)
  949. *tr2 = *tr
  950. tr2.mu = new(sync.RWMutex)
  951. tr2.cow = new(cow)
  952. return tr2
  953. }
  954. func (tr *bTree) lock() bool {
  955. if tr.locks {
  956. tr.mu.Lock()
  957. }
  958. return tr.locks
  959. }
  960. func (tr *bTree) unlock() {
  961. tr.mu.Unlock()
  962. }
  963. func (tr *bTree) rlock() bool {
  964. if tr.locks {
  965. tr.mu.RLock()
  966. }
  967. return tr.locks
  968. }
  969. func (tr *bTree) runlock() {
  970. tr.mu.RUnlock()
  971. }
  972. // Iter represents an iterator
  973. type bIter struct {
  974. tr *bTree
  975. locked bool
  976. seeked bool
  977. atstart bool
  978. atend bool
  979. stack []iterStackItem
  980. item kind
  981. }
  982. type iterStackItem struct {
  983. n *node
  984. i int
  985. }
  986. // Iter returns a read-only iterator.
  987. // The Release method must be called finished with iterator.
  988. func (tr *bTree) Iter() bIter {
  989. var iter bIter
  990. iter.tr = tr
  991. iter.locked = tr.rlock()
  992. return iter
  993. }
  994. // Seek to item greater-or-equal-to key.
  995. // Returns false if there was no item found.
  996. func (iter *bIter) Seek(key kind) bool {
  997. if iter.tr == nil {
  998. return false
  999. }
  1000. iter.seeked = true
  1001. iter.stack = iter.stack[:0]
  1002. if iter.tr.root == nil {
  1003. return false
  1004. }
  1005. n := iter.tr.root
  1006. for {
  1007. i, found := iter.tr.find(n, key, nil, 0)
  1008. iter.stack = append(iter.stack, iterStackItem{n, i})
  1009. if found {
  1010. return true
  1011. }
  1012. if n.leaf() {
  1013. if i == len(n.items) {
  1014. iter.stack = iter.stack[:0]
  1015. return false
  1016. }
  1017. return true
  1018. }
  1019. n = (*n.children)[i]
  1020. }
  1021. }
  1022. // First moves iterator to first item in tree.
  1023. // Returns false if the tree is empty.
  1024. func (iter *bIter) First() bool {
  1025. if iter.tr == nil {
  1026. return false
  1027. }
  1028. iter.atend = false
  1029. iter.atstart = false
  1030. iter.seeked = true
  1031. iter.stack = iter.stack[:0]
  1032. if iter.tr.root == nil {
  1033. return false
  1034. }
  1035. n := iter.tr.root
  1036. for {
  1037. iter.stack = append(iter.stack, iterStackItem{n, 0})
  1038. if n.leaf() {
  1039. break
  1040. }
  1041. n = (*n.children)[0]
  1042. }
  1043. s := &iter.stack[len(iter.stack)-1]
  1044. iter.item = s.n.items[s.i]
  1045. return true
  1046. }
  1047. // Last moves iterator to last item in tree.
  1048. // Returns false if the tree is empty.
  1049. func (iter *bIter) Last() bool {
  1050. if iter.tr == nil {
  1051. return false
  1052. }
  1053. iter.seeked = true
  1054. iter.stack = iter.stack[:0]
  1055. if iter.tr.root == nil {
  1056. return false
  1057. }
  1058. n := iter.tr.root
  1059. for {
  1060. iter.stack = append(iter.stack, iterStackItem{n, len(n.items)})
  1061. if n.leaf() {
  1062. iter.stack[len(iter.stack)-1].i--
  1063. break
  1064. }
  1065. n = (*n.children)[len(n.items)]
  1066. }
  1067. s := &iter.stack[len(iter.stack)-1]
  1068. iter.item = s.n.items[s.i]
  1069. return true
  1070. }
  1071. // First moves iterator to first item in tree.
  1072. // Returns false if the tree is empty.
  1073. func (iter *bIter) Release() {
  1074. if iter.tr == nil {
  1075. return
  1076. }
  1077. if iter.locked {
  1078. iter.tr.runlock()
  1079. iter.locked = false
  1080. }
  1081. iter.stack = nil
  1082. iter.tr = nil
  1083. }
  1084. // Next moves iterator to the next item in iterator.
  1085. // Returns false if the tree is empty or the iterator is at the end of
  1086. // the tree.
  1087. func (iter *bIter) Next() bool {
  1088. if iter.tr == nil {
  1089. return false
  1090. }
  1091. if !iter.seeked {
  1092. return iter.First()
  1093. }
  1094. if len(iter.stack) == 0 {
  1095. if iter.atstart {
  1096. return iter.First() && iter.Next()
  1097. }
  1098. return false
  1099. }
  1100. s := &iter.stack[len(iter.stack)-1]
  1101. s.i++
  1102. if s.n.leaf() {
  1103. if s.i == len(s.n.items) {
  1104. for {
  1105. iter.stack = iter.stack[:len(iter.stack)-1]
  1106. if len(iter.stack) == 0 {
  1107. iter.atend = true
  1108. return false
  1109. }
  1110. s = &iter.stack[len(iter.stack)-1]
  1111. if s.i < len(s.n.items) {
  1112. break
  1113. }
  1114. }
  1115. }
  1116. } else {
  1117. n := (*s.n.children)[s.i]
  1118. for {
  1119. iter.stack = append(iter.stack, iterStackItem{n, 0})
  1120. if n.leaf() {
  1121. break
  1122. }
  1123. n = (*n.children)[0]
  1124. }
  1125. }
  1126. s = &iter.stack[len(iter.stack)-1]
  1127. iter.item = s.n.items[s.i]
  1128. return true
  1129. }
  1130. // Prev moves iterator to the previous item in iterator.
  1131. // Returns false if the tree is empty or the iterator is at the beginning of
  1132. // the tree.
  1133. func (iter *bIter) Prev() bool {
  1134. if iter.tr == nil {
  1135. return false
  1136. }
  1137. if !iter.seeked {
  1138. return false
  1139. }
  1140. if len(iter.stack) == 0 {
  1141. if iter.atend {
  1142. return iter.Last() && iter.Prev()
  1143. }
  1144. return false
  1145. }
  1146. s := &iter.stack[len(iter.stack)-1]
  1147. if s.n.leaf() {
  1148. s.i--
  1149. if s.i == -1 {
  1150. for {
  1151. iter.stack = iter.stack[:len(iter.stack)-1]
  1152. if len(iter.stack) == 0 {
  1153. iter.atstart = true
  1154. return false
  1155. }
  1156. s = &iter.stack[len(iter.stack)-1]
  1157. s.i--
  1158. if s.i > -1 {
  1159. break
  1160. }
  1161. }
  1162. }
  1163. } else {
  1164. n := (*s.n.children)[s.i]
  1165. for {
  1166. iter.stack = append(iter.stack, iterStackItem{n, len(n.items)})
  1167. if n.leaf() {
  1168. iter.stack[len(iter.stack)-1].i--
  1169. break
  1170. }
  1171. n = (*n.children)[len(n.items)]
  1172. }
  1173. }
  1174. s = &iter.stack[len(iter.stack)-1]
  1175. iter.item = s.n.items[s.i]
  1176. return true
  1177. }
  1178. // Item returns the current iterator item.
  1179. func (iter *bIter) Item() kind {
  1180. return iter.item
  1181. }