Add support for Trie.Walk() and Node.SetValue() (#4)

* udpate: add support for SelectOnValue, find item in the tree based on Value, using eval function

* update: renaming, cleaning tests

* update: renaming, SelectOnValue -> Walk, implement with childrenDLL to avoid non deterministic sort.

* Add support for Trie.Walk() and Node.SetValue()

* Fix err on Windows + Fix docs

---------

Co-authored-by: tophe <cvigny@artprice.com>
Co-authored-by: Shivam Mamgain <smamgain@microsoft.com>

.github/workflows/ci.yml

@@ -35,6 +35,6 @@ jobs:
       - name: Checkout code
         uses: actions/checkout@v2
       - name: Run tests
-        run: go test ./...
+        run: go test -v ./...
       - name: Run benchmarks
-        run: go test -bench=.
+        run: go test -v -bench=.

search_test.go

@@ -515,6 +515,7 @@ func getWordsTrie() *trie.Trie {
 			panic(err)
 		}
 		word = strings.TrimRight(word, "\n")
+		word = strings.TrimRight(word, "\r") // windows
 		key := strings.Split(word, "")
 		tri.Put(key, nil)
 	}

trie.go

@@ -47,6 +47,13 @@ func (n *Node) Value() interface{} {
 	return n.value
 }
 
+// SetValue sets the value for the key ending at this Node. If Node is not a terminal, value is not set.
+func (n *Node) SetValue(value interface{}) {
+	if n.isTerminal {
+		n.value = value
+	}
+}
+
 // ChildNodes returns the child-nodes of this Node.
 func (n *Node) ChildNodes() []*Node {
 	return n.childNodes()
@@ -101,7 +108,7 @@ func (t *Trie) Root() *Node {
 	return t.root
 }
 
-// Put upserts value the given key in the Trie. It returns a boolean depending on
+// Put upserts value for the given key in the Trie. It returns a boolean depending on
 // whether the key already existed or not.
 func (t *Trie) Put(key []string, value interface{}) (existed bool) {
 	node := t.root

trie_test.go

@@ -8,6 +8,24 @@ import (
 	"github.com/stretchr/testify/assert"
 )
 
+func TestNode_SetValue(t *testing.T) {
+	tri := trie.New()
+	tri.Put([]string{"a", "b"}, 1)
+	tri.Put([]string{"a", "b", "c"}, 2)
+
+	node := tri.Root()
+	node.SetValue(10)
+	assert.Equal(t, nil, node.Value())
+
+	node = tri.Root().ChildNodes()[0]
+	node.SetValue(10)
+	assert.Equal(t, nil, node.Value())
+
+	node = tri.Root().ChildNodes()[0].ChildNodes()[0]
+	node.SetValue(10)
+	assert.Equal(t, 10, node.Value())
+}
+
 func TestTrie_Put(t *testing.T) {
 	tri := trie.New()
 	existed := tri.Put([]string{"an", "umbrella"}, 2)

walk.go

@@ -0,0 +1,38 @@
+package trie
+
+type WalkFunc func(key []string, node *Node) error
+
+// Walk traverses the Trie and calls walker function. If walker function returns an error, Walk early-returns with that error.
+// Traversal follows insertion order.
+func (t *Trie) Walk(key []string, walker WalkFunc) error {
+	node := t.root
+	for _, keyPart := range key {
+		child, ok := node.children[keyPart]
+		if !ok {
+			return nil
+		}
+		node = child
+	}
+	return t.walk(node, &key, walker)
+}
+
+func (t *Trie) walk(node *Node, prefixKey *[]string, walker WalkFunc) error {
+	if node.isTerminal {
+		key := make([]string, len(*prefixKey))
+		copy(key, *prefixKey)
+		if err := walker(key, node); err != nil {
+			return err
+		}
+	}
+
+	for dllNode := node.childrenDLL.head; dllNode != nil; dllNode = dllNode.next {
+		child := dllNode.trieNode
+		*prefixKey = append(*prefixKey, child.keyPart)
+		err := t.walk(child, prefixKey, walker)
+		*prefixKey = (*prefixKey)[:len(*prefixKey)-1]
+		if err != nil {
+			return err
+		}
+	}
+	return nil
+}

walk_test.go

@@ -0,0 +1,75 @@
+package trie_test
+
+import (
+	"errors"
+	"testing"
+
+	"github.com/shivamMg/trie"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestTrie_WalkErr(t *testing.T) {
+	tri := trie.New()
+	tri.Put([]string{"d", "a", "l", "i"}, 1)
+	tri.Put([]string{"d", "a", "l", "i", "b"}, 2)
+	tri.Put([]string{"d", "a", "l", "i", "b", "e"}, 3)
+	tri.Put([]string{"d", "a", "l", "i", "b", "e", "r", "t"}, 4)
+
+	var selected []string
+	walker := func(key []string, node *trie.Node) error {
+		what := node.Value().(int)
+		if what == 3 {
+			selected = key
+			return errors.New("found")
+		}
+		return nil
+	}
+
+	err := tri.Walk(nil, walker)
+	assert.EqualError(t, err, "found")
+	assert.EqualValues(t, []string{"d", "a", "l", "i", "b", "e"}, selected)
+}
+
+func TestTrie_Walk(t *testing.T) {
+	tri := trie.New()
+	tri.Put([]string{"d", "a", "l", "i"}, []int{0, 1, 2, 4, 5})
+	tri.Put([]string{"d", "a", "l", "i", "b"}, []int{1, 2, 4, 5})
+	tri.Put([]string{"d", "a", "l", "i", "b", "e"}, []int{1, 0, 2, 4, 5, 0})
+	tri.Put([]string{"d", "a", "l", "i", "b", "e", "r", "t"}, []int{1, 2, 4, 5})
+	type KVPair struct {
+		key   []string
+		value []int
+	}
+	var selected []KVPair
+	walker := func(key []string, node *trie.Node) error {
+		what := node.Value().([]int)
+		for _, i := range what {
+			if i == 0 {
+				selected = append(selected, KVPair{key, what})
+				break
+			}
+		}
+		return nil
+	}
+
+	err := tri.Walk(nil, walker)
+	assert.NoError(t, err)
+	expected := []KVPair{
+		{[]string{"d", "a", "l", "i"}, []int{0, 1, 2, 4, 5}},
+		{[]string{"d", "a", "l", "i", "b", "e"}, []int{1, 0, 2, 4, 5, 0}},
+	}
+	assert.EqualValues(t, expected, selected)
+
+	selected = nil
+	err = tri.Walk([]string{"d", "a", "l", "i", "b"}, walker)
+	assert.NoError(t, err)
+	expected = []KVPair{
+		{[]string{"d", "a", "l", "i", "b", "e"}, []int{1, 0, 2, 4, 5, 0}},
+	}
+	assert.EqualValues(t, expected, selected)
+
+	selected = nil
+	err = tri.Walk([]string{"a", "b", "c"}, walker)
+	assert.NoError(t, err)
+	assert.Nil(t, selected)
+}