Skip to content

Commit

Permalink
Convert prefixtree to a generic implementation
Browse files Browse the repository at this point in the history
Value types are now generic instead of relying on type assertions.
  • Loading branch information
beevik committed Jul 6, 2024
1 parent 3e5091e commit 94cf616
Show file tree
Hide file tree
Showing 4 changed files with 69 additions and 66 deletions.
16 changes: 8 additions & 8 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@ The following code adds strings and associated data (in this case an integer)
to a prefix tree.

```go
tree := prefixtree.New()
tree := prefixtree.New[int]()

tree.Add("apple", 10)
tree.Add("orange", 20)
Expand Down Expand Up @@ -61,20 +61,20 @@ Output:
```
prefix value error
------ ----- -----
a <nil> prefixtree: prefix ambiguous
appl <nil> prefixtree: prefix ambiguous
a 0 prefixtree: prefix ambiguous
appl 0 prefixtree: prefix ambiguous
apple 10 <nil>
apple p 30 <nil>
apple pie 30 <nil>
apple pies <nil> prefixtree: prefix not found
apple pies 0 prefixtree: prefix not found
o 20 <nil>
orang 20 <nil>
orange 20 <nil>
oranges <nil> prefixtree: prefix not found
l <nil> prefixtree: prefix ambiguous
lemo <nil> prefixtree: prefix ambiguous
oranges 0 prefixtree: prefix not found
l 0 prefixtree: prefix ambiguous
lemo 0 prefixtree: prefix ambiguous
lemon 40 <nil>
lemon m 50 <nil>
lemon meringue 50 <nil>
pear <nil> prefixtree: prefix not found
pear 0 prefixtree: prefix not found
```
14 changes: 7 additions & 7 deletions example_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@ import (
func ExampleTree_usage() {
// Create the tree. Add 5 strings, each with an associated
// integer.
tree := prefixtree.New()
tree := prefixtree.New[int]()
for i, s := range []string{
"apple",
"orange",
Expand Down Expand Up @@ -53,19 +53,19 @@ func ExampleTree_usage() {
// Output:
// prefix value error
// ------ ----- -----
// a <nil> prefixtree: prefix ambiguous
// appl <nil> prefixtree: prefix ambiguous
// a 0 prefixtree: prefix ambiguous
// appl 0 prefixtree: prefix ambiguous
// apple 0 <nil>
// apple p 2 <nil>
// apple pie 2 <nil>
// apple pies <nil> prefixtree: prefix not found
// apple pies 0 prefixtree: prefix not found
// o 1 <nil>
// orang 1 <nil>
// orange 1 <nil>
// oranges <nil> prefixtree: prefix not found
// lemo <nil> prefixtree: prefix ambiguous
// oranges 0 prefixtree: prefix not found
// lemo 0 prefixtree: prefix ambiguous
// lemon 4 <nil>
// lemon m 3 <nil>
// lemon meringue 3 <nil>
// lemon meringues <nil> prefixtree: prefix not found
// lemon meringues 0 prefixtree: prefix not found
}
83 changes: 43 additions & 40 deletions prefixtree.go
Original file line number Diff line number Diff line change
Expand Up @@ -25,43 +25,44 @@ var (
ErrPrefixAmbiguous = errors.New("prefixtree: prefix ambiguous")
)

// A KeyValue type encapsulates a key string and its associated value.
type KeyValue struct {
// A KeyValue type encapsulates a key string and its associated value of type
// V.
type KeyValue[V any] struct {
Key string
Value any
Value V
}

// A Tree represents a prefix tree containing strings and their associated
// value data. The tree is implemented as a trie and can be searched
// value data of type V. The tree is implemented as a trie and can be searched
// efficiently for unique prefix matches.
type Tree struct {
type Tree[V any] struct {
key string
value any
links []link
value V
links []link[V]
descendants int
}

type link struct {
type link[V any] struct {
keyseg string
tree *Tree
tree *Tree[V]
}

// New returns an empty prefix tree.
func New() *Tree {
return new(Tree)
// New returns an empty prefix tree with a value type of V.
func New[V any]() *Tree[V] {
return new(Tree[V])
}

// isTerminal returns true if the tree is a terminal subtree in the
// prefix tree.
func (t *Tree) isTerminal() bool {
func (t *Tree[V]) isTerminal() bool {
return t.key != ""
}

// FindKey searches the prefix tree for a key string that uniquely matches the
// prefix. If found, the full matching key is returned. If not found,
// ErrPrefixNotFound is returned. If the prefix matches more than one key in
// the tree, ErrPrefixAmbiguous is returned.
func (t *Tree) FindKey(prefix string) (key string, err error) {
func (t *Tree[V]) FindKey(prefix string) (key string, err error) {
st, err := t.findSubtree(prefix)
if err != nil {
return "", err
Expand All @@ -74,17 +75,17 @@ func (t *Tree) FindKey(prefix string) (key string, err error) {
// value is returned. If not found, ErrPrefixNotFound is returned. If the
// prefix matches more than one key in the tree, ErrPrefixAmbiguous is
// returned.
func (t *Tree) FindKeyValue(prefix string) (kv KeyValue, err error) {
func (t *Tree[V]) FindKeyValue(prefix string) (kv KeyValue[V], err error) {
st, err := t.findSubtree(prefix)
if err != nil {
return KeyValue{}, err
return KeyValue[V]{}, err
}
return KeyValue{st.key, st.value}, nil
return KeyValue[V]{st.key, st.value}, nil
}

// FindKeys searches the prefix tree for all key strings prefixed by the
// provided prefix and returns them.
func (t *Tree) FindKeys(prefix string) (keys []string) {
func (t *Tree[V]) FindKeys(prefix string) (keys []string) {
st, err := t.findSubtree(prefix)
if err == ErrPrefixNotFound {
return []string{}
Expand All @@ -99,43 +100,44 @@ func (t *Tree) FindKeys(prefix string) (keys []string) {
// the prefix. If found, the value associated with the key is returned. If not
// found, ErrPrefixNotFound is returned. If the prefix matches more than one
// key in the tree, ErrPrefixAmbiguous is returned.
func (t *Tree) FindValue(prefix string) (value any, err error) {
func (t *Tree[V]) FindValue(prefix string) (value V, err error) {
st, err := t.findSubtree(prefix)
if err != nil {
return nil, err
var empty V
return empty, err
}
return st.value, nil
}

// FindKeyValues searches the prefix tree for all key strings prefixed by the
// provided prefix. All discovered keys and their values are returned.
func (t *Tree) FindKeyValues(prefix string) (values []KeyValue) {
func (t *Tree[V]) FindKeyValues(prefix string) (values []KeyValue[V]) {
st, err := t.findSubtree(prefix)
if err == ErrPrefixNotFound {
return []KeyValue{}
return []KeyValue[V]{}
}
if st.isTerminal() && err != ErrPrefixAmbiguous {
return []KeyValue{{st.key, st.value}}
return []KeyValue[V]{{st.key, st.value}}
}
return appendDescendantKeyValues(st, nil)
}

// FindValues searches the prefix tree for all key strings prefixed by the
// provided prefix. All associated values are returned.
func (t *Tree) FindValues(prefix string) (values []any) {
func (t *Tree[V]) FindValues(prefix string) (values []V) {
st, err := t.findSubtree(prefix)
if err == ErrPrefixNotFound {
return []any{}
return []V{}
}
if st.isTerminal() && err != ErrPrefixAmbiguous {
return []any{st.value}
return []V{st.value}
}
return appendDescendantValues(st, nil)
}

// findSubtree searches the prefix tree for the deepest subtree matching
// the prefix.
func (t *Tree) findSubtree(prefix string) (*Tree, error) {
func (t *Tree[V]) findSubtree(prefix string) (*Tree[V], error) {
outerLoop:
for {
// Ran out of prefix?
Expand Down Expand Up @@ -201,7 +203,7 @@ func matchingChars(s1, s2 string) int {

// appendDescendantKeys recursively appends a tree's descendant keys
// to an array of keys.
func appendDescendantKeys(t *Tree, keys []string) []string {
func appendDescendantKeys[V any](t *Tree[V], keys []string) []string {
if t.isTerminal() {
keys = append(keys, t.key)
}
Expand All @@ -213,9 +215,9 @@ func appendDescendantKeys(t *Tree, keys []string) []string {

// appendDescendantKeyValues recursively appends a tree's descendant keys
// to an array of key/value pairs.
func appendDescendantKeyValues(t *Tree, kv []KeyValue) []KeyValue {
func appendDescendantKeyValues[V any](t *Tree[V], kv []KeyValue[V]) []KeyValue[V] {
if t.isTerminal() {
kv = append(kv, KeyValue{t.key, t.value})
kv = append(kv, KeyValue[V]{t.key, t.value})
}
for i := 0; i < len(t.links); i++ {
kv = appendDescendantKeyValues(t.links[i].tree, kv)
Expand All @@ -225,7 +227,7 @@ func appendDescendantKeyValues(t *Tree, kv []KeyValue) []KeyValue {

// appendDescendantValues recursively appends a tree's descendant values
// to an array of values.
func appendDescendantValues(t *Tree, values []any) []any {
func appendDescendantValues[V any](t *Tree[V], values []V) []V {
if t.isTerminal() {
values = append(values, t.value)
}
Expand All @@ -236,7 +238,7 @@ func appendDescendantValues(t *Tree, values []any) []any {
}

// Add a key string and its associated value data to the prefix tree.
func (t *Tree) Add(key string, value any) {
func (t *Tree[V]) Add(key string, value V) {
k := key
outerLoop:
for {
Expand All @@ -255,7 +257,7 @@ outerLoop:

// Check the links before and after the insertion point for a matching
// prefix to see if we need to split one of them.
var splitLink *link
var splitLink *link[V]
var splitIndex int
innerLoop:
for li, lm := max(ix-1, 0), min(ix, len(t.links)-1); li <= lm; li++ {
Expand All @@ -275,24 +277,25 @@ outerLoop:

// No split necessary, so insert a new link and subtree.
if splitLink == nil {
child := &Tree{
child := &Tree[V]{
key: key,
value: value,
links: nil,
descendants: 1,
}
t.links = append(t.links[:ix],
append([]link{{k, child}}, t.links[ix:]...)...)
append([]link[V]{{k, child}}, t.links[ix:]...)...)
break outerLoop
}

// A split is necessary, so split the current link's string and insert
// a child tree.
k1, k2 := splitLink.keyseg[:splitIndex], splitLink.keyseg[splitIndex:]
child := &Tree{
var empty V
child := &Tree[V]{
key: "",
value: nil,
links: []link{{k2, splitLink.tree}},
value: empty,
links: []link[V]{{k2, splitLink.tree}},
descendants: splitLink.tree.descendants,
}
splitLink.keyseg, splitLink.tree = k1, child
Expand All @@ -302,11 +305,11 @@ outerLoop:

// Output the structure of the tree to stdout. This function exists for
// debugging purposes.
func (t *Tree) Output() {
func (t *Tree[V]) Output() {
t.outputNode(0)
}

func (t *Tree) outputNode(level int) {
func (t *Tree[V]) outputNode(level int) {
fmt.Printf("%sNode: key=\"%s\" term=%v desc=%d value=%v\n",
strings.Repeat(" ", level), t.key, t.isTerminal(), t.descendants, t.value)
for i, l := range t.links {
Expand Down
22 changes: 11 additions & 11 deletions prefixtree_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -13,21 +13,21 @@ import (

type entry struct {
key string
value any
value int
}

type testcase struct {
key string
value any
value int
err error
}

func test(t *testing.T, entries []entry, cases []testcase) *Tree {
func test(t *testing.T, entries []entry, cases []testcase) *Tree[int] {
// Run 256 iterations of build/find using random tree entry
// insertion orders.
var tree *Tree
var tree *Tree[int]
for i := 0; i < 256; i++ {
tree = New()
tree = New[int]()
for _, i := range rand.Perm(len(entries)) {
tree.Add(entries[i].key, entries[i].value)
}
Expand Down Expand Up @@ -189,7 +189,7 @@ func TestFindKeys(t *testing.T) {
{"bog", 6},
}

tree := New()
tree := New[int]()
for _, entry := range entries {
tree.Add(entry.key, entry.value)
}
Expand Down Expand Up @@ -250,7 +250,7 @@ func TestFindValues(t *testing.T) {
{"bee", 5},
}

tree := New()
tree := New[int]()
for _, entry := range entries {
tree.Add(entry.key, entry.value)
}
Expand Down Expand Up @@ -332,10 +332,10 @@ func TestDictionary(t *testing.T) {
}

// Scan all words from the dictionary into the tree.
tree := New()
tree := New[bool]()
scanner := bufio.NewScanner(file)
for scanner.Scan() {
tree.Add(scanner.Text(), nil)
tree.Add(scanner.Text(), true)
}
file.Close()

Expand Down Expand Up @@ -382,10 +382,10 @@ func BenchmarkDictionary(b *testing.B) {
}

// Scan all words from the dictionary into the tree.
tree := New()
tree := New[bool]()
scanner := bufio.NewScanner(file)
for scanner.Scan() {
tree.Add(scanner.Text(), nil)
tree.Add(scanner.Text(), true)
}
file.Close()

Expand Down

0 comments on commit 94cf616

Please sign in to comment.