map.go 28 KB

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