Long over due Trie delete implemented

2 files changed, 121 insertions, 3 deletions

diff --git a/ethutil/trie.go b/ethutil/trie.go
index e900a0a63..322f77647 100644
--- a/ethutil/trie.go
+++ b/ethutil/trie.go
@@ -148,6 +148,10 @@ func (t *Trie) Get(key string) string {
 	return c.Str()
 }
 
+func (t *Trie) Delete(key string) {
+	t.Update(key, "")
+}
+
 func (t *Trie) GetState(node interface{}, key []int) interface{} {
 	n := NewValue(node)
 	// Return the node if key is empty (= found)
@@ -202,9 +206,10 @@ func (t *Trie) UpdateState(node interface{}, key []int, value string) interface{
 		return t.InsertState(node, key, value)
 	} else {
 		// delete it
+		return t.DeleteState(node, key)
 	}
 
-	return ""
+	return t.Root
 }
 
 func (t *Trie) Put(node interface{}) interface{} {
@@ -313,6 +318,87 @@ func (t *Trie) InsertState(node interface{}, key []int, value interface{}) inter
 	return ""
 }
 
+func (t *Trie) DeleteState(node interface{}, key []int) interface{} {
+	if len(key) == 0 {
+		return ""
+	}
+
+	// New node
+	n := NewValue(node)
+	if node == nil || (n.Type() == reflect.String && (n.Str() == "" || n.Get(0).IsNil())) || n.Len() == 0 {
+		return ""
+	}
+
+	currentNode := t.GetNode(node)
+	// Check for "special" 2 slice type node
+	if currentNode.Len() == 2 {
+		// Decode the key
+		k := CompactDecode(currentNode.Get(0).Str())
+		v := currentNode.Get(1).Raw()
+
+		// Matching key pair (ie. there's already an object with this key)
+		if CompareIntSlice(k, key) {
+			return ""
+		} else if CompareIntSlice(key[:len(k)], k) {
+			hash := t.DeleteState(v, key[len(k):])
+			child := t.GetNode(hash)
+
+			var newNode []interface{}
+			if child.Len() == 2 {
+				newKey := append(k, CompactDecode(child.Get(0).Str())...)
+				newNode = []interface{}{CompactEncode(newKey), child.Get(1).Raw()}
+			} else {
+				newNode = []interface{}{currentNode.Get(0).Str(), hash}
+			}
+
+			return t.Put(newNode)
+		} else {
+			return node
+		}
+	} else {
+		// Copy the current node over to the new node and replace the first nibble in the key
+		n := EmptyStringSlice(17)
+		var newNode []interface{}
+
+		for i := 0; i < 17; i++ {
+			cpy := currentNode.Get(i).Raw()
+			if cpy != nil {
+				n[i] = cpy
+			}
+		}
+
+		n[key[0]] = t.DeleteState(n[key[0]], key[1:])
+		amount := -1
+		for i := 0; i < 17; i++ {
+			if n[i] != "" {
+				if amount == -1 {
+					amount = i
+				} else {
+					amount = -2
+				}
+			}
+		}
+		if amount == 16 {
+			newNode = []interface{}{CompactEncode([]int{16}), n[amount]}
+		} else if amount >= 0 {
+			child := t.GetNode(n[amount])
+			if child.Len() == 17 {
+				newNode = []interface{}{CompactEncode([]int{amount}), n[amount]}
+			} else if child.Len() == 2 {
+				key := append([]int{amount}, CompactDecode(child.Get(0).Str())...)
+				newNode = []interface{}{CompactEncode(key), child.Get(1).Str()}
+			}
+
+		} else {
+			newNode = n
+		}
+
+		return t.Put(newNode)
+	}
+
+	return ""
+}
+
 // Simple compare function which creates a rlp value out of the evaluated objects
 func (t *Trie) Cmp(trie *Trie) bool {
 	return NewValue(t.Root).Cmp(NewValue(trie.Root))
diff --git a/ethutil/trie_test.go b/ethutil/trie_test.go
index fa60c8cfc..9d2c8e19f 100644
--- a/ethutil/trie_test.go
+++ b/ethutil/trie_test.go
@@ -1,8 +1,7 @@
 package ethutil
 
 import (
-	_ "encoding/hex"
-	_ "fmt"
+	"reflect"
 	"testing"
 )
 
@@ -116,3 +115,36 @@ func TestTrieCmp(t *testing.T) {
 		t.Errorf("Expected tries not to be equal %x %x", trie1.Root, trie2.Root)
 	}
 }
+
+func TestTrieDelete(t *testing.T) {
+	_, trie := New()
+	trie.Update("cat", LONG_WORD)
+	exp := trie.Root
+	trie.Update("dog", LONG_WORD)
+	trie.Delete("dog")
+	if !reflect.DeepEqual(exp, trie.Root) {
+		t.Errorf("Expected tries to be equal %x : %x", exp, trie.Root)
+	}
+
+	trie.Update("dog", LONG_WORD)
+	exp = trie.Root
+	trie.Update("dude", LONG_WORD)
+	trie.Delete("dude")
+	if !reflect.DeepEqual(exp, trie.Root) {
+		t.Errorf("Expected tries to be equal %x : %x", exp, trie.Root)
+	}
+}
+
+func TestTrieDeleteWithValue(t *testing.T) {
+	_, trie := New()
+	trie.Update("c", LONG_WORD)
+	exp := trie.Root
+	trie.Update("ca", LONG_WORD)
+	trie.Update("cat", LONG_WORD)
+	trie.Delete("ca")
+	trie.Delete("cat")
+	if !reflect.DeepEqual(exp, trie.Root) {
+		t.Errorf("Expected tries to be equal %x : %x", exp, trie.Root)
+	}
+
+}