map.go 24 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056
  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 "sync/atomic"
  6. type ordered interface {
  7. ~int | ~int8 | ~int16 | ~int32 | ~int64 |
  8. ~uint | ~uint8 | ~uint16 | ~uint32 | ~uint64 | ~uintptr |
  9. ~float32 | ~float64 | ~string
  10. }
  11. type mapPair[K ordered, V any] struct {
  12. // The `value` field should be before the `key` field because doing so
  13. // allows for the Go compiler to optimize away the `value` field when
  14. // it's a `struct{}`, which is the case for `btree.Set`.
  15. value V
  16. key K
  17. }
  18. type Map[K ordered, V any] struct {
  19. cow uint64
  20. root *mapNode[K, V]
  21. count int
  22. empty mapPair[K, V]
  23. }
  24. type mapNode[K ordered, V any] struct {
  25. cow uint64
  26. count int
  27. items []mapPair[K, V]
  28. children *[]*mapNode[K, V]
  29. }
  30. // This operation should not be inlined because it's expensive and rarely
  31. // called outside of heavy copy-on-write situations. Marking it "noinline"
  32. // allows for the parent cowLoad to be inlined.
  33. // go:noinline
  34. func (tr *Map[K, V]) copy(n *mapNode[K, V]) *mapNode[K, V] {
  35. n2 := new(mapNode[K, V])
  36. n2.cow = tr.cow
  37. n2.count = n.count
  38. n2.items = make([]mapPair[K, V], len(n.items), cap(n.items))
  39. copy(n2.items, n.items)
  40. if !n.leaf() {
  41. n2.children = new([]*mapNode[K, V])
  42. *n2.children = make([]*mapNode[K, V], len(*n.children), maxItems+1)
  43. copy(*n2.children, *n.children)
  44. }
  45. return n2
  46. }
  47. // cowLoad loads the provided node and, if needed, performs a copy-on-write.
  48. func (tr *Map[K, V]) cowLoad(cn **mapNode[K, V]) *mapNode[K, V] {
  49. if (*cn).cow != tr.cow {
  50. *cn = tr.copy(*cn)
  51. }
  52. return *cn
  53. }
  54. func (tr *Map[K, V]) Copy() *Map[K, V] {
  55. tr2 := new(Map[K, V])
  56. *tr2 = *tr
  57. tr2.cow = atomic.AddUint64(&gcow, 1)
  58. tr.cow = atomic.AddUint64(&gcow, 1)
  59. return tr2
  60. }
  61. func (tr *Map[K, V]) newNode(leaf bool) *mapNode[K, V] {
  62. n := new(mapNode[K, V])
  63. n.cow = tr.cow
  64. if !leaf {
  65. n.children = new([]*mapNode[K, V])
  66. }
  67. return n
  68. }
  69. // leaf returns true if the node is a leaf.
  70. func (n *mapNode[K, V]) leaf() bool {
  71. return n.children == nil
  72. }
  73. func (tr *Map[K, V]) bsearch(n *mapNode[K, V], key K) (index int, found bool) {
  74. low, high := 0, len(n.items)
  75. for low < high {
  76. h := int(uint(low+high) >> 1)
  77. if key >= n.items[h].key {
  78. low = h + 1
  79. } else {
  80. high = h
  81. }
  82. }
  83. if low > 0 && n.items[low-1].key >= key {
  84. return low - 1, true
  85. }
  86. return low, false
  87. }
  88. // Set or replace a value for a key
  89. func (tr *Map[K, V]) Set(key K, value V) (V, bool) {
  90. item := mapPair[K, V]{key: key, value: value}
  91. if tr.root == nil {
  92. tr.root = tr.newNode(true)
  93. tr.root.items = append([]mapPair[K, V]{}, item)
  94. tr.root.count = 1
  95. tr.count = 1
  96. return tr.empty.value, false
  97. }
  98. prev, replaced, split := tr.nodeSet(&tr.root, item)
  99. if split {
  100. left := tr.root
  101. right, median := tr.nodeSplit(left)
  102. tr.root = tr.newNode(false)
  103. *tr.root.children = make([]*mapNode[K, V], 0, maxItems+1)
  104. *tr.root.children = append([]*mapNode[K, V]{}, left, right)
  105. tr.root.items = append([]mapPair[K, V]{}, median)
  106. tr.root.updateCount()
  107. return tr.Set(item.key, item.value)
  108. }
  109. if replaced {
  110. return prev, true
  111. }
  112. tr.count++
  113. return tr.empty.value, false
  114. }
  115. func (tr *Map[K, V]) nodeSplit(n *mapNode[K, V],
  116. ) (right *mapNode[K, V], median mapPair[K, V]) {
  117. i := maxItems / 2
  118. median = n.items[i]
  119. const sliceItems = true
  120. // right node
  121. right = tr.newNode(n.leaf())
  122. if sliceItems {
  123. right.items = n.items[i+1:]
  124. if !n.leaf() {
  125. *right.children = (*n.children)[i+1:]
  126. }
  127. } else {
  128. right.items = make([]mapPair[K, V], len(n.items[i+1:]), maxItems/2)
  129. copy(right.items, n.items[i+1:])
  130. if !n.leaf() {
  131. *right.children = make([]*mapNode[K, V],
  132. len((*n.children)[i+1:]), maxItems+1)
  133. copy(*right.children, (*n.children)[i+1:])
  134. }
  135. }
  136. right.updateCount()
  137. // left node
  138. if sliceItems {
  139. n.items[i] = tr.empty
  140. n.items = n.items[:i:i]
  141. if !n.leaf() {
  142. *n.children = (*n.children)[: i+1 : i+1]
  143. }
  144. } else {
  145. for j := i; j < len(n.items); j++ {
  146. n.items[j] = tr.empty
  147. }
  148. if !n.leaf() {
  149. for j := i + 1; j < len((*n.children)); j++ {
  150. (*n.children)[j] = nil
  151. }
  152. }
  153. n.items = n.items[:i]
  154. if !n.leaf() {
  155. *n.children = (*n.children)[:i+1]
  156. }
  157. }
  158. n.updateCount()
  159. return right, median
  160. }
  161. func (n *mapNode[K, V]) updateCount() {
  162. n.count = len(n.items)
  163. if !n.leaf() {
  164. for i := 0; i < len(*n.children); i++ {
  165. n.count += (*n.children)[i].count
  166. }
  167. }
  168. }
  169. func (tr *Map[K, V]) nodeSet(pn **mapNode[K, V], item mapPair[K, V],
  170. ) (prev V, replaced bool, split bool) {
  171. n := tr.cowLoad(pn)
  172. i, found := tr.bsearch(n, item.key)
  173. if found {
  174. prev = n.items[i].value
  175. n.items[i].value = item.value
  176. return prev, true, false
  177. }
  178. if n.leaf() {
  179. if len(n.items) == maxItems {
  180. return tr.empty.value, false, true
  181. }
  182. n.items = append(n.items, tr.empty)
  183. copy(n.items[i+1:], n.items[i:])
  184. n.items[i] = item
  185. n.count++
  186. return tr.empty.value, false, false
  187. }
  188. prev, replaced, split = tr.nodeSet(&(*n.children)[i], item)
  189. if split {
  190. if len(n.items) == maxItems {
  191. return tr.empty.value, false, true
  192. }
  193. right, median := tr.nodeSplit((*n.children)[i])
  194. *n.children = append(*n.children, nil)
  195. copy((*n.children)[i+1:], (*n.children)[i:])
  196. (*n.children)[i+1] = right
  197. n.items = append(n.items, tr.empty)
  198. copy(n.items[i+1:], n.items[i:])
  199. n.items[i] = median
  200. return tr.nodeSet(&n, item)
  201. }
  202. if !replaced {
  203. n.count++
  204. }
  205. return prev, replaced, false
  206. }
  207. func (tr *Map[K, V]) Scan(iter func(key K, value V) bool) {
  208. if tr.root == nil {
  209. return
  210. }
  211. tr.root.scan(iter)
  212. }
  213. func (n *mapNode[K, V]) scan(iter func(key K, value V) bool) bool {
  214. if n.leaf() {
  215. for i := 0; i < len(n.items); i++ {
  216. if !iter(n.items[i].key, n.items[i].value) {
  217. return false
  218. }
  219. }
  220. return true
  221. }
  222. for i := 0; i < len(n.items); i++ {
  223. if !(*n.children)[i].scan(iter) {
  224. return false
  225. }
  226. if !iter(n.items[i].key, n.items[i].value) {
  227. return false
  228. }
  229. }
  230. return (*n.children)[len(*n.children)-1].scan(iter)
  231. }
  232. // Get a value for key
  233. func (tr *Map[K, V]) Get(key K) (V, bool) {
  234. if tr.root == nil {
  235. return tr.empty.value, false
  236. }
  237. n := tr.root
  238. for {
  239. i, found := tr.bsearch(n, key)
  240. if found {
  241. return n.items[i].value, true
  242. }
  243. if n.leaf() {
  244. return tr.empty.value, false
  245. }
  246. n = (*n.children)[i]
  247. }
  248. }
  249. // Len returns the number of items in the tree
  250. func (tr *Map[K, V]) Len() int {
  251. return tr.count
  252. }
  253. // Delete a value for a key and returns the deleted value.
  254. // Returns false if there was no value by that key found.
  255. func (tr *Map[K, V]) Delete(key K) (V, bool) {
  256. if tr.root == nil {
  257. return tr.empty.value, false
  258. }
  259. prev, deleted := tr.delete(&tr.root, false, key)
  260. if !deleted {
  261. return tr.empty.value, false
  262. }
  263. if len(tr.root.items) == 0 && !tr.root.leaf() {
  264. tr.root = (*tr.root.children)[0]
  265. }
  266. tr.count--
  267. if tr.count == 0 {
  268. tr.root = nil
  269. }
  270. return prev.value, true
  271. }
  272. func (tr *Map[K, V]) delete(pn **mapNode[K, V], max bool, key K,
  273. ) (mapPair[K, V], bool) {
  274. n := tr.cowLoad(pn)
  275. var i int
  276. var found bool
  277. if max {
  278. i, found = len(n.items)-1, true
  279. } else {
  280. i, found = tr.bsearch(n, key)
  281. }
  282. if n.leaf() {
  283. if found {
  284. // found the items at the leaf, remove it and return.
  285. prev := n.items[i]
  286. copy(n.items[i:], n.items[i+1:])
  287. n.items[len(n.items)-1] = tr.empty
  288. n.items = n.items[:len(n.items)-1]
  289. n.count--
  290. return prev, true
  291. }
  292. return tr.empty, false
  293. }
  294. var prev mapPair[K, V]
  295. var deleted bool
  296. if found {
  297. if max {
  298. i++
  299. prev, deleted = tr.delete(&(*n.children)[i], true, tr.empty.key)
  300. } else {
  301. prev = n.items[i]
  302. maxItem, _ := tr.delete(&(*n.children)[i], true, tr.empty.key)
  303. deleted = true
  304. n.items[i] = maxItem
  305. }
  306. } else {
  307. prev, deleted = tr.delete(&(*n.children)[i], max, key)
  308. }
  309. if !deleted {
  310. return tr.empty, false
  311. }
  312. n.count--
  313. if len((*n.children)[i].items) < minItems {
  314. tr.nodeRebalance(n, i)
  315. }
  316. return prev, true
  317. }
  318. // nodeRebalance rebalances the child nodes following a delete operation.
  319. // Provide the index of the child node with the number of items that fell
  320. // below minItems.
  321. func (tr *Map[K, V]) nodeRebalance(n *mapNode[K, V], i int) {
  322. if i == len(n.items) {
  323. i--
  324. }
  325. // ensure copy-on-write
  326. left := tr.cowLoad(&(*n.children)[i])
  327. right := tr.cowLoad(&(*n.children)[i+1])
  328. if len(left.items)+len(right.items) < maxItems {
  329. // Merges the left and right children nodes together as a single node
  330. // that includes (left,item,right), and places the contents into the
  331. // existing left node. Delete the right node altogether and move the
  332. // following items and child nodes to the left by one slot.
  333. // merge (left,item,right)
  334. left.items = append(left.items, n.items[i])
  335. left.items = append(left.items, right.items...)
  336. if !left.leaf() {
  337. *left.children = append(*left.children, *right.children...)
  338. }
  339. left.count += right.count + 1
  340. // move the items over one slot
  341. copy(n.items[i:], n.items[i+1:])
  342. n.items[len(n.items)-1] = tr.empty
  343. n.items = n.items[:len(n.items)-1]
  344. // move the children over one slot
  345. copy((*n.children)[i+1:], (*n.children)[i+2:])
  346. (*n.children)[len(*n.children)-1] = nil
  347. (*n.children) = (*n.children)[:len(*n.children)-1]
  348. } else if len(left.items) > len(right.items) {
  349. // move left -> right over one slot
  350. // Move the item of the parent node at index into the right-node first
  351. // slot, and move the left-node last item into the previously moved
  352. // parent item slot.
  353. right.items = append(right.items, tr.empty)
  354. copy(right.items[1:], right.items)
  355. right.items[0] = n.items[i]
  356. right.count++
  357. n.items[i] = left.items[len(left.items)-1]
  358. left.items[len(left.items)-1] = tr.empty
  359. left.items = left.items[:len(left.items)-1]
  360. left.count--
  361. if !left.leaf() {
  362. // move the left-node last child into the right-node first slot
  363. *right.children = append(*right.children, nil)
  364. copy((*right.children)[1:], *right.children)
  365. (*right.children)[0] = (*left.children)[len(*left.children)-1]
  366. (*left.children)[len(*left.children)-1] = nil
  367. (*left.children) = (*left.children)[:len(*left.children)-1]
  368. left.count -= (*right.children)[0].count
  369. right.count += (*right.children)[0].count
  370. }
  371. } else {
  372. // move left <- right over one slot
  373. // Same as above but the other direction
  374. left.items = append(left.items, n.items[i])
  375. left.count++
  376. n.items[i] = right.items[0]
  377. copy(right.items, right.items[1:])
  378. right.items[len(right.items)-1] = tr.empty
  379. right.items = right.items[:len(right.items)-1]
  380. right.count--
  381. if !left.leaf() {
  382. *left.children = append(*left.children, (*right.children)[0])
  383. copy(*right.children, (*right.children)[1:])
  384. (*right.children)[len(*right.children)-1] = nil
  385. *right.children = (*right.children)[:len(*right.children)-1]
  386. left.count += (*left.children)[len(*left.children)-1].count
  387. right.count -= (*left.children)[len(*left.children)-1].count
  388. }
  389. }
  390. }
  391. // Ascend the tree within the range [pivot, last]
  392. // Pass nil for pivot to scan all item in ascending order
  393. // Return false to stop iterating
  394. func (tr *Map[K, V]) Ascend(pivot K, iter func(key K, value V) bool) {
  395. if tr.root == nil {
  396. return
  397. }
  398. tr.ascend(tr.root, pivot, iter)
  399. }
  400. // The return value of this function determines whether we should keep iterating
  401. // upon this functions return.
  402. func (tr *Map[K, V]) ascend(n *mapNode[K, V], pivot K,
  403. iter func(key K, value V) bool,
  404. ) bool {
  405. i, found := tr.bsearch(n, pivot)
  406. if !found {
  407. if !n.leaf() {
  408. if !tr.ascend((*n.children)[i], pivot, iter) {
  409. return false
  410. }
  411. }
  412. }
  413. // We are either in the case that
  414. // - node is found, we should iterate through it starting at `i`,
  415. // the index it was located at.
  416. // - node is not found, and TODO: fill in.
  417. for ; i < len(n.items); i++ {
  418. if !iter(n.items[i].key, n.items[i].value) {
  419. return false
  420. }
  421. if !n.leaf() {
  422. if !(*n.children)[i+1].scan(iter) {
  423. return false
  424. }
  425. }
  426. }
  427. return true
  428. }
  429. func (tr *Map[K, V]) Reverse(iter func(key K, value V) bool) {
  430. if tr.root == nil {
  431. return
  432. }
  433. tr.root.reverse(iter)
  434. }
  435. func (n *mapNode[K, V]) reverse(iter func(key K, value V) bool) bool {
  436. if n.leaf() {
  437. for i := len(n.items) - 1; i >= 0; i-- {
  438. if !iter(n.items[i].key, n.items[i].value) {
  439. return false
  440. }
  441. }
  442. return true
  443. }
  444. if !(*n.children)[len(*n.children)-1].reverse(iter) {
  445. return false
  446. }
  447. for i := len(n.items) - 1; i >= 0; i-- {
  448. if !iter(n.items[i].key, n.items[i].value) {
  449. return false
  450. }
  451. if !(*n.children)[i].reverse(iter) {
  452. return false
  453. }
  454. }
  455. return true
  456. }
  457. // Descend the tree within the range [pivot, first]
  458. // Pass nil for pivot to scan all item in descending order
  459. // Return false to stop iterating
  460. func (tr *Map[K, V]) Descend(pivot K, iter func(key K, value V) bool) {
  461. if tr.root == nil {
  462. return
  463. }
  464. tr.descend(tr.root, pivot, iter)
  465. }
  466. func (tr *Map[K, V]) descend(n *mapNode[K, V], pivot K,
  467. iter func(key K, value V) bool,
  468. ) bool {
  469. i, found := tr.bsearch(n, pivot)
  470. if !found {
  471. if !n.leaf() {
  472. if !tr.descend((*n.children)[i], pivot, iter) {
  473. return false
  474. }
  475. }
  476. i--
  477. }
  478. for ; i >= 0; i-- {
  479. if !iter(n.items[i].key, n.items[i].value) {
  480. return false
  481. }
  482. if !n.leaf() {
  483. if !(*n.children)[i].reverse(iter) {
  484. return false
  485. }
  486. }
  487. }
  488. return true
  489. }
  490. // Load is for bulk loading pre-sorted items
  491. func (tr *Map[K, V]) Load(key K, value V) (V, bool) {
  492. item := mapPair[K, V]{key: key, value: value}
  493. if tr.root == nil {
  494. return tr.Set(item.key, item.value)
  495. }
  496. n := tr.cowLoad(&tr.root)
  497. for {
  498. n.count++ // optimistically update counts
  499. if n.leaf() {
  500. if len(n.items) < maxItems {
  501. if n.items[len(n.items)-1].key < item.key {
  502. n.items = append(n.items, item)
  503. tr.count++
  504. return tr.empty.value, false
  505. }
  506. }
  507. break
  508. }
  509. n = tr.cowLoad(&(*n.children)[len(*n.children)-1])
  510. }
  511. // revert the counts
  512. n = tr.root
  513. for {
  514. n.count--
  515. if n.leaf() {
  516. break
  517. }
  518. n = (*n.children)[len(*n.children)-1]
  519. }
  520. return tr.Set(item.key, item.value)
  521. }
  522. // Min returns the minimum item in tree.
  523. // Returns nil if the treex has no items.
  524. func (tr *Map[K, V]) Min() (K, V, bool) {
  525. if tr.root == nil {
  526. return tr.empty.key, tr.empty.value, false
  527. }
  528. n := tr.root
  529. for {
  530. if n.leaf() {
  531. item := n.items[0]
  532. return item.key, item.value, true
  533. }
  534. n = (*n.children)[0]
  535. }
  536. }
  537. // Max returns the maximum item in tree.
  538. // Returns nil if the tree has no items.
  539. func (tr *Map[K, V]) Max() (K, V, bool) {
  540. if tr.root == nil {
  541. return tr.empty.key, tr.empty.value, false
  542. }
  543. n := tr.root
  544. for {
  545. if n.leaf() {
  546. item := n.items[len(n.items)-1]
  547. return item.key, item.value, true
  548. }
  549. n = (*n.children)[len(*n.children)-1]
  550. }
  551. }
  552. // PopMin removes the minimum item in tree and returns it.
  553. // Returns nil if the tree has no items.
  554. func (tr *Map[K, V]) PopMin() (K, V, bool) {
  555. if tr.root == nil {
  556. return tr.empty.key, tr.empty.value, false
  557. }
  558. n := tr.cowLoad(&tr.root)
  559. var item mapPair[K, V]
  560. for {
  561. n.count-- // optimistically update counts
  562. if n.leaf() {
  563. item = n.items[0]
  564. if len(n.items) == minItems {
  565. break
  566. }
  567. copy(n.items[:], n.items[1:])
  568. n.items[len(n.items)-1] = tr.empty
  569. n.items = n.items[:len(n.items)-1]
  570. tr.count--
  571. if tr.count == 0 {
  572. tr.root = nil
  573. }
  574. return item.key, item.value, true
  575. }
  576. n = tr.cowLoad(&(*n.children)[0])
  577. }
  578. // revert the counts
  579. n = tr.root
  580. for {
  581. n.count++
  582. if n.leaf() {
  583. break
  584. }
  585. n = (*n.children)[0]
  586. }
  587. value, deleted := tr.Delete(item.key)
  588. if deleted {
  589. return item.key, value, true
  590. }
  591. return tr.empty.key, tr.empty.value, false
  592. }
  593. // PopMax removes the maximum item in tree and returns it.
  594. // Returns nil if the tree has no items.
  595. func (tr *Map[K, V]) PopMax() (K, V, bool) {
  596. if tr.root == nil {
  597. return tr.empty.key, tr.empty.value, false
  598. }
  599. n := tr.cowLoad(&tr.root)
  600. var item mapPair[K, V]
  601. for {
  602. n.count-- // optimistically update counts
  603. if n.leaf() {
  604. item = n.items[len(n.items)-1]
  605. if len(n.items) == minItems {
  606. break
  607. }
  608. n.items[len(n.items)-1] = tr.empty
  609. n.items = n.items[:len(n.items)-1]
  610. tr.count--
  611. if tr.count == 0 {
  612. tr.root = nil
  613. }
  614. return item.key, item.value, true
  615. }
  616. n = tr.cowLoad(&(*n.children)[len(*n.children)-1])
  617. }
  618. // revert the counts
  619. n = tr.root
  620. for {
  621. n.count++
  622. if n.leaf() {
  623. break
  624. }
  625. n = (*n.children)[len(*n.children)-1]
  626. }
  627. value, deleted := tr.Delete(item.key)
  628. if deleted {
  629. return item.key, value, true
  630. }
  631. return tr.empty.key, tr.empty.value, false
  632. }
  633. // GetAt returns the value at index.
  634. // Return nil if the tree is empty or the index is out of bounds.
  635. func (tr *Map[K, V]) GetAt(index int) (K, V, bool) {
  636. if tr.root == nil || index < 0 || index >= tr.count {
  637. return tr.empty.key, tr.empty.value, false
  638. }
  639. n := tr.root
  640. for {
  641. if n.leaf() {
  642. return n.items[index].key, n.items[index].value, true
  643. }
  644. i := 0
  645. for ; i < len(n.items); i++ {
  646. if index < (*n.children)[i].count {
  647. break
  648. } else if index == (*n.children)[i].count {
  649. return n.items[i].key, n.items[i].value, true
  650. }
  651. index -= (*n.children)[i].count + 1
  652. }
  653. n = (*n.children)[i]
  654. }
  655. }
  656. // DeleteAt deletes the item at index.
  657. // Return nil if the tree is empty or the index is out of bounds.
  658. func (tr *Map[K, V]) DeleteAt(index int) (K, V, bool) {
  659. if tr.root == nil || index < 0 || index >= tr.count {
  660. return tr.empty.key, tr.empty.value, false
  661. }
  662. var pathbuf [8]uint8 // track the path
  663. path := pathbuf[:0]
  664. var item mapPair[K, V]
  665. n := tr.cowLoad(&tr.root)
  666. outer:
  667. for {
  668. n.count-- // optimistically update counts
  669. if n.leaf() {
  670. // the index is the item position
  671. item = n.items[index]
  672. if len(n.items) == minItems {
  673. path = append(path, uint8(index))
  674. break outer
  675. }
  676. copy(n.items[index:], n.items[index+1:])
  677. n.items[len(n.items)-1] = tr.empty
  678. n.items = n.items[:len(n.items)-1]
  679. tr.count--
  680. if tr.count == 0 {
  681. tr.root = nil
  682. }
  683. return item.key, item.value, true
  684. }
  685. i := 0
  686. for ; i < len(n.items); i++ {
  687. if index < (*n.children)[i].count {
  688. break
  689. } else if index == (*n.children)[i].count {
  690. item = n.items[i]
  691. path = append(path, uint8(i))
  692. break outer
  693. }
  694. index -= (*n.children)[i].count + 1
  695. }
  696. path = append(path, uint8(i))
  697. n = tr.cowLoad(&(*n.children)[i])
  698. }
  699. // revert the counts
  700. n = tr.root
  701. for i := 0; i < len(path); i++ {
  702. n.count++
  703. if !n.leaf() {
  704. n = (*n.children)[uint8(path[i])]
  705. }
  706. }
  707. value, deleted := tr.Delete(item.key)
  708. if deleted {
  709. return item.key, value, true
  710. }
  711. return tr.empty.key, tr.empty.value, false
  712. }
  713. // Height returns the height of the tree.
  714. // Returns zero if tree has no items.
  715. func (tr *Map[K, V]) Height() int {
  716. var height int
  717. if tr.root != nil {
  718. n := tr.root
  719. for {
  720. height++
  721. if n.leaf() {
  722. break
  723. }
  724. n = (*n.children)[0]
  725. }
  726. }
  727. return height
  728. }
  729. // MapIter represents an iterator for btree.Map
  730. type MapIter[K ordered, V any] struct {
  731. tr *Map[K, V]
  732. seeked bool
  733. atstart bool
  734. atend bool
  735. stack []mapIterStackItem[K, V]
  736. item mapPair[K, V]
  737. }
  738. type mapIterStackItem[K ordered, V any] struct {
  739. n *mapNode[K, V]
  740. i int
  741. }
  742. // Iter returns a read-only iterator.
  743. func (tr *Map[K, V]) Iter() MapIter[K, V] {
  744. var iter MapIter[K, V]
  745. iter.tr = tr
  746. return iter
  747. }
  748. // Seek to item greater-or-equal-to key.
  749. // Returns false if there was no item found.
  750. func (iter *MapIter[K, V]) Seek(key K) bool {
  751. if iter.tr == nil {
  752. return false
  753. }
  754. iter.seeked = true
  755. iter.stack = iter.stack[:0]
  756. if iter.tr.root == nil {
  757. return false
  758. }
  759. n := iter.tr.root
  760. for {
  761. i, found := iter.tr.bsearch(n, key)
  762. iter.stack = append(iter.stack, mapIterStackItem[K, V]{n, i})
  763. if found {
  764. iter.item = n.items[i]
  765. return true
  766. }
  767. if n.leaf() {
  768. iter.stack[len(iter.stack)-1].i--
  769. return iter.Next()
  770. }
  771. n = (*n.children)[i]
  772. }
  773. }
  774. // First moves iterator to first item in tree.
  775. // Returns false if the tree is empty.
  776. func (iter *MapIter[K, V]) First() bool {
  777. if iter.tr == nil {
  778. return false
  779. }
  780. iter.atend = false
  781. iter.atstart = false
  782. iter.seeked = true
  783. iter.stack = iter.stack[:0]
  784. if iter.tr.root == nil {
  785. return false
  786. }
  787. n := iter.tr.root
  788. for {
  789. iter.stack = append(iter.stack, mapIterStackItem[K, V]{n, 0})
  790. if n.leaf() {
  791. break
  792. }
  793. n = (*n.children)[0]
  794. }
  795. s := &iter.stack[len(iter.stack)-1]
  796. iter.item = s.n.items[s.i]
  797. return true
  798. }
  799. // Last moves iterator to last item in tree.
  800. // Returns false if the tree is empty.
  801. func (iter *MapIter[K, V]) Last() bool {
  802. if iter.tr == nil {
  803. return false
  804. }
  805. iter.seeked = true
  806. iter.stack = iter.stack[:0]
  807. if iter.tr.root == nil {
  808. return false
  809. }
  810. n := iter.tr.root
  811. for {
  812. iter.stack = append(iter.stack, mapIterStackItem[K, V]{n, len(n.items)})
  813. if n.leaf() {
  814. iter.stack[len(iter.stack)-1].i--
  815. break
  816. }
  817. n = (*n.children)[len(n.items)]
  818. }
  819. s := &iter.stack[len(iter.stack)-1]
  820. iter.item = s.n.items[s.i]
  821. return true
  822. }
  823. // Next moves iterator to the next item in iterator.
  824. // Returns false if the tree is empty or the iterator is at the end of
  825. // the tree.
  826. func (iter *MapIter[K, V]) Next() bool {
  827. if iter.tr == nil {
  828. return false
  829. }
  830. if !iter.seeked {
  831. return iter.First()
  832. }
  833. if len(iter.stack) == 0 {
  834. if iter.atstart {
  835. return iter.First() && iter.Next()
  836. }
  837. return false
  838. }
  839. s := &iter.stack[len(iter.stack)-1]
  840. s.i++
  841. if s.n.leaf() {
  842. if s.i == len(s.n.items) {
  843. for {
  844. iter.stack = iter.stack[:len(iter.stack)-1]
  845. if len(iter.stack) == 0 {
  846. iter.atend = true
  847. return false
  848. }
  849. s = &iter.stack[len(iter.stack)-1]
  850. if s.i < len(s.n.items) {
  851. break
  852. }
  853. }
  854. }
  855. } else {
  856. n := (*s.n.children)[s.i]
  857. for {
  858. iter.stack = append(iter.stack, mapIterStackItem[K, V]{n, 0})
  859. if n.leaf() {
  860. break
  861. }
  862. n = (*n.children)[0]
  863. }
  864. }
  865. s = &iter.stack[len(iter.stack)-1]
  866. iter.item = s.n.items[s.i]
  867. return true
  868. }
  869. // Prev moves iterator to the previous item in iterator.
  870. // Returns false if the tree is empty or the iterator is at the beginning of
  871. // the tree.
  872. func (iter *MapIter[K, V]) Prev() bool {
  873. if iter.tr == nil {
  874. return false
  875. }
  876. if !iter.seeked {
  877. return false
  878. }
  879. if len(iter.stack) == 0 {
  880. if iter.atend {
  881. return iter.Last() && iter.Prev()
  882. }
  883. return false
  884. }
  885. s := &iter.stack[len(iter.stack)-1]
  886. if s.n.leaf() {
  887. s.i--
  888. if s.i == -1 {
  889. for {
  890. iter.stack = iter.stack[:len(iter.stack)-1]
  891. if len(iter.stack) == 0 {
  892. iter.atstart = true
  893. return false
  894. }
  895. s = &iter.stack[len(iter.stack)-1]
  896. s.i--
  897. if s.i > -1 {
  898. break
  899. }
  900. }
  901. }
  902. } else {
  903. n := (*s.n.children)[s.i]
  904. for {
  905. iter.stack = append(iter.stack,
  906. mapIterStackItem[K, V]{n, len(n.items)})
  907. if n.leaf() {
  908. iter.stack[len(iter.stack)-1].i--
  909. break
  910. }
  911. n = (*n.children)[len(n.items)]
  912. }
  913. }
  914. s = &iter.stack[len(iter.stack)-1]
  915. iter.item = s.n.items[s.i]
  916. return true
  917. }
  918. // Key returns the current iterator item key.
  919. func (iter *MapIter[K, V]) Key() K {
  920. return iter.item.key
  921. }
  922. // Value returns the current iterator item value.
  923. func (iter *MapIter[K, V]) Value() V {
  924. return iter.item.value
  925. }
  926. // Values returns all the values in order.
  927. func (tr *Map[K, V]) Values() []V {
  928. values := make([]V, 0, tr.Len())
  929. if tr.root != nil {
  930. values = tr.root.values(values)
  931. }
  932. return values
  933. }
  934. func (n *mapNode[K, V]) values(values []V) []V {
  935. if n.leaf() {
  936. for i := 0; i < len(n.items); i++ {
  937. values = append(values, n.items[i].value)
  938. }
  939. return values
  940. }
  941. for i := 0; i < len(n.items); i++ {
  942. values = (*n.children)[i].values(values)
  943. values = append(values, n.items[i].value)
  944. }
  945. return (*n.children)[len(*n.children)-1].values(values)
  946. }
  947. // Keys returns all the keys in order.
  948. func (tr *Map[K, V]) Keys() []K {
  949. keys := make([]K, 0, tr.Len())
  950. if tr.root != nil {
  951. keys = tr.root.keys(keys)
  952. }
  953. return keys
  954. }
  955. func (n *mapNode[K, V]) keys(keys []K) []K {
  956. if n.leaf() {
  957. for i := 0; i < len(n.items); i++ {
  958. keys = append(keys, n.items[i].key)
  959. }
  960. return keys
  961. }
  962. for i := 0; i < len(n.items); i++ {
  963. keys = (*n.children)[i].keys(keys)
  964. keys = append(keys, n.items[i].key)
  965. }
  966. return (*n.children)[len(*n.children)-1].keys(keys)
  967. }
  968. // KeyValues returns all the keys and values in order.
  969. func (tr *Map[K, V]) KeyValues() ([]K, []V) {
  970. keys := make([]K, 0, tr.Len())
  971. values := make([]V, 0, tr.Len())
  972. if tr.root != nil {
  973. keys, values = tr.root.keyValues(keys, values)
  974. }
  975. return keys, values
  976. }
  977. func (n *mapNode[K, V]) keyValues(keys []K, values []V) ([]K, []V) {
  978. if n.leaf() {
  979. for i := 0; i < len(n.items); i++ {
  980. keys = append(keys, n.items[i].key)
  981. values = append(values, n.items[i].value)
  982. }
  983. return keys, values
  984. }
  985. for i := 0; i < len(n.items); i++ {
  986. keys, values = (*n.children)[i].keyValues(keys, values)
  987. keys = append(keys, n.items[i].key)
  988. values = append(values, n.items[i].value)
  989. }
  990. return (*n.children)[len(*n.children)-1].keyValues(keys, values)
  991. }