diff --git a/README.md b/README.md index b7376e1..d57c6e3 100644 --- a/README.md +++ b/README.md @@ -6,19 +6,21 @@ [![codecov](https://codecov.io/gh/gammazero/radixtree/branch/master/graph/badge.svg)](https://codecov.io/gh/gammazero/radixtree) [![License](https://img.shields.io/badge/License-MIT-blue.svg)](LICENSE) -Package `radixtree` implements an Adaptive [Radix Tree](https://en.wikipedia.org/wiki/Radix_tree), aka compressed [trie](https://en.wikipedia.org/wiki/Trie) or compact prefix tree. This data structure is useful to quickly lookup data by key, find values whose keys have a common prefix, or find values whose keys are a prefix (i.e. found along the way) of a search key. +Package `radixtree` implements an Adaptive [Radix Tree](https://en.wikipedia.org/wiki/Radix_tree), aka compressed [trie](https://en.wikipedia.org/wiki/Trie) or compact prefix tree. This data structure is useful to quickly lookup data by key, find values whose keys have a common prefix, or find values whose keys are a prefix (i.e. found along the way) of a search key. It is adaptive in the sense that nodes are not constant size, having only as many children, up to the maximum, as needed to branch to all subtrees. This package implements a radix-256 tree where each key symbol (radix) is a byte, allowing up to 256 possible branches to traverse to the next node. -The implementation is optimized for Get performance and allocate 0 bytes of heap memory for any read operation (Get, Walk, WalkPath, etc.); therefore no garbage to collect. Once a radix tree is built, it can be repeatedly searched quickly. Concurrent searches are safe since these do not modify the data structure. Access is not synchronized (not concurrent safe with writes), allowing the caller to synchronize, if needed, in whatever manner works best for the application. +The implementation is optimized for Get performance and avoids allocate of heap memory for any read operation (Get, Iter, IterPath, etc.). Once a radix tree is built, it can be repeatedly searched quickly. Concurrent searches are safe since these do not modify the data structure. Access is not synchronized (not concurrent safe with writes), allowing the caller to synchronize, if needed, in whatever manner works best for the application. This radix tree offers the following features: - Efficient: Operations are O(k). Zero memory allocation for all read operations. -- Ordered iteration: Walking and iterating the tree is done in lexical order, making the output deterministic. +- Ordered iteration: Iterating the tree is done in lexical order, making the output deterministic. - Store `nil` values: Read operations differentiate between missing and `nil` values. - Compact: When values are stored using keys that have a common prefix, the common part of the key is only stored once. Consider this when keys are similar to a timestamp, OID, filepath, geohash, network address, etc. Only the minimum number of nodes are kept to branch at the points where keys differ. -- Iterators: An `Iterator` type walks every key-value pair stored in the tree. A `Stepper` type of iterator traverses the tree one specified byte at a time, and is useful for incremental lookup. A Stepper can be copied in order to branch a search and iterate the copies concurrently. +- Iterators: Go 1.23 iterators allow ranging over key-value pairs stored in the tree. Iterators can traverse all key-value pairss, pairs with a key having specified prefix, or pairs along a key-path from root to a specified key. +- A `Stepper` type of iterator traverses the tree one specified byte at a time, and is useful for incremental lookup. A Stepper can be copied in order to branch a search and iterate the copies concurrently. +- Generics: The tree stores the specified type of value without needing to do interface type assertion. ## Install @@ -49,21 +51,19 @@ func main() { } // Output: Found TOM - // Find all items whose keys start with "tom" - rt.Walk("tom", func(key string, value any) bool { - fmt.Println(value) - return false - }) + // Find all items whose keys start with "tom". + for key, value := range rt.IterAt("tom") { + fmt.Println(key, "=", value) + } // Output: - // TOM - // TOMATO - // TOMMY + // tom = TOM + // tomato = TOMATO + // tommy = TOMMY // Find all items whose keys are a prefix of "tomato" - rt.WalkPath("tomato", func(key string, value any) bool { + for _, value := range rt.IterPath("tomato") { fmt.Println(value) - return false - }) + } // Output: // TOM // TOMATO diff --git a/bench_test.go b/bench_test.go index 67a8439..cb6921c 100644 --- a/bench_test.go +++ b/bench_test.go @@ -44,23 +44,23 @@ func BenchmarkPut(b *testing.B) { }) } -func BenchmarkWalk(b *testing.B) { +func BenchmarkIter(b *testing.B) { b.Run("Words", func(b *testing.B) { - benchmarkWalk(b, web2Path) + benchmarkIter(b, web2Path) }) b.Run("Web2a", func(b *testing.B) { - benchmarkWalk(b, web2aPath) + benchmarkIter(b, web2aPath) }) } -func BenchmarkWalkPath(b *testing.B) { +func BenchmarkIterPath(b *testing.B) { b.Run("Words", func(b *testing.B) { - benchmarkWalkPath(b, web2Path) + benchmarkIterPath(b, web2Path) }) b.Run("Web2a", func(b *testing.B) { - benchmarkWalkPath(b, web2aPath) + benchmarkIterPath(b, web2aPath) }) } @@ -69,7 +69,7 @@ func benchmarkGet(b *testing.B, filePath string) { if err != nil { b.Skip(err.Error()) } - tree := new(Tree) + tree := new(Tree[string]) for _, w := range words { tree.Put(w, w) } @@ -92,19 +92,19 @@ func benchmarkPut(b *testing.B, filePath string) { b.ResetTimer() b.ReportAllocs() for n := 0; n < b.N; n++ { - tree := new(Tree) + tree := new(Tree[string]) for _, w := range words { tree.Put(w, w) } } } -func benchmarkWalk(b *testing.B, filePath string) { +func benchmarkIter(b *testing.B, filePath string) { words, err := loadWords(filePath) if err != nil { b.Skip(err.Error()) } - tree := new(Tree) + tree := new(Tree[string]) for _, w := range words { tree.Put(w, w) } @@ -113,22 +113,21 @@ func benchmarkWalk(b *testing.B, filePath string) { var count int for n := 0; n < b.N; n++ { count = 0 - tree.Walk("", func(k string, value any) bool { + for range tree.Iter() { count++ - return false - }) + } } if count != len(words) { - b.Fatalf("Walk wrong count, expected %d got %d", len(words), count) + b.Fatalf("iter wrong count, expected %d got %d", len(words), count) } } -func benchmarkWalkPath(b *testing.B, filePath string) { +func benchmarkIterPath(b *testing.B, filePath string) { words, err := loadWords(filePath) if err != nil { b.Skip(err.Error()) } - tree := new(Tree) + tree := new(Tree[string]) for _, w := range words { tree.Put(w, w) } @@ -137,13 +136,12 @@ func benchmarkWalkPath(b *testing.B, filePath string) { for n := 0; n < b.N; n++ { found := false for _, w := range words { - tree.WalkPath(w, func(key string, value any) bool { + for range tree.IterPath(w) { found = true - return false - }) + } } if !found { - b.Fatal("Walk did not find word") + b.Fatal("IterPath did not find word") } } } diff --git a/doc.go b/doc.go index 8f09390..f21f0ef 100644 --- a/doc.go +++ b/doc.go @@ -1,19 +1,17 @@ -/* -Package radixtree implements an Adaptive Radix Tree, aka compressed trie or -compact prefix tree. It is adaptive in the sense that nodes are not constant -size, having as few or many children as needed to branch to all subtrees. - -This package implements a radix-256 tree where each key symbol (radix) is a -byte, allowing up to 256 possible branches to traverse to the next node. - -The implementation is optimized for Get performance and allocates 0 bytes of -heap memory per Get; therefore no garbage to collect. Once the radix tree is -built, it can be repeatedly searched quickly. Concurrent searches are safe -since these do not modify the radixtree. Access is not synchronized (not -concurrent safe with writes), allowing the caller to synchronize, if needed, in -whatever manner works best for the application. - -The API uses string keys, since strings are immutable and therefore it is not -necessary make a copy of the key provided to the radix tree. -*/ +// Package radixtree implements an Adaptive Radix Tree, aka compressed trie or +// compact prefix tree. It is adaptive in the sense that nodes are not constant +// size, having as few or many children as needed to branch to all subtrees. +// +// This package implements a radix-256 tree where each key symbol (radix) is a +// byte, allowing up to 256 possible branches to traverse to the next node. +// +// The implementation is optimized for Get performance and avoids allocates 0 +// bytes of heap memory per Get. Once the radix tree is built, it can be +// repeatedly searched quickly. Concurrent searches are safe since these do not +// modify the radixtree. Access is not synchronized (not concurrent safe with +// writes), allowing the caller to synchronize, if needed, in whatever manner +// works best for the application. +// +// The API uses string keys, since strings are immutable and therefore it is +// not necessary make a copy of the key provided to the radix tree. package radixtree diff --git a/doc_test.go b/doc_test.go index 9f47886..2a47d53 100644 --- a/doc_test.go +++ b/doc_test.go @@ -6,74 +6,71 @@ import ( "github.com/gammazero/radixtree" ) -func ExampleTree_Walk() { - rt := radixtree.New() - rt.Put("tomato", "TOMATO") - rt.Put("tom", "TOM") - rt.Put("tommy", "TOMMY") - rt.Put("tornado", "TORNADO") +func ExampleTree_Iter() { + rt := radixtree.New[int]() + rt.Put("mercury", 1) + rt.Put("venus", 2) + rt.Put("earth", 3) + rt.Put("mars", 4) - // Find all items whose keys start with "tom" - rt.Walk("tom", func(key string, value any) bool { - fmt.Println(value) - return false - }) + // Find all items that that have a key that is a prefix of "tomato". + for key, value := range rt.Iter() { + fmt.Println(key, "=", value) + } + // Output: + // earth = 3 + // mars = 4 + // mercury = 1 + // venus = 2 } -func ExampleTree_WalkPath() { - rt := radixtree.New() +func ExampleTree_IterAt() { + rt := radixtree.New[string]() rt.Put("tomato", "TOMATO") rt.Put("tom", "TOM") rt.Put("tommy", "TOMMY") rt.Put("tornado", "TORNADO") - // Find all items that are a prefix of "tomato" - rt.WalkPath("tomato", func(key string, value any) bool { + // Find all items whose keys start with "tom". + for _, value := range rt.IterAt("tom") { fmt.Println(value) - return false - }) + } // Output: // TOM // TOMATO + // TOMMY } -func ExampleTree_NewIterator() { - rt := radixtree.New() +func ExampleTree_IterPath() { + rt := radixtree.New[string]() rt.Put("tomato", "TOMATO") rt.Put("tom", "TOM") rt.Put("tommy", "TOMMY") rt.Put("tornado", "TORNADO") - iter := rt.NewIterator() - for { - key, val, done := iter.Next() - if done { - break - } - fmt.Println(key, "=", val) + // Find all items that that have a key that is a prefix of "tomato". + for key, value := range rt.IterPath("tomato") { + fmt.Println(key, "=", value) } - // Output: // tom = TOM // tomato = TOMATO - // tommy = TOMMY - // tornado = TORNADO } func ExampleTree_NewStepper() { - rt := radixtree.New() + rt := radixtree.New[string]() rt.Put("tomato", "TOMATO") rt.Put("tom", "TOM") rt.Put("tommy", "TOMMY") rt.Put("tornado", "TORNADO") - iter := rt.NewStepper() + stepper := rt.NewStepper() word := "tomato" for i := range word { - if !iter.Next(word[i]) { + if !stepper.Next(word[i]) { break } - if val, ok := iter.Value(); ok { + if val, ok := stepper.Value(); ok { fmt.Println(val) } } diff --git a/go.mod b/go.mod index 8b32901..7aece52 100644 --- a/go.mod +++ b/go.mod @@ -1,3 +1,3 @@ module github.com/gammazero/radixtree -go 1.21 +go 1.23 diff --git a/iterator.go b/iterator.go deleted file mode 100644 index 79df767..0000000 --- a/iterator.go +++ /dev/null @@ -1,47 +0,0 @@ -package radixtree - -// Iterator iterates all keys and values in the radix tree. -// -// Is is safe to use different Iterator instances concurrently. Any -// modification to the Tree that the Iterator was created from invalidates the -// Iterator instance. -type Iterator struct { - nodes []*radixNode -} - -// NewIterator returns a new Iterator. -func (t *Tree) NewIterator() *Iterator { - return &Iterator{ - nodes: []*radixNode{&t.root}, - } -} - -// Copy creates a new Iterator at this iterator's state of iteration. -func (it *Iterator) Copy() *Iterator { - nodes := make([]*radixNode, len(it.nodes)) - copy(nodes, it.nodes) - return &Iterator{ - nodes: nodes, - } -} - -// Next returns the next key and value stored in the Tree, and true when -// iteration is complete. -func (it *Iterator) Next() (key string, value any, done bool) { - for { - if len(it.nodes) == 0 { - break - } - node := it.nodes[len(it.nodes)-1] - it.nodes = it.nodes[:len(it.nodes)-1] - - for i := len(node.edges) - 1; i >= 0; i-- { - it.nodes = append(it.nodes, node.edges[i].node) - } - - if node.leaf != nil { - return node.leaf.key, node.leaf.value, false - } - } - return "", nil, true -} diff --git a/stepper.go b/stepper.go index e57c646..0ca1c2a 100644 --- a/stepper.go +++ b/stepper.go @@ -3,15 +3,15 @@ package radixtree // Stepper traverses a Tree one byte at a time. // // Any modification to the tree invalidates the Stepper. -type Stepper struct { +type Stepper[T any] struct { p int - node *radixNode + node *radixNode[T] } // NewStepper returns a new Stepper instance that begins at the root of the // tree. -func (t *Tree) NewStepper() *Stepper { - return &Stepper{ +func (t *Tree[T]) NewStepper() *Stepper[T] { + return &Stepper[T]{ node: &t.root, } } @@ -19,8 +19,8 @@ func (t *Tree) NewStepper() *Stepper { // Copy makes a copy of the current Stepper. This allows branching a Stepper // into two that can take separate paths. These Steppers do not affect each // other and can be used concurrently. -func (s *Stepper) Copy() *Stepper { - return &Stepper{ +func (s *Stepper[T]) Copy() *Stepper[T] { + return &Stepper[T]{ p: s.p, node: s.node, } @@ -33,7 +33,7 @@ func (s *Stepper) Copy() *Stepper { // // When false is returned the Stepper is not modified. This allows different // values to be used in subsequent calls to Next. -func (s *Stepper) Next(radix byte) bool { +func (s *Stepper[T]) Next(radix byte) bool { // The tree.prefix represents single-edge parents without values that were // compressed out of the tree. Let prefix consume key symbols. if s.p < len(s.node.prefix) { @@ -58,7 +58,7 @@ func (s *Stepper) Next(radix byte) bool { // Item returns an Item containing the key and value at the current Stepper // position, or returns nil if no value is present at the position. -func (s *Stepper) Item() *Item { +func (s *Stepper[T]) Item() *Item[T] { // Only return item if all of this node's prefix was matched. Otherwise, // have not fully traversed into this node (edge not completely traversed). if s.p == len(s.node.prefix) { @@ -69,10 +69,11 @@ func (s *Stepper) Item() *Item { // Value returns the value at the current Stepper position, and true or false // to indicate if a value is present at the position. -func (s *Stepper) Value() (any, bool) { +func (s *Stepper[T]) Value() (T, bool) { item := s.Item() if item == nil { - return nil, false + var zero T + return zero, false } return item.value, true } diff --git a/stepper_test.go b/stepper_test.go index 44384f0..12d66df 100644 --- a/stepper_test.go +++ b/stepper_test.go @@ -5,7 +5,7 @@ import ( ) func TestStepper(t *testing.T) { - rt := new(Tree) + rt := new(Tree[string]) rt.Put("tom", "TOM") rt.Put("tomato", "TOMATO") rt.Put("torn", "TORN") @@ -21,7 +21,7 @@ func TestStepper(t *testing.T) { t.Fatal("'t' should have advanced iterator") } val, ok := iter.Value() - if ok || val != nil { + if ok || val != "" { t.Fatal("should not have value at 't'") } item := iter.Item() diff --git a/tree.go b/tree.go index 52adfb7..f861ddc 100644 --- a/tree.go +++ b/tree.go @@ -1,65 +1,59 @@ package radixtree import ( + "iter" "strings" ) // Tree is a radix tree of bytes keys and any values. -type Tree struct { - root radixNode +type Tree[T any] struct { + root radixNode[T] size int } // New creates a new bytes-based radix tree -func New() *Tree { - return new(Tree) +func New[T any]() *Tree[T] { + return new(Tree[T]) } -type radixNode struct { +type radixNode[T any] struct { // prefix is the edge label between this node and the parent, minus the key // segment used in the parent to index this child. prefix string - edges []edge - leaf *Item + edges []edge[T] + leaf *Item[T] } -// WalkFunc is the type of the function called for each value visited by Walk -// or WalkPath. The key argument contains the elements of the key at which the -// value is stored. -// -// If the function returns true Walk stops immediately and returns. This -// applies to WalkPath as well. -type WalkFunc func(key string, value any) bool - // InspectFunc is the type of the function called for each node visited by // Inspect. The key argument contains the key at which the node is located, the // depth is the distance from the root of the tree, and children is the number // of children the node has. // // If the function returns true Inspect stops immediately and returns. -type InspectFunc func(link, prefix, key string, depth, children int, hasValue bool, value any) bool +type InspectFunc[T any] func(link, prefix, key string, depth, children int, hasValue bool, value T) bool -type Item struct { +type Item[T any] struct { key string - value any + value T } -func (kv *Item) Key() string { return kv.key } -func (kv *Item) Value() any { return kv.value } +func (kv *Item[T]) Key() string { return kv.key } +func (kv *Item[T]) Value() T { return kv.value } -type edge struct { +type edge[T any] struct { radix byte - node *radixNode + node *radixNode[T] } // Len returns the number of values stored in the tree. -func (t *Tree) Len() int { +func (t *Tree[T]) Len() int { return t.size } // Get returns the value stored at the given key. Returns false if there is no // value present for the key. -func (t *Tree) Get(key string) (any, bool) { +func (t *Tree[T]) Get(key string) (T, bool) { + var zero T node := &t.root // Consume key data while mathcing edge and prefix; return if remaining key // data matches nothing. @@ -67,30 +61,30 @@ func (t *Tree) Get(key string) (any, bool) { // Find edge for radix. node = node.getEdge(key[0]) if node == nil { - return nil, false + return zero, false } // Consume key data. key = key[1:] if !strings.HasPrefix(key, node.prefix) { - return nil, false + return zero, false } key = key[len(node.prefix):] } if node.leaf != nil { return node.leaf.value, true } - return nil, false + return zero, false } // Put inserts the value into the tree at the given key, replacing any existing // items. It returns true if it adds a new value, false if it replaces an // existing value. -func (t *Tree) Put(key string, value any) bool { +func (t *Tree[T]) Put(key string, value T) bool { var ( p int isNewValue bool - newEdge edge + newEdge edge[T] hasNewEdge bool ) node := &t.root @@ -110,8 +104,8 @@ func (t *Tree) Put(key string, value any) bool { // Descended as far as prefixes and edges match key, and still have key // data, so add child that has a prefix of the unmatched key data and // set its value to the new value. - newChild := &radixNode{ - leaf: &Item{ + newChild := &radixNode[T]{ + leaf: &Item[T]{ key: key, value: value, }, @@ -119,7 +113,7 @@ func (t *Tree) Put(key string, value any) bool { if i < len(key)-1 { newChild.prefix = key[i+1:] } - newEdge = edge{radix, newChild} + newEdge = edge[T]{radix, newChild} hasNewEdge = true break } @@ -142,7 +136,7 @@ func (t *Tree) Put(key string, value any) bool { isNewValue = true t.size++ } - node.leaf = &Item{ + node.leaf = &Item[T]{ key: key, value: value, } @@ -154,10 +148,10 @@ func (t *Tree) Put(key string, value any) bool { // Delete removes the value associated with the given key. Returns true if // there was a value stored for the key. If the node or any of its ancestors // becomes childless as a result, they are removed from the tree. -func (t *Tree) Delete(key string) bool { +func (t *Tree[T]) Delete(key string) bool { node := &t.root var ( - parents []*radixNode + parents []*radixNode[T] links []byte ) for len(key) != 0 { @@ -200,10 +194,10 @@ func (t *Tree) Delete(key string) bool { // DeletePrefix removes all values whose key is prefixed by the given prefix. // Returns true if any values were removed. -func (t *Tree) DeletePrefix(prefix string) bool { +func (t *Tree[T]) DeletePrefix(prefix string) bool { node := &t.root var ( - parents []*radixNode + parents []*radixNode[T] links []byte ) for len(prefix) != 0 { @@ -231,10 +225,9 @@ func (t *Tree) DeletePrefix(prefix string) bool { if node.edges != nil { var count int - node.walk(func(k string, _ any) bool { + for range node.iter(1) { count++ - return false - }) + } t.size -= count node.edges = nil } else { @@ -253,16 +246,26 @@ func (t *Tree) DeletePrefix(prefix string) bool { return true } -// Walk visits all nodes whose keys match or are prefixed by the specified key, -// calling walkFn for each value found. If walkFn returns true, Walk returns. -// Use empty key "" to visit all nodes starting from the root or the Tree. +// IterAt visits all nodes in the tree, yielding the key and value of each. +// +// The tree is traversed in lexical order, making the output deterministic. +func (t *Tree[T]) Iter() iter.Seq2[string, T] { + return t.root.iter(t.size) +} + +// IterAt visits all nodes whose keys match or are prefixed by the specified +// key, yielding the key and value of each. An empty key "" to visits all +// nodes, and is the same as calling Iter. // // The tree is traversed in lexical order, making the output deterministic. -func (t *Tree) Walk(key string, walkFn WalkFunc) { +func (t *Tree[T]) IterAt(key string) iter.Seq2[string, T] { + nothing := func(yield func(string, T) bool) { return } + + // Find the subtree with a matching prefix. node := &t.root for len(key) != 0 { if node = node.getEdge(key[0]); node == nil { - return + return nothing } // Consume key data @@ -271,40 +274,66 @@ func (t *Tree) Walk(key string, walkFn WalkFunc) { if strings.HasPrefix(node.prefix, key) { break } - return + return nothing } key = key[len(node.prefix):] } + // Iterate the subtree. + return node.iter(1) +} - // Walk down tree starting at node located at key. - node.walk(walkFn) +func (node *radixNode[T]) iter(size int) iter.Seq2[string, T] { + return func(yield func(string, T) bool) { + // Iterate nodes of subtree. + nodes := make([]*radixNode[T], 1, size) + nodes[0] = node + + for len(nodes) != 0 { + // pop next node + node := nodes[len(nodes)-1] + nodes = nodes[:len(nodes)-1] + + // Append all the nodes edges to the nodes list + for i := len(node.edges) - 1; i >= 0; i-- { + nodes = append(nodes, node.edges[i].node) + } + + // If the node is a leaf, yield its value. + if node.leaf != nil { + if !yield(node.leaf.key, node.leaf.value) { + return + } + } + } + } } -// WalkPath walks each node along the path from the root to the node at the -// given key, calling walkFn for each node that has a value. If walkFn returns -// true, WalkPath returns. +// IterPath returns an iterator that visits each node along the path from the +// root to the node at the given key. yielding the key and value of each. // // The tree is traversed in lexical order, making the output deterministic. -func (t *Tree) WalkPath(key string, walkFn WalkFunc) { - node := &t.root - for { - if node.leaf != nil && walkFn(node.leaf.key, node.leaf.value) { - return - } +func (t *Tree[T]) IterPath(key string) iter.Seq2[string, T] { + return func(yield func(string, T) bool) { + node := &t.root + for { + if node.leaf != nil && !yield(node.leaf.key, node.leaf.value) { + return + } - if len(key) == 0 { - return - } + if len(key) == 0 { + return + } - if node = node.getEdge(key[0]); node == nil { - return - } + if node = node.getEdge(key[0]); node == nil { + return + } - key = key[1:] - if !strings.HasPrefix(key, node.prefix) { - return + key = key[1:] + if !strings.HasPrefix(key, node.prefix) { + return + } + key = key[len(node.prefix):] } - key = key[len(node.prefix):] } } @@ -316,7 +345,7 @@ func (t *Tree) WalkPath(key string, walkFn WalkFunc) { // If inspectFn returns false, the traversal is stopped and Inspect returns. // // The tree is traversed in lexical order, making the output deterministic. -func (t *Tree) Inspect(inspectFn InspectFunc) { +func (t *Tree[T]) Inspect(inspectFn InspectFunc[T]) { t.root.inspect("", "", 0, inspectFn) } @@ -327,8 +356,8 @@ func (t *Tree) Inspect(inspectFn InspectFunc) { // is split into parent branching node, and a child leaf node: // // ("pre", nil, edges[f])--->("ix", leaf, edges[]) -func (node *radixNode) split(p int) { - split := &radixNode{ +func (node *radixNode[T]) split(p int) { + split := &radixNode[T]{ edges: node.edges, leaf: node.leaf, } @@ -336,7 +365,7 @@ func (node *radixNode) split(p int) { split.prefix = node.prefix[p+1:] } node.edges = nil - node.addEdge(edge{node.prefix[p], split}) + node.addEdge(edge[T]{node.prefix[p], split}) if p == 0 { node.prefix = "" } else { @@ -345,7 +374,7 @@ func (node *radixNode) split(p int) { node.leaf = nil } -func (node *radixNode) prune(parents []*radixNode, links []byte) *radixNode { +func (node *radixNode[T]) prune(parents []*radixNode[T], links []byte) *radixNode[T] { if node.edges != nil { return node } @@ -366,7 +395,7 @@ func (node *radixNode) prune(parents []*radixNode, links []byte) *radixNode { return node } -func (node *radixNode) compress() { +func (node *radixNode[T]) compress() { if len(node.edges) != 1 || node.leaf != nil { return } @@ -381,21 +410,9 @@ func (node *radixNode) compress() { node.edges = edge.node.edges } -func (node *radixNode) walk(walkFn WalkFunc) bool { - if node.leaf != nil && walkFn(node.leaf.key, node.leaf.value) { - return true - } - for _, edge := range node.edges { - if edge.node.walk(walkFn) { - return true - } - } - return false -} - -func (node *radixNode) inspect(link, key string, depth int, inspectFn InspectFunc) bool { +func (node *radixNode[T]) inspect(link, key string, depth int, inspectFn InspectFunc[T]) bool { key += link + node.prefix - var val any + var val T var hasVal bool if node.leaf != nil { val = node.leaf.value @@ -415,7 +432,7 @@ func (node *radixNode) inspect(link, key string, depth int, inspectFn InspectFun // indexEdge binary searches for the edge index. // // This is faster then going through sort.Interface for repeated searches. -func (node *radixNode) indexEdge(radix byte) int { +func (node *radixNode[T]) indexEdge(radix byte) int { n := len(node.edges) i, j := 0, n for i < j { @@ -430,7 +447,7 @@ func (node *radixNode) indexEdge(radix byte) int { } // getEdge binary searches for edge. -func (node *radixNode) getEdge(radix byte) *radixNode { +func (node *radixNode[T]) getEdge(radix byte) *radixNode[T] { idx := node.indexEdge(radix) if idx < len(node.edges) && node.edges[idx].radix == radix { return node.edges[idx].node @@ -439,19 +456,19 @@ func (node *radixNode) getEdge(radix byte) *radixNode { } // addEdge binary searches to find where to insert edge, and inserts at. -func (node *radixNode) addEdge(e edge) { +func (node *radixNode[T]) addEdge(e edge[T]) { idx := node.indexEdge(e.radix) - node.edges = append(node.edges, edge{}) + node.edges = append(node.edges, edge[T]{}) copy(node.edges[idx+1:], node.edges[idx:]) node.edges[idx] = e } // delEdge binary searches for edge and removes it. -func (node *radixNode) delEdge(radix byte) { +func (node *radixNode[T]) delEdge(radix byte) { idx := node.indexEdge(radix) if idx < len(node.edges) && node.edges[idx].radix == radix { copy(node.edges[idx:], node.edges[idx+1:]) - node.edges[len(node.edges)-1] = edge{} + node.edges[len(node.edges)-1] = edge[T]{} node.edges = node.edges[:len(node.edges)-1] } } diff --git a/tree_test.go b/tree_test.go index 3cb7141..cefeb5d 100644 --- a/tree_test.go +++ b/tree_test.go @@ -8,7 +8,7 @@ import ( ) func TestAddEnd(t *testing.T) { - rt := new(Tree) + rt := new(Tree[string]) rt.Put("tomato", "TOMATO") if len(rt.root.edges) != 1 { t.Fatal("root should have 1 child") @@ -74,7 +74,7 @@ func TestAddEnd(t *testing.T) { } func TestAddFront(t *testing.T) { - rt := new(Tree) + rt := new(Tree[string]) rt.Put("tom", "TOM") t.Log(dump(rt)) // (root) t-> ("om", TOM) @@ -121,7 +121,7 @@ func TestAddFront(t *testing.T) { } func TestAddBranch(t *testing.T) { - rt := new(Tree) + rt := new(Tree[string]) rt.Put("tom", "TOM") rt.Put("tomato", "TOMATO") @@ -200,7 +200,7 @@ func TestAddBranch(t *testing.T) { } func TestAddBranchToBranch(t *testing.T) { - rt := new(Tree) + rt := new(Tree[string]) rt.Put("tom", "TOM") rt.Put("tomato", "TOMATO") rt.Put("torn", "TORN") @@ -256,7 +256,7 @@ func TestAddBranchToBranch(t *testing.T) { } func TestAddExisting(t *testing.T) { - rt := new(Tree) + rt := new(Tree[string]) rt.Put("tom", "TOM") rt.Put("tomato", "TOMATO") rt.Put("torn", "TORN") @@ -318,7 +318,7 @@ func TestAddExisting(t *testing.T) { } func TestDelete(t *testing.T) { - rt := new(Tree) + rt := new(Tree[string]) rt.Put("tom", "TOM") rt.Put("tomato", "TOMATO") rt.Put("torn", "TORN") @@ -358,7 +358,7 @@ func TestDelete(t *testing.T) { } func TestDeletePrefix(t *testing.T) { - rt := new(Tree) + rt := new(Tree[string]) rt.Put("tom", "TOM") rt.Put("tomato", "TOMATO") rt.Put("torn", "TORN") @@ -397,7 +397,7 @@ func TestDeletePrefix(t *testing.T) { } func TestBuildEdgeCases(t *testing.T) { - tree := new(Tree) + tree := new(Tree[int]) tree.Put("ABCD", 1) t.Log(dump(tree)) @@ -407,7 +407,7 @@ func TestBuildEdgeCases(t *testing.T) { t.Log(dump(tree)) val, ok := tree.Get("ABCE") - if ok || val != nil { + if ok || val != 0 { t.Fatal("expected no value") } @@ -534,69 +534,63 @@ func TestBuildEdgeCases(t *testing.T) { t.Log(dump(tree)) } -func TestSimpleWalk(t *testing.T) { - rt := New() +func TestSimpleIterAt(t *testing.T) { + rt := New[string]() rt.Put("tomato", "TOMATO") rt.Put("tom", "TOM") rt.Put("tornado", "TORNADO") count := 0 - rt.Walk("tomato", func(key string, value any) bool { + for range rt.IterAt("tomato") { count++ - return false - }) + } if count != 1 { t.Errorf("expected to visit 1 key, visited %d", count) } count = 0 - rt.Walk("t", func(key string, value any) bool { + for range rt.IterAt("t") { count++ - return false - }) + } if count != 3 { t.Errorf("expected to visit 3 keys, visited %d", count) } count = 0 - rt.Walk("to", func(key string, value any) bool { + for range rt.IterAt("to") { count++ - return false - }) + } if count != 3 { t.Errorf("expected to visit 3 keys, visited %d", count) } count = 0 - rt.Walk("tom", func(key string, value any) bool { + for range rt.IterAt("tom") { count++ - return false - }) + } if count != 2 { t.Errorf("expected to visit 2 keys, visited %d", count) } count = 0 - rt.Walk("tomx", func(key string, value any) bool { + for range rt.IterAt("tomx") { count++ - return false - }) + } if count != 0 { t.Errorf("expected to visit 0 keys, visited %d", count) } count = 0 - rt.Walk("torn", func(key string, value any) bool { + for range rt.IterAt("torn") { count++ - return false - }) + } if count != 1 { t.Errorf("expected to visit 1 key, visited %d", count) } } func TestTree(t *testing.T) { - tree := New() + tree := New[string]() keys := []string{ "bird", @@ -640,38 +634,39 @@ func TestTree(t *testing.T) { // get for _, key := range keys { val, _ := tree.Get(key) - if val.(string) != strings.ToUpper(key) { + if val != strings.ToUpper(key) { t.Errorf("expected key %s to have value %v, got %v", key, strings.ToUpper(key), val) } } - var wvals []any - kvMap := map[string]any{} - walkFn := func(key string, value any) bool { - kvMap[key] = value - wvals = append(wvals, value) - return false - } + var wvals []string + kvMap := map[string]string{} - // walk path + // iter path t.Log(dump(tree)) key := "bad/key" - tree.WalkPath(key, walkFn) + for key, val := range tree.IterPath(key) { + kvMap[key] = val + wvals = append(wvals, val) + } if len(kvMap) != 0 { t.Error("should not have returned values, got ", kvMap) } lastKey := keys[len(keys)-1] - var expectVals []any + var expectVals []string for _, key := range keys { // If key is a prefix of lastKey, then expect value. if strings.HasPrefix(lastKey, key) { expectVals = append(expectVals, strings.ToUpper(key)) } } - kvMap = map[string]any{} + kvMap = map[string]string{} wvals = nil - tree.WalkPath(lastKey, walkFn) - if kvMap[lastKey] == nil { + for key, val := range tree.IterPath(lastKey) { + kvMap[key] = val + wvals = append(wvals, val) + } + if kvMap[lastKey] == "" { t.Fatalf("expected value for %s", lastKey) } if len(kvMap) != len(expectVals) { @@ -725,10 +720,12 @@ func TestTree(t *testing.T) { } func TestNilGet(t *testing.T) { - tree := New() + tree := New[*int]() - tree.Put("/rat", 1) - tree.Put("/ratatattat", 2) + one := 1 + tree.Put("/rat", &one) + two := 2 + tree.Put("/ratatattat", &two) tree.Put("/ratatouille", nil) val, ok := tree.Get("/ratatouille") @@ -751,7 +748,7 @@ func TestNilGet(t *testing.T) { } func TestRoot(t *testing.T) { - tree := New() + tree := New[string]() val, ok := tree.Get("") if ok { @@ -782,8 +779,8 @@ func TestRoot(t *testing.T) { } } -func TestWalk(t *testing.T) { - tree := New() +func TestIter(t *testing.T) { + tree := New[string]() keys := []string{ "bird", @@ -820,30 +817,18 @@ func TestWalk(t *testing.T) { } } - var err error - walkFn := func(key string, value any) bool { - // value for each walked key is correct - if value != strings.ToUpper(key) { - err = fmt.Errorf("expected key %s to have value %v, got %v", key, strings.ToUpper(key), value) - return true - } - count := visited[key] - visited[key] = count + 1 - return false - } - for _, notKey := range notKeys { - tree.Walk(notKey, walkFn) - if err != nil { - t.Error(err) + for key, val := range tree.IterAt(notKey) { + if val != strings.ToUpper(key) { + t.Fatalf("expected key %s to have value %v, got %v", key, strings.ToUpper(key), val) + } + visited[key]++ } - } t.Log(dump(tree)) for _, notKey := range notKeys { - _, ok := visited[notKey] - if ok { + if _, ok := visited[notKey]; ok { t.Fatalf("%s should not have been visited", notKey) } } @@ -866,75 +851,16 @@ func TestWalk(t *testing.T) { visited = make(map[string]int, len(keys)) - // Walk from root - tree.Walk("", walkFn) - if err != nil { - t.Error(err) - } - - if len(visited) != len(keys) { - t.Error("wrong number of iterm iterated") - } - - // each key/value visited exactly once - for key, visitedCount := range visited { - if visitedCount != 1 { - t.Errorf("expected key %s to be visited exactly once, got %v", key, visitedCount) - } - } - - visited = make(map[string]int, len(keys)) - - var iterCopy *Iterator - iter := tree.NewIterator() - for { - key, val, done := iter.Next() - if done { - break - } - t.Log("Iterated key:", key) - if walkFn(key, val) { - if err != nil { - t.Fatal(err) - } - break - } - if key == "rat/winks/wisely/once" { - iterCopy = iter.Copy() + // Iter from root. + for key, val := range tree.Iter() { + if val != strings.ToUpper(key) { + t.Fatalf("expected key %s to have value %v, got %v", key, strings.ToUpper(key), val) } + visited[key]++ } - if len(visited) != len(keys) { t.Error("wrong number of iterm iterated") } - - // each key/value visited exactly once - for key, visitedCount := range visited { - if visitedCount != 1 { - t.Errorf("expected key %s to be visited exactly once, got %v", key, visitedCount) - } - } - - visited = make(map[string]int, 4) - - for { - key, val, done := iterCopy.Next() - if done { - break - } - t.Log("Iterated key:", key) - if walkFn(key, val) { - if err != nil { - t.Fatal(err) - } - break - } - } - - if len(visited) != 4 { - t.Error("Wrong iteration count of copied iterator:", len(visited)) - } - // each key/value visited exactly once for key, visitedCount := range visited { if visitedCount != 1 { @@ -942,11 +868,9 @@ func TestWalk(t *testing.T) { } } - visited = make(map[string]int, len(keys)) - - tree.Walk("rat", walkFn) - if err != nil { - t.Errorf("expected error nil, got %v", err) + clear(visited) + for key, _ := range tree.IterAt("rat") { + visited[key]++ } if visited[keys[0]] != 0 { t.Error(keys[0], "should not have been visited") @@ -967,14 +891,9 @@ func TestWalk(t *testing.T) { t.Error(keys[7], "should have been visited") } - // Reset visited counts - for _, k := range keys { - visited[k] = 0 - } - - tree.Walk("rat/whis/kers", walkFn) - if err != nil { - t.Errorf("expected error nil, got %v", err) + clear(visited) + for key, _ := range tree.IterAt("rat/whis/kers") { + visited[key]++ } for _, key := range keys { if key == "rat/whis/kers" { @@ -988,66 +907,47 @@ func TestWalk(t *testing.T) { } } - // Reset visited counts - for _, k := range keys { - visited[k] = 0 - } - + clear(visited) testKey := "rat/winks/wryly/once" keys = append(keys, testKey) tree.Put(testKey, strings.ToUpper(testKey)) - walkPFn := func(key string, value any) bool { - // value for each walked key is correct + for key, val := range tree.IterPath(testKey) { v := strings.ToUpper(key) - if value != v { - err = fmt.Errorf("expected key %s to have value %v, got %v", key, v, value) - return true + if val != v { + t.Fatalf("expected key %s to have value %v, got %v", key, v, val) } visited[key]++ - return false } - - tree.WalkPath(testKey, walkPFn) - if err != nil { - t.Errorf("expected error nil, got %v", err) - } - err = checkVisited(visited, "rat", "rat/winks/wryly", testKey) + err := checkVisited(visited, "rat", "rat/winks/wryly", testKey) if err != nil { t.Error(err) } - // Reset visited counts - for _, k := range keys { - visited[k] = 0 - } - tree.WalkPath(testKey, func(key string, value any) bool { + clear(visited) + for key, _ := range tree.IterPath(testKey) { pfx := "rat/winks/wryly" if strings.HasPrefix(key, pfx) && len(key) > len(pfx) { - return false + continue } visited[key]++ - return false - }) + } err = checkVisited(visited, "rat", "rat/winks/wryly") if err != nil { t.Error(err) } - // Reset visited counts - for _, k := range keys { - visited[k] = 0 - } - tree.WalkPath(testKey, func(key string, value any) bool { + var found bool + clear(visited) + for key, _ := range tree.IterPath(testKey) { visited[key]++ if key == "rat/winks/wryly" { - err = fmt.Errorf("error at key %s", key) - return true + found = true + break } - return false - }) - if err == nil { - t.Errorf("expected error") + } + if !found { + t.Fatal("expected found to be true") } err = checkVisited(visited, "rat", "rat/winks/wryly") if err != nil { @@ -1056,53 +956,20 @@ func TestWalk(t *testing.T) { var foundRoot bool tree.Put("", "ROOT") - tree.WalkPath(testKey, func(key string, value any) bool { - if key == "" && value == "ROOT" { + for key, val := range tree.IterPath(testKey) { + if key == "" && val == "ROOT" { foundRoot = true + break } - return false - }) + } if !foundRoot { t.Error("did not find root") } - for _, k := range keys { - visited[k] = 0 - } - - tree.WalkPath(testKey, func(key string, value any) bool { - if key == "" && value == "ROOT" { - return false - } - return true - }) - for k, count := range visited { - if count != 0 { - t.Error("should not have visited ", k) - } - } - - tree.WalkPath(testKey, func(key string, value any) bool { - if key == "" && value == "ROOT" { - err = errors.New("error at root") - return true - } - return false - }) - if err == nil { - t.Errorf("expected error") - } - for k, count := range visited { - if count != 0 { - t.Error("should not have visited ", k) - } - } - var lastKey string - tree.WalkPath("rat/winks/wisely/x/y/z/w", func(key string, value any) bool { + for key, _ := range tree.IterPath("rat/winks/wisely/x/y/z/w") { lastKey = key - return false - }) + } if lastKey != "rat/winks/wisely/x/y/z" { t.Error("did not get expected last key") } @@ -1127,8 +994,8 @@ func checkVisited(visited map[string]int, expectVisited ...string) error { return nil } -func TestWalkStop(t *testing.T) { - tree := New() +func TestIterStop(t *testing.T) { + tree := New[int]() table := []struct { key string @@ -1145,20 +1012,18 @@ func TestWalkStop(t *testing.T) { tree.Put(table[i].key, table[i].val) } - walkErr := errors.New("walk error") + iterErr := errors.New("iter error") var walked int var err error - walkFn := func(k string, value any) bool { - if value == 999 { - err = walkErr - return true + for _, val := range tree.Iter() { + if val == 999 { + err = iterErr + break } walked++ - return false } - tree.Walk("", walkFn) - if err != walkErr { - t.Errorf("expected error %v, got %v", walkErr, err) + if err != iterErr { + t.Fatalf("expected error %v, got %v", iterErr, err) } if len(table) == walked { t.Errorf("expected nodes walked < %d, got %d", len(table), walked) @@ -1166,7 +1031,7 @@ func TestWalkStop(t *testing.T) { } func TestInspectStop(t *testing.T) { - tree := New() + tree := New[int]() table := []struct { key string @@ -1183,7 +1048,7 @@ func TestInspectStop(t *testing.T) { tree.Put(table[i].key, table[i].val) } var keys []string - inspectFn := func(link, prefix, key string, depth, children int, hasValue bool, value any) bool { + inspectFn := func(link, prefix, key string, depth, children int, hasValue bool, value int) bool { if !hasValue { // Do not count internal nodes return false @@ -1200,12 +1065,12 @@ func TestInspectStop(t *testing.T) { } tree.Inspect(inspectFn) if len(keys) != len(table)-2 { - t.Errorf("expected nodes walked to be %d, got %d: %v", len(table)-2, len(keys), keys) + t.Errorf("expected nodes iterated to be %d, got %d: %v", len(table)-2, len(keys), keys) } } func TestGetAfterDelete(t *testing.T) { - tree := New() + tree := New[string]() keys := []string{ "bird", @@ -1233,7 +1098,7 @@ func TestGetAfterDelete(t *testing.T) { } func TestStringConvert(t *testing.T) { - tree := New() + tree := New[string]() for _, w := range []string{"Bart", `Bartók`, `AbónXw`, `AbónYz`} { ok := tree.Put(w, w) if !ok { @@ -1241,43 +1106,33 @@ func TestStringConvert(t *testing.T) { } v, _ := tree.Get(w) - if v == nil { + if v == "" { t.Log(dump(tree)) t.Fatal("nil value returned getting", w) } - s, ok := v.(string) - if !ok { - t.Fatal("value is not a string") - } - if s != w { - t.Fatalf("returned wrong value - expected %q got %q", w, s) + if v != w { + t.Fatalf("returned wrong value - expected %q got %q", w, v) } } - tree.Walk("", func(key string, val any) bool { + for key, val := range tree.Iter() { t.Log("Key:", key) - s, ok := val.(string) - if !ok { - t.Log(dump(tree)) - t.Fatal("value is not a string") - } - t.Log("Val:", s) - if key != s { + t.Log("Val:", val) + if key != val { t.Log(dump(tree)) t.Fatal("Key and value do not match") } - return false - }) + } } // Use the Inspect functionality to create a function to dump the tree. -func dump(tree *Tree) string { +func dump[T any](tree *Tree[T]) string { var b strings.Builder - tree.Inspect(func(link, prefix, key string, depth, children int, hasValue bool, value any) bool { + tree.Inspect(func(link, prefix, key string, depth, children int, hasValue bool, value T) bool { for ; depth > 0; depth-- { b.WriteString(" ") } if hasValue { - b.WriteString(fmt.Sprintf("%s-> (%q, [%s: %q]) children: %d\n", link, prefix, key, value, children)) + b.WriteString(fmt.Sprintf("%s-> (%q, [%s: %v]) children: %d\n", link, prefix, key, value, children)) } else { b.WriteString(fmt.Sprintf("%s-> (%q, [%s])] children: %d\n", link, prefix, key, children)) }