pax_global_header 0000666 0000000 0000000 00000000064 15044500710 0014506 g ustar 00root root 0000000 0000000 52 comment=490165c8137d1cb1f669d07554305665b4c2067a
btree-1.8.1/ 0000775 0000000 0000000 00000000000 15044500710 0012616 5 ustar 00root root 0000000 0000000 btree-1.8.1/.github/ 0000775 0000000 0000000 00000000000 15044500710 0014156 5 ustar 00root root 0000000 0000000 btree-1.8.1/.github/workflows/ 0000775 0000000 0000000 00000000000 15044500710 0016213 5 ustar 00root root 0000000 0000000 btree-1.8.1/.github/workflows/go.yml 0000664 0000000 0000000 00000000742 15044500710 0017346 0 ustar 00root root 0000000 0000000 name: Go
on:
workflow_dispatch:
push:
branches: [ master ]
pull_request:
branches: [ master ]
jobs:
build:
strategy:
matrix:
os: [macos-latest, windows-latest, ubuntu-latest]
runs-on: ${{ matrix.os }}
steps:
- uses: actions/checkout@v4
- name: Set up Go
uses: actions/setup-go@v4
with:
go-version-file: './go.mod'
- name: Build
run: go build -v ./...
- name: Test
run: go test -v ./...
btree-1.8.1/LICENSE 0000664 0000000 0000000 00000002036 15044500710 0013624 0 ustar 00root root 0000000 0000000 Copyright (c) 2020 Josh Baker
Permission is hereby granted, free of charge, to any person obtaining a copy of
this software and associated documentation files (the "Software"), to deal in
the Software without restriction, including without limitation the rights to
use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
the Software, and to permit persons to whom the Software is furnished to do so,
subject to the following conditions:
The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
btree-1.8.1/PATH_HINT.md 0000664 0000000 0000000 00000006474 15044500710 0014531 0 ustar 00root root 0000000 0000000 # B-tree Path Hints
I use a thing I call path hints in my B-tree [C](https://github.com/tidwall/btree.c) and [Go](https://github.com/tidwall/btree) implementations. It's a search optimization.
## The B-tree
A standard [B-tree](https://en.wikipedia.org/wiki/B-tree) is an ordered tree-based data structure that stores its items in nodes. The B-tree has a single root node, which may have children nodes, and those children nodes may also have children nodes.
Searching for items in a B-tree is fast. [O(log N)](https://en.wikipedia.org/wiki/Big_O_notation) to be exact.
This is because the [binary search algorithm](https://en.wikipedia.org/wiki/Binary_search_algorithm) is used.
A binary search works by first comparing the item at the middle-most index of the root node with the target item.
If the middle item is greater than the target item, then it divides the node in two and does the binary search on the left part of the node. If the middle is less, it searches the right part. And so on. If the target item is found, then the search stop. If the item is not found, then the search is passed to the child node at the appropriate index. This traversal terminates when item is found or there are no more child nodes.
## The Path
Each index is a component of the path to the item (or where the item should be stored, if it does not exist in the tree).
Take the first example image. The item 9 is at path “1/0”. The item 16 is at path “1”. The item 21 is at path “2/1”. The item 5 is at path “0/2”.
## The Path Hint
A Path Hint is a predefined path that is provided to B-tree operations. It’s just a hint that says, “Hey B-tree, instead of starting your binary search with the middle index, start with what I provide you. My path may be wrong, and if so please provide me with the correct path so I get it right the next time.”
I’ve found using path hints can lead to a little performance increase of 150% - 300%. This is because in real-world cases the items that I’m working with are usually nearby each other in the tree.
Take for example inserting a group of timeseries points. They may often be received as chucks of near-contiguous items.
Or, I'm sequentially inserting an ordered group of rows somewhere in the middle of a table.
Or, I have a Redis-style key/value store, where the keys look have the common structure “user:98512:name”, “user:98512:email”, and I want to update a bunch of values for specified user.
Using a path hint may help to avoid the unnecessary binary searching in each of these examples.
While I may see a 3x boost in when the path hint is right on, I'll only see around 5% decrease when the path hint is totally wrong.
## Using a Path Hint
All of the functions that take in a path hint argument mutate the path hint argument.
For single-threaded programs, it’s possible to use one shared path hint per B-tree for the life of the program.
For multi-threaded programs, I find it best to use one path hint per B-tree , per thread.
For server-client programs, one path hint per B-tree, per client should suffice.
btree-1.8.1/README.md 0000664 0000000 0000000 00000024321 15044500710 0014077 0 ustar 00root root 0000000 0000000 # btree
[](https://godoc.org/github.com/tidwall/btree)
An efficient [B-tree](https://en.wikipedia.org/wiki/B-tree) implementation in Go.
## Features
- Support for [Generics](#generics) (Go 1.18+).
- `Map` and `Set` types for ordered key-value maps and sets,
- Fast bulk loading for pre-ordered data using the `Load()` method.
- `Copy()` method with copy-on-write support.
- [Path hinting](PATH_HINT.md) optimization for operations with nearby keys.
- Allows for array-like operations. ([Counted B-tree](https://www.chiark.greenend.org.uk/~sgtatham/algorithms/cbtree.html))
## Using
To start using this package, install Go and run:
```sh
$ go get github.com/tidwall/btree
```
## B-tree types
This package includes the following types of B-trees:
- [`btree.Map`](#btreemap):
A fast B-tree for storing ordered key value pairs.
- [`btree.Set`](#btreeset):
Like `Map`, but only for storing keys.
- [`btree.BTreeG`](#btreebtreeg):
A feature-rich B-tree for storing data using a custom comparator. Thread-safe.
- [`btree.BTree`](#btreebtree):
Like `BTreeG` but uses the `interface{}` type for data. Backwards compatible. Thread-safe.
### btree.Map
```go
// Basic
Set(key, value) // insert or replace an item
Get(key, value) // get an existing item
Delete(key) // delete an item
Len() // return the number of items in the map
// Iteration
Scan(iter) // scan items in ascending order
Reverse(iter) // scan items in descending order
Ascend(key, iter) // scan items in ascending order that are >= to key
Descend(key, iter) // scan items in descending order that are <= to key.
Iter() // returns a read-only iterator for for-loops.
// Array-like operations
GetAt(index) // returns the item at index
DeleteAt(index) // deletes the item at index
// Bulk-loading
Load(key, value) // load presorted items into tree
```
#### Example
```go
package main
import (
"fmt"
"github.com/tidwall/btree"
)
func main() {
// create a map
var users btree.Map[string, string]
// add some users
users.Set("user:4", "Andrea")
users.Set("user:6", "Andy")
users.Set("user:2", "Andy")
users.Set("user:1", "Jane")
users.Set("user:5", "Janet")
users.Set("user:3", "Steve")
// Iterate over the maps and print each user
users.Scan(func(key, value string) bool {
fmt.Printf("%s %s\n", key, value)
return true
})
fmt.Printf("\n")
// Delete a couple
users.Delete("user:5")
users.Delete("user:1")
// print the map again
users.Scan(func(key, value string) bool {
fmt.Printf("%s %s\n", key, value)
return true
})
fmt.Printf("\n")
// Output:
// user:1 Jane
// user:2 Andy
// user:3 Steve
// user:4 Andrea
// user:5 Janet
// user:6 Andy
//
// user:2 Andy
// user:3 Steve
// user:4 Andrea
// user:6 Andy
}
```
### btree.Set
```go
// Basic
Insert(key) // insert an item
Contains(key) // test if item exists
Delete(key) // delete an item
Len() // return the number of items in the set
// Iteration
Scan(iter) // scan items in ascending order
Reverse(iter) // scan items in descending order
Ascend(key, iter) // scan items in ascending order that are >= to key
Descend(key, iter) // scan items in descending order that are <= to key.
Iter() // returns a read-only iterator for for-loops.
// Array-like operations
GetAt(index) // returns the item at index
DeleteAt(index) // deletes the item at index
// Bulk-loading
Load(key) // load presorted item into tree
```
#### Example
```go
package main
import (
"fmt"
"github.com/tidwall/btree"
)
func main() {
// create a set
var names btree.Set[string]
// add some names
names.Insert("Jane")
names.Insert("Andrea")
names.Insert("Steve")
names.Insert("Andy")
names.Insert("Janet")
names.Insert("Andy")
// Iterate over the maps and print each user
names.Scan(func(key string) bool {
fmt.Printf("%s\n", key)
return true
})
fmt.Printf("\n")
// Delete a couple
names.Delete("Steve")
names.Delete("Andy")
// print the map again
names.Scan(func(key string) bool {
fmt.Printf("%s\n", key)
return true
})
fmt.Printf("\n")
// Output:
// Andrea
// Andy
// Jane
// Janet
// Steve
//
// Andrea
// Jane
// Janet
}
```
### btree.BTreeG
```go
// Basic
Set(item) // insert or replace an item
Get(item) // get an existing item
Delete(item) // delete an item
Len() // return the number of items in the btree
// Iteration
Scan(iter) // scan items in ascending order
Reverse(iter) // scan items in descending order
Ascend(key, iter) // scan items in ascending order that are >= to key
Descend(key, iter) // scan items in descending order that are <= to key.
Iter() // returns a read-only iterator for for-loops.
// Array-like operations
GetAt(index) // returns the item at index
DeleteAt(index) // deletes the item at index
// Bulk-loading
Load(item) // load presorted items into tree
// Path hinting
SetHint(item, *hint) // insert or replace an existing item
GetHint(item, *hint) // get an existing item
DeleteHint(item, *hint) // delete an item
AscendHint(key, iter, *hint)
DescendHint(key, iter, *hint)
SeekHint(key, iter, *hint)
// Copy-on-write
Copy() // copy the btree
```
#### Example
```go
package main
import (
"fmt"
"github.com/tidwall/btree"
)
type Item struct {
Key, Val string
}
// byKeys is a comparison function that compares item keys and returns true
// when a is less than b.
func byKeys(a, b Item) bool {
return a.Key < b.Key
}
// byVals is a comparison function that compares item values and returns true
// when a is less than b.
func byVals(a, b Item) bool {
if a.Val < b.Val {
return true
}
if a.Val > b.Val {
return false
}
// Both vals are equal so we should fall though
// and let the key comparison take over.
return byKeys(a, b)
}
func main() {
// Create a tree for keys and a tree for values.
// The "keys" tree will be sorted on the Keys field.
// The "values" tree will be sorted on the Values field.
keys := btree.NewBTreeG[Item](byKeys)
vals := btree.NewBTreeG[Item](byVals)
// Create some items.
users := []Item{
Item{Key: "user:1", Val: "Jane"},
Item{Key: "user:2", Val: "Andy"},
Item{Key: "user:3", Val: "Steve"},
Item{Key: "user:4", Val: "Andrea"},
Item{Key: "user:5", Val: "Janet"},
Item{Key: "user:6", Val: "Andy"},
}
// Insert each user into both trees
for _, user := range users {
keys.Set(user)
vals.Set(user)
}
// Iterate over each user in the key tree
keys.Scan(func(item Item) bool {
fmt.Printf("%s %s\n", item.Key, item.Val)
return true
})
fmt.Printf("\n")
// Iterate over each user in the val tree
vals.Scan(func(item Item) bool {
fmt.Printf("%s %s\n", item.Key, item.Val)
return true
})
// Output:
// user:1 Jane
// user:2 Andy
// user:3 Steve
// user:4 Andrea
// user:5 Janet
// user:6 Andy
//
// user:4 Andrea
// user:2 Andy
// user:6 Andy
// user:1 Jane
// user:5 Janet
// user:3 Steve
}
```
### btree.BTree
```go
// Basic
Set(item) // insert or replace an item
Get(item) // get an existing item
Delete(item) // delete an item
Len() // return the number of items in the btree
// Iteration
Scan(iter) // scan items in ascending order
Reverse(iter) // scan items in descending order
Ascend(key, iter) // scan items in ascending order that are >= to key
Descend(key, iter) // scan items in descending order that are <= to key.
Iter() // returns a read-only iterator for for-loops.
// Array-like operations
GetAt(index) // returns the item at index
DeleteAt(index) // deletes the item at index
// Bulk-loading
Load(item) // load presorted items into tree
// Path hinting
SetHint(item, *hint) // insert or replace an existing item
GetHint(item, *hint) // get an existing item
DeleteHint(item, *hint) // delete an item
AscendHint(key, iter, *hint)
DescendHint(key, iter, *hint)
SeekHint(key, iter, *hint)
// Copy-on-write
Copy() // copy the btree
```
#### Example
```go
package main
import (
"fmt"
"github.com/tidwall/btree"
)
type Item struct {
Key, Val string
}
// byKeys is a comparison function that compares item keys and returns true
// when a is less than b.
func byKeys(a, b interface{}) bool {
i1, i2 := a.(*Item), b.(*Item)
return i1.Key < i2.Key
}
// byVals is a comparison function that compares item values and returns true
// when a is less than b.
func byVals(a, b interface{}) bool {
i1, i2 := a.(*Item), b.(*Item)
if i1.Val < i2.Val {
return true
}
if i1.Val > i2.Val {
return false
}
// Both vals are equal so we should fall though
// and let the key comparison take over.
return byKeys(a, b)
}
func main() {
// Create a tree for keys and a tree for values.
// The "keys" tree will be sorted on the Keys field.
// The "values" tree will be sorted on the Values field.
keys := btree.New(byKeys)
vals := btree.New(byVals)
// Create some items.
users := []*Item{
&Item{Key: "user:1", Val: "Jane"},
&Item{Key: "user:2", Val: "Andy"},
&Item{Key: "user:3", Val: "Steve"},
&Item{Key: "user:4", Val: "Andrea"},
&Item{Key: "user:5", Val: "Janet"},
&Item{Key: "user:6", Val: "Andy"},
}
// Insert each user into both trees
for _, user := range users {
keys.Set(user)
vals.Set(user)
}
// Iterate over each user in the key tree
keys.Ascend(nil, func(item interface{}) bool {
kvi := item.(*Item)
fmt.Printf("%s %s\n", kvi.Key, kvi.Val)
return true
})
fmt.Printf("\n")
// Iterate over each user in the val tree
vals.Ascend(nil, func(item interface{}) bool {
kvi := item.(*Item)
fmt.Printf("%s %s\n", kvi.Key, kvi.Val)
return true
})
// Output:
// user:1 Jane
// user:2 Andy
// user:3 Steve
// user:4 Andrea
// user:5 Janet
// user:6 Andy
//
// user:4 Andrea
// user:2 Andy
// user:6 Andy
// user:1 Jane
// user:5 Janet
// user:3 Steve
}
```
## Performance
See [tidwall/btree-benchmark](https://github.com/tidwall/btree-benchmark) for benchmark numbers.
## Contact
Josh Baker [@tidwall](http://twitter.com/tidwall)
## License
Source code is available under the MIT [License](/LICENSE).
btree-1.8.1/btree.go 0000664 0000000 0000000 00000021206 15044500710 0014247 0 ustar 00root root 0000000 0000000 // 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
type BTree struct {
base *BTreeG[any]
}
// New returns a new BTree
func New(less func(a, b any) bool) *BTree {
if less == nil {
panic("nil less")
}
return &BTree{base: NewBTreeG(less)}
}
// NewNonConcurrent returns a new BTree which is not safe for concurrent
// write operations by multiple goroutines.
//
// This is useful for when you do not need the BTree to manage the locking,
// but would rather do it yourself.
//
// Deprecated: use NewOptions
func NewNonConcurrent(less func(a, b any) bool) *BTree {
if less == nil {
panic("nil less")
}
return &BTree{base: NewBTreeGOptions(less, Options{NoLocks: true})}
}
// NewOptions returns a new BTree
func NewOptions(less func(a, b any) bool, opts Options) *BTree {
if less == nil {
panic("nil less")
}
return &BTree{base: NewBTreeGOptions(less, opts)}
}
// Less is a convenience function that performs a comparison of two items
// using the same "less" function provided to New.
func (tr *BTree) Less(a, b any) bool {
return tr.base.Less(a, b)
}
// Set or replace a value for a key
// Returns the value for the replaced item or nil if the key was not found.
func (tr *BTree) Set(item any) (prev any) {
return tr.SetHint(item, nil)
}
// SetHint sets or replace a value for a key using a path hint
// Returns the value for the replaced item or nil if the key was not found.
func (tr *BTree) SetHint(item any, hint *PathHint) (prev any) {
if item == nil {
panic("nil item")
}
v, ok := tr.base.SetHint(item, hint)
if !ok {
return nil
}
return v
}
// Get a value for key.
// Returns nil if the key was not found.
func (tr *BTree) Get(key any) any {
return tr.getHintMut(key, nil, false)
}
func (tr *BTree) GetMut(key any) any {
return tr.getHintMut(key, nil, true)
}
func (tr *BTree) GetHint(key any, hint *PathHint) any {
return tr.getHintMut(key, hint, false)
}
func (tr *BTree) GetHintMut(key any, hint *PathHint) any {
return tr.getHintMut(key, hint, true)
}
// GetHint gets a value for key using a path hint.
// Returns nil if the item was not found.
func (tr *BTree) getHintMut(key any, hint *PathHint, mut bool) (value any) {
if key == nil {
return nil
}
var v any
var ok bool
if mut {
v, ok = tr.base.GetHintMut(key, hint)
} else {
v, ok = tr.base.GetHint(key, hint)
}
if !ok {
return nil
}
return v
}
// Len returns the number of items in the tree
func (tr *BTree) Len() int {
return tr.base.Len()
}
// Delete an item for a key.
// Returns the deleted value or nil if the key was not found.
func (tr *BTree) Delete(key any) (prev any) {
return tr.DeleteHint(key, nil)
}
// DeleteHint deletes a value for a key using a path hint
// Returns the deleted value or nil if the key was not found.
func (tr *BTree) DeleteHint(key any, hint *PathHint) (prev any) {
if key == nil {
return nil
}
v, ok := tr.base.DeleteHint(key, hint)
if !ok {
return nil
}
return v
}
// 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 *BTree) Ascend(pivot any, iter func(item any) bool) {
if pivot == nil {
tr.base.Scan(iter)
} else {
tr.base.Ascend(pivot, iter)
}
}
func (tr *BTree) AscendMut(pivot any, iter func(item any) bool) {
if pivot == nil {
tr.base.ScanMut(iter)
} else {
tr.base.AscendMut(pivot, iter)
}
}
func (tr *BTree) AscendHint(pivot any, iter func(item any) bool,
hint *PathHint,
) {
if pivot == nil {
tr.base.Scan(iter)
} else {
tr.base.AscendHint(pivot, iter, hint)
}
}
func (tr *BTree) AscendHintMut(pivot any, iter func(item any) bool,
hint *PathHint,
) {
if pivot == nil {
tr.base.ScanMut(iter)
} else {
tr.base.AscendHintMut(pivot, iter, hint)
}
}
// 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 *BTree) Descend(pivot any, iter func(item any) bool) {
if pivot == nil {
tr.base.Reverse(iter)
} else {
tr.base.Descend(pivot, iter)
}
}
func (tr *BTree) DescendMut(pivot any, iter func(item any) bool) {
if pivot == nil {
tr.base.ReverseMut(iter)
} else {
tr.base.DescendMut(pivot, iter)
}
}
func (tr *BTree) DescendHint(pivot any, iter func(item any) bool,
hint *PathHint,
) {
if pivot == nil {
tr.base.Reverse(iter)
} else {
tr.base.DescendHint(pivot, iter, hint)
}
}
func (tr *BTree) DescendHintMut(pivot any, iter func(item any) bool,
hint *PathHint,
) {
if pivot == nil {
tr.base.ReverseMut(iter)
} else {
tr.base.DescendHintMut(pivot, iter, hint)
}
}
// Load is for bulk loading pre-sorted items
// If the load replaces and existing item then the value for the replaced item
// is returned.
func (tr *BTree) Load(item any) (prev any) {
if item == nil {
panic("nil item")
}
v, ok := tr.base.Load(item)
if !ok {
return nil
}
return v
}
// Min returns the minimum item in tree.
// Returns nil if the tree has no items.
func (tr *BTree) Min() any {
v, ok := tr.base.Min()
if !ok {
return nil
}
return v
}
func (tr *BTree) MinMut() any {
v, ok := tr.base.MinMut()
if !ok {
return nil
}
return v
}
// Max returns the maximum item in tree.
// Returns nil if the tree has no items.
func (tr *BTree) Max() any {
v, ok := tr.base.Max()
if !ok {
return nil
}
return v
}
func (tr *BTree) MaxMut() any {
v, ok := tr.base.Max()
if !ok {
return nil
}
return v
}
// PopMin removes the minimum item in tree and returns it.
// Returns nil if the tree has no items.
func (tr *BTree) PopMin() any {
v, ok := tr.base.PopMin()
if !ok {
return nil
}
return v
}
// PopMax removes the maximum item in tree and returns it.
// Returns nil if the tree has no items.
func (tr *BTree) PopMax() any {
v, ok := tr.base.PopMax()
if !ok {
return nil
}
return v
}
// GetAt returns the value at index.
// Return nil if the tree is empty or the index is out of bounds.
func (tr *BTree) GetAt(index int) any {
v, ok := tr.base.GetAt(index)
if !ok {
return nil
}
return v
}
func (tr *BTree) GetAtMut(index int) any {
v, ok := tr.base.GetAtMut(index)
if !ok {
return nil
}
return v
}
// DeleteAt deletes the item at index.
// Return nil if the tree is empty or the index is out of bounds.
func (tr *BTree) DeleteAt(index int) any {
v, ok := tr.base.DeleteAt(index)
if !ok {
return nil
}
return v
}
// Height returns the height of the tree.
// Returns zero if tree has no items.
func (tr *BTree) Height() int {
return tr.base.Height()
}
// Walk iterates over all items in tree, in order.
// The items param will contain one or more items.
func (tr *BTree) Walk(iter func(items []any)) {
tr.base.Walk(func(items []any) bool {
iter(items)
return true
})
}
func (tr *BTree) WalkMut(iter func(items []any)) {
tr.base.WalkMut(func(items []any) bool {
iter(items)
return true
})
}
// Copy the tree. This is a copy-on-write operation and is very fast because
// it only performs a shadowed copy.
func (tr *BTree) Copy() *BTree {
return &BTree{base: tr.base.Copy()}
}
func (tr *BTree) IsoCopy() *BTree {
return &BTree{base: tr.base.IsoCopy()}
}
// Clear will delete all items.
func (tr *BTree) Clear() {
tr.base.Clear()
}
// Iter is an iterator for
type Iter struct {
base IterG[any]
}
// Iter returns a read-only iterator.
// The Release method must be called finished with iterator.
func (tr *BTree) Iter() Iter {
return Iter{tr.base.Iter()}
}
func (tr *BTree) IterMut() Iter {
return Iter{tr.base.IterMut()}
}
// Seek to item greater-or-equal-to key.
// Returns false if there was no item found.
func (iter *Iter) Seek(key any) bool {
return iter.base.Seek(key)
}
func (iter *Iter) SeekHint(key any, hint *PathHint) bool {
return iter.base.SeekHint(key, hint)
}
// First moves iterator to first item in tree.
// Returns false if the tree is empty.
func (iter *Iter) First() bool {
return iter.base.First()
}
// Last moves iterator to last item in tree.
// Returns false if the tree is empty.
func (iter *Iter) Last() bool {
return iter.base.Last()
}
// First moves iterator to first item in tree.
// Returns false if the tree is empty.
func (iter *Iter) Release() {
iter.base.Release()
}
// 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 *Iter) Next() bool {
return iter.base.Next()
}
// 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 *Iter) Prev() bool {
return iter.base.Prev()
}
// Item returns the current iterator item.
func (iter *Iter) Item() any {
return iter.base.Item()
}
btree-1.8.1/btree_test.go 0000664 0000000 0000000 00000011467 15044500710 0015316 0 ustar 00root root 0000000 0000000 package btree
import (
"sync"
"testing"
)
func assert(x bool) {
if !x {
panic("assert failed")
}
}
func intLess(a, b interface{}) bool {
return a.(int) < b.(int)
}
func TestBTree(t *testing.T) {
func() {
defer func() {
msg, ok := recover().(string)
assert(ok && msg == "nil less")
}()
New(nil)
}()
func() {
defer func() {
msg, ok := recover().(string)
assert(ok && msg == "nil less")
}()
NewNonConcurrent(nil)
}()
N := 1_000_000
for j := 0; j < 2; j++ {
var tr *BTree
if j == 0 {
tr = New(intLess)
} else {
tr = NewNonConcurrent(intLess)
}
for i := 0; i < N; i++ {
assert(tr.Load(i) == nil)
}
assert(tr.Len() == N)
for i := 0; i < N; i++ {
assert(tr.Get(i) == i)
}
count := 0
tr.Ascend(nil, func(_ interface{}) bool {
count++
return true
})
assert(count == N)
count = 0
tr.Ascend(N/2, func(_ interface{}) bool {
count++
return true
})
assert(count == N/2)
count = 0
tr.Descend(nil, func(_ interface{}) bool {
count++
return true
})
assert(count == N)
count = 0
tr.Descend(N/2, func(_ interface{}) bool {
count++
return true
})
assert(count == N/2+1)
for i := 0; i < N; i++ {
assert(tr.Delete(i) == i)
}
assert(tr.Len() == 0)
assert(tr.Min() == nil)
assert(tr.Max() == nil)
assert(tr.PopMin() == nil)
assert(tr.PopMax() == nil)
for i := 0; i < N; i++ {
assert(tr.Get(i) == nil)
}
for i := 0; i < N; i++ {
assert(tr.Set(i) == nil)
}
assert(tr.Len() == N)
var hint PathHint
for i := 0; i < N; i++ {
assert(tr.SetHint(i, &hint) == i)
}
assert(tr.Len() == N)
for i := 0; i < N; i++ {
assert(tr.Load(i) == i)
}
assert(tr.Len() == N)
assert(tr.Min() == 0)
assert(tr.Max() == N-1)
assert(tr.PopMin() == 0)
assert(tr.PopMax() == N-1)
assert(tr.Set(0) == nil)
assert(tr.Set(N-1) == nil)
assert(tr.GetAt(0) == 0)
assert(tr.GetAt(N) == nil)
assert(tr.Set(N-1) == N-1)
assert(tr.Height() > 0)
assert(tr.DeleteAt(0) == 0)
assert(tr.Set(0) == nil)
assert(tr.DeleteAt(N-1) == N-1)
assert(tr.DeleteAt(N) == nil)
var wg sync.WaitGroup
wg.Add(1)
go func(tr *BTree) {
wg.Wait()
count := 0
tr.Walk(func(items []interface{}) {
count += len(items)
})
assert(count == N-1)
}(tr.Copy())
for i := 0; i < N/2; i++ {
tr.Delete(i)
}
for i := 0; i < N; i++ {
tr.Set(i)
}
wg.Done()
count = 0
tr.Walk(func(items []interface{}) {
count += len(items)
})
assert(count == N)
func() {
defer func() {
msg, ok := recover().(string)
assert(ok && msg == "nil item")
}()
tr := NewNonConcurrent(intLess)
tr.Set(nil)
}()
func() {
defer func() {
msg, ok := recover().(string)
assert(ok && msg == "nil item")
}()
tr := NewNonConcurrent(intLess)
tr.Load(nil)
}()
assert(tr.Get(nil) == nil)
assert(tr.Delete(nil) == nil)
for i := 0; i < N; i++ {
assert(tr.GetHint(i, &hint) == i)
}
for i := 0; i < N; i++ {
assert(tr.DeleteHint(i, &hint) == i)
}
for i := 0; i < N; i++ {
assert(tr.GetHint(i, &hint) == nil)
}
for i := 0; i < N; i++ {
assert(tr.DeleteHint(i, &hint) == nil)
}
assert(tr.Less(1, 2))
assert(tr.Less(2, 10))
}
}
func TestClear(t *testing.T) {
tr := New(intLess)
for i := 0; i < 100; i++ {
tr.Set(i)
}
assert(tr.Len() == 100)
tr.Clear()
assert(tr.Len() == 0)
for i := 0; i < 100; i++ {
tr.Set(i)
}
assert(tr.Len() == 100)
}
func TestIter(t *testing.T) {
N := 100_000
lt := func(a, b interface{}) bool { return a.(int) < b.(int) }
eq := func(a, b interface{}) bool { return !lt(a, b) && !lt(b, a) }
tr := New(lt)
var all []int
for i := 0; i < N; i++ {
tr.Load(i)
all = append(all, i)
}
var count int
var i int
iter := tr.Iter()
for ok := iter.First(); ok; ok = iter.Next() {
if !eq(all[i], iter.Item()) {
panic("!")
}
count++
i++
}
if count != N {
t.Fatalf("expected %v, got %v", N, count)
}
iter.Release()
count = 0
i = len(all) - 1
iter = tr.Iter()
for ok := iter.Last(); ok; ok = iter.Prev() {
if !eq(all[i], iter.Item()) {
panic("!")
}
i--
count++
}
if count != N {
t.Fatalf("expected %v, got %v", N, count)
}
iter.Release()
i = 0
iter = tr.Iter()
for ok := iter.First(); ok; ok = iter.Next() {
if !eq(all[i], iter.Item()) {
panic("!")
}
i++
}
i--
for ok := iter.Prev(); ok; ok = iter.Prev() {
i--
if !eq(all[i], iter.Item()) {
panic("!")
}
}
if i != 0 {
panic("!")
}
i++
for ok := iter.Next(); ok; ok = iter.Next() {
if !eq(all[i], iter.Item()) {
panic("!")
}
i++
}
if i != N {
panic("!")
}
i = 0
for ok := iter.First(); ok; ok = iter.Next() {
if !eq(all[i], iter.Item()) {
panic("!")
}
if eq(iter.Item(), N/2) {
for ok = iter.Prev(); ok; ok = iter.Prev() {
i--
if !eq(all[i], iter.Item()) {
panic("!")
}
}
break
}
i++
}
iter.Release()
}
btree-1.8.1/btreeg.go 0000664 0000000 0000000 00000122577 15044500710 0014433 0 ustar 00root root 0000000 0000000 // 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"
type BTreeG[T any] struct {
isoid uint64
mu *sync.RWMutex
root *node[T]
count int
locks bool
copyItems bool
isoCopyItems bool
less func(a, b T) bool
empty T
max int
min int
}
type node[T any] struct {
isoid uint64
count int
items []T
children *[]*node[T]
}
// PathHint is a utility type used with the *Hint() functions. Hints provide
// faster operations for clustered keys.
type PathHint struct {
used [8]bool
path [8]uint8
}
// Options for passing to New when creating a new BTree.
type Options struct {
// Degree is used to define how many items and children each internal node
// can contain before it must branch. For example, a degree of 2 will
// create a 2-3-4 tree, where each node may contains 1-3 items and
// 2-4 children. See https://en.wikipedia.org/wiki/2–3–4_tree.
// Default is 32
Degree int
// NoLocks will disable locking. Otherwide a sync.RWMutex is used to
// ensure all operations are safe across multiple goroutines.
NoLocks bool
}
// New returns a new BTree
func NewBTreeG[T any](less func(a, b T) bool) *BTreeG[T] {
return NewBTreeGOptions(less, Options{})
}
func NewBTreeGOptions[T any](less func(a, b T) bool, opts Options) *BTreeG[T] {
tr := new(BTreeG[T])
tr.isoid = newIsoID()
tr.locks = !opts.NoLocks
if tr.locks {
tr.mu = new(sync.RWMutex)
}
tr.less = less
tr.init(opts.Degree)
return tr
}
func (tr *BTreeG[T]) init(degree int) {
if tr.min != 0 {
return
}
tr.min, tr.max = degreeToMinMax(degree)
_, tr.copyItems = ((interface{})(tr.empty)).(copier[T])
if !tr.copyItems {
_, tr.isoCopyItems = ((interface{})(tr.empty)).(isoCopier[T])
}
}
// Less is a convenience function that performs a comparison of two items
// using the same "less" function provided to New.
func (tr *BTreeG[T]) Less(a, b T) bool {
return tr.less(a, b)
}
func (tr *BTreeG[T]) newNode(leaf bool) *node[T] {
n := &node[T]{isoid: tr.isoid}
if !leaf {
n.children = new([]*node[T])
}
return n
}
// leaf returns true if the node is a leaf.
func (n *node[T]) leaf() bool {
return n.children == nil
}
func (tr *BTreeG[T]) bsearch(n *node[T], key T) (index int, found bool) {
low, high := 0, len(n.items)
for low < high {
h := int(uint(low+high) >> 1) // avoid overflow when computing h
if !tr.less(key, n.items[h]) {
low = h + 1
} else {
high = h
}
}
if low > 0 && !tr.less(n.items[low-1], key) {
return low - 1, true
}
return low, false
}
func (tr *BTreeG[T]) find(n *node[T], key T, hint *PathHint, depth int,
) (index int, found bool) {
if hint == nil {
return tr.bsearch(n, key)
}
return tr.hintsearch(n, key, hint, depth)
}
func (tr *BTreeG[T]) hintsearch(n *node[T], key T, hint *PathHint, depth int,
) (index int, found bool) {
// Best case finds the exact match, updates the hint and returns.
// Worst case, updates the low and high bounds to binary search between.
low := 0
high := len(n.items) - 1
if depth < 8 && hint.used[depth] {
index = int(hint.path[depth])
if index >= len(n.items) {
// tail item
if tr.Less(n.items[len(n.items)-1], key) {
index = len(n.items)
goto path_match
}
index = len(n.items) - 1
}
if tr.Less(key, n.items[index]) {
if index == 0 || tr.Less(n.items[index-1], key) {
goto path_match
}
high = index - 1
} else if tr.Less(n.items[index], key) {
low = index + 1
} else {
found = true
goto path_match
}
}
// Do a binary search between low and high
// keep on going until low > high, where the guarantee on low is that
// key >= items[low - 1]
for low <= high {
mid := low + ((high+1)-low)/2
// if key >= n.items[mid], low = mid + 1
// which implies that key >= everything below low
if !tr.Less(key, n.items[mid]) {
low = mid + 1
} else {
high = mid - 1
}
}
// if low > 0, n.items[low - 1] >= key,
// we have from before that key >= n.items[low - 1]
// therefore key = n.items[low - 1],
// and we have found the entry for key.
// Otherwise we must keep searching for the key in index `low`.
if low > 0 && !tr.Less(n.items[low-1], key) {
index = low - 1
found = true
} else {
index = low
found = false
}
path_match:
if depth < 8 {
hint.used[depth] = true
var pathIndex uint8
if n.leaf() && found {
pathIndex = uint8(index + 1)
} else {
pathIndex = uint8(index)
}
if pathIndex != hint.path[depth] {
hint.path[depth] = pathIndex
for i := depth + 1; i < 8; i++ {
hint.used[i] = false
}
}
}
return index, found
}
// SetHint sets or replace a value for a key using a path hint
func (tr *BTreeG[T]) SetHint(item T, hint *PathHint) (prev T, replaced bool) {
if tr.locks {
tr.mu.Lock()
prev, replaced = tr.setHint(item, hint)
tr.mu.Unlock()
} else {
prev, replaced = tr.setHint(item, hint)
}
return prev, replaced
}
func (tr *BTreeG[T]) setHint(item T, hint *PathHint) (prev T, replaced bool) {
if tr.root == nil {
tr.init(0)
tr.root = tr.newNode(true)
tr.root.items = append([]T{}, item)
tr.root.count = 1
tr.count = 1
return tr.empty, false
}
prev, replaced, split := tr.nodeSet(&tr.root, item, hint, 0)
if split {
left := tr.isoLoad(&tr.root, true)
right, median := tr.nodeSplit(left)
tr.root = tr.newNode(false)
*tr.root.children = make([]*node[T], 0, tr.max+1)
*tr.root.children = append([]*node[T]{}, left, right)
tr.root.items = append([]T{}, median)
tr.root.updateCount()
return tr.setHint(item, hint)
}
if replaced {
return prev, true
}
tr.count++
return tr.empty, false
}
// Set or replace a value for a key
func (tr *BTreeG[T]) Set(item T) (T, bool) {
return tr.SetHint(item, nil)
}
func (tr *BTreeG[T]) nodeSplit(n *node[T]) (right *node[T], median T) {
i := tr.max / 2
median = n.items[i]
// right node
right = tr.newNode(n.leaf())
right.items = n.items[i+1:]
if !n.leaf() {
*right.children = (*n.children)[i+1:]
}
right.updateCount()
// left node
n.items[i] = tr.empty
n.items = n.items[:i:i]
if !n.leaf() {
*n.children = (*n.children)[: i+1 : i+1]
}
n.updateCount()
return right, median
}
func (n *node[T]) updateCount() {
n.count = len(n.items)
if !n.leaf() {
for i := 0; i < len(*n.children); i++ {
n.count += (*n.children)[i].count
}
}
}
// Copy the node for safe isolation.
func (tr *BTreeG[T]) copy(n *node[T]) *node[T] {
n2 := new(node[T])
n2.isoid = tr.isoid
n2.count = n.count
n2.items = make([]T, len(n.items), cap(n.items))
copy(n2.items, n.items)
if tr.copyItems {
for i := 0; i < len(n2.items); i++ {
n2.items[i] = ((interface{})(n2.items[i])).(copier[T]).Copy()
}
} else if tr.isoCopyItems {
for i := 0; i < len(n2.items); i++ {
n2.items[i] = ((interface{})(n2.items[i])).(isoCopier[T]).IsoCopy()
}
}
if !n.leaf() {
n2.children = new([]*node[T])
*n2.children = make([]*node[T], len(*n.children), tr.max+1)
copy(*n2.children, *n.children)
}
return n2
}
// isoLoad loads the provided node and, if needed, performs a copy-on-write.
func (tr *BTreeG[T]) isoLoad(cn **node[T], mut bool) *node[T] {
if mut && (*cn).isoid != tr.isoid {
*cn = tr.copy(*cn)
}
return *cn
}
func (tr *BTreeG[T]) nodeSet(cn **node[T], item T,
hint *PathHint, depth int,
) (prev T, replaced bool, split bool) {
if (*cn).isoid != tr.isoid {
*cn = tr.copy(*cn)
}
n := *cn
var i int
var found bool
if hint == nil {
i, found = tr.bsearch(n, item)
} else {
i, found = tr.hintsearch(n, item, hint, depth)
}
if found {
prev = n.items[i]
n.items[i] = item
return prev, true, false
}
if n.leaf() {
if len(n.items) == tr.max {
return tr.empty, 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, false, false
}
prev, replaced, split = tr.nodeSet(&(*n.children)[i], item, hint, depth+1)
if split {
if len(n.items) == tr.max {
return tr.empty, 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, hint, depth)
}
if !replaced {
n.count++
}
return prev, replaced, false
}
func (tr *BTreeG[T]) Scan(iter func(item T) bool) {
tr.scan(iter, false)
}
func (tr *BTreeG[T]) ScanMut(iter func(item T) bool) {
tr.scan(iter, true)
}
func (tr *BTreeG[T]) scan(iter func(item T) bool, mut bool) {
if tr.lock(mut) {
defer tr.unlock(mut)
}
if tr.root == nil {
return
}
tr.nodeScan(&tr.root, iter, mut)
}
func (tr *BTreeG[T]) nodeScan(cn **node[T], iter func(item T) bool, mut bool,
) bool {
n := tr.isoLoad(cn, mut)
if n.leaf() {
for i := 0; i < len(n.items); i++ {
if !iter(n.items[i]) {
return false
}
}
return true
}
for i := 0; i < len(n.items); i++ {
if !tr.nodeScan(&(*n.children)[i], iter, mut) {
return false
}
if !iter(n.items[i]) {
return false
}
}
return tr.nodeScan(&(*n.children)[len(*n.children)-1], iter, mut)
}
// Get a value for key
func (tr *BTreeG[T]) Get(key T) (T, bool) {
return tr.getHint(key, nil, false)
}
func (tr *BTreeG[T]) GetMut(key T) (T, bool) {
return tr.getHint(key, nil, true)
}
// GetHint gets a value for key using a path hint
func (tr *BTreeG[T]) GetHint(key T, hint *PathHint) (value T, ok bool) {
return tr.getHint(key, hint, false)
}
func (tr *BTreeG[T]) GetHintMut(key T, hint *PathHint) (value T, ok bool) {
return tr.getHint(key, hint, true)
}
// GetHint gets a value for key using a path hint
func (tr *BTreeG[T]) getHint(key T, hint *PathHint, mut bool) (T, bool) {
if tr.lock(mut) {
defer tr.unlock(mut)
}
if tr.root == nil {
return tr.empty, false
}
n := tr.isoLoad(&tr.root, mut)
depth := 0
for {
i, found := tr.find(n, key, hint, depth)
if found {
return n.items[i], true
}
if n.children == nil {
return tr.empty, false
}
n = tr.isoLoad(&(*n.children)[i], mut)
depth++
}
}
// Action for DeleteAscend
type Action int
const (
Stop Action = iota
Keep
Delete
)
// DeleteAscend ascends over tree within the range [pivot, last].
// Return Delete to delete the item.
// Return Keep to keep the item, avoiding deletion.
// Return Stop to stop iterating
func (tr *BTreeG[T]) DeleteAscend(pivot T, iter func(item T) Action) {
if tr.lock(true) {
defer tr.unlock(true)
}
var hint PathHint
type stackItem struct {
node *node[T]
index int
}
var stack []stackItem
restart:
if tr.root == nil {
return
}
n := tr.isoLoad(&tr.root, true)
stack = append(stack, stackItem{n, 0})
for {
i, found := tr.find(n, pivot, &hint, len(stack)-1)
next:
if n.children != nil {
if found {
act := iter(n.items[i])
switch act {
case Delete:
pivot = n.items[i]
tr.deleteHint(n.items[i], &hint)
stack = stack[:0]
goto restart
case Keep:
n = tr.isoLoad(&(*n.children)[i+1], true)
stack = append(stack, stackItem{n, i + 1})
i = 0
for !n.leaf() {
n = tr.isoLoad(&(*n.children)[i], true)
stack = append(stack, stackItem{n, i})
}
default:
// user stop requested
return
}
} else {
// at branch, continue to leaf
n = tr.isoLoad(&(*n.children)[i], true)
stack = append(stack, stackItem{n, i})
continue
}
}
// at leaf
maxbulk := len(n.items) - i - tr.min
if maxbulk > 1 {
var act Action
j := 0
for ; j < maxbulk; j++ {
act = iter(n.items[i+j])
if act != Delete {
break
}
}
copy(n.items[i:], n.items[i+j:])
for k := len(n.items) - j; k < len(n.items); k++ {
n.items[k] = tr.empty
}
n.items = n.items[:len(n.items)-j]
for k := 0; k < len(stack); k++ {
stack[k].node.count -= j
}
tr.count -= j
if act == Stop {
return
}
if act == Keep {
i++
}
}
for ; i < len(n.items); i++ {
act := iter(n.items[i])
switch act {
case Delete:
if len(n.items) > tr.min {
copy(n.items[i:], n.items[i+1:])
n.items[len(n.items)-1] = tr.empty
n.items = n.items[:len(n.items)-1]
for j := 0; j < len(stack); j++ {
stack[j].node.count--
}
tr.count--
i--
} else {
pivot = n.items[i]
tr.deleteHint(n.items[i], &hint)
stack = stack[:0]
goto restart
}
case Keep:
default:
// user stop requested
return
}
}
// end of leaf items. traverse upwards
for {
i = stack[len(stack)-1].index
stack = stack[:len(stack)-1]
if len(stack) == 0 {
// end of tree
return
}
n = stack[len(stack)-1].node
if i < len(n.items) {
found = true
goto next
}
}
}
}
type eitem[T any] struct {
item T
node *node[T]
}
// List of deleted items from the DeleteRange function.
type List[T any] struct {
less func(a, b T) bool // comparator
count int // total number of items
np int // number of pop eitems
q []eitem[T] // priority queue or reversed/popped array
}
func (e *List[T]) push(eitem eitem[T]) {
e.q = append(e.q, eitem)
i := len(e.q) - 1
parent := (i - 1) / 2
for i != 0 && e.less(e.q[i].item, e.q[parent].item) {
e.q[parent], e.q[i] = e.q[i], e.q[parent]
i = parent
parent = (i - 1) / 2
}
}
func (e *List[T]) pop() (eitem[T], bool) {
if len(e.q) == 0 {
return eitem[T]{}, false
}
var n eitem[T]
n, e.q[0] = e.q[0], e.q[len(e.q)-1]
e.q = e.q[:len(e.q)-1]
i := 0
for {
smallest := i
left := i*2 + 1
right := i*2 + 2
if left < len(e.q) && !e.less(e.q[smallest].item, e.q[left].item) {
smallest = left
}
if right < len(e.q) && !e.less(e.q[smallest].item, e.q[right].item) {
smallest = right
}
if smallest == i {
break
}
e.q[smallest], e.q[i] = e.q[i], e.q[smallest]
i = smallest
}
return n, true
}
func (e *List[T]) append(item T, n *node[T]) {
e.count++
if n != nil {
e.count += n.count
}
e.push(eitem[T]{item, n})
}
func extractNodeScan[T any](n *node[T], iter func(item T) bool) bool {
if n.leaf() {
for i := 0; i < len(n.items); i++ {
if !iter(n.items[i]) {
return false
}
}
return true
}
for i := 0; i < len(n.items); i++ {
if !extractNodeScan((*n.children)[i], iter) {
return false
}
if !iter(n.items[i]) {
return false
}
}
return extractNodeScan((*n.children)[len(*n.children)-1], iter)
}
// Len returns the number of items in the list
func (e *List[T]) Len() int {
return e.count
}
// Scan performs an ordered scan of all items in the list.
func (e *List[T]) Scan(iter func(item T) bool) {
pitems := e.q[:cap(e.q)]
for i := 0; i < e.np; i++ {
eitem := pitems[cap(pitems)-i-1]
if !iter(eitem.item) {
return
}
if eitem.node != nil {
if !extractNodeScan(eitem.node, iter) {
return
}
}
}
for {
// pop the next item
eitem, ok := e.pop()
if !ok {
return
}
// Place item in reverse order for the extract can be scanned again.
e.q[:cap(e.q)][cap(e.q)-e.np-1] = eitem
e.np++
if !iter(eitem.item) {
return
}
if eitem.node != nil {
if !extractNodeScan(eitem.node, iter) {
return
}
}
}
}
func (e *List[T]) Clear() {
e.less = nil
e.q = e.q[:0]
e.count = 0
e.np = 0
}
// Options for passing to the DeleteRange function
type DeleteRangeOptions struct {
// Do not return the deleted items.
// Default false
NoReturn bool
// Make the max inclusive.
// Default false
MaxInclusive bool
}
// DeleteRange will delete all items within the provided min (inclusive) and
// max (exclusive) sub-range.
// Returns the deleted items as an ordered list that can be iterated over using
// the list.Scan() method.
func (tr *BTreeG[T]) DeleteRange(min, max T, opts *DeleteRangeOptions) (deleted List[T]) {
return tr.DeleteRangeReuse(min, max, opts, List[T]{})
}
// DeleteRangeReuse is the same as DeleteRange, but it takes a List as an argument to
// avoid allocating/growing a new List on each call to DeleteRange. It is unsafe to use
// the same List across concurrent calls to DeleteRange.
func (tr *BTreeG[T]) DeleteRangeReuse(min, max T, opts *DeleteRangeOptions, deleted List[T]) List[T] {
if tr.lock(true) {
defer tr.unlock(true)
}
extract := opts == nil || !opts.NoReturn
maxincl := opts != nil && opts.MaxInclusive
// Clear just in case it hasn't been cleared yet.
deleted.Clear()
maxstop := func(item T) bool {
if maxincl {
return tr.less(max, item)
}
return !tr.less(item, max)
}
if extract {
deleted.less = tr.less
}
var hint PathHint
type stackItem struct {
node *node[T]
index int
}
pivot := min
var stack []stackItem
restart:
if tr.root == nil {
return deleted
}
n := tr.isoLoad(&tr.root, true)
stack = append(stack, stackItem{n, 0})
for {
i, found := tr.find(n, pivot, &hint, len(stack)-1)
next:
if n.children != nil {
// Can the item and child to the right be completely deleted?
if i < len(n.items)-1 && (len(n.items) > tr.min ||
(len(n.items) > 1 && len(stack) == 1)) &&
tr.less(n.items[i+1], max) {
// Yes, delete both without traversing
ditem := 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]
dnode := (*n.children)[i+1]
copy((*n.children)[i+1:], (*n.children)[i+2:])
(*n.children)[len(*n.children)-1] = nil
*n.children = (*n.children)[:len(*n.children)-1]
for k := 0; k < len(stack); k++ {
stack[k].node.count -= dnode.count + 1
}
tr.count -= dnode.count + 1
if extract {
deleted.append(ditem, dnode)
}
if found {
goto next
}
continue
}
if found {
if maxstop(n.items[i]) {
// stop
return deleted
}
pivot = n.items[i]
if extract {
deleted.append(n.items[i], nil)
}
tr.deleteHint(n.items[i], &hint)
stack = stack[:0]
goto restart
}
// at branch, continue to leaf
n = tr.isoLoad(&(*n.children)[i], true)
stack = append(stack, stackItem{n, i})
continue
}
// at leaf
maxbulk := len(n.items) - i - tr.min
if maxbulk > 1 {
var stop bool
j := 0
for ; j < maxbulk; j++ {
stop := maxstop(n.items[i+j])
if stop {
break
}
if extract {
deleted.append(n.items[i+j], nil)
}
}
copy(n.items[i:], n.items[i+j:])
for k := len(n.items) - j; k < len(n.items); k++ {
n.items[k] = tr.empty
}
n.items = n.items[:len(n.items)-j]
for k := 0; k < len(stack); k++ {
stack[k].node.count -= j
}
tr.count -= j
if stop {
return deleted
}
}
for ; i < len(n.items); i++ {
if maxstop(n.items[i]) {
// stop
return deleted
}
if len(n.items) > tr.min {
if extract {
deleted.append(n.items[i], nil)
}
copy(n.items[i:], n.items[i+1:])
n.items[len(n.items)-1] = tr.empty
n.items = n.items[:len(n.items)-1]
for j := 0; j < len(stack); j++ {
stack[j].node.count--
}
tr.count--
i--
} else {
pivot = n.items[i]
if extract {
deleted.append(n.items[i], nil)
}
tr.deleteHint(n.items[i], &hint)
stack = stack[:0]
goto restart
}
}
// end of leaf items. traverse upwards
for {
i = stack[len(stack)-1].index
stack = stack[:len(stack)-1]
if len(stack) == 0 {
// end of tree
return deleted
}
n = stack[len(stack)-1].node
if i < len(n.items) {
found = true
goto next
}
}
}
}
// Len returns the number of items in the tree
func (tr *BTreeG[T]) 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 *BTreeG[T]) Delete(key T) (T, bool) {
return tr.DeleteHint(key, nil)
}
// DeleteHint deletes a value for a key using a path hint and returns the
// deleted value.
// Returns false if there was no value by that key found.
func (tr *BTreeG[T]) DeleteHint(key T, hint *PathHint) (T, bool) {
if tr.lock(true) {
defer tr.unlock(true)
}
return tr.deleteHint(key, hint)
}
func (tr *BTreeG[T]) deleteHint(key T, hint *PathHint) (T, bool) {
if tr.root == nil {
return tr.empty, false
}
prev, deleted := tr.delete(&tr.root, false, key, hint, 0)
if !deleted {
return tr.empty, 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, true
}
func (tr *BTreeG[T]) delete(cn **node[T], max bool, key T,
hint *PathHint, depth int,
) (T, bool) {
n := tr.isoLoad(cn, true)
var i int
var found bool
if max {
i, found = len(n.items)-1, true
} else {
i, found = tr.find(n, key, hint, depth)
}
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 T
var deleted bool
if found {
if max {
i++
prev, deleted = tr.delete(&(*n.children)[i], true, tr.empty, nil, 0)
} else {
prev = n.items[i]
maxItem, _ := tr.delete(&(*n.children)[i], true, tr.empty, nil, 0)
deleted = true
n.items[i] = maxItem
}
} else {
prev, deleted = tr.delete(&(*n.children)[i], max, key, hint, depth+1)
}
if !deleted {
return tr.empty, false
}
n.count--
if len((*n.children)[i].items) < tr.min {
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 *BTreeG[T]) nodeRebalance(n *node[T], i int) {
if i == len(n.items) {
i--
}
// ensure copy-on-write
left := tr.isoLoad(&(*n.children)[i], true)
right := tr.isoLoad(&(*n.children)[i+1], true)
if len(left.items)+len(right.items) < tr.max {
// 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 *BTreeG[T]) Ascend(pivot T, iter func(item T) bool) {
tr.ascend(pivot, iter, false, nil)
}
func (tr *BTreeG[T]) AscendMut(pivot T, iter func(item T) bool) {
tr.ascend(pivot, iter, true, nil)
}
func (tr *BTreeG[T]) ascend(pivot T, iter func(item T) bool, mut bool,
hint *PathHint,
) {
if tr.lock(mut) {
defer tr.unlock(mut)
}
if tr.root == nil {
return
}
tr.nodeAscend(&tr.root, pivot, hint, 0, iter, mut)
}
func (tr *BTreeG[T]) AscendHint(pivot T, iter func(item T) bool, hint *PathHint,
) {
tr.ascend(pivot, iter, false, hint)
}
func (tr *BTreeG[T]) AscendHintMut(pivot T, iter func(item T) bool,
hint *PathHint,
) {
tr.ascend(pivot, iter, true, hint)
}
// The return value of this function determines whether we should keep iterating
// upon this functions return.
func (tr *BTreeG[T]) nodeAscend(cn **node[T], pivot T, hint *PathHint,
depth int, iter func(item T) bool, mut bool,
) bool {
n := tr.isoLoad(cn, mut)
i, found := tr.find(n, pivot, hint, depth)
if !found {
if !n.leaf() {
if !tr.nodeAscend(&(*n.children)[i], pivot, hint, depth+1, iter,
mut) {
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]) {
return false
}
if !n.leaf() {
if !tr.nodeScan(&(*n.children)[i+1], iter, mut) {
return false
}
}
}
return true
}
func (tr *BTreeG[T]) Reverse(iter func(item T) bool) {
tr.reverse(iter, false)
}
func (tr *BTreeG[T]) ReverseMut(iter func(item T) bool) {
tr.reverse(iter, true)
}
func (tr *BTreeG[T]) reverse(iter func(item T) bool, mut bool) {
if tr.lock(mut) {
defer tr.unlock(mut)
}
if tr.root == nil {
return
}
tr.nodeReverse(&tr.root, iter, mut)
}
func (tr *BTreeG[T]) nodeReverse(cn **node[T], iter func(item T) bool, mut bool,
) bool {
n := tr.isoLoad(cn, mut)
if n.leaf() {
for i := len(n.items) - 1; i >= 0; i-- {
if !iter(n.items[i]) {
return false
}
}
return true
}
if !tr.nodeReverse(&(*n.children)[len(*n.children)-1], iter, mut) {
return false
}
for i := len(n.items) - 1; i >= 0; i-- {
if !iter(n.items[i]) {
return false
}
if !tr.nodeReverse(&(*n.children)[i], iter, mut) {
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 *BTreeG[T]) Descend(pivot T, iter func(item T) bool) {
tr.descend(pivot, iter, false, nil)
}
func (tr *BTreeG[T]) DescendMut(pivot T, iter func(item T) bool) {
tr.descend(pivot, iter, true, nil)
}
func (tr *BTreeG[T]) descend(pivot T, iter func(item T) bool, mut bool,
hint *PathHint,
) {
if tr.lock(mut) {
defer tr.unlock(mut)
}
if tr.root == nil {
return
}
tr.nodeDescend(&tr.root, pivot, hint, 0, iter, mut)
}
func (tr *BTreeG[T]) DescendHint(pivot T, iter func(item T) bool,
hint *PathHint,
) {
tr.descend(pivot, iter, false, hint)
}
func (tr *BTreeG[T]) DescendHintMut(pivot T, iter func(item T) bool,
hint *PathHint,
) {
tr.descend(pivot, iter, true, hint)
}
func (tr *BTreeG[T]) nodeDescend(cn **node[T], pivot T, hint *PathHint,
depth int, iter func(item T) bool, mut bool,
) bool {
n := tr.isoLoad(cn, mut)
i, found := tr.find(n, pivot, hint, depth)
if !found {
if !n.leaf() {
if !tr.nodeDescend(&(*n.children)[i], pivot, hint, depth+1, iter,
mut) {
return false
}
}
i--
}
for ; i >= 0; i-- {
if !iter(n.items[i]) {
return false
}
if !n.leaf() {
if !tr.nodeReverse(&(*n.children)[i], iter, mut) {
return false
}
}
}
return true
}
// Load is for bulk loading pre-sorted items
func (tr *BTreeG[T]) Load(item T) (T, bool) {
if tr.lock(true) {
defer tr.unlock(true)
}
if tr.root == nil {
return tr.setHint(item, nil)
}
n := tr.isoLoad(&tr.root, true)
for {
n.count++ // optimistically update counts
if n.leaf() {
if len(n.items) < tr.max {
if tr.Less(n.items[len(n.items)-1], item) {
n.items = append(n.items, item)
tr.count++
return tr.empty, false
}
}
break
}
n = tr.isoLoad(&(*n.children)[len(*n.children)-1], true)
}
// revert the counts
n = tr.root
for {
n.count--
if n.leaf() {
break
}
n = (*n.children)[len(*n.children)-1]
}
return tr.setHint(item, nil)
}
// Min returns the minimum item in tree.
// Returns nil if the treex has no items.
func (tr *BTreeG[T]) Min() (T, bool) {
return tr.minMut(false)
}
func (tr *BTreeG[T]) MinMut() (T, bool) {
return tr.minMut(true)
}
func (tr *BTreeG[T]) minMut(mut bool) (T, bool) {
if tr.lock(mut) {
defer tr.unlock(mut)
}
if tr.root == nil {
return tr.empty, false
}
n := tr.isoLoad(&tr.root, mut)
for {
if n.leaf() {
return n.items[0], true
}
n = tr.isoLoad(&(*n.children)[0], mut)
}
}
// Max returns the maximum item in tree.
// Returns nil if the tree has no items.
func (tr *BTreeG[T]) Max() (T, bool) {
return tr.maxMut(false)
}
func (tr *BTreeG[T]) MaxMut() (T, bool) {
return tr.maxMut(true)
}
func (tr *BTreeG[T]) maxMut(mut bool) (T, bool) {
if tr.lock(mut) {
defer tr.unlock(mut)
}
if tr.root == nil {
return tr.empty, false
}
n := tr.isoLoad(&tr.root, mut)
for {
if n.leaf() {
return n.items[len(n.items)-1], true
}
n = tr.isoLoad(&(*n.children)[len(*n.children)-1], mut)
}
}
// PopMin removes the minimum item in tree and returns it.
// Returns nil if the tree has no items.
func (tr *BTreeG[T]) PopMin() (T, bool) {
if tr.lock(true) {
defer tr.unlock(true)
}
if tr.root == nil {
return tr.empty, false
}
n := tr.isoLoad(&tr.root, true)
var item T
for {
n.count-- // optimistically update counts
if n.leaf() {
item = n.items[0]
if len(n.items) == tr.min {
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, true
}
n = tr.isoLoad(&(*n.children)[0], true)
}
// revert the counts
n = tr.root
for {
n.count++
if n.leaf() {
break
}
n = (*n.children)[0]
}
return tr.deleteHint(item, nil)
}
// PopMax removes the maximum item in tree and returns it.
// Returns nil if the tree has no items.
func (tr *BTreeG[T]) PopMax() (T, bool) {
if tr.lock(true) {
defer tr.unlock(true)
}
if tr.root == nil {
return tr.empty, false
}
n := tr.isoLoad(&tr.root, true)
var item T
for {
n.count-- // optimistically update counts
if n.leaf() {
item = n.items[len(n.items)-1]
if len(n.items) == tr.min {
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, true
}
n = tr.isoLoad(&(*n.children)[len(*n.children)-1], true)
}
// revert the counts
n = tr.root
for {
n.count++
if n.leaf() {
break
}
n = (*n.children)[len(*n.children)-1]
}
return tr.deleteHint(item, nil)
}
// GetAt returns the value at index.
// Return nil if the tree is empty or the index is out of bounds.
func (tr *BTreeG[T]) GetAt(index int) (T, bool) {
return tr.getAt(index, false)
}
func (tr *BTreeG[T]) GetAtMut(index int) (T, bool) {
return tr.getAt(index, true)
}
func (tr *BTreeG[T]) getAt(index int, mut bool) (T, bool) {
if tr.lock(mut) {
defer tr.unlock(mut)
}
if tr.root == nil || index < 0 || index >= tr.count {
return tr.empty, false
}
n := tr.isoLoad(&tr.root, mut)
for {
if n.leaf() {
return n.items[index], 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], true
}
index -= (*n.children)[i].count + 1
}
n = tr.isoLoad(&(*n.children)[i], mut)
}
}
// DeleteAt deletes the item at index.
// Return nil if the tree is empty or the index is out of bounds.
func (tr *BTreeG[T]) DeleteAt(index int) (T, bool) {
if tr.lock(true) {
defer tr.unlock(true)
}
if tr.root == nil || index < 0 || index >= tr.count {
return tr.empty, false
}
var pathbuf [8]uint8 // track the path
path := pathbuf[:0]
var item T
n := tr.isoLoad(&tr.root, true)
outer:
for {
n.count-- // optimistically update counts
if n.leaf() {
// the index is the item position
item = n.items[index]
if len(n.items) == tr.min {
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, 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.isoLoad(&(*n.children)[i], true)
}
// revert the counts
var hint PathHint
n = tr.root
for i := 0; i < len(path); i++ {
if i < len(hint.path) {
hint.path[i] = uint8(path[i])
hint.used[i] = true
}
n.count++
if !n.leaf() {
n = (*n.children)[uint8(path[i])]
}
}
return tr.deleteHint(item, &hint)
}
// Height returns the height of the tree.
// Returns zero if tree has no items.
func (tr *BTreeG[T]) Height() int {
if tr.lock(false) {
defer tr.unlock(false)
}
var height int
if tr.root != nil {
n := tr.root
for {
height++
if n.leaf() {
break
}
n = (*n.children)[0]
}
}
return height
}
// Walk iterates over all items in tree, in order.
// The items param will contain one or more items.
func (tr *BTreeG[T]) Walk(iter func(item []T) bool) {
tr.walk(iter, false)
}
func (tr *BTreeG[T]) WalkMut(iter func(item []T) bool) {
tr.walk(iter, true)
}
func (tr *BTreeG[T]) walk(iter func(item []T) bool, mut bool) {
if tr.lock(mut) {
defer tr.unlock(mut)
}
if tr.root == nil {
return
}
tr.nodeWalk(&tr.root, iter, mut)
}
func (tr *BTreeG[T]) nodeWalk(cn **node[T], iter func(item []T) bool, mut bool,
) bool {
n := tr.isoLoad(cn, mut)
if n.leaf() {
if !iter(n.items) {
return false
}
} else {
for i := 0; i < len(n.items); i++ {
if !tr.nodeWalk(&(*n.children)[i], iter, mut) {
return false
}
if !iter(n.items[i : i+1]) {
return false
}
}
if !tr.nodeWalk(&(*n.children)[len(n.items)], iter, mut) {
return false
}
}
return true
}
// Copy the tree. This is a copy-on-write operation and is very fast because
// it only performs a shadowed copy.
func (tr *BTreeG[T]) Copy() *BTreeG[T] {
return tr.IsoCopy()
}
func (tr *BTreeG[T]) IsoCopy() *BTreeG[T] {
var mu *sync.RWMutex
if tr.lock(true) {
mu = new(sync.RWMutex)
defer tr.unlock(true)
}
tr.isoid = newIsoID()
tr2 := new(BTreeG[T])
*tr2 = *tr
tr2.mu = mu
tr2.isoid = newIsoID()
return tr2
}
func (tr *BTreeG[T]) lock(write bool) bool {
if tr.locks {
if write {
tr.mu.Lock()
} else {
tr.mu.RLock()
}
}
return tr.locks
}
func (tr *BTreeG[T]) unlock(write bool) {
if write {
tr.mu.Unlock()
} else {
tr.mu.RUnlock()
}
}
// Iter represents an iterator
type IterG[T any] struct {
tr *BTreeG[T]
mut bool
locked bool
seeked bool
atstart bool
atend bool
stack0 [4]iterStackItemG[T]
stack []iterStackItemG[T]
item T
}
type iterStackItemG[T any] struct {
n *node[T]
i int
}
// Iter returns a read-only iterator.
// The Release method must be called finished with iterator.
func (tr *BTreeG[T]) Iter() IterG[T] {
return tr.iter(false)
}
func (tr *BTreeG[T]) IterMut() IterG[T] {
return tr.iter(true)
}
func (tr *BTreeG[T]) iter(mut bool) IterG[T] {
var iter IterG[T]
iter.tr = tr
iter.mut = mut
iter.locked = tr.lock(iter.mut)
iter.stack = iter.stack0[:0]
return iter
}
// Seek to item greater-or-equal-to key.
// Returns false if there was no item found.
func (iter *IterG[T]) Seek(key T) bool {
return iter.seek(key, nil)
}
func (iter *IterG[T]) SeekHint(key T, hint *PathHint) bool {
return iter.seek(key, hint)
}
func (iter *IterG[T]) seek(key T, hint *PathHint) 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.isoLoad(&iter.tr.root, iter.mut)
var depth int
for {
i, found := iter.tr.find(n, key, hint, depth)
iter.stack = append(iter.stack, iterStackItemG[T]{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 = iter.tr.isoLoad(&(*n.children)[i], iter.mut)
depth++
}
}
// First moves iterator to first item in tree.
// Returns false if the tree is empty.
func (iter *IterG[T]) 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.isoLoad(&iter.tr.root, iter.mut)
for {
iter.stack = append(iter.stack, iterStackItemG[T]{n, 0})
if n.leaf() {
break
}
n = iter.tr.isoLoad(&(*n.children)[0], iter.mut)
}
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 *IterG[T]) 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.isoLoad(&iter.tr.root, iter.mut)
for {
iter.stack = append(iter.stack, iterStackItemG[T]{n, len(n.items)})
if n.leaf() {
iter.stack[len(iter.stack)-1].i--
break
}
n = iter.tr.isoLoad(&(*n.children)[len(n.items)], iter.mut)
}
s := &iter.stack[len(iter.stack)-1]
iter.item = s.n.items[s.i]
return true
}
// Release the iterator.
func (iter *IterG[T]) Release() {
if iter.tr == nil {
return
}
if iter.locked {
iter.tr.unlock(iter.mut)
iter.locked = false
}
iter.stack = nil
iter.tr = nil
}
// 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 *IterG[T]) 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 := iter.tr.isoLoad(&(*s.n.children)[s.i], iter.mut)
for {
iter.stack = append(iter.stack, iterStackItemG[T]{n, 0})
if n.leaf() {
break
}
n = iter.tr.isoLoad(&(*n.children)[0], iter.mut)
}
}
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 *IterG[T]) 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 := iter.tr.isoLoad(&(*s.n.children)[s.i], iter.mut)
for {
iter.stack = append(iter.stack, iterStackItemG[T]{n, len(n.items)})
if n.leaf() {
iter.stack[len(iter.stack)-1].i--
break
}
n = iter.tr.isoLoad(&(*n.children)[len(n.items)], iter.mut)
}
}
s = &iter.stack[len(iter.stack)-1]
iter.item = s.n.items[s.i]
return true
}
// Item returns the current iterator item.
func (iter *IterG[T]) Item() T {
return iter.item
}
// Items returns all the items in order.
func (tr *BTreeG[T]) Items() []T {
return tr.items(false)
}
func (tr *BTreeG[T]) ItemsMut() []T {
return tr.items(true)
}
func (tr *BTreeG[T]) items(mut bool) []T {
if tr.lock(mut) {
defer tr.unlock(mut)
}
items := make([]T, 0, tr.Len())
if tr.root != nil {
items = tr.nodeItems(&tr.root, items, mut)
}
return items
}
func (tr *BTreeG[T]) nodeItems(cn **node[T], items []T, mut bool) []T {
n := tr.isoLoad(cn, mut)
if n.leaf() {
return append(items, n.items...)
}
for i := 0; i < len(n.items); i++ {
items = tr.nodeItems(&(*n.children)[i], items, mut)
items = append(items, n.items[i])
}
return tr.nodeItems(&(*n.children)[len(*n.children)-1], items, mut)
}
// Clear will delete all items.
func (tr *BTreeG[T]) Clear() {
if tr.lock(true) {
defer tr.unlock(true)
}
tr.root = nil
tr.count = 0
}
// Generic BTree
//
// Deprecated: use BTreeG
type Generic[T any] struct {
*BTreeG[T]
}
// NewGeneric returns a generic BTree
//
// Deprecated: use NewBTreeG
func NewGeneric[T any](less func(a, b T) bool) *Generic[T] {
return &Generic[T]{NewBTreeGOptions(less, Options{})}
}
// NewGenericOptions returns a generic BTree
//
// Deprecated: use NewBTreeGOptions
func NewGenericOptions[T any](less func(a, b T) bool, opts Options,
) *Generic[T] {
return &Generic[T]{NewBTreeGOptions(less, opts)}
}
func (tr *Generic[T]) Copy() *Generic[T] {
return &Generic[T]{tr.BTreeG.Copy()}
}
btree-1.8.1/btreeg_test.go 0000664 0000000 0000000 00000114266 15044500710 0015466 0 ustar 00root root 0000000 0000000 package btree
import (
"fmt"
"math/rand"
"os"
"runtime"
"sort"
"strconv"
"sync/atomic"
"testing"
"time"
)
func init() {
seed, err := strconv.ParseInt(os.Getenv("SEED"), 10, 64)
if err != nil {
seed = time.Now().UnixNano()
}
fmt.Printf("seed: %d\n", seed)
rand.Seed(seed)
}
// testKind is the item type.
// It's important to use the equal symbol, which tells Go to create an alias of
// the type, rather than creating an entirely new type.
type testKind = int
func testLess(a, b testKind) bool {
return a < b
}
// The functions below, which begin with "test*", are required by the
// btree_test.go file. If you choose not use include the btree_test.go file in
// your project then these functions may be omitted.
// testMakeItem must return a valid item for testing.
// It's required that the returned item maintains equal order as the
// provided int, such that:
//
// testMakeItem(0) < testMakeItem(1) < testMakeItem(2) < testMakeItem(10)
func testMakeItem(x int) (item testKind) {
return x
}
// testNewBTree must return an operational btree for testing.
func testNewBTree() *BTreeG[testKind] {
return NewBTreeG(testLess)
}
func randKeys(N int) (keys []testKind) {
keys = make([]testKind, N)
for _, i := range rand.Perm(N) {
keys[i] = testMakeItem(i)
}
return keys
}
func (tr *BTreeG[T]) lt(a, b T) bool { return tr.less(a, b) }
func (tr *BTreeG[T]) eq(a, b T) bool { return !(tr.lt(a, b) || tr.lt(b, a)) }
func (tr *BTreeG[T]) lte(a, b T) bool { return tr.lt(a, b) || tr.eq(a, b) }
func (tr *BTreeG[T]) gt(a, b T) bool { return tr.lt(b, a) }
func (tr *BTreeG[T]) gte(a, b T) bool { return tr.gt(a, b) || tr.eq(a, b) }
func kindsAreEqual(a, b []testKind) bool {
tr := NewBTreeG(testLess)
if len(a) != len(b) {
return false
}
for i := 0; i < len(a); i++ {
if !tr.eq(a[i], b[i]) {
return false
}
}
return true
}
func TestGenericMakeItemOrder(t *testing.T) {
tr := testNewBTree()
ints := []int{0, 1, 2, 3, 4, 10, 20, 30, 40, 100, 200, 300, 400}
for i := 0; i < len(ints)-1; i++ {
a := testMakeItem(ints[i])
b := testMakeItem(ints[i+1])
if !tr.lt(a, b) {
t.Fatalf("bad ordering for testMakeItem: '%v' !< '%v'", a, b)
}
}
}
func TestGenericDescend(t *testing.T) {
tr := testNewBTree()
var count int
tr.Descend(testMakeItem(rand.Int()), func(item testKind) bool {
count++
return true
})
if count > 0 {
t.Fatalf("expected 0, got %v", count)
}
var keys []testKind
for i := 0; i < 1000; i += 10 {
keys = append(keys, testMakeItem(i))
tr.Set(keys[len(keys)-1])
}
var exp []testKind
tr.Reverse(func(item testKind) bool {
exp = append(exp, item)
return true
})
for i := 999; i >= 0; i-- {
key := testMakeItem(i)
var all []testKind
tr.Descend(key, func(item testKind) bool {
all = append(all, item)
return true
})
for len(exp) > 0 && tr.Less(key, exp[0]) {
exp = exp[1:]
}
var count int
tr.Descend(key, func(item testKind) bool {
if count == (i+1)%tr.max {
return false
}
count++
return true
})
if count > len(exp) {
t.Fatalf("expected 1, got %v", count)
}
if !kindsAreEqual(exp, all) {
fmt.Printf("exp: %v\n", exp)
fmt.Printf("all: %v\n", all)
t.Fatal("mismatch")
}
for j := 0; j < tr.Len(); j++ {
count = 0
tr.Descend(key, func(item testKind) bool {
if count == j {
return false
}
count++
return true
})
}
}
}
func TestGenericDescendHint(t *testing.T) {
tr := testNewBTree()
var count int
var hint PathHint
tr.DescendHint(testMakeItem(rand.Int()), func(item testKind) bool {
count++
return true
}, &hint)
if count > 0 {
t.Fatalf("expected 0, got %v", count)
}
var keys []testKind
for i := 0; i < 1000; i += 10 {
keys = append(keys, testMakeItem(i))
tr.Set(keys[len(keys)-1])
}
var exp []testKind
tr.Reverse(func(item testKind) bool {
exp = append(exp, item)
return true
})
for i := 999; i >= 0; i-- {
key := testMakeItem(i)
var all []testKind
tr.DescendHint(key, func(item testKind) bool {
all = append(all, item)
return true
}, &hint)
for len(exp) > 0 && tr.Less(key, exp[0]) {
exp = exp[1:]
}
var count int
tr.DescendHint(key, func(item testKind) bool {
if count == (i+1)%tr.max {
return false
}
count++
return true
}, &hint)
if count > len(exp) {
t.Fatalf("expected 1, got %v", count)
}
if !kindsAreEqual(exp, all) {
fmt.Printf("exp: %v\n", exp)
fmt.Printf("all: %v\n", all)
t.Fatal("mismatch")
}
for j := 0; j < tr.Len(); j++ {
count = 0
tr.DescendHint(key, func(item testKind) bool {
if count == j {
return false
}
count++
return true
}, &hint)
}
}
}
func TestGenericAscend(t *testing.T) {
tr := testNewBTree()
var count int
tr.Ascend(testMakeItem(1), func(item testKind) bool {
count++
return true
})
if count > 0 {
t.Fatalf("expected 0, got %v", count)
}
var keys []testKind
for i := 0; i < 1000; i += 10 {
keys = append(keys, testMakeItem(i))
tr.Set(keys[len(keys)-1])
tr.sane()
}
exp := keys
for i := -1; i < 1000; i++ {
key := testMakeItem(i)
var all []testKind
tr.Ascend(key, func(item testKind) bool {
all = append(all, item)
return true
})
for len(exp) > 0 && tr.Less(exp[0], key) {
exp = exp[1:]
}
var count int
tr.Ascend(key, func(item testKind) bool {
if count == (i+1)%tr.max {
return false
}
count++
return true
})
if count > len(exp) {
t.Fatalf("expected 1, got %v", count)
}
if !kindsAreEqual(exp, all) {
t.Fatal("mismatch")
}
}
}
func TestGenericAscendHint(t *testing.T) {
tr := testNewBTree()
var hint PathHint
var count int
tr.AscendHint(testMakeItem(1), func(item testKind) bool {
count++
return true
}, &hint)
if count > 0 {
t.Fatalf("expected 0, got %v", count)
}
var keys []testKind
for i := 0; i < 1000; i += 10 {
keys = append(keys, testMakeItem(i))
tr.Set(keys[len(keys)-1])
tr.sane()
}
exp := keys
for i := -1; i < 1000; i++ {
key := testMakeItem(i)
var all []testKind
tr.AscendHint(key, func(item testKind) bool {
all = append(all, item)
return true
}, &hint)
for len(exp) > 0 && tr.Less(exp[0], key) {
exp = exp[1:]
}
var count int
tr.AscendHint(key, func(item testKind) bool {
if count == (i+1)%tr.max {
return false
}
count++
return true
}, &hint)
if count > len(exp) {
t.Fatalf("expected 1, got %v", count)
}
if !kindsAreEqual(exp, all) {
t.Fatal("mismatch")
}
}
}
func TestGenericItems(t *testing.T) {
tr := testNewBTree()
if len(tr.Items()) != 0 {
t.Fatalf("expected 0, got %v", len(tr.Items()))
}
var keys []testKind
for i := 0; i < 100000; i += 10 {
keys = append(keys, testMakeItem(i))
tr.Set(keys[len(keys)-1])
tr.sane()
}
keys2 := tr.Items()
if !kindsAreEqual(keys, keys2) {
t.Fatal("mismatch")
}
}
func TestGenericSimpleRandom(t *testing.T) {
start := time.Now()
for time.Since(start) < time.Second*2 {
N := 100_000
items := randKeys(N)
tr := testNewBTree()
tr.sane()
for i := 0; i < len(items); i++ {
if v, ok := tr.Get(items[i]); ok || !tr.eq(v, tr.empty) {
panic("!")
}
if v, ok := tr.Set(items[i]); ok || !tr.eq(v, tr.empty) {
panic("!")
}
if v, ok := tr.Get(items[i]); !ok || !tr.eq(v, items[i]) {
panic("!")
}
}
tr.sane()
for i := 0; i < len(items); i++ {
if v, ok := tr.Set(items[i]); !ok || !tr.eq(v, items[i]) {
panic("!")
}
}
pivot := items[len(items)/2]
tr.Ascend(pivot, func(item testKind) bool {
if tr.Less(item, pivot) {
panic("!")
}
return true
})
var min testKind
index := 0
tr.Scan(func(item testKind) bool {
if index == len(items)/2 {
return false
}
if index > 0 {
if tr.Less(item, min) {
panic("!")
}
}
min = item
index++
return true
})
tr.sane()
for i := 0; i < len(items); i++ {
if v, ok := tr.Delete(items[i]); !ok || !tr.eq(v, items[i]) {
panic("!")
}
if i%97 == 0 {
tr.sane()
}
if v, ok := tr.Delete(items[i]); ok || !tr.eq(v, tr.empty) {
panic("!")
}
}
if tr.Len() != 0 {
panic("!")
}
tr.sane()
for i := 0; i < len(items); i++ {
if v, ok := tr.Delete(items[i]); ok || !tr.eq(v, tr.empty) {
panic("!")
}
}
tr.sane()
tr.Scan(func(item testKind) bool {
panic("!")
})
}
}
func TestGenericBTree(t *testing.T) {
N := 10000
tr := testNewBTree()
tr.sane()
keys := randKeys(N)
// insert all items
for _, key := range keys {
if v, ok := tr.Set(key); ok || !tr.eq(v, tr.empty) {
t.Fatal("expected false")
}
tr.sane()
}
// check length
if tr.Len() != len(keys) {
t.Fatalf("expected %v, got %v", len(keys), tr.Len())
}
// get each value
for _, key := range keys {
if v, ok := tr.Get(key); !ok || !tr.eq(v, key) {
t.Fatalf("expected '%v', got '%v'", key, v)
}
}
// scan all items
var prev testKind
var count int
tr.Scan(func(item testKind) bool {
if count > 0 {
if tr.lte(item, prev) {
t.Fatal("out of order")
}
}
prev = item
count++
return true
})
if count != len(keys) {
t.Fatalf("expected '%v', got '%v'", len(keys), count)
}
// reverse all items
count = 0
tr.Reverse(func(item testKind) bool {
if count > 0 {
if tr.gte(item, prev) {
t.Fatal("out of order")
}
}
prev = item
count++
return true
})
if count != len(keys) {
t.Fatalf("expected '%v', got '%v'", len(keys), count)
}
// try to get an invalid item
if v, ok := tr.Get(testMakeItem(-1)); ok || !tr.eq(v, tr.empty) {
t.Fatal("expected nil")
}
// scan and quit at various steps
for i := 0; i < 100; i++ {
var j int
tr.Scan(func(item testKind) bool {
if j == i {
return false
}
j++
return true
})
}
// reverse and quit at various steps
for i := 0; i < 100; i++ {
var j int
tr.Reverse(func(item testKind) bool {
if j == i {
return false
}
j++
return true
})
}
// delete half the items
for _, key := range keys[:len(keys)/2] {
if v, ok := tr.Delete(key); !ok || !tr.eq(v, key) {
t.Fatalf("expected '%v', got '%v'", key, v)
}
}
// check length
if tr.Len() != len(keys)/2 {
t.Fatalf("expected %v, got %v", len(keys)/2, tr.Len())
}
// try delete half again
for _, key := range keys[:len(keys)/2] {
if v, ok := tr.Delete(key); ok || !tr.eq(v, tr.empty) {
t.Fatal("expected false")
}
tr.sane()
}
// check length
if tr.Len() != len(keys)/2 {
t.Fatalf("expected %v, got %v", len(keys)/2, tr.Len())
}
// scan items
count = 0
tr.Scan(func(item testKind) bool {
if count > 0 {
if tr.lte(item, prev) {
t.Fatal("out of order")
}
}
prev = item
count++
return true
})
if count != len(keys)/2 {
t.Fatalf("expected '%v', got '%v'", len(keys), count)
}
// replace second half
for _, key := range keys[len(keys)/2:] {
if v, ok := tr.Set(key); !ok || !tr.eq(v, key) {
t.Fatalf("expected '%v', got '%v'", key, v)
}
tr.sane()
}
// delete next half the items
for _, key := range keys[len(keys)/2:] {
if v, ok := tr.Delete(key); !ok || !tr.eq(v, key) {
t.Fatalf("expected '%v', got '%v'", key, v)
}
tr.sane()
}
// check length
if tr.Len() != 0 {
t.Fatalf("expected %v, got %v", 0, tr.Len())
}
// do some stuff on an empty tree
if v, ok := tr.Get(keys[0]); ok || !tr.eq(v, tr.empty) {
t.Fatal("expected nil")
}
tr.Scan(func(item testKind) bool {
t.Fatal("should not be reached")
return true
})
tr.Reverse(func(item testKind) bool {
t.Fatal("should not be reached")
return true
})
if v, ok := tr.Delete(testMakeItem(-1)); ok || !tr.eq(v, tr.empty) {
t.Fatal("expected nil")
}
tr.sane()
}
func TestGenericBTreeOne(t *testing.T) {
tr := testNewBTree()
tr.Set(testMakeItem(1))
tr.Delete(testMakeItem(1))
tr.Set(testMakeItem(1))
tr.Delete(testMakeItem(1))
tr.Set(testMakeItem(1))
tr.Delete(testMakeItem(1))
if tr.Len() != 0 {
panic("!")
}
tr.sane()
}
func TestGenericBTree256(t *testing.T) {
tr := testNewBTree()
var n int
for j := 0; j < 2; j++ {
for _, i := range rand.Perm(256) {
tr.Set(testMakeItem(i))
n++
if tr.Len() != n {
t.Fatalf("expected 256, got %d", n)
}
}
for _, i := range rand.Perm(256) {
if v, ok := tr.Get(testMakeItem(i)); !ok || !tr.eq(v, testMakeItem(i)) {
t.Fatalf("expected %v, got %v", i, v)
}
}
for _, i := range rand.Perm(256) {
tr.Delete(testMakeItem(i))
n--
if tr.Len() != n {
t.Fatalf("expected 256, got %d", n)
}
}
for _, i := range rand.Perm(256) {
if v, ok := tr.Get(testMakeItem(i)); ok || !tr.eq(v, tr.empty) {
t.Fatal("expected nil")
}
}
}
}
func shuffleItems(keys []testKind) {
for i := range keys {
j := rand.Intn(i + 1)
keys[i], keys[j] = keys[j], keys[i]
}
}
func sortItems(keys []testKind) {
tr := testNewBTree()
sort.Slice(keys, func(i, j int) bool {
return tr.lt(keys[i], keys[j])
})
}
func TestGenericRandom(t *testing.T) {
N := 200000
keys := randKeys(N)
tr := testNewBTree()
tr.sane()
if v, ok := tr.Min(); ok || !tr.eq(v, tr.empty) {
t.Fatalf("expected nil")
}
if v, ok := tr.Max(); ok || !tr.eq(v, tr.empty) {
t.Fatalf("expected nil")
}
if v, ok := tr.PopMin(); ok || !tr.eq(v, tr.empty) {
t.Fatalf("expected nil")
}
if v, ok := tr.PopMax(); ok || !tr.eq(v, tr.empty) {
t.Fatalf("expected nil")
}
if tr.Height() != 0 {
t.Fatalf("expected 0, got %d", tr.Height())
}
tr.sane()
shuffleItems(keys)
for i := 0; i < len(keys); i++ {
if v, ok := tr.Set(keys[i]); ok || !tr.eq(v, tr.empty) {
t.Fatalf("expected nil")
}
if i%123 == 0 {
tr.sane()
}
}
tr.sane()
sortItems(keys)
var n int
tr.Scan(func(item testKind) bool {
n++
return false
})
if n != 1 {
t.Fatalf("expected 1, got %d", n)
}
n = 0
tr.Scan(func(item testKind) bool {
if !tr.eq(item, keys[n]) {
t.Fatalf("expected %v, got %v", keys[n], item)
}
n++
return true
})
if n != len(keys) {
t.Fatalf("expected %d, got %d", len(keys), n)
}
if tr.Len() != len(keys) {
t.Fatalf("expected %d, got %d", tr.Len(), len(keys))
}
for i := 0; i < tr.Len(); i++ {
if v, ok := tr.GetAt(i); !ok || !tr.eq(v, keys[i]) {
t.Fatalf("expected %v, got %v", keys[i], v)
}
}
n = 0
tr.Reverse(func(item testKind) bool {
n++
return false
})
if n != 1 {
t.Fatalf("expected 1, got %d", n)
}
n = 0
tr.Reverse(func(item testKind) bool {
if !tr.eq(item, keys[len(keys)-n-1]) {
t.Fatalf("expected %v, got %v", keys[len(keys)-n-1], item)
}
n++
return true
})
if n != len(keys) {
t.Fatalf("expected %d, got %d", len(keys), n)
}
if tr.Len() != len(keys) {
t.Fatalf("expected %d, got %d", tr.Len(), len(keys))
}
tr.sane()
n = 0
for i := 0; i < 1000; i++ {
n := 0
tr.Scan(func(item testKind) bool {
if n == i {
return false
}
n++
return true
})
if n != i {
t.Fatalf("expected %d, got %d", i, n)
}
}
n = 0
for i := 0; i < 1000; i++ {
n = 0
tr.Reverse(func(item testKind) bool {
if n == i {
return false
}
n++
return true
})
if n != i {
t.Fatalf("expected %d, got %d", i, n)
}
}
sortItems(keys)
for i := 0; i < len(keys); i++ {
var res testKind
var j int
tr.Ascend(keys[i], func(item testKind) bool {
if j == 0 {
res = item
}
j++
return j == i%500
})
if !tr.eq(res, keys[i]) {
t.Fatal("not equal")
}
}
for i := len(keys) - 1; i >= 0; i-- {
var res testKind
var j int
tr.Descend(keys[i], func(item testKind) bool {
if j == 0 {
res = item
}
j++
return j == i%500
})
if !tr.eq(res, keys[i]) {
t.Fatal("not equal")
}
}
if tr.Height() == 0 {
t.Fatalf("expected non-zero")
}
if v, ok := tr.Min(); !ok || !tr.eq(v, keys[0]) {
t.Fatalf("expected '%v', got '%v'", keys[0], v)
}
if v, ok := tr.Max(); !ok || !tr.eq(v, keys[len(keys)-1]) {
t.Fatalf("expected '%v', got '%v'", keys[len(keys)-1], v)
}
if v, ok := tr.PopMin(); !ok || !tr.eq(v, keys[0]) {
t.Fatalf("expected '%v', got '%v'", keys[0], v)
}
tr.sane()
if v, ok := tr.PopMax(); !ok || !tr.eq(v, keys[len(keys)-1]) {
t.Fatalf("expected '%v', got '%v'", keys[len(keys)-1], v)
}
tr.sane()
tr.Set(keys[0])
tr.Set(keys[len(keys)-1])
shuffleItems(keys)
var hint PathHint
for i := 0; i < len(keys); i++ {
if v, ok := tr.Get(keys[i]); !ok || !tr.eq(v, keys[i]) {
t.Fatalf("expected '%v', got '%v'", keys[i], v)
}
if v, ok := tr.GetHint(keys[i], &hint); !ok || !tr.eq(v, keys[i]) {
t.Fatalf("expected '%v', got '%v'", keys[i], v)
}
}
sortItems(keys)
for i := 0; i < len(keys); i++ {
if v, ok := tr.PopMin(); !ok || !tr.eq(v, keys[i]) {
t.Fatalf("expected '%v', got '%v'", keys[i], v)
}
}
for i := 0; i < len(keys); i++ {
if v, ok := tr.Set(keys[i]); ok || !tr.eq(v, tr.empty) {
t.Fatalf("expected nil")
}
}
for i := len(keys) - 1; i >= 0; i-- {
if v, ok := tr.PopMax(); !ok || !tr.eq(v, keys[i]) {
t.Fatalf("expected '%v', got '%v'", keys[i], v)
}
}
for i := 0; i < len(keys); i++ {
if v, ok := tr.Set(keys[i]); ok || !tr.eq(v, tr.empty) {
t.Fatalf("expected nil")
}
}
if v, ok := tr.Delete(testMakeItem(-1)); ok || !tr.eq(v, tr.empty) {
t.Fatal("expected nil")
}
tr.sane()
shuffleItems(keys)
if v, ok := tr.Delete(keys[len(keys)/2]); !ok || !tr.eq(v, keys[len(keys)/2]) {
t.Fatalf("expected '%v', got '%v'", keys[len(keys)/2], v)
}
tr.sane()
if v, ok := tr.Delete(keys[len(keys)/2]); ok || !tr.eq(v, tr.empty) {
t.Fatalf("expected '%v', got '%v'", tr.empty, v)
}
tr.sane()
tr.Set(keys[len(keys)/2])
tr.sane()
for i := 0; i < len(keys); i++ {
if v, ok := tr.Delete(keys[i]); !ok || !tr.eq(v, keys[i]) {
t.Fatalf("expected '%v', got '%v'", keys[i], v)
}
if v, ok := tr.Get(keys[i]); ok || !tr.eq(v, tr.empty) {
t.Fatalf("expected nil")
}
if v, ok := tr.GetHint(keys[i], &hint); ok || !tr.eq(v, tr.empty) {
t.Fatalf("expected nil")
}
if i%97 == 0 {
tr.sane()
}
}
if tr.Height() != 0 {
t.Fatalf("expected 0, got %d", tr.Height())
}
shuffleItems(keys)
for i := 0; i < len(keys); i++ {
if v, ok := tr.Load(keys[i]); ok || !tr.eq(v, tr.empty) {
t.Fatalf("expected nil")
}
if i%97 == 0 {
tr.sane()
}
}
for i := 0; i < len(keys); i++ {
if v, ok := tr.Get(keys[i]); !ok || !tr.eq(v, keys[i]) {
t.Fatalf("expected '%v', got '%v'", keys[i], v)
}
}
shuffleItems(keys)
for i := 0; i < len(keys); i++ {
if v, ok := tr.Delete(keys[i]); !ok || !tr.eq(v, keys[i]) {
t.Fatalf("expected '%v', got '%v'", keys[i], v)
}
if v, ok := tr.Get(keys[i]); ok || !tr.eq(v, tr.empty) {
t.Fatalf("expected nil")
}
}
sortItems(keys)
for i := 0; i < len(keys); i++ {
if v, ok := tr.Load(keys[i]); ok || !tr.eq(v, tr.empty) {
t.Fatalf("expected nil")
}
if i%97 == 0 {
tr.sane()
}
}
shuffleItems(keys)
if v, ok := tr.Load(keys[0]); !ok || !tr.eq(v, keys[0]) {
t.Fatalf("expected '%v', got '%v'", keys[0], v)
}
tr.sane()
}
func TestGenericLess(t *testing.T) {
tr := testNewBTree()
if !tr.Less(testMakeItem(1), testMakeItem(2)) {
panic("invalid")
}
if tr.Less(testMakeItem(2), testMakeItem(1)) {
panic("invalid")
}
if tr.Less(testMakeItem(1), testMakeItem(1)) {
panic("invalid")
}
}
func TestGenericDeleteRandom(t *testing.T) {
N := 2_000_000
tr := testNewBTree()
for i := 0; i < N; i++ {
tr.Load(testMakeItem(i))
}
tr.sane()
for tr.Len() > 0 {
var item testKind
var ok bool
switch rand.Int() % 3 {
case 0:
item, ok = tr.GetAt(tr.Len() / 2)
case 1:
item, ok = tr.Min()
case 2:
item, ok = tr.Max()
}
if !ok {
panic("!")
}
v, ok := tr.Delete(item)
if !ok || !tr.eq(v, item) {
panic("!")
}
}
}
func TestGenericDeleteAt(t *testing.T) {
N := 10_000
tr := testNewBTree()
keys := randKeys(N)
for _, key := range keys {
tr.Set(key)
}
tr.sane()
for tr.Len() > 0 {
index := rand.Intn(tr.Len())
item1, ok1 := tr.GetAt(index)
item2, ok2 := tr.DeleteAt(index)
if !ok1 || !ok2 || !tr.eq(item1, item2) {
panic("mismatch")
}
tr.sane()
}
}
func TestGenericCopy(t *testing.T) {
items := randKeys(100000)
itemsM := testNewBTree()
for i := 0; i < len(items); i++ {
itemsM.Set(items[i])
}
tr := testNewBTree()
for i := 0; i < len(items); i++ {
tr.Set(items[i])
}
var wait atomic.Int32
var testCopyDeep func(tr *BTreeG[testKind], parent bool)
testCopyDeep = func(tr *BTreeG[testKind], parent bool) {
defer func() { wait.Add(-1) }()
if parent {
// 2 grandchildren
for i := 0; i < 2; i++ {
wait.Add(1)
go testCopyDeep(tr.Copy(), false)
}
}
items2 := make([]testKind, 10000)
for i := 0; i < len(items2); i++ {
x := testMakeItem(rand.Int())
_, ok := itemsM.Get(x)
for ok {
x = testMakeItem(rand.Int())
_, ok = itemsM.Get(x)
}
items2[i] = x
}
for i := 0; i < len(items2); i++ {
if v, ok := tr.Set(items2[i]); ok || !tr.eq(v, tr.empty) {
panic("!")
}
}
tr.sane()
if tr.Len() != len(items)+len(items2) {
panic("!")
}
for i := 0; i < len(items); i++ {
if v, ok := tr.Get(items[i]); !ok || !tr.eq(v, items[i]) {
panic("!")
}
}
for i := 0; i < len(items2); i++ {
if v, ok := tr.Get(items2[i]); !ok || !tr.eq(v, items2[i]) {
panic("!")
}
}
for i := 0; i < len(items); i++ {
if v, ok := tr.Delete(items[i]); !ok || !tr.eq(v, items[i]) {
panic("!")
}
}
tr.sane()
if tr.Len() != len(items2) {
panic("!")
}
for i := 0; i < len(items2); i++ {
if v, ok := tr.Get(items2[i]); !ok || !tr.eq(v, items2[i]) {
panic("!")
}
}
sortItems(items2)
var i int
for len(items2) > 0 {
if i%2 == 0 {
if v, ok := tr.PopMin(); !ok || !tr.eq(v, items2[0]) {
panic("!")
}
items2 = items2[1:]
} else {
if v, ok := tr.PopMax(); !ok || !tr.eq(v, items2[len(items2)-1]) {
panic("!")
}
items2 = items2[:len(items2)-1]
}
if i%123 == 0 {
tr.sane()
if tr.Len() != len(items2) {
panic("!")
}
for i := 0; i < len(items2); i++ {
if v, ok := tr.Get(items2[i]); !ok || !tr.eq(v, items2[i]) {
panic("!")
}
}
}
i++
}
tr.sane()
if tr.Len() != len(items2) {
panic("!")
}
}
// 10 children
for i := 0; i < 10; i++ {
wait.Add(1)
go testCopyDeep(tr.Copy(), true)
}
for wait.Load() > 0 {
tr.sane()
if tr.Len() != len(items) {
panic("!")
}
for i := 0; i < len(items); i++ {
if v, ok := tr.Get(items[i]); !ok || !tr.eq(v, items[i]) {
panic("!")
}
}
runtime.Gosched()
}
}
func TestGenericVarious(t *testing.T) {
N := 1_000_000
tr := testNewBTree()
var hint PathHint
for _, i := range randKeys(N) {
if v, ok := tr.SetHint(i, &hint); ok || !tr.eq(v, tr.empty) {
panic("!")
}
}
for _, i := range randKeys(N) {
if v, ok := tr.GetHint(i, &hint); !ok || !tr.eq(v, i) {
panic("!")
}
}
for _, i := range randKeys(N) {
if v, ok := tr.DeleteHint(i, &hint); !ok || !tr.eq(v, i) {
panic("!")
}
}
if v, ok := tr.DeleteAt(0); ok || !tr.eq(v, tr.empty) {
panic("!")
}
if v, ok := tr.GetAt(0); ok || !tr.eq(v, tr.empty) {
panic("!")
}
for i := 0; i < N; i++ {
item := testMakeItem(i)
if v, ok := tr.SetHint(item, &hint); ok || !tr.eq(v, tr.empty) {
panic("!")
}
item = testMakeItem(i)
if v, ok := tr.SetHint(item, &hint); !ok || !tr.eq(v, item) {
panic("!")
}
item = testMakeItem(i)
if v, ok := tr.SetHint(item, &hint); !ok || !tr.eq(v, item) {
panic("!")
}
}
for i := 0; i < N; i++ {
item := testMakeItem(i)
if v, ok := tr.GetHint(item, &hint); !ok || !tr.eq(v, item) {
panic("!")
}
}
for i := 0; i < 100; i++ {
var count int
tr.Walk(func(_ []testKind) bool {
if count == i {
return false
}
count++
return true
})
}
for i := 0; i < N; i++ {
item := testMakeItem(i)
if v, ok := tr.DeleteHint(item, &hint); !ok || !tr.eq(v, item) {
panic("!")
}
}
}
func (tr *BTreeG[T]) sane() {
if err := tr.Sane(); err != nil {
panic(err)
}
}
type saneError string
func (err saneError) Error() string {
return string(err)
}
// btree_sane returns true if the entire btree and every node are valid.
// - height of all leaves are the equal to the btree height.
// - deep count matches the btree count.
// - all nodes have the correct number of items and counts.
// - all items are in order.
func (tr *BTreeG[T]) Sane() error {
if tr == nil {
return nil
}
if !tr.saneheight() {
return saneError("!sane-height")
}
if tr.Len() != tr.count || tr.deepcount() != tr.count {
return saneError("!sane-count")
}
if !tr.saneprops() {
return saneError("!sane-props")
}
if !tr.saneorder() {
return saneError("!sane-order")
}
if !tr.sanenils() {
return saneError("!sane-nils")
}
return nil
}
// btree_saneheight returns true if the height of all leaves match the height
// of the btree.
func (tr *BTreeG[T]) saneheight() bool {
height := tr.Height()
if tr.root != nil {
if height == 0 {
return false
}
return tr.root.saneheight(1, height)
}
return height == 0
}
func (n *node[T]) saneheight(height, maxheight int) bool {
if n.leaf() {
if height != maxheight {
return false
}
} else {
i := 0
for ; i < len(n.items); i++ {
if !(*n.children)[i].saneheight(height+1, maxheight) {
return false
}
}
if !(*n.children)[i].saneheight(height+1, maxheight) {
return false
}
}
return true
}
// btree_deepcount returns the number of items in the btree.
func (tr *BTreeG[T]) deepcount() int {
if tr.root != nil {
return tr.root.deepcount()
}
return 0
}
func (n *node[T]) deepcount() int {
count := len(n.items)
if !n.leaf() {
for i := 0; i <= len(n.items); i++ {
count += (*n.children)[i].deepcount()
}
}
if n.count != count {
return -1
}
return count
}
func (tr *BTreeG[T]) nodesaneprops(n *node[T], height int) bool {
if height == 1 {
if len(n.items) < 1 || len(n.items) > tr.max {
println(len(n.items) < 1)
return false
}
} else {
if len(n.items) < tr.min || len(n.items) > tr.max {
println(2)
return false
}
}
if !n.leaf() {
if len(*n.children) != len(n.items)+1 {
println(3)
return false
}
for i := 0; i < len(n.items); i++ {
if !tr.nodesaneprops((*n.children)[i], height+1) {
println(4)
return false
}
}
if !tr.nodesaneprops((*n.children)[len(n.items)], height+1) {
println(5)
return false
}
}
return true
}
func (tr *BTreeG[T]) saneprops() bool {
if tr.root != nil {
return tr.nodesaneprops(tr.root, 1)
}
return true
}
func (tr *BTreeG[T]) sanenilsnode(n *node[T]) bool {
items := n.items[:cap(n.items):cap(n.items)]
for i := len(n.items); i < len(items); i++ {
if !tr.eq(items[i], tr.empty) {
return false
}
}
if !n.leaf() {
for i := 0; i < len(*n.children); i++ {
if (*n.children)[i] == nil {
return false
}
}
children := (*n.children)[:cap(*n.children):cap(*n.children)]
for i := len(*n.children); i < len(children); i++ {
if children[i] != nil {
return false
}
}
for i := 0; i < len(*n.children); i++ {
if !tr.sanenilsnode((*n.children)[i]) {
return false
}
}
}
return true
}
// sanenils checks that all the slots in the item slice that are not used,
//
// n.items[len(n.items):cap(n.items):cap(n.items)]
//
// are equal to the empty value of the kind.
func (tr *BTreeG[T]) sanenils() bool {
if tr.root != nil {
return tr.sanenilsnode(tr.root)
}
return true
}
func (tr *BTreeG[T]) saneorder() bool {
var last T
var count int
var bad bool
tr.Walk(func(items []T) bool {
for _, item := range items {
if count > 0 {
if !tr.Less(last, item) {
bad = true
return false
}
}
last = item
count++
}
return true
})
return !bad && count == tr.count
}
func TestGenericIter(t *testing.T) {
N := 100_000
tr := testNewBTree()
var all []testKind
for i := 0; i < N; i++ {
tr.Load(testMakeItem(i))
all = append(all, testMakeItem(i))
}
var count int
var i int
iter := tr.Iter()
for ok := iter.First(); ok; ok = iter.Next() {
if !tr.eq(all[i], iter.Item()) {
panic("!")
}
count++
i++
}
if count != N {
t.Fatalf("expected %v, got %v", N, count)
}
iter.Release()
count = 0
i = len(all) - 1
iter = tr.Iter()
for ok := iter.Last(); ok; ok = iter.Prev() {
if !tr.eq(all[i], iter.Item()) {
panic("!")
}
i--
count++
}
if count != N {
t.Fatalf("expected %v, got %v", N, count)
}
iter.Release()
i = 0
iter = tr.Iter()
for ok := iter.First(); ok; ok = iter.Next() {
if !tr.eq(all[i], iter.Item()) {
panic("!")
}
i++
}
i--
for ok := iter.Prev(); ok; ok = iter.Prev() {
i--
if !tr.eq(all[i], iter.Item()) {
panic("!")
}
}
if i != 0 {
panic("!")
}
i++
for ok := iter.Next(); ok; ok = iter.Next() {
if !tr.eq(all[i], iter.Item()) {
panic("!")
}
i++
}
if i != N {
panic("!")
}
i = 0
for ok := iter.First(); ok; ok = iter.Next() {
if !tr.eq(all[i], iter.Item()) {
panic("!")
}
if tr.eq(iter.Item(), testMakeItem(N/2)) {
for ok = iter.Prev(); ok; ok = iter.Prev() {
i--
if !tr.eq(all[i], iter.Item()) {
panic("!")
}
}
break
}
i++
}
iter.Release()
}
func TestGenericIterSeek(t *testing.T) {
tr := NewBTreeG(func(a, b int) bool {
return a < b
})
var all []int
for i := 0; i < 10000; i++ {
tr.Set(i * 2)
all = append(all, i)
}
_ = all
{
iter := tr.Iter()
var vals []int
for ok := iter.Seek(501); ok; ok = iter.Next() {
vals = append(vals, iter.Item())
}
iter.Release()
assert(vals[0] == 502 && vals[1] == 504)
}
{
iter := tr.Iter()
var vals []int
for ok := iter.Seek(501); ok; ok = iter.Prev() {
vals = append(vals, iter.Item())
}
iter.Release()
assert(vals[0] == 502 && vals[1] == 500)
}
}
func TestGenericIterSeekHint(t *testing.T) {
tr := NewBTreeG(func(a, b int) bool {
return a < b
})
var all []int
for i := 0; i < 10000; i++ {
tr.Set(i * 2)
all = append(all, i)
}
var hint PathHint
_ = all
{
iter := tr.Iter()
var vals []int
for ok := iter.SeekHint(501, &hint); ok; ok = iter.Next() {
vals = append(vals, iter.Item())
}
iter.Release()
assert(vals[0] == 502 && vals[1] == 504)
}
{
iter := tr.Iter()
var vals []int
for ok := iter.SeekHint(501, &hint); ok; ok = iter.Prev() {
vals = append(vals, iter.Item())
}
iter.Release()
assert(vals[0] == 502 && vals[1] == 500)
}
}
func TestGenericIterSeekPrefix(t *testing.T) {
tr := NewBTreeG(func(a, b int) bool {
return a < b
})
count := 10_000
for i := 0; i < count; i++ {
tr.Set(i * 2)
}
for i := 0; i < count; i++ {
iter := tr.Iter()
ret := iter.Seek(i*2 - 1)
assert(ret == true)
iter.Release()
}
}
func testContiguousDelete(t *testing.T, bench bool, withskips bool, sorted bool,
delop func(tr *BTreeG[int], window, skips []int),
checkop func(tr *BTreeG[int], window, skips []int),
) {
const maxBTreeSize = 200000 // maximum number of items in tree
const deleteTarget = 200000 // stop running when this num items deleted
const maxWindowSize = 0.50 // percent: number of items in tree to delete
var runs int
var count int
var total time.Duration
if bench {
fmt.Printf("%-35s ", t.Name())
}
limit := deleteTarget
if withskips {
limit /= 2
}
for count < limit {
items := rand.Perm(rand.Int() % maxBTreeSize)
for i := 0; i < len(items); i++ {
items[i] *= 10
}
if sorted {
sort.Ints(items)
}
opts := Options{
Degree: (rand.Int() % 16) + 2,
}
tr := NewBTreeGOptions(func(a, b int) bool {
return a < b
}, opts)
for i := 0; i < len(items); i++ {
if sorted {
tr.Load(items[i])
} else {
tr.Set(items[i])
}
}
// take half the items from some random location
min := int(float64(len(items)) * 1.10 * (rand.Float64() - 0.05))
max := min + int(float64(len(items))*maxWindowSize*rand.Float64())
if max < min {
max, min = min, max
}
var window, skips []int
tr.Ascend(min, func(item int) bool {
if item > max {
return false
}
if withskips && rand.Int()%2 == 0 {
skips = append(skips, item)
} else {
window = append(window, item)
}
return true
})
len1 := tr.Len()
start := time.Now()
delop(tr, window, skips)
total += time.Since(start)
len2 := tr.Len()
count += len1 - len2
tr.sane()
for i := 0; i < len(items); i++ {
_, ok := tr.Get(items[i])
if items[i] < min || items[i] > max {
assert(ok)
}
}
for i := 0; i < len(window); i++ {
_, ok := tr.Get(window[i])
assert(!ok)
}
for i := 0; i < len(skips); i++ {
_, ok := tr.Get(skips[i])
assert(ok)
}
checkop(tr, window, skips)
runs++
}
if bench {
fmt.Printf("deleted %d items in %.4f secs, %.0f ns/item, %.0f/sec\n",
count,
total.Seconds(),
float64(total.Nanoseconds())/float64(count),
float64(count)/total.Seconds())
}
}
func TestContiguousDelete(t *testing.T) {
testContiguousDelete(t, true, false, false,
func(tr *BTreeG[int], window, skips []int) {
// DELETE
for i := 0; i < len(window); i++ {
tr.Delete(window[i])
}
},
func(tr *BTreeG[int], window, skips []int) {
// CHECK
},
)
}
func TestContiguousDeleteLoad(t *testing.T) {
if os.Getenv("LOADBENCH") == "" {
t.Skip()
}
testContiguousDelete(t, true, false, true,
func(tr *BTreeG[int], window, skips []int) {
// DELETE
for i := 0; i < len(window); i++ {
tr.Delete(window[i])
}
},
func(tr *BTreeG[int], window, skips []int) {
// CHECK
},
)
}
func TestContiguousDeleteHint(t *testing.T) {
var hint PathHint
testContiguousDelete(t, true, false, false,
func(tr *BTreeG[int], window, skips []int) {
// DELETE
for i := 0; i < len(window); i++ {
tr.DeleteHint(window[i], &hint)
}
},
func(tr *BTreeG[int], window, skips []int) {
// CHECK
},
)
}
func TestContiguousDeleteHintLoad(t *testing.T) {
if os.Getenv("LOADBENCH") == "" {
t.Skip()
}
var hint PathHint
testContiguousDelete(t, true, false, true,
func(tr *BTreeG[int], window, skips []int) {
// DELETE
for i := 0; i < len(window); i++ {
tr.DeleteHint(window[i], &hint)
}
},
func(tr *BTreeG[int], window, skips []int) {
// CHECK
},
)
}
func TestContiguousDeleteAscend(t *testing.T) {
testContiguousDelete(t, true, false, false,
func(tr *BTreeG[int], window, skips []int) {
// DELETE
if len(window) == 0 {
return
}
min := window[0]
max := window[len(window)-1]
tr.DeleteAscend(min, func(item int) Action {
if item > max {
return Stop
}
return Delete
})
},
func(tr *BTreeG[int], window, skips []int) {
// CHECK
},
)
}
func TestContiguousDeleteAscendLoad(t *testing.T) {
if os.Getenv("LOADBENCH") == "" {
t.Skip()
}
testContiguousDelete(t, true, false, true,
func(tr *BTreeG[int], window, skips []int) {
// DELETE
if len(window) == 0 {
return
}
min := window[0]
max := window[len(window)-1]
tr.DeleteAscend(min, func(item int) Action {
if item > max {
return Stop
}
return Delete
})
},
func(tr *BTreeG[int], window, skips []int) {
// CHECK
},
)
}
func TestContiguousDeleteAscendSkip(t *testing.T) {
testContiguousDelete(t, false, true, false,
func(tr *BTreeG[int], window, skips []int) {
// DELETE
if len(window) == 0 {
return
}
min := window[0]
max := window[len(window)-1]
if len(skips) > 0 {
if skips[0] < min {
min = skips[0]
}
if skips[len(skips)-1] > max {
max = skips[len(skips)-1]
}
}
tr.DeleteAscend(min, func(item int) Action {
if item > max {
return Stop
}
if len(skips) > 0 && item == skips[0] {
skips = skips[1:]
return Keep
}
return Delete
})
},
func(tr *BTreeG[int], window, skips []int) {
// CHECK
},
)
}
func TestContiguousDeleteAscendSkipLoad(t *testing.T) {
if os.Getenv("LOADBENCH") == "" {
t.Skip()
}
testContiguousDelete(t, false, true, true,
func(tr *BTreeG[int], window, skips []int) {
// DELETE
if len(window) == 0 {
return
}
min := window[0]
max := window[len(window)-1]
if len(skips) > 0 {
if skips[0] < min {
min = skips[0]
}
if skips[len(skips)-1] > max {
max = skips[len(skips)-1]
}
}
tr.DeleteAscend(min, func(item int) Action {
if item > max {
return Stop
}
if len(skips) > 0 && item == skips[0] {
skips = skips[1:]
return Keep
}
return Delete
})
},
func(tr *BTreeG[int], window, skips []int) {
// CHECK
},
)
}
func TestContiguousDeleteRange(t *testing.T) {
var deleted List[int]
testContiguousDelete(t, true, false, false,
func(tr *BTreeG[int], window, skips []int) {
// DELETE
deleted = List[int]{}
if len(window) == 0 {
return
}
min := window[0]
max := window[len(window)-1]
deleted = tr.DeleteRange(min, max, &DeleteRangeOptions{
NoReturn: false,
MaxInclusive: true,
})
},
func(tr *BTreeG[int], window, skips []int) {
// CHECK
assert(deleted.Len() == len(window))
var items []int
deleted.Scan(func(item int) bool {
items = append(items, item)
return true
})
assert(len(items) == len(window))
for i := 0; i < len(items); i++ {
assert(items[i] == window[i])
}
items = nil
deleted.Scan(func(item int) bool {
items = append(items, item)
return true
})
assert(len(items) == len(window))
for i := 0; i < len(items); i++ {
assert(items[i] == window[i])
}
},
)
}
func TestContiguousDeleteRangeNoExtract(t *testing.T) {
var deleted List[int]
testContiguousDelete(t, true, false, false,
func(tr *BTreeG[int], window, skips []int) {
// DELETE
deleted = List[int]{}
if len(window) == 0 {
return
}
min := window[0]
max := window[len(window)-1]
deleted = tr.DeleteRange(min, max, &DeleteRangeOptions{
NoReturn: true,
MaxInclusive: true,
})
},
func(tr *BTreeG[int], window, skips []int) {
assert(deleted.Len() == 0)
},
)
}
func TestContiguousDeleteRangeExclusive(t *testing.T) {
var deleted List[int]
var x int
var ok bool
testContiguousDelete(t, false, false, false,
func(tr *BTreeG[int], window, skips []int) {
// DELETE
deleted = List[int]{}
x = 0
ok = false
if len(window) == 0 {
return
}
min := window[0]
max := window[len(window)-1]
deleted = tr.DeleteRange(min, max, &DeleteRangeOptions{
NoReturn: false,
MaxInclusive: false,
})
// delete the max item manually to ensure the window is cleared.
x, ok = tr.Delete(max)
},
func(tr *BTreeG[int], window, skips []int) {
if len(window) > 0 {
assert(deleted.Len() == len(window)-1)
assert(ok)
assert(x == window[len(window)-1])
}
},
)
}
func TestContiguousDeleteRangeLoad(t *testing.T) {
if os.Getenv("LOADBENCH") == "" {
t.Skip()
}
testContiguousDelete(t, true, false, true,
func(tr *BTreeG[int], window, skips []int) {
// DELETE
if len(window) == 0 {
return
}
min := window[0]
max := window[len(window)-1]
tr.DeleteRange(min, max, &DeleteRangeOptions{
NoReturn: true,
MaxInclusive: true,
})
},
func(tr *BTreeG[int], window, skips []int) {
// CHECK
},
)
}
btree-1.8.1/go.mod 0000664 0000000 0000000 00000000051 15044500710 0013720 0 ustar 00root root 0000000 0000000 module github.com/tidwall/btree
go 1.20
btree-1.8.1/map.go 0000664 0000000 0000000 00000070276 15044500710 0013736 0 ustar 00root root 0000000 0000000 // 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 copier[T any] interface {
Copy() T
}
type isoCopier[T any] interface {
IsoCopy() T
}
func degreeToMinMax(deg int) (min, max int) {
if deg <= 0 {
deg = 32
} else if deg == 1 {
deg = 2 // must have at least 2
}
max = deg*2 - 1 // max items per node. max children is +1
min = max / 2
return min, max
}
var gisoid uint64
func newIsoID() uint64 {
return atomic.AddUint64(&gisoid, 1)
}
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 {
isoid uint64
root *mapNode[K, V]
count int
empty mapPair[K, V]
min int // min items
max int // max items
copyValues bool
isoCopyValues bool
}
func NewMap[K ordered, V any](degree int) *Map[K, V] {
m := new(Map[K, V])
m.init(degree)
return m
}
type mapNode[K ordered, V any] struct {
isoid uint64
count int
items []mapPair[K, V]
children *[]*mapNode[K, V]
}
// Copy the node for safe isolation.
func (tr *Map[K, V]) copy(n *mapNode[K, V]) *mapNode[K, V] {
n2 := new(mapNode[K, V])
n2.isoid = tr.isoid
n2.count = n.count
n2.items = make([]mapPair[K, V], len(n.items), cap(n.items))
copy(n2.items, n.items)
if tr.copyValues {
for i := 0; i < len(n2.items); i++ {
n2.items[i].value =
((interface{})(n2.items[i].value)).(copier[V]).Copy()
}
} else if tr.isoCopyValues {
for i := 0; i < len(n2.items); i++ {
n2.items[i].value =
((interface{})(n2.items[i].value)).(isoCopier[V]).IsoCopy()
}
}
if !n.leaf() {
n2.children = new([]*mapNode[K, V])
*n2.children = make([]*mapNode[K, V], len(*n.children), tr.max+1)
copy(*n2.children, *n.children)
}
return n2
}
// isoLoad loads the provided node and, if needed, performs a copy-on-write.
func (tr *Map[K, V]) isoLoad(cn **mapNode[K, V], mut bool) *mapNode[K, V] {
if mut && (*cn).isoid != tr.isoid {
*cn = tr.copy(*cn)
}
return *cn
}
func (tr *Map[K, V]) Copy() *Map[K, V] {
return tr.IsoCopy()
}
func (tr *Map[K, V]) IsoCopy() *Map[K, V] {
tr2 := new(Map[K, V])
*tr2 = *tr
tr2.isoid = newIsoID()
tr.isoid = newIsoID()
return tr2
}
func (tr *Map[K, V]) newNode(leaf bool) *mapNode[K, V] {
n := new(mapNode[K, V])
n.isoid = tr.isoid
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]) search(n *mapNode[K, V], key K) (index int, found bool) {
low, high := 0, len(n.items)
for low < high {
h := (low + high) / 2
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
}
func (tr *Map[K, V]) init(degree int) {
if tr.min != 0 {
return
}
tr.min, tr.max = degreeToMinMax(degree)
_, tr.copyValues = ((interface{})(tr.empty.value)).(copier[V])
if !tr.copyValues {
_, tr.isoCopyValues = ((interface{})(tr.empty.value)).(isoCopier[V])
}
}
// 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.init(0)
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, tr.max+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 := tr.max / 2
median = n.items[i]
// right node
right = tr.newNode(n.leaf())
right.items = n.items[i+1:]
if !n.leaf() {
*right.children = (*n.children)[i+1:]
}
right.updateCount()
// left node
n.items[i] = tr.empty
n.items = n.items[:i:i]
if !n.leaf() {
*n.children = (*n.children)[: i+1 : 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.isoLoad(pn, true)
i, found := tr.search(n, item.key)
if found {
prev = n.items[i].value
n.items[i] = item
return prev, true, false
}
if n.leaf() {
if len(n.items) == tr.max {
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) == tr.max {
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) {
tr.scan(iter, false)
}
func (tr *Map[K, V]) ScanMut(iter func(key K, value V) bool) {
tr.scan(iter, true)
}
func (tr *Map[K, V]) scan(iter func(key K, value V) bool, mut bool) {
if tr.root == nil {
return
}
tr.nodeScan(&tr.root, iter, mut)
}
func (tr *Map[K, V]) nodeScan(cn **mapNode[K, V],
iter func(key K, value V) bool, mut bool,
) bool {
n := tr.isoLoad(cn, mut)
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 !tr.nodeScan(&(*n.children)[i], iter, mut) {
return false
}
if !iter(n.items[i].key, n.items[i].value) {
return false
}
}
return tr.nodeScan(&(*n.children)[len(*n.children)-1], iter, mut)
}
// Get a value for key.
func (tr *Map[K, V]) Get(key K) (V, bool) {
return tr.get(key, false)
}
// GetMut gets a value for key.
// If needed, this may perform a copy the resulting value before returning.
//
// Mut methods are only useful when all of the following are true:
// - The interior data of the value requires changes.
// - The value is a pointer type.
// - The BTree has been copied using `Copy()` or `IsoCopy()`.
// - The value itself has a `Copy()` or `IsoCopy()` method.
//
// Mut methods may modify the tree structure and should have the same
// considerations as other mutable operations like Set, Delete, Clear, etc.
func (tr *Map[K, V]) GetMut(key K) (V, bool) {
return tr.get(key, true)
}
func (tr *Map[K, V]) get(key K, mut bool) (V, bool) {
if tr.root == nil {
return tr.empty.value, false
}
n := tr.isoLoad(&tr.root, mut)
for {
i, found := tr.search(n, key)
if found {
return n.items[i].value, true
}
if n.leaf() {
return tr.empty.value, false
}
n = tr.isoLoad(&(*n.children)[i], mut)
}
}
// 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.isoLoad(pn, true)
var i int
var found bool
if max {
i, found = len(n.items)-1, true
} else {
i, found = tr.search(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) < tr.min {
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.isoLoad(&(*n.children)[i], true)
right := tr.isoLoad(&(*n.children)[i+1], true)
if len(left.items)+len(right.items) < tr.max {
// 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) {
tr.ascend(pivot, iter, false)
}
func (tr *Map[K, V]) AscendMut(pivot K, iter func(key K, value V) bool) {
tr.ascend(pivot, iter, true)
}
func (tr *Map[K, V]) ascend(pivot K, iter func(key K, value V) bool, mut bool) {
if tr.root == nil {
return
}
tr.nodeAscend(&tr.root, pivot, iter, mut)
}
// The return value of this function determines whether we should keep iterating
// upon this functions return.
func (tr *Map[K, V]) nodeAscend(cn **mapNode[K, V], pivot K,
iter func(key K, value V) bool, mut bool,
) bool {
n := tr.isoLoad(cn, mut)
i, found := tr.search(n, pivot)
if !found {
if !n.leaf() {
if !tr.nodeAscend(&(*n.children)[i], pivot, iter, mut) {
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 !tr.nodeScan(&(*n.children)[i+1], iter, mut) {
return false
}
}
}
return true
}
func (tr *Map[K, V]) Reverse(iter func(key K, value V) bool) {
tr.reverse(iter, false)
}
func (tr *Map[K, V]) ReverseMut(iter func(key K, value V) bool) {
tr.reverse(iter, true)
}
func (tr *Map[K, V]) reverse(iter func(key K, value V) bool, mut bool) {
if tr.root == nil {
return
}
tr.nodeReverse(&tr.root, iter, mut)
}
func (tr *Map[K, V]) nodeReverse(cn **mapNode[K, V],
iter func(key K, value V) bool, mut bool,
) bool {
n := tr.isoLoad(cn, mut)
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 !tr.nodeReverse(&(*n.children)[len(*n.children)-1], iter, mut) {
return false
}
for i := len(n.items) - 1; i >= 0; i-- {
if !iter(n.items[i].key, n.items[i].value) {
return false
}
if !tr.nodeReverse(&(*n.children)[i], iter, mut) {
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) {
tr.descend(pivot, iter, false)
}
func (tr *Map[K, V]) DescendMut(pivot K, iter func(key K, value V) bool) {
tr.descend(pivot, iter, true)
}
func (tr *Map[K, V]) descend(
pivot K,
iter func(key K, value V) bool,
mut bool,
) {
if tr.root == nil {
return
}
tr.nodeDescend(&tr.root, pivot, iter, mut)
}
func (tr *Map[K, V]) nodeDescend(cn **mapNode[K, V], pivot K,
iter func(key K, value V) bool, mut bool,
) bool {
n := tr.isoLoad(cn, mut)
i, found := tr.search(n, pivot)
if !found {
if !n.leaf() {
if !tr.nodeDescend(&(*n.children)[i], pivot, iter, mut) {
return false
}
}
i--
}
for ; i >= 0; i-- {
if !iter(n.items[i].key, n.items[i].value) {
return false
}
if !n.leaf() {
if !tr.nodeReverse(&(*n.children)[i], iter, mut) {
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.isoLoad(&tr.root, true)
for {
n.count++ // optimistically update counts
if n.leaf() {
if len(n.items) < tr.max {
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.isoLoad(&(*n.children)[len(*n.children)-1], true)
}
// 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) {
return tr.minMut(false)
}
func (tr *Map[K, V]) MinMut() (K, V, bool) {
return tr.minMut(true)
}
func (tr *Map[K, V]) minMut(mut bool) (key K, value V, ok bool) {
if tr.root == nil {
return key, value, false
}
n := tr.isoLoad(&tr.root, mut)
for {
if n.leaf() {
item := n.items[0]
return item.key, item.value, true
}
n = tr.isoLoad(&(*n.children)[0], mut)
}
}
// Max returns the maximum item in tree.
// Returns nil if the tree has no items.
func (tr *Map[K, V]) Max() (K, V, bool) {
return tr.maxMut(false)
}
func (tr *Map[K, V]) MaxMut() (K, V, bool) {
return tr.maxMut(true)
}
func (tr *Map[K, V]) maxMut(mut bool) (K, V, bool) {
if tr.root == nil {
return tr.empty.key, tr.empty.value, false
}
n := tr.isoLoad(&tr.root, mut)
for {
if n.leaf() {
item := n.items[len(n.items)-1]
return item.key, item.value, true
}
n = tr.isoLoad(&(*n.children)[len(*n.children)-1], mut)
}
}
// 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.isoLoad(&tr.root, true)
var item mapPair[K, V]
for {
n.count-- // optimistically update counts
if n.leaf() {
item = n.items[0]
if len(n.items) == tr.min {
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.isoLoad(&(*n.children)[0], true)
}
// 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.isoLoad(&tr.root, true)
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) == tr.min {
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.isoLoad(&(*n.children)[len(*n.children)-1], true)
}
// 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) {
return tr.getAt(index, false)
}
func (tr *Map[K, V]) GetAtMut(index int) (K, V, bool) {
return tr.getAt(index, true)
}
func (tr *Map[K, V]) getAt(index int, mut bool) (K, V, bool) {
if tr.root == nil || index < 0 || index >= tr.count {
return tr.empty.key, tr.empty.value, false
}
n := tr.isoLoad(&tr.root, mut)
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 = tr.isoLoad(&(*n.children)[i], mut)
}
}
// 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.isoLoad(&tr.root, true)
outer:
for {
n.count-- // optimistically update counts
if n.leaf() {
// the index is the item position
item = n.items[index]
if len(n.items) == tr.min {
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.isoLoad(&(*n.children)[i], true)
}
// 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]
mut bool
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] {
return tr.iter(false)
}
func (tr *Map[K, V]) IterMut() MapIter[K, V] {
return tr.iter(true)
}
func (tr *Map[K, V]) iter(mut bool) MapIter[K, V] {
var iter MapIter[K, V]
iter.tr = tr
iter.mut = mut
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.isoLoad(&iter.tr.root, iter.mut)
for {
i, found := iter.tr.search(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 = iter.tr.isoLoad(&(*n.children)[i], iter.mut)
}
}
// 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.isoLoad(&iter.tr.root, iter.mut)
for {
iter.stack = append(iter.stack, mapIterStackItem[K, V]{n, 0})
if n.leaf() {
break
}
n = iter.tr.isoLoad(&(*n.children)[0], iter.mut)
}
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.isoLoad(&iter.tr.root, iter.mut)
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 = iter.tr.isoLoad(&(*n.children)[len(n.items)], iter.mut)
}
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 := iter.tr.isoLoad(&(*s.n.children)[s.i], iter.mut)
for {
iter.stack = append(iter.stack, mapIterStackItem[K, V]{n, 0})
if n.leaf() {
break
}
n = iter.tr.isoLoad(&(*n.children)[0], iter.mut)
}
}
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 := iter.tr.isoLoad(&(*s.n.children)[s.i], iter.mut)
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 = iter.tr.isoLoad(&(*n.children)[len(n.items)], iter.mut)
}
}
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 {
return tr.values(false)
}
func (tr *Map[K, V]) ValuesMut() []V {
return tr.values(true)
}
func (tr *Map[K, V]) values(mut bool) []V {
values := make([]V, 0, tr.Len())
if tr.root != nil {
values = tr.nodeValues(&tr.root, values, mut)
}
return values
}
func (tr *Map[K, V]) nodeValues(cn **mapNode[K, V], values []V, mut bool) []V {
n := tr.isoLoad(cn, mut)
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 = tr.nodeValues(&(*n.children)[i], values, mut)
values = append(values, n.items[i].value)
}
return tr.nodeValues(&(*n.children)[len(*n.children)-1], values, mut)
}
// 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) {
return tr.keyValues(false)
}
func (tr *Map[K, V]) KeyValuesMut() ([]K, []V) {
return tr.keyValues(true)
}
func (tr *Map[K, V]) keyValues(mut bool) ([]K, []V) {
keys := make([]K, 0, tr.Len())
values := make([]V, 0, tr.Len())
if tr.root != nil {
keys, values = tr.nodeKeyValues(&tr.root, keys, values, mut)
}
return keys, values
}
func (tr *Map[K, V]) nodeKeyValues(cn **mapNode[K, V], keys []K, values []V,
mut bool,
) ([]K, []V) {
n := tr.isoLoad(cn, mut)
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 = tr.nodeKeyValues(&(*n.children)[i], keys, values, mut)
keys = append(keys, n.items[i].key)
values = append(values, n.items[i].value)
}
return tr.nodeKeyValues(&(*n.children)[len(*n.children)-1], keys, values,
mut)
}
// Clear will delete all items.
func (tr *Map[K, V]) Clear() {
tr.count = 0
tr.root = nil
}
btree-1.8.1/map_test.go 0000664 0000000 0000000 00000076012 15044500710 0014767 0 ustar 00root root 0000000 0000000 package btree
import (
"fmt"
"math/rand"
"reflect"
"sort"
"sync"
"testing"
"time"
)
type testMapKind = int
func testMapMakeItem(x int) (item testMapKind) {
return x
}
// testNewBTree must return an operational btree for testing.
func testMapNewBTree() *Map[testMapKind, testMapKind] {
return new(Map[testMapKind, testMapKind])
}
func testMapNewBTreeDegrees(degree int) *Map[testMapKind, testMapKind] {
return NewMap[testMapKind, testMapKind](degree)
}
func randMapKeys(N int) (keys []testMapKind) {
keys = make([]testMapKind, N)
for _, i := range rand.Perm(N) {
keys[i] = testMapMakeItem(i)
}
return keys
}
func (tr *Map[K, V]) lt(a, b K) bool { return a < b }
func (tr *Map[K, V]) eq(a, b K) bool { return !(tr.lt(a, b) || tr.lt(b, a)) }
func (tr *Map[K, V]) lte(a, b K) bool { return tr.lt(a, b) || tr.eq(a, b) }
func (tr *Map[K, V]) gt(a, b K) bool { return tr.lt(b, a) }
func (tr *Map[K, V]) gte(a, b K) bool { return tr.gt(a, b) || tr.eq(a, b) }
func mapKindsAreEqual(a, b []testMapKind) bool {
var tr Map[testMapKind, testMapKind]
if len(a) != len(b) {
return false
}
for i := 0; i < len(a); i++ {
if !tr.eq(a[i], b[i]) {
return false
}
}
return true
}
func TestMapMakeItemOrder(t *testing.T) {
tr := testMapNewBTree()
ints := []int{0, 1, 2, 3, 4, 10, 20, 30, 40, 100, 200, 300, 400}
for i := 0; i < len(ints)-1; i++ {
a := testMapMakeItem(ints[i])
b := testMapMakeItem(ints[i+1])
if !tr.lt(a, b) {
t.Fatalf("bad ordering for testMakeItem: '%v' !< '%v'", a, b)
}
}
}
func TestMapDescend(t *testing.T) {
tr := testMapNewBTree()
var count int
tr.Descend(testMapMakeItem(rand.Int()), func(item, value testMapKind) bool {
count++
return true
})
if count > 0 {
t.Fatalf("expected 0, got %v", count)
}
var keys []testMapKind
for i := 0; i < 1000; i += 10 {
keys = append(keys, testMapMakeItem(i))
tr.Set(keys[len(keys)-1], keys[len(keys)-1])
}
var exp []testMapKind
tr.Reverse(func(item, value testMapKind) bool {
exp = append(exp, item)
return true
})
for i := 999; i >= 0; i-- {
key := testMapMakeItem(i)
var all []testMapKind
tr.Descend(key, func(item, value testMapKind) bool {
all = append(all, item)
return true
})
for len(exp) > 0 && tr.lt(key, exp[0]) {
exp = exp[1:]
}
var count int
tr.Descend(key, func(item, value testMapKind) bool {
if count == (i+1)%tr.max {
return false
}
count++
return true
})
if count > len(exp) {
t.Fatalf("expected 1, got %v", count)
}
if !mapKindsAreEqual(exp, all) {
fmt.Printf("exp: %v\n", exp)
fmt.Printf("all: %v\n", all)
t.Fatal("mismatch")
}
for j := 0; j < tr.Len(); j++ {
count = 0
tr.Descend(key, func(item, value testMapKind) bool {
if count == j {
return false
}
count++
return true
})
}
}
}
func TestMapAscend(t *testing.T) {
tr := testMapNewBTree()
var count int
tr.Ascend(testMapMakeItem(1), func(item, value testMapKind) bool {
count++
return true
})
if count > 0 {
t.Fatalf("expected 0, got %v", count)
}
var keys []testMapKind
for i := 0; i < 1000; i += 10 {
keys = append(keys, testMapMakeItem(i))
tr.Set(keys[len(keys)-1], keys[len(keys)-1])
tr.sane()
}
exp := keys
for i := -1; i < 1000; i++ {
key := testMapMakeItem(i)
var all []testMapKind
tr.Ascend(key, func(item, value testMapKind) bool {
all = append(all, item)
return true
})
for len(exp) > 0 && tr.lt(exp[0], key) {
exp = exp[1:]
}
var count int
tr.Ascend(key, func(item, value testMapKind) bool {
if count == (i+1)%tr.max {
return false
}
count++
return true
})
if count > len(exp) {
t.Fatalf("expected 1, got %v", count)
}
if !mapKindsAreEqual(exp, all) {
t.Fatal("mismatch")
}
}
}
func TestMapKeyValues(t *testing.T) {
tr := testMapNewBTree()
if len(tr.Keys()) != 0 {
t.Fatalf("expected 0, got %v", len(tr.Keys()))
}
if len(tr.Values()) != 0 {
t.Fatalf("expected 0, got %v", len(tr.Values()))
}
keys, values := tr.KeyValues()
if len(keys) != 0 {
t.Fatalf("expected 0, got %v", len(keys))
}
if len(values) != 0 {
t.Fatalf("expected 0, got %v", len(values))
}
keys = nil
values = nil
for i := 0; i < 100000; i += 10 {
keys = append(keys, testMapMakeItem(i))
values = append(values, testMapMakeItem(i)*10)
tr.Set(keys[len(keys)-1], values[len(values)-1])
tr.sane()
}
keys2 := tr.Keys()
values2 := tr.Values()
if !kindsAreEqual(keys, keys2) {
t.Fatalf("not equal")
}
if !kindsAreEqual(values, values2) {
t.Fatalf("not equal")
}
keys2, values2 = tr.KeyValues()
if !kindsAreEqual(keys, keys2) {
t.Fatalf("not equal")
}
if !kindsAreEqual(values, values2) {
t.Fatalf("not equal")
}
}
func TestMapSimpleRandom(t *testing.T) {
start := time.Now()
for time.Since(start) < time.Second*2 {
N := 100_000
items := randMapKeys(N)
tr := testMapNewBTree()
tr.sane()
for i := 0; i < len(items); i++ {
if v, ok := tr.Get(items[i]); ok || !tr.eq(v, tr.empty.value) {
panic("!")
}
if v, ok := tr.Set(items[i], items[i]); ok || !tr.eq(v, tr.empty.value) {
panic("!")
}
if v, ok := tr.Get(items[i]); !ok || !tr.eq(v, items[i]) {
panic("!")
}
}
tr.sane()
for i := 0; i < len(items); i++ {
if v, ok := tr.Set(items[i], items[i]); !ok || !tr.eq(v, items[i]) {
panic("!")
}
}
pivot := items[len(items)/2]
tr.Ascend(pivot, func(item, value testMapKind) bool {
if tr.lt(item, pivot) {
panic("!")
}
return true
})
var min testMapKind
index := 0
tr.Scan(func(item, value testMapKind) bool {
if index == len(items)/2 {
return false
}
if index > 0 {
if tr.lt(item, min) {
panic("!")
}
}
min = item
index++
return true
})
tr.sane()
for i := 0; i < len(items); i++ {
if v, ok := tr.Delete(items[i]); !ok || !tr.eq(v, items[i]) {
panic("!")
}
if i%97 == 0 {
tr.sane()
}
if v, ok := tr.Delete(items[i]); ok || !tr.eq(v, tr.empty.value) {
panic("!")
}
}
if tr.Len() != 0 {
panic("!")
}
tr.sane()
for i := 0; i < len(items); i++ {
if v, ok := tr.Delete(items[i]); ok || !tr.eq(v, tr.empty.value) {
panic("!")
}
}
tr.sane()
tr.Scan(func(item, value testMapKind) bool {
panic("!")
})
}
}
func TestMapBTree(t *testing.T) {
N := 10000
tr := testMapNewBTree()
tr.sane()
keys := randMapKeys(N)
// insert all items
for _, key := range keys {
if v, ok := tr.Set(key, key); ok || !tr.eq(v, tr.empty.value) {
t.Fatal("expected false")
}
tr.sane()
}
// check length
if tr.Len() != len(keys) {
t.Fatalf("expected %v, got %v", len(keys), tr.Len())
}
// get each value
for _, key := range keys {
if v, ok := tr.Get(key); !ok || !tr.eq(v, key) {
t.Fatalf("expected '%v', got '%v'", key, v)
}
}
// scan all items
var prev testMapKind
var count int
tr.Scan(func(item, value testMapKind) bool {
if count > 0 {
if tr.lte(item, prev) {
t.Fatal("out of order")
}
}
prev = item
count++
return true
})
if count != len(keys) {
t.Fatalf("expected '%v', got '%v'", len(keys), count)
}
// reverse all items
count = 0
tr.Reverse(func(item, value testMapKind) bool {
if count > 0 {
if tr.gte(item, prev) {
t.Fatal("out of order")
}
}
prev = item
count++
return true
})
if count != len(keys) {
t.Fatalf("expected '%v', got '%v'", len(keys), count)
}
// try to get an invalid item
if v, ok := tr.Get(testMapMakeItem(-1)); ok || !tr.eq(v, tr.empty.value) {
t.Fatal("expected nil")
}
// scan and quit at various steps
for i := 0; i < 100; i++ {
var j int
tr.Scan(func(item, value testMapKind) bool {
if j == i {
return false
}
j++
return true
})
}
// reverse and quit at various steps
for i := 0; i < 100; i++ {
var j int
tr.Reverse(func(item, value testMapKind) bool {
if j == i {
return false
}
j++
return true
})
}
// delete half the items
for _, key := range keys[:len(keys)/2] {
if v, ok := tr.Delete(key); !ok || !tr.eq(v, key) {
t.Fatalf("expected '%v', got '%v'", key, v)
}
}
// check length
if tr.Len() != len(keys)/2 {
t.Fatalf("expected %v, got %v", len(keys)/2, tr.Len())
}
// try delete half again
for _, key := range keys[:len(keys)/2] {
if v, ok := tr.Delete(key); ok || !tr.eq(v, tr.empty.value) {
t.Fatal("expected false")
}
tr.sane()
}
// check length
if tr.Len() != len(keys)/2 {
t.Fatalf("expected %v, got %v", len(keys)/2, tr.Len())
}
// scan items
count = 0
tr.Scan(func(item, value testMapKind) bool {
if count > 0 {
if tr.lte(item, prev) {
t.Fatal("out of order")
}
}
prev = item
count++
return true
})
if count != len(keys)/2 {
t.Fatalf("expected '%v', got '%v'", len(keys), count)
}
// replace second half
for _, key := range keys[len(keys)/2:] {
if v, ok := tr.Set(key, key); !ok || !tr.eq(v, key) {
t.Fatalf("expected '%v', got '%v'", key, v)
}
tr.sane()
}
// delete next half the items
for _, key := range keys[len(keys)/2:] {
if v, ok := tr.Delete(key); !ok || !tr.eq(v, key) {
t.Fatalf("expected '%v', got '%v'", key, v)
}
tr.sane()
}
// check length
if tr.Len() != 0 {
t.Fatalf("expected %v, got %v", 0, tr.Len())
}
// do some stuff on an empty tree
if v, ok := tr.Get(keys[0]); ok || !tr.eq(v, tr.empty.value) {
t.Fatal("expected nil")
}
tr.Scan(func(item, value testMapKind) bool {
t.Fatal("should not be reached")
return true
})
tr.Reverse(func(item, value testMapKind) bool {
t.Fatal("should not be reached")
return true
})
if v, ok := tr.Delete(testMapMakeItem(-1)); ok || !tr.eq(v, tr.empty.value) {
t.Fatal("expected nil")
}
tr.sane()
}
func TestMapBTreeOne(t *testing.T) {
tr := testMapNewBTree()
tr.Set(testMapMakeItem(1), testMapMakeItem(1))
tr.Delete(testMapMakeItem(1))
tr.Set(testMapMakeItem(1), testMapMakeItem(1))
tr.Delete(testMapMakeItem(1))
tr.Set(testMapMakeItem(1), testMapMakeItem(1))
tr.Delete(testMapMakeItem(1))
if tr.Len() != 0 {
panic("!")
}
tr.sane()
}
func TestMapBTree256(t *testing.T) {
for degree := -1; degree < 256; degree++ {
tr := testMapNewBTreeDegrees(degree)
var n int
for j := 0; j < 2; j++ {
for _, i := range rand.Perm(256) {
tr.Set(testMapMakeItem(i), testMapMakeItem(i))
n++
if tr.Len() != n {
t.Fatalf("expected 256, got %d", n)
}
}
for _, i := range rand.Perm(256) {
if v, ok := tr.Get(testMapMakeItem(i)); !ok || !tr.eq(v, testMapMakeItem(i)) {
t.Fatalf("expected %v, got %v", i, v)
}
}
for _, i := range rand.Perm(256) {
tr.Delete(testMapMakeItem(i))
n--
if tr.Len() != n {
t.Fatalf("expected 256, got %d", n)
}
}
for _, i := range rand.Perm(256) {
if v, ok := tr.Get(testMapMakeItem(i)); ok || !tr.eq(v, tr.empty.value) {
t.Fatal("expected nil")
}
}
}
}
}
func shuffleMapItems(keys []testMapKind) {
for i := range keys {
j := rand.Intn(i + 1)
keys[i], keys[j] = keys[j], keys[i]
}
}
func sortMapItems(keys []testMapKind) {
tr := testMapNewBTree()
sort.Slice(keys, func(i, j int) bool {
return tr.lt(keys[i], keys[j])
})
}
func TestMapRandom(t *testing.T) {
N := 200000
keys := randMapKeys(N)
tr := testMapNewBTree()
tr.sane()
if _, v, ok := tr.Min(); ok || !tr.eq(v, tr.empty.value) {
t.Fatalf("expected nil")
}
if _, v, ok := tr.Max(); ok || !tr.eq(v, tr.empty.value) {
t.Fatalf("expected nil")
}
if _, v, ok := tr.PopMin(); ok || !tr.eq(v, tr.empty.value) {
t.Fatalf("expected nil")
}
if _, v, ok := tr.PopMax(); ok || !tr.eq(v, tr.empty.value) {
t.Fatalf("expected nil")
}
if tr.Height() != 0 {
t.Fatalf("expected 0, got %d", tr.Height())
}
tr.sane()
shuffleMapItems(keys)
for i := 0; i < len(keys); i++ {
if v, ok := tr.Set(keys[i], keys[i]); ok || !tr.eq(v, tr.empty.value) {
t.Fatalf("expected nil")
}
if i%123 == 0 {
tr.sane()
}
}
tr.sane()
sortMapItems(keys)
var n int
tr.Scan(func(item, value testMapKind) bool {
n++
return false
})
if n != 1 {
t.Fatalf("expected 1, got %d", n)
}
n = 0
tr.Scan(func(item, value testMapKind) bool {
if !tr.eq(item, keys[n]) {
t.Fatalf("expected %v, got %v", keys[n], item)
}
n++
return true
})
if n != len(keys) {
t.Fatalf("expected %d, got %d", len(keys), n)
}
if tr.Len() != len(keys) {
t.Fatalf("expected %d, got %d", tr.Len(), len(keys))
}
for i := 0; i < tr.Len(); i++ {
if _, v, ok := tr.GetAt(i); !ok || !tr.eq(v, keys[i]) {
t.Fatalf("expected %v, got %v", keys[i], v)
}
}
n = 0
tr.Reverse(func(item, value testMapKind) bool {
n++
return false
})
if n != 1 {
t.Fatalf("expected 1, got %d", n)
}
n = 0
tr.Reverse(func(item, value testMapKind) bool {
if !tr.eq(item, keys[len(keys)-n-1]) {
t.Fatalf("expected %v, got %v", keys[len(keys)-n-1], item)
}
n++
return true
})
if n != len(keys) {
t.Fatalf("expected %d, got %d", len(keys), n)
}
if tr.Len() != len(keys) {
t.Fatalf("expected %d, got %d", tr.Len(), len(keys))
}
tr.sane()
n = 0
for i := 0; i < 1000; i++ {
n := 0
tr.Scan(func(item, value testMapKind) bool {
if n == i {
return false
}
n++
return true
})
if n != i {
t.Fatalf("expected %d, got %d", i, n)
}
}
n = 0
for i := 0; i < 1000; i++ {
n = 0
tr.Reverse(func(item, value testMapKind) bool {
if n == i {
return false
}
n++
return true
})
if n != i {
t.Fatalf("expected %d, got %d", i, n)
}
}
sortMapItems(keys)
for i := 0; i < len(keys); i++ {
var res testMapKind
var j int
tr.Ascend(keys[i], func(item, value testMapKind) bool {
if j == 0 {
res = item
}
j++
return j == i%500
})
if !tr.eq(res, keys[i]) {
t.Fatal("not equal")
}
}
for i := len(keys) - 1; i >= 0; i-- {
var res testMapKind
var j int
tr.Descend(keys[i], func(item, value testMapKind) bool {
if j == 0 {
res = item
}
j++
return j == i%500
})
if !tr.eq(res, keys[i]) {
t.Fatal("not equal")
}
}
if tr.Height() == 0 {
t.Fatalf("expected non-zero")
}
if _, v, ok := tr.Min(); !ok || !tr.eq(v, keys[0]) {
t.Fatalf("expected '%v', got '%v'", keys[0], v)
}
if _, v, ok := tr.Max(); !ok || !tr.eq(v, keys[len(keys)-1]) {
t.Fatalf("expected '%v', got '%v'", keys[len(keys)-1], v)
}
if _, v, ok := tr.PopMin(); !ok || !tr.eq(v, keys[0]) {
t.Fatalf("expected '%v', got '%v'", keys[0], v)
}
tr.sane()
if _, v, ok := tr.PopMax(); !ok || !tr.eq(v, keys[len(keys)-1]) {
t.Fatalf("expected '%v', got '%v'", keys[len(keys)-1], v)
}
tr.sane()
tr.Set(keys[0], keys[0])
tr.Set(keys[len(keys)-1], keys[len(keys)-1])
shuffleMapItems(keys)
for i := 0; i < len(keys); i++ {
if v, ok := tr.Get(keys[i]); !ok || !tr.eq(v, keys[i]) {
t.Fatalf("expected '%v', got '%v'", keys[i], v)
}
if v, ok := tr.Get(keys[i]); !ok || !tr.eq(v, keys[i]) {
t.Fatalf("expected '%v', got '%v'", keys[i], v)
}
}
sortMapItems(keys)
for i := 0; i < len(keys); i++ {
if _, v, ok := tr.PopMin(); !ok || !tr.eq(v, keys[i]) {
t.Fatalf("expected '%v', got '%v'", keys[i], v)
}
}
for i := 0; i < len(keys); i++ {
if v, ok := tr.Set(keys[i], keys[i]); ok || !tr.eq(v, tr.empty.value) {
t.Fatalf("expected nil")
}
}
for i := len(keys) - 1; i >= 0; i-- {
if _, v, ok := tr.PopMax(); !ok || !tr.eq(v, keys[i]) {
t.Fatalf("expected '%v', got '%v'", keys[i], v)
}
}
for i := 0; i < len(keys); i++ {
if v, ok := tr.Set(keys[i], keys[i]); ok || !tr.eq(v, tr.empty.value) {
t.Fatalf("expected nil")
}
}
if v, ok := tr.Delete(testMapMakeItem(-1)); ok || !tr.eq(v, tr.empty.value) {
t.Fatal("expected nil")
}
tr.sane()
shuffleMapItems(keys)
if v, ok := tr.Delete(keys[len(keys)/2]); !ok || !tr.eq(v, keys[len(keys)/2]) {
t.Fatalf("expected '%v', got '%v'", keys[len(keys)/2], v)
}
tr.sane()
if v, ok := tr.Delete(keys[len(keys)/2]); ok || !tr.eq(v, tr.empty.value) {
t.Fatalf("expected '%v', got '%v'", tr.empty.value, v)
}
tr.sane()
tr.Set(keys[len(keys)/2], keys[len(keys)/2])
tr.sane()
for i := 0; i < len(keys); i++ {
if v, ok := tr.Delete(keys[i]); !ok || !tr.eq(v, keys[i]) {
t.Fatalf("expected '%v', got '%v'", keys[i], v)
}
if v, ok := tr.Get(keys[i]); ok || !tr.eq(v, tr.empty.value) {
t.Fatalf("expected nil")
}
if v, ok := tr.Get(keys[i]); ok || !tr.eq(v, tr.empty.value) {
t.Fatalf("expected nil")
}
if i%97 == 0 {
tr.sane()
}
}
if tr.Height() != 0 {
t.Fatalf("expected 0, got %d", tr.Height())
}
shuffleMapItems(keys)
for i := 0; i < len(keys); i++ {
if v, ok := tr.Load(keys[i], keys[i]); ok || !tr.eq(v, tr.empty.value) {
t.Fatalf("expected nil")
}
if i%97 == 0 {
tr.sane()
}
}
for i := 0; i < len(keys); i++ {
if v, ok := tr.Get(keys[i]); !ok || !tr.eq(v, keys[i]) {
t.Fatalf("expected '%v', got '%v'", keys[i], v)
}
}
shuffleMapItems(keys)
for i := 0; i < len(keys); i++ {
if v, ok := tr.Delete(keys[i]); !ok || !tr.eq(v, keys[i]) {
t.Fatalf("expected '%v', got '%v'", keys[i], v)
}
if v, ok := tr.Get(keys[i]); ok || !tr.eq(v, tr.empty.value) {
t.Fatalf("expected nil")
}
}
sortMapItems(keys)
for i := 0; i < len(keys); i++ {
if v, ok := tr.Load(keys[i], keys[i]); ok || !tr.eq(v, tr.empty.value) {
t.Fatalf("expected nil")
}
if i%97 == 0 {
tr.sane()
}
}
shuffleMapItems(keys)
if v, ok := tr.Load(keys[0], keys[0]); !ok || !tr.eq(v, keys[0]) {
t.Fatalf("expected '%v', got '%v'", keys[0], v)
}
tr.sane()
}
func TestMapLess(t *testing.T) {
tr := testMapNewBTree()
if !tr.lt(testMapMakeItem(1), testMapMakeItem(2)) {
panic("invalid")
}
if tr.lt(testMapMakeItem(2), testMapMakeItem(1)) {
panic("invalid")
}
if tr.lt(testMapMakeItem(1), testMapMakeItem(1)) {
panic("invalid")
}
}
func TestMapDeleteRandom(t *testing.T) {
N := 2_000_000
tr := testMapNewBTree()
for i := 0; i < N; i++ {
tr.Load(testMapMakeItem(i), testMapMakeItem(i))
}
tr.sane()
for tr.Len() > 0 {
var item testMapKind
var ok bool
switch rand.Int() % 3 {
case 0:
_, item, ok = tr.GetAt(tr.Len() / 2)
case 1:
_, item, ok = tr.Min()
case 2:
_, item, ok = tr.Max()
}
if !ok {
panic("!")
}
v, ok := tr.Delete(item)
if !ok || !tr.eq(v, item) {
panic("!")
}
}
}
func TestMapDeleteAt(t *testing.T) {
N := 10_000
tr := testMapNewBTree()
keys := randMapKeys(N)
for _, key := range keys {
tr.Set(key, key)
}
tr.sane()
for tr.Len() > 0 {
index := rand.Intn(tr.Len())
_, item1, ok1 := tr.GetAt(index)
_, item2, ok2 := tr.DeleteAt(index)
if !ok1 || !ok2 || !tr.eq(item1, item2) {
panic("mismatch")
}
tr.sane()
}
}
func TestMapVarious(t *testing.T) {
N := 1_000_000
tr := testMapNewBTree()
for _, i := range randMapKeys(N) {
if v, ok := tr.Set(i, i); ok || !tr.eq(v, tr.empty.value) {
panic("!")
}
}
for _, i := range randMapKeys(N) {
if v, ok := tr.Get(i); !ok || !tr.eq(v, i) {
panic("!")
}
}
for _, i := range randMapKeys(N) {
if v, ok := tr.Delete(i); !ok || !tr.eq(v, i) {
panic("!")
}
}
if _, v, ok := tr.DeleteAt(0); ok || !tr.eq(v, tr.empty.value) {
panic("!")
}
if _, v, ok := tr.GetAt(0); ok || !tr.eq(v, tr.empty.value) {
panic("!")
}
for i := 0; i < N; i++ {
item := testMapMakeItem(i)
if v, ok := tr.Set(item, item); ok || !tr.eq(v, tr.empty.value) {
panic("!")
}
item = testMapMakeItem(i)
if v, ok := tr.Set(item, item); !ok || !tr.eq(v, item) {
panic("!")
}
item = testMapMakeItem(i)
if v, ok := tr.Set(item, item); !ok || !tr.eq(v, item) {
panic("!")
}
}
for i := 0; i < N; i++ {
item := testMapMakeItem(i)
if v, ok := tr.Get(item); !ok || !tr.eq(v, item) {
panic("!")
}
}
for i := 0; i < 100; i++ {
var count int
tr.Scan(func(_, _ testMapKind) bool {
if count == i {
return false
}
count++
return true
})
}
for i := 0; i < N; i++ {
item := testMapMakeItem(i)
if v, ok := tr.Delete(item); !ok || !tr.eq(v, item) {
panic("!")
}
}
}
func (tr *Map[K, V]) sane() {
if err := tr.Sane(); err != nil {
panic(err)
}
}
type saneMapError string
func (err saneMapError) Error() string {
return string(err)
}
// btree_sane returns true if the entire btree and every node are valid.
// - height of all leaves are the equal to the btree height.
// - deep count matches the btree count.
// - all nodes have the correct number of items and counts.
// - all items are in order.
func (tr *Map[K, V]) Sane() error {
if tr == nil {
return nil
}
if !tr.saneheight() {
return saneMapError("!sane-height")
}
if tr.Len() != tr.count || tr.deepcount() != tr.count {
return saneMapError("!sane-count")
}
if !tr.saneprops() {
return saneMapError("!sane-props")
}
if !tr.saneorder() {
return saneMapError("!sane-order")
}
if !tr.sanenils() {
return saneMapError("!sane-nils")
}
return nil
}
// btree_saneheight returns true if the height of all leaves match the height
// of the btree.
func (tr *Map[K, V]) saneheight() bool {
height := tr.Height()
if tr.root != nil {
if height == 0 {
return false
}
return tr.root.saneheight(1, height)
}
return height == 0
}
func (n *mapNode[K, V]) saneheight(height, maxheight int) bool {
if n.leaf() {
if height != maxheight {
return false
}
} else {
i := 0
for ; i < len(n.items); i++ {
if !(*n.children)[i].saneheight(height+1, maxheight) {
return false
}
}
if !(*n.children)[i].saneheight(height+1, maxheight) {
return false
}
}
return true
}
// btree_deepcount returns the number of items in the btree.
func (tr *Map[K, V]) deepcount() int {
if tr.root != nil {
return tr.root.deepcount()
}
return 0
}
func (n *mapNode[K, V]) deepcount() int {
count := len(n.items)
if !n.leaf() {
for i := 0; i <= len(n.items); i++ {
count += (*n.children)[i].deepcount()
}
}
if n.count != count {
return -1
}
return count
}
func (tr *Map[K, V]) nodesaneprops(n *mapNode[K, V], height int) bool {
if height == 1 {
if len(n.items) < 1 || len(n.items) > tr.max {
println(len(n.items) < 1)
return false
}
} else {
if len(n.items) < tr.min || len(n.items) > tr.max {
println(2)
return false
}
}
if !n.leaf() {
if len(*n.children) != len(n.items)+1 {
println(3)
return false
}
for i := 0; i < len(n.items); i++ {
if !tr.nodesaneprops((*n.children)[i], height+1) {
println(4)
return false
}
}
if !tr.nodesaneprops((*n.children)[len(n.items)], height+1) {
println(5)
return false
}
}
return true
}
func (tr *Map[K, V]) saneprops() bool {
if tr.root != nil {
return tr.nodesaneprops(tr.root, 1)
}
return true
}
func (tr *Map[K, V]) sanenilsnode(n *mapNode[K, V]) bool {
items := n.items[:cap(n.items):cap(n.items)]
for i := len(n.items); i < len(items); i++ {
if !tr.eq(items[i].key, tr.empty.key) {
return false
}
}
if !n.leaf() {
for i := 0; i < len(*n.children); i++ {
if (*n.children)[i] == nil {
return false
}
}
children := (*n.children)[:cap(*n.children):cap(*n.children)]
for i := len(*n.children); i < len(children); i++ {
if children[i] != nil {
return false
}
}
for i := 0; i < len(*n.children); i++ {
if !tr.sanenilsnode((*n.children)[i]) {
return false
}
}
}
return true
}
// sanenils checks that all the slots in the item slice that are not used,
//
// n.items[len(n.items):cap(n.items):cap(n.items)]
//
// are equal to the empty value of the kind.
func (tr *Map[K, V]) sanenils() bool {
if tr.root != nil {
return tr.sanenilsnode(tr.root)
}
return true
}
func (tr *Map[K, V]) saneorder() bool {
var last K
var count int
var bad bool
tr.Scan(func(key K, value V) bool {
if count > 0 {
if !tr.lt(last, key) {
bad = true
return false
}
}
last = key
count++
return true
})
return !bad && count == tr.count
}
func TestMapIter(t *testing.T) {
N := 100_000
tr := testMapNewBTree()
var all []testMapKind
for i := 0; i < N; i++ {
tr.Load(testMapMakeItem(i), testMapMakeItem(i))
all = append(all, testMapMakeItem(i))
}
var count int
var i int
iter := tr.Iter()
for ok := iter.First(); ok; ok = iter.Next() {
if !tr.eq(all[i], iter.Key()) {
panic("!")
}
count++
i++
}
if count != N {
t.Fatalf("expected %v, got %v", N, count)
}
count = 0
i = len(all) - 1
iter = tr.Iter()
for ok := iter.Last(); ok; ok = iter.Prev() {
if !tr.eq(all[i], iter.Key()) {
panic("!")
}
i--
count++
}
if count != N {
t.Fatalf("expected %v, got %v", N, count)
}
i = 0
iter = tr.Iter()
for ok := iter.First(); ok; ok = iter.Next() {
if !tr.eq(all[i], iter.Key()) {
panic("!")
}
i++
}
i--
for ok := iter.Prev(); ok; ok = iter.Prev() {
i--
if !tr.eq(all[i], iter.Key()) {
panic("!")
}
}
if i != 0 {
panic("!")
}
i++
for ok := iter.Next(); ok; ok = iter.Next() {
if !tr.eq(all[i], iter.Key()) {
panic("!")
}
i++
}
if i != N {
panic("!")
}
i = 0
for ok := iter.First(); ok; ok = iter.Next() {
if !tr.eq(all[i], iter.Key()) {
panic("!")
}
if tr.eq(iter.Key(), testMapMakeItem(N/2)) {
for ok = iter.Prev(); ok; ok = iter.Prev() {
i--
if !tr.eq(all[i], iter.Key()) {
panic("!")
}
}
break
}
i++
}
}
func TestMapIterSeek(t *testing.T) {
var tr Map[int, struct{}]
var all []int
for i := 0; i < 10000; i++ {
tr.Set(i*2, struct{}{})
all = append(all, i)
}
_ = all
{
iter := tr.Iter()
var vals []int
for ok := iter.Seek(501); ok; ok = iter.Next() {
vals = append(vals, iter.Key())
}
assert(vals[0] == 502 && vals[1] == 504)
}
{
iter := tr.Iter()
var vals []int
for ok := iter.Seek(501); ok; ok = iter.Prev() {
vals = append(vals, iter.Key())
}
assert(vals[0] == 502 && vals[1] == 500)
}
}
func TestMapIterSeekPrefix(t *testing.T) {
var tr Map[int, struct{}]
count := 10_000
for i := 0; i < count; i++ {
tr.Set(i*2, struct{}{})
}
for i := 0; i < count; i++ {
iter := tr.Iter()
ret := iter.Seek(i*2 - 1)
assert(ret == true)
}
}
func copyMapEntries(m *Map[int, int]) []mapPair[int, int] {
all := make([]mapPair[int, int], m.Len())
keys := m.Keys()
vals := m.Values()
for i := 0; i < len(keys); i++ {
all[i].key = keys[i]
all[i].value = vals[i]
}
sort.Slice(all, func(i, j int) bool {
return all[i].key < all[j].key
})
return all
}
func mapEntriesEqual(a, b []mapPair[int, int]) bool {
return reflect.DeepEqual(a, b)
}
func copyMapTest(N int, m1 *Map[int, int], e11 []mapPair[int, int], deep bool) {
e12 := copyMapEntries(m1)
if !mapEntriesEqual(e11, e12) {
panic("!")
}
// Make a copy and compare the values
m2 := m1.Copy()
e21 := copyMapEntries(m1)
if !mapEntriesEqual(e21, e12) {
panic("!")
}
// Delete every other key
var e22 []mapPair[int, int]
for i, j := range rand.Perm(N) {
if i&1 == 0 {
e22 = append(e22, e21[j])
} else {
prev, deleted := m2.Delete(e21[j].key)
if !deleted {
panic("!")
}
if prev != e21[j].value {
panic("!")
}
}
}
if m2.Len() != N/2 {
panic("!")
}
sort.Slice(e22, func(i, j int) bool {
return e22[i].key < e22[j].key
})
e23 := copyMapEntries(m2)
if !mapEntriesEqual(e23, e22) {
panic("!")
}
if !deep {
var wg sync.WaitGroup
wg.Add(2)
go func() {
defer wg.Done()
copyMapTest(N/2, m2, e23, true)
}()
go func() {
defer wg.Done()
copyMapTest(N/2, m2, e23, true)
}()
wg.Wait()
}
e24 := copyMapEntries(m2)
if !mapEntriesEqual(e24, e23) {
panic("!")
}
}
func TestMapCopy(t *testing.T) {
N := 1_000
// create the initial map
m1 := new(Map[int, int])
for m1.Len() < N {
m1.Set(rand.Int(), rand.Int())
}
e11 := copyMapEntries(m1)
dur := time.Second * 2
var wg sync.WaitGroup
for i := 0; i < 16; i++ {
wg.Add(1)
go func() {
defer wg.Done()
start := time.Now()
for time.Since(start) < dur {
copyMapTest(N, m1, e11, false)
}
}()
}
wg.Wait()
e12 := copyMapEntries(m1)
if !mapEntriesEqual(e11, e12) {
panic("!")
}
}
type testNonCopyItem struct {
data string
}
func newTestNonCopyItem(data string) *testNonCopyItem {
return &testNonCopyItem{data: data}
}
type testCopyItem struct {
data string
}
func newTestCopyItem(data string) *testCopyItem {
return &testCopyItem{data: data}
}
func (item *testCopyItem) Copy() *testCopyItem {
return &testCopyItem{data: item.data}
}
type testIsoCopyItem struct {
data string
}
func newTestIsoCopyItem(data string) *testIsoCopyItem {
return &testIsoCopyItem{data: data}
}
func (item *testIsoCopyItem) Copy() *testIsoCopyItem {
return &testIsoCopyItem{data: item.data}
}
func TestMapValueCopy(t *testing.T) {
t.Run("without-copy", func(t *testing.T) {
var m Map[string, *testNonCopyItem]
m.Set("hello", newTestNonCopyItem("world"))
m2 := m.Copy()
v, _ := m.Get("hello")
assert(v.data == "world")
v, _ = m2.Get("hello")
assert(v.data == "world")
// now get and change the value, mutable, this will affect both trees
v, _ = m.GetMut("hello")
v.data = "planet"
v, _ = m.Get("hello")
assert(v.data == "planet")
v, _ = m2.Get("hello")
assert(v.data == "planet")
})
t.Run("with-copy", func(t *testing.T) {
var m Map[string, *testCopyItem]
m.Set("hello", newTestCopyItem("world"))
m2 := m.Copy()
v, _ := m.Get("hello")
assert(v.data == "world")
v, _ = m2.Get("hello")
assert(v.data == "world")
// now get and change the value, mutable, this will only affect the
// first tree.
v, _ = m.GetMut("hello")
v.data = "planet"
v, _ = m.Get("hello")
assert(v.data == "planet")
v, _ = m2.Get("hello")
assert(v.data == "world")
})
t.Run("with-isocopy", func(t *testing.T) {
var m Map[string, *testIsoCopyItem]
m.Set("hello", newTestIsoCopyItem("world"))
m2 := m.Copy()
v, _ := m.Get("hello")
assert(v.data == "world")
v, _ = m2.Get("hello")
assert(v.data == "world")
// now get and change the value, mutable, this will only affect the
// first tree.
v, _ = m.GetMut("hello")
v.data = "planet"
v, _ = m.Get("hello")
assert(v.data == "planet")
v, _ = m2.Get("hello")
assert(v.data == "world")
})
}
func TestMapDeepCopy(t *testing.T) {
Ncols := 1000
Nvals := 1000
// Create a collection of maps that are each a collection of key/value
// pairs of string.
cols1 := NewMap[string, *Map[string, string]](4)
for i := 0; i < Ncols; i++ {
col := NewMap[string, string](4)
for j := 0; j < Nvals; j++ {
col.Set(fmt.Sprintf("key:%d", j), fmt.Sprintf("val:%d", j))
}
cols1.Set(fmt.Sprintf("col:%d", i), col)
}
// Copy the root tree
cols2 := cols1.Copy()
// Update the second cols2 by deleting half the entries
for i := 0; i < Ncols; i++ {
col, _ := cols2.GetMut(fmt.Sprintf("col:%d", i))
for j := 0; j < Nvals; j += 2 {
col.Delete(fmt.Sprintf("key:%d", j))
}
}
// Now Count the total of keys in all collections
var count1 int
for i := 0; i < Ncols; i++ {
col, _ := cols1.Get(fmt.Sprintf("col:%d", i))
count1 += col.Len()
}
var count2 int
for i := 0; i < Ncols; i++ {
col, _ := cols2.Get(fmt.Sprintf("col:%d", i))
count2 += col.Len()
}
assert(count1 == Ncols*Nvals)
assert(count2 == Ncols*Nvals/2)
// Copy again, but this time use Get instead of GetMut
cols2 = cols1.Copy()
// Update the second cols2 by deleting half the entries
for i := 0; i < Ncols; i++ {
col, _ := cols2.Get(fmt.Sprintf("col:%d", i))
for j := 0; j < Nvals; j += 2 {
col.Delete(fmt.Sprintf("key:%d", j))
}
}
// Now Count the total of keys in all collections
count1 = 0
for i := 0; i < Ncols; i++ {
col, _ := cols1.Get(fmt.Sprintf("col:%d", i))
count1 += col.Len()
}
count2 = 0
for i := 0; i < Ncols; i++ {
col, _ := cols2.Get(fmt.Sprintf("col:%d", i))
count2 += col.Len()
}
assert(count1 == Ncols*Nvals/2)
assert(count2 == Ncols*Nvals/2)
}
btree-1.8.1/set.go 0000664 0000000 0000000 00000010225 15044500710 0013740 0 ustar 00root root 0000000 0000000 package btree
type Set[K ordered] struct {
base Map[K, struct{}]
}
// Copy
func (tr *Set[K]) Copy() *Set[K] {
tr2 := new(Set[K])
tr2.base = *tr.base.Copy()
return tr2
}
func (tr *Set[K]) IsoCopy() *Set[K] {
tr2 := new(Set[K])
tr2.base = *tr.base.IsoCopy()
return tr2
}
// Insert an item
func (tr *Set[K]) Insert(key K) {
tr.base.Set(key, struct{}{})
}
func (tr *Set[K]) Scan(iter func(key K) bool) {
tr.base.Scan(func(key K, value struct{}) bool {
return iter(key)
})
}
// Get a value for key
func (tr *Set[K]) Contains(key K) bool {
_, ok := tr.base.Get(key)
return ok
}
// Len returns the number of items in the tree
func (tr *Set[K]) Len() int {
return tr.base.Len()
}
// Delete an item
func (tr *Set[K]) Delete(key K) {
tr.base.Delete(key)
}
// 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 *Set[K]) Ascend(pivot K, iter func(key K) bool) {
tr.base.Ascend(pivot, func(key K, value struct{}) bool {
return iter(key)
})
}
func (tr *Set[K]) Reverse(iter func(key K) bool) {
tr.base.Reverse(func(key K, value struct{}) bool {
return iter(key)
})
}
// 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 *Set[K]) Descend(pivot K, iter func(key K) bool) {
tr.base.Descend(pivot, func(key K, value struct{}) bool {
return iter(key)
})
}
// Load is for bulk loading pre-sorted items
func (tr *Set[K]) Load(key K) {
tr.base.Load(key, struct{}{})
}
// Min returns the minimum item in tree.
// Returns nil if the treex has no items.
func (tr *Set[K]) Min() (K, bool) {
key, _, ok := tr.base.Min()
return key, ok
}
// Max returns the maximum item in tree.
// Returns nil if the tree has no items.
func (tr *Set[K]) Max() (K, bool) {
key, _, ok := tr.base.Max()
return key, ok
}
// PopMin removes the minimum item in tree and returns it.
// Returns nil if the tree has no items.
func (tr *Set[K]) PopMin() (K, bool) {
key, _, ok := tr.base.PopMin()
return key, ok
}
// PopMax removes the maximum item in tree and returns it.
// Returns nil if the tree has no items.
func (tr *Set[K]) PopMax() (K, bool) {
key, _, ok := tr.base.PopMax()
return key, ok
}
// GetAt returns the value at index.
// Return nil if the tree is empty or the index is out of bounds.
func (tr *Set[K]) GetAt(index int) (K, bool) {
key, _, ok := tr.base.GetAt(index)
return key, ok
}
// DeleteAt deletes the item at index.
// Return nil if the tree is empty or the index is out of bounds.
func (tr *Set[K]) DeleteAt(index int) (K, bool) {
key, _, ok := tr.base.DeleteAt(index)
return key, ok
}
// Height returns the height of the tree.
// Returns zero if tree has no items.
func (tr *Set[K]) Height() int {
return tr.base.Height()
}
// SetIter represents an iterator for btree.Set
type SetIter[K ordered] struct {
base MapIter[K, struct{}]
}
// Iter returns a read-only iterator.
func (tr *Set[K]) Iter() SetIter[K] {
return SetIter[K]{tr.base.Iter()}
}
// Seek to item greater-or-equal-to key.
// Returns false if there was no item found.
func (iter *SetIter[K]) Seek(key K) bool {
return iter.base.Seek(key)
}
// First moves iterator to first item in tree.
// Returns false if the tree is empty.
func (iter *SetIter[K]) First() bool {
return iter.base.First()
}
// Last moves iterator to last item in tree.
// Returns false if the tree is empty.
func (iter *SetIter[K]) Last() bool {
return iter.base.Last()
}
// 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 *SetIter[K]) Next() bool {
return iter.base.Next()
}
// 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 *SetIter[K]) Prev() bool {
return iter.base.Prev()
}
// Key returns the current iterator item key.
func (iter *SetIter[K]) Key() K {
return iter.base.Key()
}
// Keys returns all the keys in order.
func (tr *Set[K]) Keys() []K {
return tr.base.Keys()
}
// Clear will delete all items.
func (tr *Set[K]) Clear() {
tr.base.Clear()
}
btree-1.8.1/set_test.go 0000664 0000000 0000000 00000012330 15044500710 0014776 0 ustar 00root root 0000000 0000000 package btree
import (
"math/rand"
"reflect"
"sort"
"sync"
"testing"
"time"
)
func TestSet(t *testing.T) {
N := 1_000_000
var tr Set[int]
for i := 0; i < N; i++ {
tr.Load(i)
}
assert(tr.Len() == N)
for i := 0; i < N; i++ {
assert(tr.Contains(i))
}
count := 0
tr.Scan(func(_ int) bool {
count++
return true
})
assert(count == N)
count = 0
tr.Ascend(N/2, func(_ int) bool {
count++
return true
})
assert(count == N/2)
count = 0
tr.Reverse(func(_ int) bool {
count++
return true
})
assert(count == N)
count = 0
tr.Descend(N/2, func(_ int) bool {
count++
return true
})
assert(count == N/2+1)
for i := 0; i < N; i++ {
tr.Delete(i)
}
dotup := func(v int, ok bool) interface{} {
if !ok {
return nil
}
return v
}
assert(tr.Len() == 0)
assert(dotup(tr.Min()) == nil)
assert(dotup(tr.Max()) == nil)
assert(dotup(tr.PopMin()) == nil)
assert(dotup(tr.PopMax()) == nil)
for i := 0; i < N; i++ {
assert(!tr.Contains(i))
}
for i := 0; i < N; i++ {
tr.Insert(i)
}
assert(tr.Len() == N)
for i := 0; i < N; i++ {
tr.Insert(i)
}
assert(tr.Len() == N)
for i := 0; i < N; i++ {
tr.Load(i)
}
assert(tr.Len() == N)
assert(dotup(tr.Min()) == 0)
assert(dotup(tr.Max()) == N-1)
assert(dotup(tr.PopMin()) == 0)
assert(dotup(tr.PopMax()) == N-1)
tr.Insert(0)
tr.Insert(N - 1)
assert(dotup(tr.GetAt(0)) == 0)
assert(dotup(tr.GetAt(N)) == nil)
tr.Insert(N - 1)
assert(tr.Height() > 0)
assert(dotup(tr.DeleteAt(0)) == 0)
tr.Insert(0)
assert(dotup(tr.DeleteAt(N-1)) == N-1)
assert(dotup(tr.DeleteAt(N)) == nil)
tr.Insert(N - 1)
count = 0
tr.Scan(func(item int) bool {
count++
return true
})
assert(count == N)
func() {
defer func() {
msg, ok := recover().(string)
assert(ok && msg == "nil item")
}()
tr := NewNonConcurrent(intLess)
tr.Set(nil)
}()
func() {
defer func() {
msg, ok := recover().(string)
assert(ok && msg == "nil item")
}()
tr := NewNonConcurrent(intLess)
tr.Load(nil)
}()
for i := 0; i < N; i++ {
assert(tr.Contains(i))
}
for i := 0; i < N; i++ {
tr.Delete(i)
}
for i := 0; i < N; i++ {
assert(!tr.Contains(i))
}
assert(tr.base.lt(1, 2))
assert(tr.base.lt(2, 10))
}
func TestSetClear(t *testing.T) {
var tr Set[int]
for i := 0; i < 100; i++ {
tr.Insert(i)
}
assert(tr.Len() == 100)
tr.Clear()
assert(tr.Len() == 0)
for i := 0; i < 100; i++ {
tr.Insert(i)
}
assert(tr.Len() == 100)
}
func TestSetIter(t *testing.T) {
N := 100_000
lt := func(a, b int) bool { return a < b }
eq := func(a, b int) bool { return !lt(a, b) && !lt(b, a) }
var tr Set[int]
var all []int
for i := 0; i < N; i++ {
tr.Load(i)
all = append(all, i)
}
var count int
var i int
iter := tr.Iter()
for ok := iter.First(); ok; ok = iter.Next() {
if !eq(all[i], iter.Key()) {
panic("!")
}
count++
i++
}
if count != N {
t.Fatalf("expected %v, got %v", N, count)
}
count = 0
i = len(all) - 1
iter = tr.Iter()
for ok := iter.Last(); ok; ok = iter.Prev() {
if !eq(all[i], iter.Key()) {
panic("!")
}
i--
count++
}
if count != N {
t.Fatalf("expected %v, got %v", N, count)
}
i = 0
iter = tr.Iter()
for ok := iter.First(); ok; ok = iter.Next() {
if !eq(all[i], iter.Key()) {
panic("!")
}
i++
}
i--
for ok := iter.Prev(); ok; ok = iter.Prev() {
i--
if !eq(all[i], iter.Key()) {
panic("!")
}
}
if i != 0 {
panic("!")
}
i++
for ok := iter.Next(); ok; ok = iter.Next() {
if !eq(all[i], iter.Key()) {
panic("!")
}
i++
}
if i != N {
panic("!")
}
i = 0
for ok := iter.First(); ok; ok = iter.Next() {
if !eq(all[i], iter.Key()) {
panic("!")
}
if eq(iter.Key(), N/2) {
for ok = iter.Prev(); ok; ok = iter.Prev() {
i--
if !eq(all[i], iter.Key()) {
panic("!")
}
}
break
}
i++
}
}
func copySetEntries(m *Set[int]) []int {
all := m.Keys()
sort.Ints(all)
return all
}
func setEntriesEqual(a, b []int) bool {
return reflect.DeepEqual(a, b)
}
func copySetTest(N int, s1 *Set[int], e11 []int, deep bool) {
e12 := copySetEntries(s1)
if !setEntriesEqual(e11, e12) {
panic("!")
}
// Make a copy and compare the values
s2 := s1.Copy()
e21 := copySetEntries(s1)
if !setEntriesEqual(e21, e12) {
panic("!")
}
// Delete every other key
var e22 []int
for i, j := range rand.Perm(N) {
if i&1 == 0 {
e22 = append(e22, e21[j])
} else {
s2.Delete(e21[j])
}
}
if s2.Len() != N/2 {
panic("!")
}
sort.Ints(e22)
e23 := copySetEntries(s2)
if !setEntriesEqual(e23, e22) {
panic("!")
}
if !deep {
var wg sync.WaitGroup
wg.Add(2)
go func() {
defer wg.Done()
copySetTest(N/2, s2, e23, true)
}()
go func() {
defer wg.Done()
copySetTest(N/2, s2, e23, true)
}()
wg.Wait()
}
e24 := copySetEntries(s2)
if !setEntriesEqual(e24, e23) {
panic("!")
}
}
func TestSetCopy(t *testing.T) {
N := 1_000
// create the initial map
s1 := new(Set[int])
for s1.Len() < N {
s1.Insert(rand.Int())
}
e11 := copySetEntries(s1)
dur := time.Second * 2
var wg sync.WaitGroup
for i := 0; i < 16; i++ {
wg.Add(1)
go func() {
defer wg.Done()
start := time.Now()
for time.Since(start) < dur {
copySetTest(N, s1, e11, false)
}
}()
}
wg.Wait()
e12 := copySetEntries(s1)
if !setEntriesEqual(e11, e12) {
panic("!")
}
}