123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056 |
- // Copyright 2020 Joshua J Baker. All rights reserved.
- // Use of this source code is governed by an MIT-style
- // license that can be found in the LICENSE file.
- package btree
- import "sync/atomic"
- type ordered interface {
- ~int | ~int8 | ~int16 | ~int32 | ~int64 |
- ~uint | ~uint8 | ~uint16 | ~uint32 | ~uint64 | ~uintptr |
- ~float32 | ~float64 | ~string
- }
- type mapPair[K ordered, V any] struct {
- // The `value` field should be before the `key` field because doing so
- // allows for the Go compiler to optimize away the `value` field when
- // it's a `struct{}`, which is the case for `btree.Set`.
- value V
- key K
- }
- type Map[K ordered, V any] struct {
- cow uint64
- root *mapNode[K, V]
- count int
- empty mapPair[K, V]
- }
- type mapNode[K ordered, V any] struct {
- cow uint64
- count int
- items []mapPair[K, V]
- children *[]*mapNode[K, V]
- }
- // This operation should not be inlined because it's expensive and rarely
- // called outside of heavy copy-on-write situations. Marking it "noinline"
- // allows for the parent cowLoad to be inlined.
- // go:noinline
- func (tr *Map[K, V]) copy(n *mapNode[K, V]) *mapNode[K, V] {
- n2 := new(mapNode[K, V])
- n2.cow = tr.cow
- n2.count = n.count
- n2.items = make([]mapPair[K, V], len(n.items), cap(n.items))
- copy(n2.items, n.items)
- if !n.leaf() {
- n2.children = new([]*mapNode[K, V])
- *n2.children = make([]*mapNode[K, V], len(*n.children), maxItems+1)
- copy(*n2.children, *n.children)
- }
- return n2
- }
- // cowLoad loads the provided node and, if needed, performs a copy-on-write.
- func (tr *Map[K, V]) cowLoad(cn **mapNode[K, V]) *mapNode[K, V] {
- if (*cn).cow != tr.cow {
- *cn = tr.copy(*cn)
- }
- return *cn
- }
- func (tr *Map[K, V]) Copy() *Map[K, V] {
- tr2 := new(Map[K, V])
- *tr2 = *tr
- tr2.cow = atomic.AddUint64(&gcow, 1)
- tr.cow = atomic.AddUint64(&gcow, 1)
- return tr2
- }
- func (tr *Map[K, V]) newNode(leaf bool) *mapNode[K, V] {
- n := new(mapNode[K, V])
- n.cow = tr.cow
- if !leaf {
- n.children = new([]*mapNode[K, V])
- }
- return n
- }
- // leaf returns true if the node is a leaf.
- func (n *mapNode[K, V]) leaf() bool {
- return n.children == nil
- }
- func (tr *Map[K, V]) bsearch(n *mapNode[K, V], key K) (index int, found bool) {
- low, high := 0, len(n.items)
- for low < high {
- h := int(uint(low+high) >> 1)
- if key >= n.items[h].key {
- low = h + 1
- } else {
- high = h
- }
- }
- if low > 0 && n.items[low-1].key >= key {
- return low - 1, true
- }
- return low, false
- }
- // Set or replace a value for a key
- func (tr *Map[K, V]) Set(key K, value V) (V, bool) {
- item := mapPair[K, V]{key: key, value: value}
- if tr.root == nil {
- tr.root = tr.newNode(true)
- tr.root.items = append([]mapPair[K, V]{}, item)
- tr.root.count = 1
- tr.count = 1
- return tr.empty.value, false
- }
- prev, replaced, split := tr.nodeSet(&tr.root, item)
- if split {
- left := tr.root
- right, median := tr.nodeSplit(left)
- tr.root = tr.newNode(false)
- *tr.root.children = make([]*mapNode[K, V], 0, maxItems+1)
- *tr.root.children = append([]*mapNode[K, V]{}, left, right)
- tr.root.items = append([]mapPair[K, V]{}, median)
- tr.root.updateCount()
- return tr.Set(item.key, item.value)
- }
- if replaced {
- return prev, true
- }
- tr.count++
- return tr.empty.value, false
- }
- func (tr *Map[K, V]) nodeSplit(n *mapNode[K, V],
- ) (right *mapNode[K, V], median mapPair[K, V]) {
- i := maxItems / 2
- median = n.items[i]
- const sliceItems = true
- // right node
- right = tr.newNode(n.leaf())
- if sliceItems {
- right.items = n.items[i+1:]
- if !n.leaf() {
- *right.children = (*n.children)[i+1:]
- }
- } else {
- right.items = make([]mapPair[K, V], len(n.items[i+1:]), maxItems/2)
- copy(right.items, n.items[i+1:])
- if !n.leaf() {
- *right.children = make([]*mapNode[K, V],
- len((*n.children)[i+1:]), maxItems+1)
- copy(*right.children, (*n.children)[i+1:])
- }
- }
- right.updateCount()
- // left node
- if sliceItems {
- n.items[i] = tr.empty
- n.items = n.items[:i:i]
- if !n.leaf() {
- *n.children = (*n.children)[: i+1 : i+1]
- }
- } else {
- for j := i; j < len(n.items); j++ {
- n.items[j] = tr.empty
- }
- if !n.leaf() {
- for j := i + 1; j < len((*n.children)); j++ {
- (*n.children)[j] = nil
- }
- }
- n.items = n.items[:i]
- if !n.leaf() {
- *n.children = (*n.children)[:i+1]
- }
- }
- n.updateCount()
- return right, median
- }
- func (n *mapNode[K, V]) updateCount() {
- n.count = len(n.items)
- if !n.leaf() {
- for i := 0; i < len(*n.children); i++ {
- n.count += (*n.children)[i].count
- }
- }
- }
- func (tr *Map[K, V]) nodeSet(pn **mapNode[K, V], item mapPair[K, V],
- ) (prev V, replaced bool, split bool) {
- n := tr.cowLoad(pn)
- i, found := tr.bsearch(n, item.key)
- if found {
- prev = n.items[i].value
- n.items[i].value = item.value
- return prev, true, false
- }
- if n.leaf() {
- if len(n.items) == maxItems {
- return tr.empty.value, false, true
- }
- n.items = append(n.items, tr.empty)
- copy(n.items[i+1:], n.items[i:])
- n.items[i] = item
- n.count++
- return tr.empty.value, false, false
- }
- prev, replaced, split = tr.nodeSet(&(*n.children)[i], item)
- if split {
- if len(n.items) == maxItems {
- return tr.empty.value, false, true
- }
- right, median := tr.nodeSplit((*n.children)[i])
- *n.children = append(*n.children, nil)
- copy((*n.children)[i+1:], (*n.children)[i:])
- (*n.children)[i+1] = right
- n.items = append(n.items, tr.empty)
- copy(n.items[i+1:], n.items[i:])
- n.items[i] = median
- return tr.nodeSet(&n, item)
- }
- if !replaced {
- n.count++
- }
- return prev, replaced, false
- }
- func (tr *Map[K, V]) Scan(iter func(key K, value V) bool) {
- if tr.root == nil {
- return
- }
- tr.root.scan(iter)
- }
- func (n *mapNode[K, V]) scan(iter func(key K, value V) bool) bool {
- if n.leaf() {
- for i := 0; i < len(n.items); i++ {
- if !iter(n.items[i].key, n.items[i].value) {
- return false
- }
- }
- return true
- }
- for i := 0; i < len(n.items); i++ {
- if !(*n.children)[i].scan(iter) {
- return false
- }
- if !iter(n.items[i].key, n.items[i].value) {
- return false
- }
- }
- return (*n.children)[len(*n.children)-1].scan(iter)
- }
- // Get a value for key
- func (tr *Map[K, V]) Get(key K) (V, bool) {
- if tr.root == nil {
- return tr.empty.value, false
- }
- n := tr.root
- for {
- i, found := tr.bsearch(n, key)
- if found {
- return n.items[i].value, true
- }
- if n.leaf() {
- return tr.empty.value, false
- }
- n = (*n.children)[i]
- }
- }
- // Len returns the number of items in the tree
- func (tr *Map[K, V]) Len() int {
- return tr.count
- }
- // Delete a value for a key and returns the deleted value.
- // Returns false if there was no value by that key found.
- func (tr *Map[K, V]) Delete(key K) (V, bool) {
- if tr.root == nil {
- return tr.empty.value, false
- }
- prev, deleted := tr.delete(&tr.root, false, key)
- if !deleted {
- return tr.empty.value, false
- }
- if len(tr.root.items) == 0 && !tr.root.leaf() {
- tr.root = (*tr.root.children)[0]
- }
- tr.count--
- if tr.count == 0 {
- tr.root = nil
- }
- return prev.value, true
- }
- func (tr *Map[K, V]) delete(pn **mapNode[K, V], max bool, key K,
- ) (mapPair[K, V], bool) {
- n := tr.cowLoad(pn)
- var i int
- var found bool
- if max {
- i, found = len(n.items)-1, true
- } else {
- i, found = tr.bsearch(n, key)
- }
- if n.leaf() {
- if found {
- // found the items at the leaf, remove it and return.
- prev := n.items[i]
- copy(n.items[i:], n.items[i+1:])
- n.items[len(n.items)-1] = tr.empty
- n.items = n.items[:len(n.items)-1]
- n.count--
- return prev, true
- }
- return tr.empty, false
- }
- var prev mapPair[K, V]
- var deleted bool
- if found {
- if max {
- i++
- prev, deleted = tr.delete(&(*n.children)[i], true, tr.empty.key)
- } else {
- prev = n.items[i]
- maxItem, _ := tr.delete(&(*n.children)[i], true, tr.empty.key)
- deleted = true
- n.items[i] = maxItem
- }
- } else {
- prev, deleted = tr.delete(&(*n.children)[i], max, key)
- }
- if !deleted {
- return tr.empty, false
- }
- n.count--
- if len((*n.children)[i].items) < minItems {
- tr.nodeRebalance(n, i)
- }
- return prev, true
- }
- // nodeRebalance rebalances the child nodes following a delete operation.
- // Provide the index of the child node with the number of items that fell
- // below minItems.
- func (tr *Map[K, V]) nodeRebalance(n *mapNode[K, V], i int) {
- if i == len(n.items) {
- i--
- }
- // ensure copy-on-write
- left := tr.cowLoad(&(*n.children)[i])
- right := tr.cowLoad(&(*n.children)[i+1])
- if len(left.items)+len(right.items) < maxItems {
- // Merges the left and right children nodes together as a single node
- // that includes (left,item,right), and places the contents into the
- // existing left node. Delete the right node altogether and move the
- // following items and child nodes to the left by one slot.
- // merge (left,item,right)
- left.items = append(left.items, n.items[i])
- left.items = append(left.items, right.items...)
- if !left.leaf() {
- *left.children = append(*left.children, *right.children...)
- }
- left.count += right.count + 1
- // move the items over one slot
- copy(n.items[i:], n.items[i+1:])
- n.items[len(n.items)-1] = tr.empty
- n.items = n.items[:len(n.items)-1]
- // move the children over one slot
- copy((*n.children)[i+1:], (*n.children)[i+2:])
- (*n.children)[len(*n.children)-1] = nil
- (*n.children) = (*n.children)[:len(*n.children)-1]
- } else if len(left.items) > len(right.items) {
- // move left -> right over one slot
- // Move the item of the parent node at index into the right-node first
- // slot, and move the left-node last item into the previously moved
- // parent item slot.
- right.items = append(right.items, tr.empty)
- copy(right.items[1:], right.items)
- right.items[0] = n.items[i]
- right.count++
- n.items[i] = left.items[len(left.items)-1]
- left.items[len(left.items)-1] = tr.empty
- left.items = left.items[:len(left.items)-1]
- left.count--
- if !left.leaf() {
- // move the left-node last child into the right-node first slot
- *right.children = append(*right.children, nil)
- copy((*right.children)[1:], *right.children)
- (*right.children)[0] = (*left.children)[len(*left.children)-1]
- (*left.children)[len(*left.children)-1] = nil
- (*left.children) = (*left.children)[:len(*left.children)-1]
- left.count -= (*right.children)[0].count
- right.count += (*right.children)[0].count
- }
- } else {
- // move left <- right over one slot
- // Same as above but the other direction
- left.items = append(left.items, n.items[i])
- left.count++
- n.items[i] = right.items[0]
- copy(right.items, right.items[1:])
- right.items[len(right.items)-1] = tr.empty
- right.items = right.items[:len(right.items)-1]
- right.count--
- if !left.leaf() {
- *left.children = append(*left.children, (*right.children)[0])
- copy(*right.children, (*right.children)[1:])
- (*right.children)[len(*right.children)-1] = nil
- *right.children = (*right.children)[:len(*right.children)-1]
- left.count += (*left.children)[len(*left.children)-1].count
- right.count -= (*left.children)[len(*left.children)-1].count
- }
- }
- }
- // Ascend the tree within the range [pivot, last]
- // Pass nil for pivot to scan all item in ascending order
- // Return false to stop iterating
- func (tr *Map[K, V]) Ascend(pivot K, iter func(key K, value V) bool) {
- if tr.root == nil {
- return
- }
- tr.ascend(tr.root, pivot, iter)
- }
- // The return value of this function determines whether we should keep iterating
- // upon this functions return.
- func (tr *Map[K, V]) ascend(n *mapNode[K, V], pivot K,
- iter func(key K, value V) bool,
- ) bool {
- i, found := tr.bsearch(n, pivot)
- if !found {
- if !n.leaf() {
- if !tr.ascend((*n.children)[i], pivot, iter) {
- return false
- }
- }
- }
- // We are either in the case that
- // - node is found, we should iterate through it starting at `i`,
- // the index it was located at.
- // - node is not found, and TODO: fill in.
- for ; i < len(n.items); i++ {
- if !iter(n.items[i].key, n.items[i].value) {
- return false
- }
- if !n.leaf() {
- if !(*n.children)[i+1].scan(iter) {
- return false
- }
- }
- }
- return true
- }
- func (tr *Map[K, V]) Reverse(iter func(key K, value V) bool) {
- if tr.root == nil {
- return
- }
- tr.root.reverse(iter)
- }
- func (n *mapNode[K, V]) reverse(iter func(key K, value V) bool) bool {
- if n.leaf() {
- for i := len(n.items) - 1; i >= 0; i-- {
- if !iter(n.items[i].key, n.items[i].value) {
- return false
- }
- }
- return true
- }
- if !(*n.children)[len(*n.children)-1].reverse(iter) {
- return false
- }
- for i := len(n.items) - 1; i >= 0; i-- {
- if !iter(n.items[i].key, n.items[i].value) {
- return false
- }
- if !(*n.children)[i].reverse(iter) {
- return false
- }
- }
- return true
- }
- // Descend the tree within the range [pivot, first]
- // Pass nil for pivot to scan all item in descending order
- // Return false to stop iterating
- func (tr *Map[K, V]) Descend(pivot K, iter func(key K, value V) bool) {
- if tr.root == nil {
- return
- }
- tr.descend(tr.root, pivot, iter)
- }
- func (tr *Map[K, V]) descend(n *mapNode[K, V], pivot K,
- iter func(key K, value V) bool,
- ) bool {
- i, found := tr.bsearch(n, pivot)
- if !found {
- if !n.leaf() {
- if !tr.descend((*n.children)[i], pivot, iter) {
- return false
- }
- }
- i--
- }
- for ; i >= 0; i-- {
- if !iter(n.items[i].key, n.items[i].value) {
- return false
- }
- if !n.leaf() {
- if !(*n.children)[i].reverse(iter) {
- return false
- }
- }
- }
- return true
- }
- // Load is for bulk loading pre-sorted items
- func (tr *Map[K, V]) Load(key K, value V) (V, bool) {
- item := mapPair[K, V]{key: key, value: value}
- if tr.root == nil {
- return tr.Set(item.key, item.value)
- }
- n := tr.cowLoad(&tr.root)
- for {
- n.count++ // optimistically update counts
- if n.leaf() {
- if len(n.items) < maxItems {
- if n.items[len(n.items)-1].key < item.key {
- n.items = append(n.items, item)
- tr.count++
- return tr.empty.value, false
- }
- }
- break
- }
- n = tr.cowLoad(&(*n.children)[len(*n.children)-1])
- }
- // revert the counts
- n = tr.root
- for {
- n.count--
- if n.leaf() {
- break
- }
- n = (*n.children)[len(*n.children)-1]
- }
- return tr.Set(item.key, item.value)
- }
- // Min returns the minimum item in tree.
- // Returns nil if the treex has no items.
- func (tr *Map[K, V]) Min() (K, V, bool) {
- if tr.root == nil {
- return tr.empty.key, tr.empty.value, false
- }
- n := tr.root
- for {
- if n.leaf() {
- item := n.items[0]
- return item.key, item.value, true
- }
- n = (*n.children)[0]
- }
- }
- // Max returns the maximum item in tree.
- // Returns nil if the tree has no items.
- func (tr *Map[K, V]) Max() (K, V, bool) {
- if tr.root == nil {
- return tr.empty.key, tr.empty.value, false
- }
- n := tr.root
- for {
- if n.leaf() {
- item := n.items[len(n.items)-1]
- return item.key, item.value, true
- }
- n = (*n.children)[len(*n.children)-1]
- }
- }
- // PopMin removes the minimum item in tree and returns it.
- // Returns nil if the tree has no items.
- func (tr *Map[K, V]) PopMin() (K, V, bool) {
- if tr.root == nil {
- return tr.empty.key, tr.empty.value, false
- }
- n := tr.cowLoad(&tr.root)
- var item mapPair[K, V]
- for {
- n.count-- // optimistically update counts
- if n.leaf() {
- item = n.items[0]
- if len(n.items) == minItems {
- break
- }
- copy(n.items[:], n.items[1:])
- n.items[len(n.items)-1] = tr.empty
- n.items = n.items[:len(n.items)-1]
- tr.count--
- if tr.count == 0 {
- tr.root = nil
- }
- return item.key, item.value, true
- }
- n = tr.cowLoad(&(*n.children)[0])
- }
- // revert the counts
- n = tr.root
- for {
- n.count++
- if n.leaf() {
- break
- }
- n = (*n.children)[0]
- }
- value, deleted := tr.Delete(item.key)
- if deleted {
- return item.key, value, true
- }
- return tr.empty.key, tr.empty.value, false
- }
- // PopMax removes the maximum item in tree and returns it.
- // Returns nil if the tree has no items.
- func (tr *Map[K, V]) PopMax() (K, V, bool) {
- if tr.root == nil {
- return tr.empty.key, tr.empty.value, false
- }
- n := tr.cowLoad(&tr.root)
- var item mapPair[K, V]
- for {
- n.count-- // optimistically update counts
- if n.leaf() {
- item = n.items[len(n.items)-1]
- if len(n.items) == minItems {
- break
- }
- n.items[len(n.items)-1] = tr.empty
- n.items = n.items[:len(n.items)-1]
- tr.count--
- if tr.count == 0 {
- tr.root = nil
- }
- return item.key, item.value, true
- }
- n = tr.cowLoad(&(*n.children)[len(*n.children)-1])
- }
- // revert the counts
- n = tr.root
- for {
- n.count++
- if n.leaf() {
- break
- }
- n = (*n.children)[len(*n.children)-1]
- }
- value, deleted := tr.Delete(item.key)
- if deleted {
- return item.key, value, true
- }
- return tr.empty.key, tr.empty.value, false
- }
- // GetAt returns the value at index.
- // Return nil if the tree is empty or the index is out of bounds.
- func (tr *Map[K, V]) GetAt(index int) (K, V, bool) {
- if tr.root == nil || index < 0 || index >= tr.count {
- return tr.empty.key, tr.empty.value, false
- }
- n := tr.root
- for {
- if n.leaf() {
- return n.items[index].key, n.items[index].value, true
- }
- i := 0
- for ; i < len(n.items); i++ {
- if index < (*n.children)[i].count {
- break
- } else if index == (*n.children)[i].count {
- return n.items[i].key, n.items[i].value, true
- }
- index -= (*n.children)[i].count + 1
- }
- n = (*n.children)[i]
- }
- }
- // DeleteAt deletes the item at index.
- // Return nil if the tree is empty or the index is out of bounds.
- func (tr *Map[K, V]) DeleteAt(index int) (K, V, bool) {
- if tr.root == nil || index < 0 || index >= tr.count {
- return tr.empty.key, tr.empty.value, false
- }
- var pathbuf [8]uint8 // track the path
- path := pathbuf[:0]
- var item mapPair[K, V]
- n := tr.cowLoad(&tr.root)
- outer:
- for {
- n.count-- // optimistically update counts
- if n.leaf() {
- // the index is the item position
- item = n.items[index]
- if len(n.items) == minItems {
- path = append(path, uint8(index))
- break outer
- }
- copy(n.items[index:], n.items[index+1:])
- n.items[len(n.items)-1] = tr.empty
- n.items = n.items[:len(n.items)-1]
- tr.count--
- if tr.count == 0 {
- tr.root = nil
- }
- return item.key, item.value, true
- }
- i := 0
- for ; i < len(n.items); i++ {
- if index < (*n.children)[i].count {
- break
- } else if index == (*n.children)[i].count {
- item = n.items[i]
- path = append(path, uint8(i))
- break outer
- }
- index -= (*n.children)[i].count + 1
- }
- path = append(path, uint8(i))
- n = tr.cowLoad(&(*n.children)[i])
- }
- // revert the counts
- n = tr.root
- for i := 0; i < len(path); i++ {
- n.count++
- if !n.leaf() {
- n = (*n.children)[uint8(path[i])]
- }
- }
- value, deleted := tr.Delete(item.key)
- if deleted {
- return item.key, value, true
- }
- return tr.empty.key, tr.empty.value, false
- }
- // Height returns the height of the tree.
- // Returns zero if tree has no items.
- func (tr *Map[K, V]) Height() int {
- var height int
- if tr.root != nil {
- n := tr.root
- for {
- height++
- if n.leaf() {
- break
- }
- n = (*n.children)[0]
- }
- }
- return height
- }
- // MapIter represents an iterator for btree.Map
- type MapIter[K ordered, V any] struct {
- tr *Map[K, V]
- seeked bool
- atstart bool
- atend bool
- stack []mapIterStackItem[K, V]
- item mapPair[K, V]
- }
- type mapIterStackItem[K ordered, V any] struct {
- n *mapNode[K, V]
- i int
- }
- // Iter returns a read-only iterator.
- func (tr *Map[K, V]) Iter() MapIter[K, V] {
- var iter MapIter[K, V]
- iter.tr = tr
- return iter
- }
- // Seek to item greater-or-equal-to key.
- // Returns false if there was no item found.
- func (iter *MapIter[K, V]) Seek(key K) bool {
- if iter.tr == nil {
- return false
- }
- iter.seeked = true
- iter.stack = iter.stack[:0]
- if iter.tr.root == nil {
- return false
- }
- n := iter.tr.root
- for {
- i, found := iter.tr.bsearch(n, key)
- iter.stack = append(iter.stack, mapIterStackItem[K, V]{n, i})
- if found {
- iter.item = n.items[i]
- return true
- }
- if n.leaf() {
- iter.stack[len(iter.stack)-1].i--
- return iter.Next()
- }
- n = (*n.children)[i]
- }
- }
- // First moves iterator to first item in tree.
- // Returns false if the tree is empty.
- func (iter *MapIter[K, V]) First() bool {
- if iter.tr == nil {
- return false
- }
- iter.atend = false
- iter.atstart = false
- iter.seeked = true
- iter.stack = iter.stack[:0]
- if iter.tr.root == nil {
- return false
- }
- n := iter.tr.root
- for {
- iter.stack = append(iter.stack, mapIterStackItem[K, V]{n, 0})
- if n.leaf() {
- break
- }
- n = (*n.children)[0]
- }
- s := &iter.stack[len(iter.stack)-1]
- iter.item = s.n.items[s.i]
- return true
- }
- // Last moves iterator to last item in tree.
- // Returns false if the tree is empty.
- func (iter *MapIter[K, V]) Last() bool {
- if iter.tr == nil {
- return false
- }
- iter.seeked = true
- iter.stack = iter.stack[:0]
- if iter.tr.root == nil {
- return false
- }
- n := iter.tr.root
- for {
- iter.stack = append(iter.stack, mapIterStackItem[K, V]{n, len(n.items)})
- if n.leaf() {
- iter.stack[len(iter.stack)-1].i--
- break
- }
- n = (*n.children)[len(n.items)]
- }
- s := &iter.stack[len(iter.stack)-1]
- iter.item = s.n.items[s.i]
- return true
- }
- // Next moves iterator to the next item in iterator.
- // Returns false if the tree is empty or the iterator is at the end of
- // the tree.
- func (iter *MapIter[K, V]) Next() bool {
- if iter.tr == nil {
- return false
- }
- if !iter.seeked {
- return iter.First()
- }
- if len(iter.stack) == 0 {
- if iter.atstart {
- return iter.First() && iter.Next()
- }
- return false
- }
- s := &iter.stack[len(iter.stack)-1]
- s.i++
- if s.n.leaf() {
- if s.i == len(s.n.items) {
- for {
- iter.stack = iter.stack[:len(iter.stack)-1]
- if len(iter.stack) == 0 {
- iter.atend = true
- return false
- }
- s = &iter.stack[len(iter.stack)-1]
- if s.i < len(s.n.items) {
- break
- }
- }
- }
- } else {
- n := (*s.n.children)[s.i]
- for {
- iter.stack = append(iter.stack, mapIterStackItem[K, V]{n, 0})
- if n.leaf() {
- break
- }
- n = (*n.children)[0]
- }
- }
- s = &iter.stack[len(iter.stack)-1]
- iter.item = s.n.items[s.i]
- return true
- }
- // Prev moves iterator to the previous item in iterator.
- // Returns false if the tree is empty or the iterator is at the beginning of
- // the tree.
- func (iter *MapIter[K, V]) Prev() bool {
- if iter.tr == nil {
- return false
- }
- if !iter.seeked {
- return false
- }
- if len(iter.stack) == 0 {
- if iter.atend {
- return iter.Last() && iter.Prev()
- }
- return false
- }
- s := &iter.stack[len(iter.stack)-1]
- if s.n.leaf() {
- s.i--
- if s.i == -1 {
- for {
- iter.stack = iter.stack[:len(iter.stack)-1]
- if len(iter.stack) == 0 {
- iter.atstart = true
- return false
- }
- s = &iter.stack[len(iter.stack)-1]
- s.i--
- if s.i > -1 {
- break
- }
- }
- }
- } else {
- n := (*s.n.children)[s.i]
- for {
- iter.stack = append(iter.stack,
- mapIterStackItem[K, V]{n, len(n.items)})
- if n.leaf() {
- iter.stack[len(iter.stack)-1].i--
- break
- }
- n = (*n.children)[len(n.items)]
- }
- }
- s = &iter.stack[len(iter.stack)-1]
- iter.item = s.n.items[s.i]
- return true
- }
- // Key returns the current iterator item key.
- func (iter *MapIter[K, V]) Key() K {
- return iter.item.key
- }
- // Value returns the current iterator item value.
- func (iter *MapIter[K, V]) Value() V {
- return iter.item.value
- }
- // Values returns all the values in order.
- func (tr *Map[K, V]) Values() []V {
- values := make([]V, 0, tr.Len())
- if tr.root != nil {
- values = tr.root.values(values)
- }
- return values
- }
- func (n *mapNode[K, V]) values(values []V) []V {
- if n.leaf() {
- for i := 0; i < len(n.items); i++ {
- values = append(values, n.items[i].value)
- }
- return values
- }
- for i := 0; i < len(n.items); i++ {
- values = (*n.children)[i].values(values)
- values = append(values, n.items[i].value)
- }
- return (*n.children)[len(*n.children)-1].values(values)
- }
- // Keys returns all the keys in order.
- func (tr *Map[K, V]) Keys() []K {
- keys := make([]K, 0, tr.Len())
- if tr.root != nil {
- keys = tr.root.keys(keys)
- }
- return keys
- }
- func (n *mapNode[K, V]) keys(keys []K) []K {
- if n.leaf() {
- for i := 0; i < len(n.items); i++ {
- keys = append(keys, n.items[i].key)
- }
- return keys
- }
- for i := 0; i < len(n.items); i++ {
- keys = (*n.children)[i].keys(keys)
- keys = append(keys, n.items[i].key)
- }
- return (*n.children)[len(*n.children)-1].keys(keys)
- }
- // KeyValues returns all the keys and values in order.
- func (tr *Map[K, V]) KeyValues() ([]K, []V) {
- keys := make([]K, 0, tr.Len())
- values := make([]V, 0, tr.Len())
- if tr.root != nil {
- keys, values = tr.root.keyValues(keys, values)
- }
- return keys, values
- }
- func (n *mapNode[K, V]) keyValues(keys []K, values []V) ([]K, []V) {
- if n.leaf() {
- for i := 0; i < len(n.items); i++ {
- keys = append(keys, n.items[i].key)
- values = append(values, n.items[i].value)
- }
- return keys, values
- }
- for i := 0; i < len(n.items); i++ {
- keys, values = (*n.children)[i].keyValues(keys, values)
- keys = append(keys, n.items[i].key)
- values = append(values, n.items[i].value)
- }
- return (*n.children)[len(*n.children)-1].keyValues(keys, values)
- }
|