generic.go 28 KB

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