From 8facbdf809dd6e60baf8c81872f3ff2a1f306888 Mon Sep 17 00:00:00 2001
From: Jae Kwon <jkwon.work@gmail.com>
Date: Mon, 6 Oct 2014 13:55:56 -0700
Subject: [PATCH] TypedTree.

---
 merkle/iavl_test.go | 45 ++++++++++++++++++++++++++
 merkle/iavl_tree.go | 78 ++++++++++++++++++++++++++++++++++++++++++---
 merkle/types.go     |  6 ++--
 3 files changed, 121 insertions(+), 8 deletions(-)

diff --git a/merkle/iavl_test.go b/merkle/iavl_test.go
index a8d86352c..f759771c5 100644
--- a/merkle/iavl_test.go
+++ b/merkle/iavl_test.go
@@ -4,7 +4,9 @@ import (
 	"bytes"
 	"crypto/sha256"
 	"fmt"
+	"time"
 
+	. "github.com/tendermint/tendermint/binary"
 	. "github.com/tendermint/tendermint/common"
 	"github.com/tendermint/tendermint/db"
 
@@ -232,6 +234,49 @@ func TestPersistence(t *testing.T) {
 	}
 }
 
+func TestTypedTree(t *testing.T) {
+	db := db.NewMemDB()
+
+	// Construct some tree and save it
+	t1 := NewTypedTree(NewIAVLTree(db), BasicCodec, BasicCodec)
+	t1.Set(uint8(1), "uint8(1)")
+	t1.Set(uint16(1), "uint16(1)")
+	t1.Set(uint32(1), "uint32(1)")
+	t1.Set(uint64(1), "uint64(1)")
+	t1.Set("byteslice01", []byte{byte(0x00), byte(0x01)})
+	t1.Set("byteslice23", []byte{byte(0x02), byte(0x03)})
+	t1.Set("time", time.Unix(123, 0))
+	t1.Set("nil", nil)
+	t1Hash := t1.Tree.Save()
+
+	// Reconstruct tree
+	t2 := NewTypedTree(LoadIAVLTreeFromHash(db, t1Hash), BasicCodec, BasicCodec)
+	if t2.Get(uint8(1)).(string) != "uint8(1)" {
+		t.Errorf("Expected string uint8(1)")
+	}
+	if t2.Get(uint16(1)).(string) != "uint16(1)" {
+		t.Errorf("Expected string uint16(1)")
+	}
+	if t2.Get(uint32(1)).(string) != "uint32(1)" {
+		t.Errorf("Expected string uint32(1)")
+	}
+	if t2.Get(uint64(1)).(string) != "uint64(1)" {
+		t.Errorf("Expected string uint64(1)")
+	}
+	if !bytes.Equal(t2.Get("byteslice01").([]byte), []byte{byte(0x00), byte(0x01)}) {
+		t.Errorf("Expected byteslice 0x00 0x01")
+	}
+	if !bytes.Equal(t2.Get("byteslice23").([]byte), []byte{byte(0x02), byte(0x03)}) {
+		t.Errorf("Expected byteslice 0x02 0x03")
+	}
+	if t2.Get("time").(time.Time).Unix() != 123 {
+		t.Errorf("Expected time 123")
+	}
+	if t2.Get("nil") != nil {
+		t.Errorf("Expected nil")
+	}
+}
+
 func BenchmarkHash(b *testing.B) {
 	b.StopTimer()
 
diff --git a/merkle/iavl_tree.go b/merkle/iavl_tree.go
index 907022b52..e0c8e1d65 100644
--- a/merkle/iavl_tree.go
+++ b/merkle/iavl_tree.go
@@ -3,12 +3,12 @@ package merkle
 import (
 	"bytes"
 	"container/list"
+	. "github.com/tendermint/tendermint/binary"
+	. "github.com/tendermint/tendermint/common"
 )
 
 const defaultCacheCapacity = 1000 // TODO make configurable.
 
-// XXX Make Codec tree.
-
 /*
 Immutable AVL Tree (wraps the Node root)
 
@@ -81,11 +81,11 @@ func (t *IAVLTree) HashWithCount() ([]byte, uint64) {
 	return t.root.HashWithCount()
 }
 
-func (t *IAVLTree) Save() {
+func (t *IAVLTree) Save() []byte {
 	if t.root == nil {
-		return
+		return nil
 	}
-	t.root.Save(t.ndb)
+	return t.root.Save(t.ndb)
 }
 
 func (t *IAVLTree) Get(key []byte) (value []byte) {
@@ -117,6 +117,74 @@ func (t *IAVLTree) Copy() Tree {
 
 //-----------------------------------------------------------------------------
 
+type TypedTree struct {
+	Tree       Tree
+	keyCodec   Codec
+	valueCodec Codec
+}
+
+func NewTypedTree(tree Tree, keyCodec, valueCodec Codec) *TypedTree {
+	return &TypedTree{
+		Tree:       tree,
+		keyCodec:   keyCodec,
+		valueCodec: valueCodec,
+	}
+}
+
+func (t *TypedTree) Has(key interface{}) bool {
+	bytes, err := t.keyCodec.Write(key)
+	if err != nil {
+		Panicf("Error from keyCodec: %v", err)
+	}
+	return t.Tree.Has(bytes)
+}
+
+func (t *TypedTree) Get(key interface{}) interface{} {
+	keyBytes, err := t.keyCodec.Write(key)
+	if err != nil {
+		Panicf("Error from keyCodec: %v", err)
+	}
+	valueBytes := t.Tree.Get(keyBytes)
+	if valueBytes == nil {
+		return nil
+	}
+	value, err := t.valueCodec.Read(valueBytes)
+	if err != nil {
+		Panicf("Error from valueCodec: %v", err)
+	}
+	return value
+}
+
+func (t *TypedTree) Set(key interface{}, value interface{}) bool {
+	keyBytes, err := t.keyCodec.Write(key)
+	if err != nil {
+		Panicf("Error from keyCodec: %v", err)
+	}
+	valueBytes, err := t.valueCodec.Write(value)
+	if err != nil {
+		Panicf("Error from valueCodec: %v", err)
+	}
+	return t.Tree.Set(keyBytes, valueBytes)
+}
+
+func (t *TypedTree) Remove(key interface{}) (interface{}, error) {
+	keyBytes, err := t.keyCodec.Write(key)
+	if err != nil {
+		Panicf("Error from keyCodec: %v", err)
+	}
+	valueBytes, err := t.Tree.Remove(keyBytes)
+	if valueBytes == nil {
+		return nil, err
+	}
+	value, err_ := t.valueCodec.Read(valueBytes)
+	if err_ != nil {
+		Panicf("Error from valueCodec: %v", err)
+	}
+	return value, err
+}
+
+//-----------------------------------------------------------------------------
+
 type nodeElement struct {
 	node *IAVLNode
 	elem *list.Element
diff --git a/merkle/types.go b/merkle/types.go
index 117bfa536..5538197a6 100644
--- a/merkle/types.go
+++ b/merkle/types.go
@@ -14,11 +14,11 @@ type Tree interface {
 	Height() uint8
 	Has(key []byte) bool
 	Get(key []byte) []byte
+	Set(key []byte, value []byte) bool
+	Remove(key []byte) ([]byte, error)
 	HashWithCount() ([]byte, uint64)
 	Hash() []byte
-	Save()
-	Set(key []byte, vlaue []byte) bool
-	Remove(key []byte) ([]byte, error)
+	Save() []byte
 	Copy() Tree
 }