Initial v2 code

Author: Daisuke Maki

README.md

@@ -4,32 +4,3 @@ This trie is implemented such that generic Key types can be used.
 Most other trie implementations are optimized for string based keys, but my use
 case is to match certain numeric opcodes to arbitrary data.
 
-Within this library Keys are treated as sequence of Labels.
-For example, a string can be thought of as Key that is comprised of a sequence
-of rune Labels.
-
-Each Key need to be able to break down to Labels via the `Iterate` method.
-Each Label in turn becomes the local key in a trie node.
-Each Label need to implement a `UniqueID` method to identify itself.
-
-# SYNOPSIS
-
-```go
-ctx, cancel := context.WithCancel(context.Background())
-defer cancel()
-
-t := trie.New()
-
-t.Put(ctx, trie.StringKey("foo"), 1)
-v, ok := t.Get(ctx, trie.StringKey("foo"))
-ok := t.Delete(ctx, trie.StringKey("foo"))
-for p := range t.Walk(ctx) {
-	// p.Labels
-	// p.Value
-}
-```
-
-# REFERENCES
-
-Originally based on https://github.com/koron/go-trie
-Much code stolen from https://github.com/dghubble/trie

go.mod

@@ -1,5 +1,11 @@
-module github.com/lestrrat-go/trie
+module github.com/lestrrat-go/trie/v2
 
-go 1.16
+go 1.23
 
-require github.com/stretchr/testify v1.7.0
+require github.com/stretchr/testify v1.9.0
+
+require (
+	github.com/davecgh/go-spew v1.1.1 // indirect
+	github.com/pmezard/go-difflib v1.0.0 // indirect
+	gopkg.in/yaml.v3 v3.0.1 // indirect
+)

go.sum

@@ -1,11 +1,10 @@
-github.com/davecgh/go-spew v1.1.0 h1:ZDRjVQ15GmhC3fiQ8ni8+OwkZQO4DARzQgrnXU1Liz8=
-github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
+github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
-github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
-github.com/stretchr/testify v1.7.0 h1:nwc3DEeHmmLAfoZucVR881uASk0Mfjw8xYJ99tb5CcY=
-github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
+github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg=
+github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
-gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c h1:dUUwHk2QECo/6vqA44rthZ8ie2QXMNeKRTHCNY2nXvo=
-gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
+gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
+gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=

interface.go

@@ -0,0 +1,18 @@
+package trie
+
+import (
+	"iter"
+)
+
+// String returns a Tokenizer that tokenizes a string into individual runes.
+func String() Tokenizer[string, rune] {
+	return TokenizeFunc[string, rune](func(s string) (iter.Seq[rune], error) {
+		return func(yield func(rune) bool) {
+			for _, r := range s {
+				if !yield(r) {
+					break
+				}
+			}
+		}, nil
+	})
+}

string.go

@@ -4,163 +4,280 @@
 package trie
 
 import (
-	"context"
+	"cmp"
+	"fmt"
+	"iter"
+	"slices"
+	"sort"
+	"strings"
 	"sync"
 )
 
-// Trie is a simple trie that accepts arbitrary Key types as its input.
-type Trie struct {
-	children map[interface{}]*Trie
-	hasValue bool
-	mu       sync.RWMutex
-	label    Label
-	value    interface{}
+// Tokenizer is an object that tokenize a L into individual keys.
+// For example, a string tokenizer would split a string into individual runes.
+type Tokenizer[L any, K cmp.Ordered] interface {
+	Tokenize(L) (iter.Seq[K], error)
+}
+
+// TokenizeFunc is a function that implements the Tokenizer interface
+type TokenizeFunc[L any, K cmp.Ordered] func(L) (iter.Seq[K], error)
+
+func (f TokenizeFunc[L, K]) Tokenize(in L) (iter.Seq[K], error) {
+	return f(in)
 }
 
-// New creates a new Trie
-func New() *Trie {
-	return newTrie(nil)
+// Trie is a trie that accepts arbitrary Key types as its input.
+//
+// L represents the "label", the input that is used to Get/Set/Delete
+// a value from the trie.
+//
+// K represents the "key", the individual components that are associated
+// with the nodes in the trie.
+//
+// V represents the "value", the data that is stored in the trie.
+// Data is stored at the leaf nodes of the trie.
+type Trie[L any, K cmp.Ordered, V any] struct {
+	mu        sync.RWMutex
+	root      *node[K, V]
+	tokenizer Tokenizer[L, K]
 }
 
-func newTrie(l Label) *Trie {
-	return &Trie{
-		label:    l,
-		children: make(map[interface{}]*Trie),
+// Node represents an individual node in the trie.
+type Node[K cmp.Ordered, V any] interface {
+	Key() K
+	Value() V
+	Children() iter.Seq[Node[K, V]]
+	AddChild(Node[K, V])
+}
+
+// New creates a new Trie object.
+func New[L any, K cmp.Ordered, V any](tokenizer Tokenizer[L, K]) *Trie[L, K, V] {
+	return &Trie[L, K, V]{
+		root:      newNode[K, V](),
+		tokenizer: tokenizer,
 	}
 }
 
 // Get returns the value associated with `key`. The second return value
-// is true if the value exists, false otherwise
-func (t *Trie) Get(key Key) (interface{}, bool) {
+// indicates if the value was found.
+func (t *Trie[L, K, V]) Get(key L) (V, bool) {
+	var zero V
+	iter, err := t.tokenizer.Tokenize(key)
+	if err != nil {
+		return zero, false
+	}
+
 	t.mu.RLock()
 	defer t.mu.RUnlock()
-	node := t
-	for iter := key.Labels(); iter.Next(); {
-		l := iter.Label()
-		node = node.children[l.UniqueID()]
-		if node == nil {
-			return nil, false
+	var tokens []K
+	for x := range iter {
+		tokens = append(tokens, x)
+	}
+	return get(t.root, tokens)
+}
+
+func get[K cmp.Ordered, V any](root Node[K, V], tokens []K) (V, bool) {
+	if len(tokens) > 0 {
+		for child := range root.Children() {
+			if child.Key() == tokens[0] {
+				// found the current token in the children.
+				if len(tokens) == 1 {
+					// this is the node we're looking for
+					return child.Value(), true
+				}
+				// we need to traverse down the trie
+				return get[K, V](child, tokens[1:])
+			}
 		}
 	}
-	return node.value, true
+
+	// if we got here, that means we couldn't find a common ancestor
+	var zero V
+	return zero, false
 }
 
-// Put sets `key` to point to data `value`. The return value is true
-// if the value was set anew. If this was an update operation, the return
-// value would be false
-func (t *Trie) Put(key Key, value interface{}) bool {
+// Delete removes data associated with `key`. It returns true if the value
+// was found and deleted, false otherwise
+func (t *Trie[L, K, V]) Delete(key L) bool {
+	iter, err := t.tokenizer.Tokenize(key)
+	if err != nil {
+		return false
+	}
+	var tokens []K
+	for x := range iter {
+		tokens = append(tokens, x)
+	}
+
 	t.mu.Lock()
 	defer t.mu.Unlock()
+	return delete[K, V](t.root, tokens)
+}
+
+func delete[K cmp.Ordered, V any](root *node[K, V], tokens []K) bool {
+	if len(tokens) <= 0 {
+		return false
+	}
+
+	for i, child := range root.children {
+		if child.Key() == tokens[0] {
+			if len(tokens) == 1 {
+				// this is the node we're looking for
+				root.children = slices.Delete(root.children, i, i+1)
+				return true
+			}
 
-	node := t
-	for iter := key.Labels(); iter.Next(); {
-		l := iter.Label()
-		child := node.children[l.UniqueID()]
-		if child == nil {
-			child = newTrie(l)
-			node.children[l.UniqueID()] = child
+			// we need to traverse down the trie
+			if delete[K, V](child, tokens[1:]) {
+				if len(child.children) == 0 {
+					root.children = slices.Delete(root.children, i, i+1)
+				}
+				return true
+			}
+			return false
 		}
-		node = child
 	}
 
-	isNewVal := node.hasValue
-	node.hasValue = true
-	node.value = value
-	return isNewVal
+	return false
 }
 
-func (t *Trie) isLeaf() bool {
-	return len(t.children) == 0
-}
+// Put sets `key` to point to data `value`.
+func (t *Trie[L, K, V]) Put(key L, value V) error {
+	iter, err := t.tokenizer.Tokenize(key)
+	if err != nil {
+		return fmt.Errorf(`failed to tokenize key: %w`, err)
+	}
+	node := t.root
+
+	var tokens []K
+	for x := range iter {
+		tokens = append(tokens, x)
+	}
 
-type ancestor struct {
-	Label Label
-	Node  *Trie
+	t.mu.Lock()
+	defer t.mu.Unlock()
+	put[K, V](node, tokens, value)
+	return nil
 }
 
-// Delete removes data associated with `key`. It returns true if the value
-// was found and deleted, false otherwise
-func (t *Trie) Delete(key Key) bool {
-	var ancestors []ancestor
-	node := t
-	for iter := key.Labels(); iter.Next(); {
-		l := iter.Label()
-		ancestors = append(ancestors, ancestor{
-			Label: l,
-			Node:  node,
-		})
-		node = node.children[l.UniqueID()]
-		if node == nil {
-			// node does not exist
-			return false
+func put[K cmp.Ordered, V any](root Node[K, V], tokens []K, value V) {
+	if len(tokens) == 0 {
+		return
+	}
+
+	for _, token := range tokens {
+		for child := range root.Children() {
+			if child.Key() == token {
+				// found the current token in the children.
+				// we need to traverse down the trie
+				put[K, V](child, tokens[1:], value)
+				return
+			}
 		}
 	}
 
-	// delete the node value
-	node.value = nil
+	// if we got here, that means we couldn't find a common ancestor
 
-	// if leaf, remove it from its parent's children map. Repeat for ancestors.
-	if !node.isLeaf() {
-		return true
+	// the first token has already been consumed, create a new node,
+	var newRoot *node[K, V]
+	var cur *node[K, V]
+	for _, token := range tokens { // duplicate token?
+		newNode := newNode[K, V]()
+		newNode.key = token
+		if cur == nil {
+			newRoot = newNode
+		} else {
+			cur.children = append(cur.children, newNode)
+		}
+		cur = newNode
 	}
-	// iterate backwards over the ancestors
-	for i := len(ancestors) - 1; i >= 0; i-- {
-		ancestor := ancestors[i]
-		parent := ancestor.Node
-		delete(parent.children, ancestor.Label.UniqueID())
+	// cur holds the last element.
+	cur.value = value
 
-		if !parent.isLeaf() {
-			// parent has other children, stop
-			break
-		}
-		parent.children = nil
-		if parent.hasValue {
-			// parent has a value, stop
-			break
+	root.AddChild(newRoot)
+}
+
+type node[K cmp.Ordered, V any] struct {
+	mu       sync.RWMutex
+	key      K
+	value    V
+	children []*node[K, V]
+}
+
+func newNode[K cmp.Ordered, V any]() *node[K, V] {
+	return &node[K, V]{}
+}
+
+func (n *node[K, V]) Key() K {
+	return n.key
+}
+
+func (n *node[K, V]) Value() V {
+	return n.value
+}
+
+func (n *node[K, V]) Children() iter.Seq[Node[K, V]] {
+	n.mu.RLock()
+	children := make([]*node[K, V], len(n.children))
+	copy(children, n.children)
+	n.mu.RUnlock()
+	return func(yield func(Node[K, V]) bool) {
+		for _, child := range children {
+			if !yield(child) {
+				break
+			}
 		}
 	}
-	return true
 }
 
-// WalkPair is what you get when you call `Walk()` on a trie.
-type WalkPair struct {
-	// Because we have a generic "Label" type, we unfortunately cannot
-	// provide a re-constructed Key object for the user to handle.
-	// Instead we provide this value as a slice of Labels
-	Labels []Label
+func (n *node[K, V]) AddChild(child Node[K, V]) {
+	n.mu.Lock()
+	// This is kind of gross, but we're only covering *node[T] with
+	// Node[T] interface because we don't want the users to instantiate
+	// their own nodes... so this type conversion is safe.
+	//nolint:forcetypeassert
+	n.children = append(n.children, child.(*node[K, V]))
+	sort.Slice(n.children, func(i, j int) bool {
+		return n.children[i].Key() < n.children[j].Key()
+	})
+	n.mu.Unlock()
+}
 
-	// Value is the value associated with the Labels
-	Value interface{}
+type VisitMetadata struct {
+	Depth int
 }
 
-// Walk returns a channel that you can read from to access all data
-// that is stored within this trie.
-func (t *Trie) Walk(ctx context.Context) <-chan WalkPair {
-	ch := make(chan WalkPair)
-	go t.walk(ctx, ch, nil)
-	return ch
+type Visitor[K cmp.Ordered, V any] interface {
+	Visit(Node[K, V], VisitMetadata) bool
 }
 
-func (t *Trie) walk(ctx context.Context, dst chan WalkPair, labels []Label) {
-	if labels == nil {
-		t.mu.RLock()
-		defer t.mu.RUnlock()
-		defer close(dst)
-	}
+func Walk[L any, K cmp.Ordered, V any](trie *Trie[L, K, V], v Visitor[K, V]) {
+	var meta VisitMetadata
+	meta.Depth = 1
+	walk(trie.root, v, meta)
+}
 
-	if t.hasValue {
-		p := WalkPair{
-			Labels: labels,
-			Value:  t.value,
-		}
-		select {
-		case <-ctx.Done():
-			return
-		case dst <- p:
+func walk[K cmp.Ordered, V any](node Node[K, V], v Visitor[K, V], meta VisitMetadata) {
+	for child := range node.Children() {
+		if !v.Visit(child, meta) {
+			break
 		}
+		walk(child, v, VisitMetadata{Depth: meta.Depth + 1})
 	}
+}
 
-	for _, child := range t.children {
-		child.walk(ctx, dst, append(labels, child.label))
+type dumper[K cmp.Ordered, V any] struct{}
+
+func (dumper[K, V]) Visit(n Node[K, V], meta VisitMetadata) bool {
+	var sb strings.Builder
+	for i := 0; i < meta.Depth; i++ {
+		sb.WriteString("  ")
 	}
+
+	fmt.Fprintf(&sb, "%v: %v", n.Key(), n.Value())
+	fmt.Println(sb.String())
+	return true
+}
+
+func Dumper[K cmp.Ordered, V any]() Visitor[K, V] {
+	return dumper[K, V]{}
 }

string_key.go

@@ -1,50 +1,41 @@
 package trie_test
 
 import (
-	"context"
 	"fmt"
 
-	"github.com/lestrrat-go/trie"
+	"github.com/lestrrat-go/trie/v2"
 )
 
-func ExampleStringKey() {
-	tree := trie.New()
+func Example() {
+	tree := trie.New[string, rune, any](trie.String())
 
 	// Put values in the trie
-	tree.Put(trie.StringKey("foo"), "one")
-	tree.Put(trie.StringKey("bar"), 2)
-	tree.Put(trie.StringKey("baz"), 3.0)
-	tree.Put(trie.StringKey("日本語"), []byte{'f', 'o', 'u', 'r'})
+	tree.Put("foo", "one")
+	tree.Put("bar", 2)
+	tree.Put("baz", 3.0)
+	tree.Put("日本語", []byte{'f', 'o', 'u', 'r'})
 
 	// Get a value from the trie
-	v, ok := tree.Get(trie.StringKey("日本語"))
+	v, ok := tree.Get("日本語")
 	if !ok {
 		fmt.Printf("failed to find key '日本語'\n")
 		return
 	}
 	_ = v
 
 	// Delete a key from the trie
-	if !tree.Delete(trie.StringKey("日本語")) {
+	if !tree.Delete("日本語") {
 		fmt.Printf("failed to delete key '日本語'\n")
 		return
 	}
 
 	// This time Get() should fail
-	v, ok = tree.Get(trie.StringKey("日本語"))
+	v, ok = tree.Get("日本語")
 	if ok {
 		fmt.Printf("key '日本語' should not exist\n")
 		return
 	}
 	_ = v
 
-	ctx := context.Background()
-
-	// Or, walk the entire trie
-	for p := range tree.Walk(ctx) {
-		// Do something with the values...
-		_ = p
-	}
-
 	// OUTPUT:
 }

trie.go

@@ -1,66 +1,41 @@
 package trie_test
 
 import (
-	"context"
-	"fmt"
 	"testing"
 
-	"github.com/lestrrat-go/trie"
-	"github.com/stretchr/testify/assert"
+	"github.com/lestrrat-go/trie/v2"
+	"github.com/stretchr/testify/require"
 )
 
 func TestTrie(t *testing.T) {
 	t.Parallel()
 
-	tree := trie.New()
-	tree.Put(trie.StringKey("foo"), 1)
-	tree.Put(trie.StringKey("bar"), 2)
-	tree.Put(trie.StringKey("baz"), 3)
-	tree.Put(trie.StringKey("日本語"), 4)
-
-	ctx, cancel := context.WithCancel(context.Background())
-	defer cancel()
-	for p := range tree.Walk(ctx) {
-		t.Logf("%#v", p)
-	}
+	tree := trie.New[string, rune, int](trie.String())
 
 	testcases := []struct {
-		Key      trie.Key
-		Expected interface{}
-		Missing  bool
+		Key   string
+		Value int
 	}{
-		{
-			Key:      trie.StringKey("foo"),
-			Expected: 1,
-		},
-		{
-			Key:      trie.StringKey("日本語"),
-			Expected: 4,
-		},
-		{
-			Key:     trie.StringKey("hoge"),
-			Missing: true,
-		},
+		{"foo", 1},
+		{"far", 2},
+		{"for", 3},
+		{"bar", 4},
+		{"baz", 5},
 	}
 
 	for _, tc := range testcases {
-		tc := tc
-		t.Run(fmt.Sprintf("%s", tc.Key), func(t *testing.T) {
-			t.Parallel()
-			v, ok := tree.Get(tc.Key)
-			if tc.Missing {
-				if !assert.False(t, ok, `tree.Get should return false`) {
-					return
-				}
-			} else {
-				if !assert.True(t, ok, `tree.Get should return true`) {
-					return
-				}
+		tree.Put(tc.Key, tc.Value)
+	}
 
-				if !assert.Equal(t, tc.Expected, v, `tree.Get should return expected value`) {
-					return
-				}
-			}
+	for _, tc := range testcases {
+		t.Run(tc.Key, func(t *testing.T) {
+			v, ok := tree.Get(tc.Key)
+			require.True(t, ok, `tree.Get should return true`)
+			require.Equal(t, tc.Value, v, `tree.Get should return expected value`)
 		})
 	}
+
+	require.True(t, tree.Delete("foo"), `tree.Delete should return true`)
+	_, ok := tree.Get("foo")
+	require.False(t, ok, `tree.Get should return false`)
 }