fix builder so that we handle account deletions properly and properly diff storage when an account is moved to a new path; update params

statediff/builder.go

@@ -22,10 +22,8 @@ package statediff
 import (
 	"bytes"
 	"fmt"
-	"math/big"
 
 	"github.com/ethereum/go-ethereum/common"
-	"github.com/ethereum/go-ethereum/core"
 	"github.com/ethereum/go-ethereum/core/state"
 	"github.com/ethereum/go-ethereum/crypto"
 	"github.com/ethereum/go-ethereum/log"
@@ -37,94 +35,125 @@ var nullNode = common.Hex2Bytes("00000000000000000000000000000000000000000000000
 
 // Builder interface exposes the method for building a state diff between two blocks
 type Builder interface {
-	BuildStateDiff(oldStateRoot, newStateRoot common.Hash, blockNumber *big.Int, blockHash common.Hash) (StateDiff, error)
+	BuildStateDiff(args Args, params Params) (StateDiff, error)
 }
 
 type builder struct {
-	config     Config
-	blockChain *core.BlockChain
 	stateCache state.Database
 }
 
 // NewBuilder is used to create a statediff builder
-func NewBuilder(blockChain *core.BlockChain, config Config) Builder {
+func NewBuilder(stateCache state.Database) Builder {
 	return &builder{
-		config:     config,
+		stateCache: stateCache, // state cache is safe for concurrent reads
-		blockChain: blockChain,
 	}
 }
 
-// BuildStateDiff builds a statediff object from two blocks
+// BuildStateDiff builds a statediff object from two blocks and the provided parameters
-func (sdb *builder) BuildStateDiff(oldStateRoot, newStateRoot common.Hash, blockNumber *big.Int, blockHash common.Hash) (StateDiff, error) {
+func (sdb *builder) BuildStateDiff(args Args, params Params) (StateDiff, error) {
+	if !params.IntermediateStateNodes || len(params.WatchedAddresses) > 0 { // if we are watching only specific accounts then we are only diffing leaf nodes
+		return sdb.buildStateDiffWithoutIntermediateStateNodes(args, params)
+	}
+	return sdb.buildStateDiffWithIntermediateStateNodes(args, params.IntermediateStorageNodes)
+}
+
+func (sdb *builder) buildStateDiffWithIntermediateStateNodes(args Args, intermediateStorageNodes bool) (StateDiff, error) {
 	// Generate tries for old and new states
-	sdb.stateCache = sdb.blockChain.StateCache()
+	oldTrie, err := sdb.stateCache.OpenTrie(args.OldStateRoot)
-	oldTrie, err := sdb.stateCache.OpenTrie(oldStateRoot)
 	if err != nil {
 		return StateDiff{}, fmt.Errorf("error creating trie for oldStateRoot: %v", err)
 	}
-	newTrie, err := sdb.stateCache.OpenTrie(newStateRoot)
+	newTrie, err := sdb.stateCache.OpenTrie(args.NewStateRoot)
 	if err != nil {
 		return StateDiff{}, fmt.Errorf("error creating trie for newStateRoot: %v", err)
 	}
 
-	// Find created accounts
+	createdOrUpdatedIntermediateNodes, diffAccountsAtB, diffPathsAtB, err := sdb.createdAndUpdatedNodes(oldTrie.NodeIterator([]byte{}), newTrie.NodeIterator([]byte{}))
-	creations, err := sdb.collectDiffNodes(oldTrie.NodeIterator([]byte{}), newTrie.NodeIterator([]byte{}))
 	if err != nil {
-		return StateDiff{}, fmt.Errorf("error collecting creation diff nodes: %v", err)
+		return StateDiff{}, fmt.Errorf("error collecting createdAndUpdatedNodes: %v", err)
 	}
 
-	// Find deleted accounts
+	deletedIntermediateNodes, diffAccountsAtA, err := sdb.deletedOrUpdatedNodes(oldTrie.NodeIterator([]byte{}), newTrie.NodeIterator([]byte{}), diffPathsAtB)
-	deletions, err := sdb.collectDiffNodes(newTrie.NodeIterator([]byte{}), oldTrie.NodeIterator([]byte{}))
 	if err != nil {
-		return StateDiff{}, fmt.Errorf("error collecting deletion diff nodes: %v", err)
+		return StateDiff{}, fmt.Errorf("error collecting deletedOrUpdatedNodes: %v", err)
 	}
 
 	// Find all the diffed keys
-	createKeys := sortKeys(creations)
+	createKeys := sortKeys(diffAccountsAtB)
-	deleteKeys := sortKeys(deletions)
+	deleteKeys := sortKeys(diffAccountsAtA)
 	updatedKeys := findIntersection(createKeys, deleteKeys)
 
 	// Build and return the statediff
-	updatedAccounts, err := sdb.buildDiffIncremental(creations, deletions, updatedKeys)
+	updatedAccounts, err := sdb.buildAccountUpdates(diffAccountsAtB, diffAccountsAtA, updatedKeys, intermediateStorageNodes)
 	if err != nil {
 		return StateDiff{}, fmt.Errorf("error building diff for updated accounts: %v", err)
 	}
-	createdAccounts, err := sdb.buildDiffCreations(creations)
+	createdAccounts, err := sdb.buildAccountCreations(diffAccountsAtB, intermediateStorageNodes)
 	if err != nil {
 		return StateDiff{}, fmt.Errorf("error building diff for created accounts: %v", err)
 	}
-	deletedAccounts, err := sdb.buildDiffDeletions(deletions)
+	deletedAccounts, err := sdb.buildAccountDeletions(diffAccountsAtA)
 	if err != nil {
 		return StateDiff{}, fmt.Errorf("error building diff for deleted accounts: %v", err)
 	}
 
 	return StateDiff{
-		BlockNumber:  blockNumber,
+		BlockNumber:       args.BlockNumber,
-		BlockHash:    blockHash,
+		BlockHash:         args.BlockHash,
-		CreatedNodes: createdAccounts,
+		LeafNodes:         append(append(updatedAccounts, createdAccounts...), deletedAccounts...),
-		DeletedNodes: deletedAccounts,
+		IntermediateNodes: append(createdOrUpdatedIntermediateNodes, deletedIntermediateNodes...),
-		UpdatedNodes: updatedAccounts,
 	}, nil
 }
 
-// isWatchedAddress is used to check if a state account corresponds to one of the addresses the builder is configured to watch
+func (sdb *builder) buildStateDiffWithoutIntermediateStateNodes(args Args, params Params) (StateDiff, error) {
-func (sdb *builder) isWatchedAddress(hashKey []byte) bool {
+	// Generate tries for old and new states
-	// If we aren't watching any specific addresses, we are watching everything
+	oldTrie, err := sdb.stateCache.OpenTrie(args.OldStateRoot)
-	if len(sdb.config.WatchedAddresses) == 0 {
+	if err != nil {
-		return true
+		return StateDiff{}, fmt.Errorf("error creating trie for oldStateRoot: %v", err)
 	}
-	for _, addrStr := range sdb.config.WatchedAddresses {
+	newTrie, err := sdb.stateCache.OpenTrie(args.NewStateRoot)
-		addr := common.HexToAddress(addrStr)
+	if err != nil {
-		addrHashKey := crypto.Keccak256(addr[:])
+		return StateDiff{}, fmt.Errorf("error creating trie for newStateRoot: %v", err)
-		if bytes.Equal(addrHashKey, hashKey) {
-			return true
 	}
+
+	diffAccountsAtB, err := sdb.collectDiffAccounts(oldTrie.NodeIterator([]byte{}), newTrie.NodeIterator([]byte{}), params.WatchedAddresses)
+	if err != nil {
+		return StateDiff{}, fmt.Errorf("error collecting createdAndUpdatedNodes: %v", err)
 	}
-	return false
+
+	diffAccountsAtA, err := sdb.collectDiffAccounts(newTrie.NodeIterator([]byte{}), oldTrie.NodeIterator([]byte{}), params.WatchedAddresses)
+	if err != nil {
+		return StateDiff{}, fmt.Errorf("error collecting deletedOrUpdatedNodes: %v", err)
+	}
+
+	// Find all the diffed keys
+	createKeys := sortKeys(diffAccountsAtB)
+	deleteKeys := sortKeys(diffAccountsAtA)
+	updatedKeys := findIntersection(createKeys, deleteKeys)
+
+	// Build and return the statediff
+	updatedAccounts, err := sdb.buildAccountUpdates(diffAccountsAtB, diffAccountsAtA, updatedKeys, params.IntermediateStorageNodes)
+	if err != nil {
+		return StateDiff{}, fmt.Errorf("error building diff for updated accounts: %v", err)
+	}
+	createdAccounts, err := sdb.buildAccountCreations(diffAccountsAtB, params.IntermediateStorageNodes)
+	if err != nil {
+		return StateDiff{}, fmt.Errorf("error building diff for created accounts: %v", err)
+	}
+	deletedAccounts, err := sdb.buildAccountDeletions(diffAccountsAtA)
+	if err != nil {
+		return StateDiff{}, fmt.Errorf("error building diff for deleted accounts: %v", err)
+	}
+
+	return StateDiff{
+		BlockNumber: args.BlockNumber,
+		BlockHash:   args.BlockHash,
+		LeafNodes:   append(append(updatedAccounts, createdAccounts...), deletedAccounts...),
+	}, nil
 }
 
-func (sdb *builder) collectDiffNodes(a, b trie.NodeIterator) (AccountsMap, error) {
+func (sdb *builder) collectDiffAccounts(a, b trie.NodeIterator, watchedAddresses []string) (AccountMap, error) {
-	diffAccounts := make(AccountsMap)
+	diffAcountsAtB := make(AccountMap)
 	it, _ := trie.NewDifferenceIterator(a, b)
 	for it.Next(true) {
 		// skip value nodes
@@ -148,9 +177,10 @@ func (sdb *builder) collectDiffNodes(a, b trie.NodeIterator) (AccountsMap, error
 		if err != nil {
 			return nil, err
 		}
-		nodePathHash := crypto.Keccak256Hash(nodePath)
 		switch ty {
 		case Leaf:
+			// created vs updated is important for leaf nodes since we need to diff their storage
+			// so we need to map all changed accounts at B to their leafkey, since account can change pathes but not leafkey
 			var account state.Account
 			if err := rlp.DecodeBytes(nodeElements[1].([]byte), &account); err != nil {
 				return nil, fmt.Errorf("error decoding account for leaf node at path %x nerror: %v", nodePath, err)
@@ -159,43 +189,208 @@ func (sdb *builder) collectDiffNodes(a, b trie.NodeIterator) (AccountsMap, error
 			valueNodePath := append(nodePath, partialPath...)
 			encodedPath := trie.HexToCompact(valueNodePath)
 			leafKey := encodedPath[1:]
-			if sdb.isWatchedAddress(leafKey) {
+			if isWatchedAddress(watchedAddresses, leafKey) {
-				aw := accountWrapper{
+				diffAcountsAtB[common.Bytes2Hex(leafKey)] = accountWrapper{
 					NodeType:  ty,
 					Path:      nodePath,
 					NodeValue: node,
 					LeafKey:   leafKey,
 					Account:   &account,
 				}
-				diffAccounts[nodePathHash] = aw
 			}
 		case Extension, Branch:
-			if sdb.config.IntermediateNodes {
+			// fall through to next iteration
-				diffAccounts[nodePathHash] = accountWrapper{
-					NodeType:  ty,
-					Path:      nodePath,
-					NodeValue: node,
-				}
-			}
 		default:
 			return nil, fmt.Errorf("unexpected node type %s", ty)
 		}
 	}
-	return diffAccounts, nil
+	return diffAcountsAtB, nil
 }
 
-func (sdb *builder) buildDiffCreations(accounts AccountsMap) ([]StateNode, error) {
+// isWatchedAddress is used to check if a state account corresponds to one of the addresses the builder is configured to watch
-	accountDiffs := make([]StateNode, 0, len(accounts))
+func isWatchedAddress(watchedAddresses []string, hashKey []byte) bool {
+	// If we aren't watching any specific addresses, we are watching everything
+	if len(watchedAddresses) == 0 {
+		return true
+	}
+	for _, addrStr := range watchedAddresses {
+		addr := common.HexToAddress(addrStr)
+		addrHashKey := crypto.Keccak256(addr[:])
+		if bytes.Equal(addrHashKey, hashKey) {
+			return true
+		}
+	}
+	return false
+}
+
+// createdAndUpdatedNodes returns a slice of all the intermediate nodes that are different in b than they are in a
+// it also returns a map of LEAFKEY to all the leaf nodes (accounts) that are different in b than they are in a
+// it also returns a list of the intermediate node paths that were touched and the leaf node PATHs that were touched
+// this function should only be called b = the NodeIterator for the state trie at  a.blockheight + 1
+// make version for leaf nodes only
+func (sdb *builder) createdAndUpdatedNodes(a, b trie.NodeIterator) ([]StateNode, AccountMap, map[string]bool, error) {
+	createdOrUpdatedIntermediateNodes := make([]StateNode, 0)
+	diffPathsAtB := make(map[string]bool)
+	diffAcountsAtB := make(AccountMap)
+	it, _ := trie.NewDifferenceIterator(a, b)
+	for it.Next(true) {
+		// skip value nodes
+		if it.Leaf() {
+			continue
+		}
+		if bytes.Equal(nullNode, it.Hash().Bytes()) {
+			continue
+		}
+		nodePath := make([]byte, len(it.Path()))
+		copy(nodePath, it.Path())
+		node, err := sdb.stateCache.TrieDB().Node(it.Hash())
+		if err != nil {
+			return nil, nil, nil, err
+		}
+		var nodeElements []interface{}
+		if err := rlp.DecodeBytes(node, &nodeElements); err != nil {
+			return nil, nil, nil, err
+		}
+		ty, err := CheckKeyType(nodeElements)
+		if err != nil {
+			return nil, nil, nil, err
+		}
+		switch ty {
+		case Leaf:
+			// created vs updated is important for leaf nodes since we need to diff their storage
+			// so we need to map all changed accounts at B to their leafkey, since account can change pathes but not leafkey
+			var account state.Account
+			if err := rlp.DecodeBytes(nodeElements[1].([]byte), &account); err != nil {
+				return nil, nil, nil, fmt.Errorf("error decoding account for leaf node at path %x nerror: %v", nodePath, err)
+			}
+			partialPath := trie.CompactToHex(nodeElements[0].([]byte))
+			valueNodePath := append(nodePath, partialPath...)
+			encodedPath := trie.HexToCompact(valueNodePath)
+			leafKey := encodedPath[1:]
+			diffAcountsAtB[common.Bytes2Hex(leafKey)] = accountWrapper{
+				NodeType:  ty,
+				Path:      nodePath,
+				NodeValue: node,
+				LeafKey:   leafKey,
+				Account:   &account,
+			}
+		case Extension, Branch:
+			// create a diff for any intermediate node that has changed at b
+			// created vs updated makes no difference for intermediate nodes since we do not need to diff storage
+			createdOrUpdatedIntermediateNodes = append(createdOrUpdatedIntermediateNodes, StateNode{
+				NodeType:  ty,
+				Path:      nodePath,
+				NodeValue: node,
+			})
+		default:
+			return nil, nil, nil, fmt.Errorf("unexpected node type %s", ty)
+		}
+		// add both intermediate and leaf node paths to the list of diffPathsAtB
+		diffPathsAtB[common.Bytes2Hex(nodePath)] = true
+	}
+	return createdOrUpdatedIntermediateNodes, diffAcountsAtB, diffPathsAtB, nil
+}
+
+func (sdb *builder) deletedOrUpdatedNodes(a, b trie.NodeIterator, diffPathsAtB map[string]bool) ([]StateNode, AccountMap, error) {
+	deletedIntermediateNodes := make([]StateNode, 0)
+	diffAccountAtA := make(AccountMap)
+	it, _ := trie.NewDifferenceIterator(b, a)
+	for it.Next(true) {
+		// skip value nodes
+		if it.Leaf() {
+			continue
+		}
+		if bytes.Equal(nullNode, it.Hash().Bytes()) {
+			continue
+		}
+		nodePath := make([]byte, len(it.Path()))
+		copy(nodePath, it.Path())
+		node, err := sdb.stateCache.TrieDB().Node(it.Hash())
+		if err != nil {
+			return nil, nil, err
+		}
+		var nodeElements []interface{}
+		if err := rlp.DecodeBytes(node, &nodeElements); err != nil {
+			return nil, nil, err
+		}
+		ty, err := CheckKeyType(nodeElements)
+		if err != nil {
+			return nil, nil, err
+		}
+		switch ty {
+		case Leaf:
+			// map all different accounts at A to their leafkey
+			var account state.Account
+			if err := rlp.DecodeBytes(nodeElements[1].([]byte), &account); err != nil {
+				return nil, nil, fmt.Errorf("error decoding account for leaf node at path %x nerror: %v", nodePath, err)
+			}
+			partialPath := trie.CompactToHex(nodeElements[0].([]byte))
+			valueNodePath := append(nodePath, partialPath...)
+			encodedPath := trie.HexToCompact(valueNodePath)
+			leafKey := encodedPath[1:]
+			diffAccountAtA[common.Bytes2Hex(leafKey)] = accountWrapper{
+				NodeType:  ty,
+				Path:      nodePath,
+				NodeValue: node,
+				LeafKey:   leafKey,
+				Account:   &account,
+			}
+		case Extension, Branch:
+			// if this nodePath did not show up in diffPathsAtB
+			// that means the node at this path at A was deleted in B
+			if _, ok := diffPathsAtB[common.Bytes2Hex(nodePath)]; !ok {
+				deletedIntermediateNodes = append(deletedIntermediateNodes, StateNode{
+					NodeType:  Removed,
+					Path:      nodePath,
+					NodeValue: []byte{},
+				})
+			}
+		default:
+			return nil, nil, fmt.Errorf("unexpected node type %s", ty)
+		}
+	}
+	return deletedIntermediateNodes, diffAccountAtA, nil
+}
+
+// needs to be called before building account creations and deletions as this mutattes
+// those account maps to remove the accounts which were updated
+func (sdb *builder) buildAccountUpdates(creations, deletions AccountMap, updatedKeys []string, intermediateStorageNodes bool) ([]StateNode, error) {
+	updatedAccounts := make([]StateNode, 0, len(updatedKeys))
 	var err error
-	for _, val := range accounts {
+	for _, key := range updatedKeys {
-		// If account is not nil, we need to process storage diffs
+		createdAcc := creations[key]
+		deletedAcc := deletions[key]
 		var storageDiffs []StorageNode
-		if val.Account != nil {
+		if deletedAcc.Account != nil && createdAcc.Account != nil {
-			storageDiffs, err = sdb.buildStorageNodesEventual(val.Account.Root)
+			oldSR := deletedAcc.Account.Root
+			newSR := createdAcc.Account.Root
+			storageDiffs, err = sdb.buildStorageNodesIncremental(oldSR, newSR, intermediateStorageNodes)
+			if err != nil {
+				return nil, fmt.Errorf("failed building incremental storage diffs for account with leafkey %s\r\nerror: %v", key, err)
+			}
+		}
+		updatedAccounts = append(updatedAccounts, StateNode{
+			NodeType:     createdAcc.NodeType,
+			Path:         createdAcc.Path,
+			NodeValue:    createdAcc.NodeValue,
+			LeafKey:      createdAcc.LeafKey,
+			StorageDiffs: storageDiffs,
+		})
+		delete(creations, key)
+		delete(deletions, key)
+	}
+
+	return updatedAccounts, nil
+}
+
+func (sdb *builder) buildAccountCreations(accounts AccountMap, intermediateStorageNodes bool) ([]StateNode, error) {
+	accountDiffs := make([]StateNode, 0, len(accounts))
+	for _, val := range accounts {
+		// For account creations, any storage node contained is a diff
+		storageDiffs, err := sdb.buildStorageNodesEventual(val.Account.Root, intermediateStorageNodes)
 		if err != nil {
 			return nil, fmt.Errorf("failed building eventual storage diffs for node %x\r\nerror: %v", val.Path, err)
 		}
-		}
 		accountDiffs = append(accountDiffs, StateNode{
 			NodeType:     val.NodeType,
 			Path:         val.Path,
@@ -208,10 +403,10 @@ func (sdb *builder) buildDiffCreations(accounts AccountsMap) ([]StateNode, error
 	return accountDiffs, nil
 }
 
-func (sdb *builder) buildDiffDeletions(accounts AccountsMap) ([]StateNode, error) {
+func (sdb *builder) buildAccountDeletions(accounts AccountMap) ([]StateNode, error) {
 	accountDiffs := make([]StateNode, 0, len(accounts))
 	for _, val := range accounts {
-		// deleted account can not have storage or it would not be deleted
+		// For account deletions, we can not have any storage or the account would not be deleted
 		accountDiffs = append(accountDiffs, StateNode{
 			NodeType:  Removed,
 			Path:      val.Path,
@@ -222,60 +417,24 @@ func (sdb *builder) buildDiffDeletions(accounts AccountsMap) ([]StateNode, error
 	return accountDiffs, nil
 }
 
-func (sdb *builder) buildDiffIncremental(creations AccountsMap, deletions AccountsMap, updatedKeys []string) ([]StateNode, error) {
+func (sdb *builder) buildStorageNodesEventual(sr common.Hash, intermediateNodes bool) ([]StorageNode, error) {
-	updatedAccounts := make([]StateNode, 0, len(updatedKeys))
-	var err error
-	for _, val := range updatedKeys {
-		hashKey := common.HexToHash(val)
-		createdAcc := creations[hashKey]
-		deletedAcc := deletions[hashKey]
-		var storageDiffs []StorageNode
-		if deletedAcc.Account != nil && createdAcc.Account != nil {
-			oldSR := deletedAcc.Account.Root
-			newSR := createdAcc.Account.Root
-			storageDiffs, err = sdb.buildStorageNodesIncremental(oldSR, newSR)
-			if err != nil {
-				return nil, fmt.Errorf("failed building incremental storage diffs for %s\r\nerror: %v", hashKey.Hex(), err)
-			}
-		}
-		updatedAccounts = append(updatedAccounts, StateNode{
-			NodeType:     createdAcc.NodeType,
-			Path:         createdAcc.Path,
-			NodeValue:    createdAcc.NodeValue,
-			LeafKey:      createdAcc.LeafKey,
-			StorageDiffs: storageDiffs,
-		})
-		delete(creations, common.HexToHash(val))
-		delete(deletions, common.HexToHash(val))
-	}
-
-	return updatedAccounts, nil
-}
-
-// if an account is moved to a new path that didn't previously have a node, it appears to have been "created"
-// but it actually has a previous state associated with, a previous storage trie, so this would produce the
-// wrong result... it would produce a node for the entire trie, not just the diffs since the last state
-func (sdb *builder) buildStorageNodesEventual(sr common.Hash) ([]StorageNode, error) {
 	log.Debug("Storage Root For Eventual Diff", "root", sr.Hex())
-	stateCache := sdb.blockChain.StateCache()
+	sTrie, err := sdb.stateCache.OpenTrie(sr)
-	sTrie, err := stateCache.OpenTrie(sr)
 	if err != nil {
 		log.Info("error in build storage diff eventual", "error", err)
 		return nil, err
 	}
 	it := sTrie.NodeIterator(make([]byte, 0))
-	return sdb.buildStorageNodesFromTrie(it)
+	return sdb.buildStorageNodesFromTrie(it, intermediateNodes)
 }
 
-func (sdb *builder) buildStorageNodesIncremental(oldSR common.Hash, newSR common.Hash) ([]StorageNode, error) {
+func (sdb *builder) buildStorageNodesIncremental(oldSR common.Hash, newSR common.Hash, intermediateNodes bool) ([]StorageNode, error) {
 	log.Debug("Storage Roots for Incremental Diff", "old", oldSR.Hex(), "new", newSR.Hex())
-	stateCache := sdb.blockChain.StateCache()
+	oldTrie, err := sdb.stateCache.OpenTrie(oldSR)
-
-	oldTrie, err := stateCache.OpenTrie(oldSR)
 	if err != nil {
 		return nil, err
 	}
-	newTrie, err := stateCache.OpenTrie(newSR)
+	newTrie, err := sdb.stateCache.OpenTrie(newSR)
 	if err != nil {
 		return nil, err
 	}
@@ -283,10 +442,10 @@ func (sdb *builder) buildStorageNodesIncremental(oldSR common.Hash, newSR common
 	oldIt := oldTrie.NodeIterator(make([]byte, 0))
 	newIt := newTrie.NodeIterator(make([]byte, 0))
 	it, _ := trie.NewDifferenceIterator(oldIt, newIt)
-	return sdb.buildStorageNodesFromTrie(it)
+	return sdb.buildStorageNodesFromTrie(it, intermediateNodes)
 }
 
-func (sdb *builder) buildStorageNodesFromTrie(it trie.NodeIterator) ([]StorageNode, error) {
+func (sdb *builder) buildStorageNodesFromTrie(it trie.NodeIterator, intermediateNodes bool) ([]StorageNode, error) {
 	storageDiffs := make([]StorageNode, 0)
 	for it.Next(true) {
 		// skip value nodes
@@ -316,15 +475,14 @@ func (sdb *builder) buildStorageNodesFromTrie(it trie.NodeIterator) ([]StorageNo
 			valueNodePath := append(nodePath, partialPath...)
 			encodedPath := trie.HexToCompact(valueNodePath)
 			leafKey := encodedPath[1:]
-			sd := StorageNode{
+			storageDiffs = append(storageDiffs, StorageNode{
 				NodeType:  ty,
 				Path:      nodePath,
 				NodeValue: node,
 				LeafKey:   leafKey,
-			}
+			})
-			storageDiffs = append(storageDiffs, sd)
 		case Extension, Branch:
-			if sdb.config.IntermediateNodes {
+			if intermediateNodes {
 				storageDiffs = append(storageDiffs, StorageNode{
 					NodeType:  ty,
 					Path:      nodePath,

statediff/config.go

@@ -1,24 +0,0 @@
-// Copyright 2019 The go-ethereum Authors
-// This file is part of the go-ethereum library.
-//
-// The go-ethereum library is free software: you can redistribute it and/or modify
-// it under the terms of the GNU Lesser General Public License as published by
-// the Free Software Foundation, either version 3 of the License, or
-// (at your option) any later version.
-//
-// The go-ethereum library is distributed in the hope that it will be useful,
-// but WITHOUT ANY WARRANTY; without even the implied warranty of
-// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
-// GNU Lesser General Public License for more details.
-//
-// You should have received a copy of the GNU Lesser General Public License
-// along with the go-ethereum library. If not, see <http://www.gnu.org/licenses/>.
-
-package statediff
-
-// Config is used to carry in parameters from CLI configuration
-type Config struct {
-	IntermediateNodes bool
-	StreamBlock       bool
-	WatchedAddresses  []string
-}

statediff/helpers.go

@@ -25,10 +25,10 @@ import (
 	"strings"
 )
 
-func sortKeys(data AccountsMap) []string {
+func sortKeys(data AccountMap) []string {
 	keys := make([]string, 0, len(data))
 	for key := range data {
-		keys = append(keys, key.Hex())
+		keys = append(keys, key)
 	}
 	sort.Strings(keys)
 
@@ -79,7 +79,7 @@ func CheckKeyType(elements []interface{}) (NodeType, error) {
 		return Branch, nil
 	}
 	if len(elements) < 2 {
-		return Unknown, nil
+		return Unknown, fmt.Errorf("node cannot be less than two elements in length")
 	}
 	switch elements[0].([]byte)[0] / 16 {
 	case '\x00':

statediff/types.go

@@ -34,6 +34,22 @@ type Subscription struct {
 	QuitChan    chan<- bool
 }
 
+// Params is used to carry in parameters from subscribing/requesting clients configuration
+type Params struct {
+	IntermediateStateNodes   bool
+	IntermediateStorageNodes bool
+	IncludeBlock             bool
+	IncludeReceipts          bool
+	IncludeTD                bool
+	WatchedAddresses         []string
+}
+
+// Args bundles the arguments for the state diff builder
+type Args struct {
+	OldStateRoot, NewStateRoot, BlockHash common.Hash
+	BlockNumber                           *big.Int
+}
+
 // Payload packages the data to send to statediff subscriptions
 type Payload struct {
 	BlockRlp        []byte   `json:"blockRlp"`
@@ -67,9 +83,8 @@ func (sd *Payload) Encode() ([]byte, error) {
 type StateDiff struct {
 	BlockNumber       *big.Int    `json:"blockNumber"     gencodec:"required"`
 	BlockHash         common.Hash `json:"blockHash"       gencodec:"required"`
-	CreatedNodes []StateNode `json:"createdAccounts" gencodec:"required"`
+	LeafNodes         []StateNode `json:"leafNodes" gencodec:"required"`
-	DeletedNodes []StateNode `json:"deletedAccounts" gencodec:"required"`
+	IntermediateNodes []StateNode `json:"intermediateNodes" gencodec:"required"`
-	UpdatedNodes []StateNode `json:"updatedAccounts" gencodec:"required"`
 
 	encoded []byte
 	err     error
@@ -92,10 +107,10 @@ type StorageNode struct {
 	LeafKey   []byte   `json:"leafKey"`
 }
 
-// AccountsMap is a mapping of keccak256(address) => accountWrapper
+// AccountMap is a mapping of hex encoded path => account wrapper
-type AccountsMap map[common.Hash]accountWrapper
+type AccountMap map[string]accountWrapper
 
-// AccountWrapper is used to temporary associate the unpacked account with its raw values
+// accountWrapper is used to temporary associate the unpacked node with its raw values
 type accountWrapper struct {
 	Account   *state.Account
 	NodeType  NodeType