From 254353fb861957a5520783eda9cec57441c3f399 Mon Sep 17 00:00:00 2001
From: Zachary Yedidia <zyedidia@gmail.com>
Date: Mon, 9 Jan 2023 11:50:20 -0800
Subject: [PATCH] Fix unicode handling in trie

Ref #34
---
 trie/trie.go      | 10 +++++-----
 trie/trie_test.go | 12 ++++++------
 2 files changed, 11 insertions(+), 11 deletions(-)

diff --git a/trie/trie.go b/trie/trie.go
index f1f3d04..8dddf67 100644
--- a/trie/trie.go
+++ b/trie/trie.go
@@ -134,7 +134,7 @@ func (t *Trie[V]) LongestPrefix(query string) string {
 
 // Keys returns all keys in the trie.
 func (t *Trie[V]) Keys() (queue []string) {
-	return t.collect(t.root, "", queue)
+	return t.collect(t.root, nil, queue)
 }
 
 // KeysWithPrefix returns all keys with prefix 'prefix'.
@@ -149,17 +149,17 @@ func (t *Trie[V]) KeysWithPrefix(prefix string) (queue []string) {
 	if x.valid {
 		queue = []string{prefix}
 	}
-	return t.collect(x.mid, prefix, queue)
+	return t.collect(x.mid, []byte(prefix), queue)
 }
 
-func (t *Trie[V]) collect(x *node[V], prefix string, queue []string) []string {
+func (t *Trie[V]) collect(x *node[V], prefix []byte, queue []string) []string {
 	if x == nil {
 		return queue
 	}
 	queue = t.collect(x.left, prefix, queue)
 	if x.valid {
-		queue = append(queue, prefix+string(x.c))
+		queue = append(queue, string(append(prefix, x.c)))
 	}
-	queue = t.collect(x.mid, prefix+string(x.c), queue)
+	queue = t.collect(x.mid, append(prefix, x.c), queue)
 	return t.collect(x.right, prefix, queue)
 }
diff --git a/trie/trie_test.go b/trie/trie_test.go
index f8277e9..832b81f 100644
--- a/trie/trie_test.go
+++ b/trie/trie_test.go
@@ -82,15 +82,15 @@ func TestKeys(t *testing.T) {
 
 func Example() {
 	tr := trie.New[int]()
-	tr.Put("foo", 1)
-	tr.Put("fo", 2)
+	tr.Put("f§oo", 1)
+	tr.Put("f§o", 2)
 	tr.Put("bar", 3)
 
-	fmt.Println(tr.Contains("f"))
+	fmt.Println(tr.Contains("§"))
 	fmt.Println(tr.KeysWithPrefix(""))
-	fmt.Println(tr.KeysWithPrefix("f"))
+	fmt.Println(tr.KeysWithPrefix("f§"))
 	// Output:
 	// false
-	// [bar fo foo]
-	// [fo foo]
+	// [bar f§o f§oo]
+	// [f§o f§oo]
 }