Add gob encoding and decoding functionality

bench_test.go

@@ -2,6 +2,8 @@ package radixtree
 
 import (
 	"bufio"
+	"bytes"
+	"encoding/gob"
 	"fmt"
 	"io"
 	"net/http"
@@ -64,6 +66,51 @@ func BenchmarkIterPath(b *testing.B) {
 	})
 }
 
+func BenchmarkRebuildVsReload(b *testing.B) {
+	err := downloadFile(web2URL, web2Path)
+	if err != nil {
+		b.Skip(err.Error())
+	}
+	words, err := loadWords(web2Path)
+	if err != nil {
+		b.Skip(err.Error())
+	}
+
+	b.Run("Reload from encoded", func(b *testing.B) {
+		tree := new(Tree[string])
+		for _, w := range words {
+			tree.Put(w, w)
+		}
+
+		var buf bytes.Buffer
+		encoder := gob.NewEncoder(&buf)
+		if err = encoder.Encode(tree); err != nil {
+			b.Fatal(err)
+		}
+		reader := bytes.NewReader(buf.Bytes())
+
+		b.ResetTimer()
+		for n := 0; n < b.N; n++ {
+			tree2 := New[string]()
+			decoder := gob.NewDecoder(reader)
+			err := decoder.Decode(tree2)
+			if err != nil {
+				b.Fatal(err)
+			}
+		}
+	})
+
+	b.Run("Rebuild from dictionary", func(b *testing.B) {
+		for n := 0; n < b.N; n++ {
+			tree := new(Tree[string])
+			for _, w := range words {
+				tree.Put(w, w)
+			}
+		}
+	})
+
+}
+
 func benchmarkGet(b *testing.B, filePath string) {
 	words, err := loadWords(filePath)
 	if err != nil {

gob.go

@@ -0,0 +1,121 @@
+package radixtree
+
+import (
+	"bytes"
+	"encoding/gob"
+)
+
+func (t *Tree[T]) GobEncode() ([]byte, error) {
+	var buf bytes.Buffer
+	encoder := gob.NewEncoder(&buf)
+	var err error
+
+	if err = encoder.Encode(&t.root); err != nil {
+		return nil, err
+	}
+	if err = encoder.Encode(&t.size); err != nil {
+		return nil, err
+	}
+
+	return buf.Bytes(), nil
+}
+
+func (t *Tree[T]) GobDecode(data []byte) error {
+	buf := bytes.NewBuffer(data)
+	decoder := gob.NewDecoder(buf)
+	var err error
+
+	if err = decoder.Decode(&t.root); err != nil {
+		return err
+	}
+	return decoder.Decode(&t.size)
+}
+
+func (node *radixNode[T]) GobEncode() ([]byte, error) {
+	var buf bytes.Buffer
+	encoder := gob.NewEncoder(&buf)
+
+	err := encoder.Encode(node.prefix)
+	if err != nil {
+		return nil, err
+	}
+	if err = encoder.Encode(node.edges); err != nil {
+		return nil, err
+	}
+	if node.leaf != nil {
+		if err = encoder.Encode(node.leaf); err != nil {
+			return nil, err
+		}
+	}
+
+	return buf.Bytes(), nil
+}
+
+func (node *radixNode[T]) GobDecode(data []byte) error {
+	buf := bytes.NewBuffer(data)
+	decoder := gob.NewDecoder(buf)
+
+	err := decoder.Decode(&node.prefix)
+	if err != nil {
+		return err
+	}
+	if err = decoder.Decode(&node.edges); err != nil {
+		return err
+	}
+	if buf.Len() != 0 { // if leaf is not nil
+		return decoder.Decode(&node.leaf)
+	}
+	return nil
+}
+
+func (kv *Item[T]) GobEncode() ([]byte, error) {
+	var buf bytes.Buffer
+	encoder := gob.NewEncoder(&buf)
+
+	err := encoder.Encode(kv.key)
+	if err != nil {
+		return nil, err
+	}
+	if err = encoder.Encode(kv.value); err != nil {
+		return nil, err
+	}
+
+	return buf.Bytes(), nil
+}
+
+func (kv *Item[T]) GobDecode(data []byte) error {
+	buf := bytes.NewBuffer(data)
+	decoder := gob.NewDecoder(buf)
+
+	err := decoder.Decode(&kv.key)
+	if err != nil {
+		return err
+	}
+	return decoder.Decode(&kv.value)
+}
+
+func (e *edge[T]) GobEncode() ([]byte, error) {
+	var buf bytes.Buffer
+	encoder := gob.NewEncoder(&buf)
+
+	err := encoder.Encode(e.radix)
+	if err != nil {
+		return nil, err
+	}
+	if err = encoder.Encode(e.node); err != nil {
+		return nil, err
+	}
+
+	return buf.Bytes(), nil
+}
+
+func (e *edge[T]) GobDecode(data []byte) error {
+	buf := bytes.NewBuffer(data)
+	decoder := gob.NewDecoder(buf)
+
+	err := decoder.Decode(&e.radix)
+	if err != nil {
+		return err
+	}
+	return decoder.Decode(&e.node)
+}

gob_test.go

@@ -0,0 +1,69 @@
+package radixtree_test
+
+import (
+	"bytes"
+	"encoding/gob"
+	"fmt"
+	"strings"
+	"testing"
+
+	"github.com/gammazero/radixtree"
+)
+
+func TestEncode(t *testing.T) {
+	rt := radixtree.New[string]()
+
+	keys := []string{
+		"bird",
+		"rat",
+		"bat",
+		"rats",
+		"ratatouille",
+		"rat/whis/key",
+		"rat/whis/kers",
+		"rat/whis/per/er",
+		"rat/winks/wisely/once",
+		"rat/winks/wisely/x/y/z",
+		"rat/winks/wryly",
+	}
+
+	for _, key := range keys {
+		rt.Put(key, strings.ToUpper(key))
+	}
+
+	var buf bytes.Buffer
+	encoder := gob.NewEncoder(&buf)
+
+	err := encoder.Encode(rt)
+	if err != nil {
+		panic(err)
+	}
+
+	fmt.Println("Encoded to", buf.Len(), "bytes")
+
+	rt2 := radixtree.New[string]()
+
+	decoder := gob.NewDecoder(&buf)
+	err = decoder.Decode(rt2)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if rt.Len() != rt2.Len() {
+		t.Fatalf("decoded tree has wrong size, expectd %d got %d", rt.Len(), rt2.Len())
+	}
+	for _, key := range keys {
+		checkItem(t, rt2, key, strings.ToUpper(key))
+	}
+	fmt.Println("Decoded tree matches original tree")
+}
+
+func checkItem(t *testing.T, rt *radixtree.Tree[string], key, value string) {
+	val, ok := rt.Get(key)
+	if !ok {
+		t.Fatalf("decoded tree missing value for key %q", key)
+	}
+	if val != value {
+		t.Fatalf("wrong value in tree, expected %s got%s", value, val)
+	}
+}