Restacktrie with pruning (#13)

* revert snaphash to processing all accounts

* use ReStackTrie in generateTrie

* Save memory by hashing a branch if no more insert will occur

cmd/geth/chaincmd.go

@@ -18,7 +18,6 @@ package main
 
 import (
 	"encoding/json"
-	"errors"
 	"fmt"
 	"github.com/ethereum/go-ethereum/core/state/snapshot"
 	"os"
@@ -617,18 +616,14 @@ func snapToHash(ctx *cli.Context) error {
 	if err != nil {
 		return fmt.Errorf("Could not create iterator for root %x: %v", root, err)
 	}
-	ollKorrekt := snapshot.CrosscheckTriehasher(it, 0, 10000)
-	//generatedRoot := snapshot.GenerateTrieRoot(it)
-	//if err := it.Error(); err != nil {
-	//	fmt.Printf("Iterator error: %v\n", it.Error())
-	//}
-	//if root != generatedRoot {
-	//	return fmt.Errorf("Wrong hash generated, expected %x, got %x", root, generatedRoot[:])
-	//}
-	if !ollKorrekt {
-		return errors.New("Computer says No, @gballet\n...come on man, fix me already!")
-	}
-	//log.Info("Generation done", "root", generatedRoot)
+	generatedRoot := snapshot.GenerateTrieRoot(it)
+	if err := it.Error(); err != nil {
+		fmt.Printf("Iterator error: %v\n", it.Error())
+	}
+	if root != generatedRoot {
+		return fmt.Errorf("Wrong hash generated, expected %x, got %x", root, generatedRoot[:])
+	}
+	log.Info("Generation done", "root", generatedRoot)
 	return nil
 }
 

core/state/snapshot/hextrie_generator.go

@@ -36,7 +36,7 @@ type trieGeneratorFn func(in chan (leaf), out chan (common.Hash))
 // GenerateTrieRoot takes an account iterator and reproduces the root hash.
 func GenerateTrieRoot(it AccountIterator) common.Hash {
 	//return generateTrieRoot(it, StackGenerate)
-	return generateTrieRoot(it, StdGenerate)
+	return generateTrieRoot(it, ReStackGenerate)
 }
 
 func CrosscheckTriehasher(it AccountIterator, begin, end int) bool {

trie/stacktrie.go

@@ -287,6 +287,7 @@ const (
 	extNode
 	leafNode
 	emptyNode
+	hashedNode
 )
 
 func (st *ReStackTrie) TryUpdate(key, value []byte) error {
@@ -318,6 +319,18 @@ func (st *ReStackTrie) insert(key, value []byte) {
 			st.children[idx] = NewReStackTrie()
 			st.children[idx].keyOffset = st.keyOffset + 1
 		}
+		for i := idx - 1; i >= 0; i-- {
+			if st.children[i] != nil {
+				if st.children[i].nodeType != hashedNode {
+					st.children[i].val = st.children[i].Hash().Bytes()
+					st.children[i].key = nil
+					st.children[i].nodeType = hashedNode
+				}
+
+				break
+			}
+
+		}
 		st.children[idx].insert(key, value)
 	case extNode: /* Ext */
 		// Compare both key chunks and see where they differ
@@ -415,13 +428,19 @@ func (st *ReStackTrie) insert(key, value []byte) {
 
 		// Create the two child leaves: the one containing the
 		// original value and the one containing the new value
+		// The child leave will be hashed directly in order to
+		// free up some memory.
 		origIdx := st.key[diffidx]
 		p.children[origIdx] = NewReStackTrie()
 		p.children[origIdx].nodeType = leafNode
 		p.children[origIdx].key = st.key[diffidx+1:]
 		p.children[origIdx].val = st.val
 		p.children[origIdx].keyOffset = p.keyOffset + 1
 
+		p.children[origIdx].val = p.children[origIdx].Hash().Bytes()
+		p.children[origIdx].nodeType = hashedNode
+		p.children[origIdx].key = nil
+
 		newIdx := key[diffidx+st.keyOffset]
 		p.children[newIdx] = NewReStackTrie()
 		p.children[newIdx].nodeType = leafNode
@@ -434,6 +453,8 @@ func (st *ReStackTrie) insert(key, value []byte) {
 		st.nodeType = leafNode
 		st.key = key[st.keyOffset:]
 		st.val = value
+	case hashedNode:
+		panic("trying to insert into hash")
 	default:
 		panic("invalid type")
 	}
@@ -539,18 +560,24 @@ func writeHPRLP(writer io.Writer, key, val []byte, leaf bool) {
 }
 
 func (st *ReStackTrie) Hash() (h common.Hash) {
+	/* Shortcut if node is already hashed */
+	if st.nodeType == hashedNode {
+		return common.BytesToHash(st.val)
+	}
+
 	d := sha3.NewLegacyKeccak256()
 	switch st.nodeType {
-	case 0:
+	case branchNode:
 		payload := [544]byte{}
 		pos := 3 // maximum header length given what we know
-		for _, v := range st.children {
+		for i, v := range st.children {
 			if v != nil {
 				// Write a 32 byte list to the sponge
 				payload[pos] = 0xa0
 				pos++
 				copy(payload[pos:pos+32], v.Hash().Bytes())
 				pos += 32
+				st.children[i] = nil // Reclaim mem from subtree
 			} else {
 				// Write an empty list to the sponge
 				payload[pos] = 0x80
@@ -579,12 +606,13 @@ func (st *ReStackTrie) Hash() (h common.Hash) {
 			start = 0
 		}
 		d.Write(payload[start:pos])
-	case 1:
+	case extNode:
 		ch := st.children[0].Hash().Bytes()
 		writeHPRLP(d, st.key, ch, false)
-	case 2:
+		st.children[0] = nil // Reclaim mem from subtree
+	case leafNode:
 		writeHPRLP(d, st.key, st.val, true)
-	case 3:
+	case emptyNode:
 	default:
 		panic("Invalid node type")
 	}