cmd/geth/config.go

@@ -155,8 +155,14 @@ func makeFullNode(ctx *cli.Context) *node.Node {
 	if ctx.GlobalIsSet(utils.ConstantinopleOverrideFlag.Name) {
 		cfg.Eth.ConstantinopleOverride = new(big.Int).SetUint64(ctx.GlobalUint64(utils.ConstantinopleOverrideFlag.Name))
 	}
+
 	utils.RegisterEthService(stack, &cfg.Eth)
 
+	if ctx.GlobalBool(utils.StateDiffFlag.Name) {
+		cfg.Eth.StateDiff = true
+		utils.RegisterStateDiffService(stack, ctx)
+	}
+
 	if ctx.GlobalBool(utils.DashboardEnabledFlag.Name) {
 		utils.RegisterDashboardService(stack, &cfg.Dashboard, gitCommit)
 	}
@@ -181,9 +187,6 @@ func makeFullNode(ctx *cli.Context) *node.Node {
 		utils.RegisterEthStatsService(stack, cfg.Ethstats.URL)
 	}
 
-	if ctx.GlobalBool(utils.StateDiffFlag.Name) {
-		utils.RegisterStateDiffService(stack, ctx)
-	}
 	return stack
 }
 

core/blockchain.go

@@ -71,10 +71,11 @@ const (
 // CacheConfig contains the configuration values for the trie caching/pruning
 // that's resident in a blockchain.
 type CacheConfig struct {
-	Disabled       bool          // Whether to disable trie write caching (archive node)
+	Disabled             bool          // Whether to disable trie write caching (archive node)
-	TrieCleanLimit int           // Memory allowance (MB) to use for caching trie nodes in memory
+	TrieCleanLimit       int           // Memory allowance (MB) to use for caching trie nodes in memory
-	TrieDirtyLimit int           // Memory limit (MB) at which to start flushing dirty trie nodes to disk
+	TrieDirtyLimit       int           // Memory limit (MB) at which to start flushing dirty trie nodes to disk
-	TrieTimeLimit  time.Duration // Time limit after which to flush the current in-memory trie to disk
+	TrieTimeLimit        time.Duration // Time limit after which to flush the current in-memory trie to disk
+	ProcessingStateDiffs bool          // Whether statediffs processing should be taken into a account before a trie is pruned
 }
 
 // BlockChain represents the canonical chain given a database with a genesis
@@ -136,6 +137,8 @@ type BlockChain struct {
 
 	badBlocks      *lru.Cache              // Bad block cache
 	shouldPreserve func(*types.Block) bool // Function used to determine whether should preserve the given block.
+
+	stateDiffsProcessed map[common.Hash]int
 }
 
 // NewBlockChain returns a fully initialised block chain using information
@@ -155,24 +158,26 @@ func NewBlockChain(db ethdb.Database, cacheConfig *CacheConfig, chainConfig *par
 	blockCache, _ := lru.New(blockCacheLimit)
 	futureBlocks, _ := lru.New(maxFutureBlocks)
 	badBlocks, _ := lru.New(badBlockLimit)
-
+	stateDiffsProcessed := make(map[common.Hash]int)
 	bc := &BlockChain{
-		chainConfig:    chainConfig,
+		chainConfig:         chainConfig,
-		cacheConfig:    cacheConfig,
+		cacheConfig:         cacheConfig,
-		db:             db,
+		db:                  db,
-		triegc:         prque.New(nil),
+		triegc:              prque.New(nil),
-		stateCache:     state.NewDatabaseWithCache(db, cacheConfig.TrieCleanLimit),
+		stateCache:          state.NewDatabaseWithCache(db, cacheConfig.TrieCleanLimit),
-		quit:           make(chan struct{}),
+		quit:                make(chan struct{}),
-		shouldPreserve: shouldPreserve,
+		shouldPreserve:      shouldPreserve,
-		bodyCache:      bodyCache,
+		bodyCache:           bodyCache,
-		bodyRLPCache:   bodyRLPCache,
+		bodyRLPCache:        bodyRLPCache,
-		receiptsCache:  receiptsCache,
+		receiptsCache:       receiptsCache,
-		blockCache:     blockCache,
+		blockCache:          blockCache,
-		futureBlocks:   futureBlocks,
+		futureBlocks:        futureBlocks,
-		engine:         engine,
+		engine:              engine,
-		vmConfig:       vmConfig,
+		vmConfig:            vmConfig,
-		badBlocks:      badBlocks,
+		badBlocks:           badBlocks,
+		stateDiffsProcessed: stateDiffsProcessed,
 	}
+
 	bc.SetValidator(NewBlockValidator(chainConfig, bc, engine))
 	bc.SetProcessor(NewStateProcessor(chainConfig, bc, engine))
 
@@ -922,6 +927,11 @@ func (bc *BlockChain) WriteBlockWithoutState(block *types.Block, td *big.Int) (e
 	return nil
 }
 
+func (bc *BlockChain) AddToStateDiffProcessedCollection(hash common.Hash) {
+	count := bc.stateDiffsProcessed[hash]
+	bc.stateDiffsProcessed[hash] = count + 1
+}
+
 // WriteBlockWithState writes the block and all associated state to the database.
 func (bc *BlockChain) WriteBlockWithState(block *types.Block, receipts []*types.Receipt, state *state.StateDB) (status WriteStatus, err error) {
 	bc.wg.Add(1)
@@ -994,6 +1004,16 @@ func (bc *BlockChain) WriteBlockWithState(block *types.Block, receipts []*types.
 					bc.triegc.Push(root, number)
 					break
 				}
+
+				if bc.cacheConfig.ProcessingStateDiffs {
+					if !bc.allowedRootToBeDereferenced(root.(common.Hash)) {
+						bc.triegc.Push(root, number)
+						break
+					} else {
+						delete(bc.stateDiffsProcessed, root.(common.Hash))
+					}
+				}
+
 				triedb.Dereference(root.(common.Hash))
 			}
 		}
@@ -1048,6 +1068,15 @@ func (bc *BlockChain) WriteBlockWithState(block *types.Block, receipts []*types.
 	return status, nil
 }
 
+// since we need the state tries of the current block and its parent in-memory
+// in order to process statediffs, we should avoid dereferencing roots until
+// its statediff and its child have been processed
+func (bc *BlockChain) allowedRootToBeDereferenced(root common.Hash) bool {
+	diffProcessedForSelfAndChildCount := 2
+	count := bc.stateDiffsProcessed[root]
+	return count >= diffProcessedForSelfAndChildCount
+}
+
 // addFutureBlock checks if the block is within the max allowed window to get
 // accepted for future processing, and returns an error if the block is too far
 // ahead and was not added.

core/blockchain_test.go

@@ -1483,3 +1483,84 @@ func BenchmarkBlockChain_1x1000Executions(b *testing.B) {
 
 	benchmarkLargeNumberOfValueToNonexisting(b, numTxs, numBlocks, recipientFn, dataFn)
 }
+
+func TestProcessingStateDiffs(t *testing.T) {
+	defaultTrieCleanCache := 256
+	defaultTrieDirtyCache := 256
+	defaultTrieTimeout := 60 * time.Minute
+	cacheConfig := &CacheConfig{
+		Disabled:             false,
+		TrieCleanLimit:       defaultTrieCleanCache,
+		TrieDirtyLimit:       defaultTrieDirtyCache,
+		TrieTimeLimit:        defaultTrieTimeout,
+		ProcessingStateDiffs: true,
+	}
+	db := ethdb.NewMemDatabase()
+	genesis := new(Genesis).MustCommit(db)
+	numberOfBlocks := triesInMemory
+	engine := ethash.NewFaker()
+	blockchain, _ := NewBlockChain(db, cacheConfig, params.AllEthashProtocolChanges, engine, vm.Config{}, nil)
+	blocks := makeBlockChain(genesis, numberOfBlocks+1, engine, db, canonicalSeed)
+	_, err := blockchain.InsertChain(blocks)
+	if err != nil {
+		t.Fatalf("failed to create pristine chain: %v", err)
+	}
+	defer blockchain.Stop()
+
+	//when adding a root hash to the collection, it will increment the count
+	firstStateRoot := blocks[0].Root()
+	blockchain.AddToStateDiffProcessedCollection(firstStateRoot)
+	value, ok := blockchain.stateDiffsProcessed[firstStateRoot]
+	if !ok {
+		t.Error("state root not found in collection")
+	}
+	if value != 1 {
+		t.Error("state root count not correct", "want", 1, "got", value)
+	}
+
+	blockchain.AddToStateDiffProcessedCollection(firstStateRoot)
+	value, ok = blockchain.stateDiffsProcessed[firstStateRoot]
+	if !ok {
+		t.Error("state root not found in collection")
+	}
+	if value != 2 {
+		t.Error("state root count not correct", "want", 2, "got", value)
+	}
+
+	moreBlocks := makeBlockChain(blocks[len(blocks)-1], 1, engine, db, canonicalSeed)
+	_, err = blockchain.InsertChain(moreBlocks)
+
+	//a root hash can be dereferenced when it's state diff and it's child's state diff have been processed
+	//(i.e. it has a count of 2 in stateDiffsProcessed)
+	nodes := blockchain.stateCache.TrieDB().Nodes()
+	if containsRootHash(nodes, firstStateRoot) {
+		t.Errorf("stateRoot %s in nodes, want: %t, got: %t", firstStateRoot.Hex(), false, true)
+	}
+
+	//a root hash should still be in the in-mem db if it's child's state diff hasn't yet been processed
+	//(i.e. it has a count of 1 stateDiffsProcessed)
+	secondStateRoot := blocks[1].Root()
+	blockchain.AddToStateDiffProcessedCollection(secondStateRoot)
+	if !containsRootHash(nodes, secondStateRoot) {
+		t.Errorf("stateRoot %s in nodes, want: %t, got: %t", secondStateRoot.Hex(), true, false)
+	}
+
+	//the stateDiffsProcessed collection is cleaned up once a hash has been dereferenced
+	_, ok = blockchain.stateDiffsProcessed[firstStateRoot]
+	if ok {
+		t.Errorf("stateRoot %s in stateDiffsProcessed collection, want: %t, got: %t",
+			firstStateRoot.Hex(),
+			false,
+			ok,
+		)
+	}
+}
+
+func containsRootHash(collection []common.Hash, hash common.Hash) bool {
+	for _, n := range collection {
+		if n == hash {
+			return true
+		}
+	}
+	return false
+}

eth/backend.go

@@ -157,10 +157,11 @@ func New(ctx *node.ServiceContext, config *Config) (*Ethereum, error) {
 			EVMInterpreter:          config.EVMInterpreter,
 		}
 		cacheConfig = &core.CacheConfig{
-			Disabled:       config.NoPruning,
+			Disabled:             config.NoPruning,
-			TrieCleanLimit: config.TrieCleanCache,
+			TrieCleanLimit:       config.TrieCleanCache,
-			TrieDirtyLimit: config.TrieDirtyCache,
+			TrieDirtyLimit:       config.TrieDirtyCache,
-			TrieTimeLimit:  config.TrieTimeout,
+			TrieTimeLimit:        config.TrieTimeout,
+			ProcessingStateDiffs: config.StateDiff,
 		}
 	)
 	eth.blockchain, err = core.NewBlockChain(chainDb, cacheConfig, eth.chainConfig, eth.engine, vmConfig, eth.shouldPreserve)

eth/config.go

@@ -59,6 +59,8 @@ var DefaultConfig = Config{
 		Blocks:     20,
 		Percentile: 60,
 	},
+
+	StateDiff: false,
 }
 
 func init() {
@@ -135,6 +137,8 @@ type Config struct {
 
 	// Constantinople block override (TODO: remove after the fork)
 	ConstantinopleOverride *big.Int
+
+	StateDiff bool
 }
 
 type configMarshaling struct {

statediff/builder/builder.go

@@ -22,6 +22,7 @@ package builder
 import (
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/common/hexutil"
+	"github.com/ethereum/go-ethereum/core"
 	"github.com/ethereum/go-ethereum/core/state"
 	"github.com/ethereum/go-ethereum/ethdb"
 	"github.com/ethereum/go-ethereum/log"
@@ -35,25 +36,27 @@ type Builder interface {
 
 type builder struct {
 	chainDB    ethdb.Database
-	trieDB     *trie.Database
+	blockChain *core.BlockChain
-	cachedTrie *trie.Trie
 }
 
-func NewBuilder(db ethdb.Database) *builder {
+type AccountsMap map[common.Hash]*state.Account
+
+func NewBuilder(db ethdb.Database, blockChain *core.BlockChain) *builder {
 	return &builder{
-		chainDB: db,
+		chainDB:    db,
-		trieDB:  trie.NewDatabase(db),
+		blockChain: blockChain,
 	}
 }
 
 func (sdb *builder) BuildStateDiff(oldStateRoot, newStateRoot common.Hash, blockNumber int64, blockHash common.Hash) (*StateDiff, error) {
 	// Generate tries for old and new states
-	oldTrie, err := trie.New(oldStateRoot, sdb.trieDB)
+	stateCache := sdb.blockChain.StateCache()
+	oldTrie, err := stateCache.OpenTrie(oldStateRoot)
 	if err != nil {
 		log.Error("Error creating trie for oldStateRoot", "error", err)
 		return nil, err
 	}
-	newTrie, err := trie.New(newStateRoot, sdb.trieDB)
+	newTrie, err := stateCache.OpenTrie(newStateRoot)
 	if err != nil {
 		log.Error("Error creating trie for newStateRoot", "error", err)
 		return nil, err
@@ -108,33 +111,27 @@ func (sdb *builder) BuildStateDiff(oldStateRoot, newStateRoot common.Hash, block
 	}, nil
 }
 
-func (sdb *builder) collectDiffNodes(a, b trie.NodeIterator) (map[common.Address]*state.Account, error) {
+func (sdb *builder) collectDiffNodes(a, b trie.NodeIterator) (AccountsMap, error) {
-	var diffAccounts = make(map[common.Address]*state.Account)
+	var diffAccounts = make(AccountsMap)
 	it, _ := trie.NewDifferenceIterator(a, b)
 
 	for {
 		log.Debug("Current Path and Hash", "path", pathToStr(it), "hashold", it.Hash())
 		if it.Leaf() {
-
+			leafKey := make([]byte, len(it.LeafKey()))
-			// lookup address
+			copy(leafKey, it.LeafKey())
-			path := make([]byte, len(it.Path())-1)
+			leafKeyHash := common.BytesToHash(leafKey)
-			copy(path, it.Path())
-			addr, err := sdb.addressByPath(path)
-			if err != nil {
-				log.Error("Error looking up address via path", "path", path, "error", err)
-				return nil, err
-			}
 
 			// lookup account state
 			var account state.Account
 			if err := rlp.DecodeBytes(it.LeafBlob(), &account); err != nil {
-				log.Error("Error looking up account via address", "address", addr, "error", err)
+				log.Error("Error looking up account via address", "address", leafKeyHash, "error", err)
 				return nil, err
 			}
 
 			// record account to diffs (creation if we are looking at new - old; deletion if old - new)
-			log.Debug("Account lookup successful", "address", addr, "account", account)
+			log.Debug("Account lookup successful", "address", leafKeyHash, "account", account)
-			diffAccounts[*addr] = &account
+			diffAccounts[leafKeyHash] = &account
 		}
 		cont := it.Next(true)
 		if !cont {
@@ -145,8 +142,8 @@ func (sdb *builder) collectDiffNodes(a, b trie.NodeIterator) (map[common.Address
 	return diffAccounts, nil
 }
 
-func (sdb *builder) buildDiffEventual(accounts map[common.Address]*state.Account) (map[common.Address]AccountDiff, error) {
+func (sdb *builder) buildDiffEventual(accounts AccountsMap) (AccountDiffsMap, error) {
-	accountDiffs := make(map[common.Address]AccountDiff)
+	accountDiffs := make(AccountDiffsMap)
 	for addr, val := range accounts {
 		sr := val.Root
 		storageDiffs, err := sdb.buildStorageDiffsEventual(sr)
@@ -172,11 +169,11 @@ func (sdb *builder) buildDiffEventual(accounts map[common.Address]*state.Account
 	return accountDiffs, nil
 }
 
-func (sdb *builder) buildDiffIncremental(creations map[common.Address]*state.Account, deletions map[common.Address]*state.Account, updatedKeys []string) (map[common.Address]AccountDiff, error) {
+func (sdb *builder) buildDiffIncremental(creations AccountsMap, deletions AccountsMap, updatedKeys []string) (AccountDiffsMap, error) {
-	updatedAccounts := make(map[common.Address]AccountDiff)
+	updatedAccounts := make(AccountDiffsMap)
 	for _, val := range updatedKeys {
-		createdAcc := creations[common.HexToAddress(val)]
+		createdAcc := creations[common.HexToHash(val)]
-		deletedAcc := deletions[common.HexToAddress(val)]
+		deletedAcc := deletions[common.HexToHash(val)]
 		oldSR := deletedAcc.Root
 		newSR := createdAcc.Root
 		if storageDiffs, err := sdb.buildStorageDiffsIncremental(oldSR, newSR); err != nil {
@@ -190,15 +187,15 @@ func (sdb *builder) buildDiffIncremental(creations map[common.Address]*state.Acc
 			nHexRoot := createdAcc.Root.Hex()
 			contractRoot := DiffString{Value: &nHexRoot}
 
-			updatedAccounts[common.HexToAddress(val)] = AccountDiff{
+			updatedAccounts[common.HexToHash(val)] = AccountDiff{
 				Nonce:        nonce,
 				Balance:      balance,
 				CodeHash:     codeHash,
 				ContractRoot: contractRoot,
 				Storage:      storageDiffs,
 			}
-			delete(creations, common.HexToAddress(val))
+			delete(creations, common.HexToHash(val))
-			delete(deletions, common.HexToAddress(val))
+			delete(deletions, common.HexToHash(val))
 		}
 	}
 	return updatedAccounts, nil
@@ -206,7 +203,8 @@ func (sdb *builder) buildDiffIncremental(creations map[common.Address]*state.Acc
 
 func (sdb *builder) buildStorageDiffsEventual(sr common.Hash) (map[string]DiffStorage, error) {
 	log.Debug("Storage Root For Eventual Diff", "root", sr.Hex())
-	sTrie, err := trie.New(sr, sdb.trieDB)
+	stateCache := sdb.blockChain.StateCache()
+	sTrie, err := stateCache.OpenTrie(sr)
 	if err != nil {
 		log.Info("error in build storage diff eventual", "error", err)
 		return nil, err
@@ -218,11 +216,13 @@ func (sdb *builder) buildStorageDiffsEventual(sr common.Hash) (map[string]DiffSt
 
 func (sdb *builder) buildStorageDiffsIncremental(oldSR common.Hash, newSR common.Hash) (map[string]DiffStorage, error) {
 	log.Debug("Storage Roots for Incremental Diff", "old", oldSR.Hex(), "new", newSR.Hex())
-	oldTrie, err := trie.New(oldSR, sdb.trieDB)
+	stateCache := sdb.blockChain.StateCache()
+
+	oldTrie, err := stateCache.OpenTrie(oldSR)
 	if err != nil {
 		return nil, err
 	}
-	newTrie, err := trie.New(newSR, sdb.trieDB)
+	newTrie, err := stateCache.OpenTrie(newSR)
 	if err != nil {
 		return nil, err
 	}

statediff/builder/builder_test.go

@@ -10,10 +10,12 @@ import (
 	"github.com/ethereum/go-ethereum/consensus/ethash"
 	"github.com/ethereum/go-ethereum/core"
 	"github.com/ethereum/go-ethereum/core/types"
+	"github.com/ethereum/go-ethereum/core/vm"
 	"github.com/ethereum/go-ethereum/crypto"
 	"github.com/ethereum/go-ethereum/ethdb"
 	"github.com/ethereum/go-ethereum/params"
 	b "github.com/ethereum/go-ethereum/statediff/builder"
+	"github.com/ethereum/go-ethereum/statediff/testhelpers"
 )
 
 var (
@@ -21,26 +23,33 @@ var (
 
 	testBankKey, _  = crypto.HexToECDSA("b71c71a67e1177ad4e901695e1b4b9ee17ae16c6668d313eac2f96dbcda3f291")
 	testBankAddress = crypto.PubkeyToAddress(testBankKey.PublicKey) //0x71562b71999873DB5b286dF957af199Ec94617F7
+	bankLeafKey     = testhelpers.AddressToLeafKey(testBankAddress)
 	testBankFunds   = big.NewInt(100000000)
 	genesis         = core.GenesisBlockForTesting(testdb, testBankAddress, testBankFunds)
 
 	account1Key, _                                 = crypto.HexToECDSA("8a1f9a8f95be41cd7ccb6168179afb4504aefe388d1e14474d32c45c72ce7b7a")
 	account2Key, _                                 = crypto.HexToECDSA("49a7b37aa6f6645917e7b807e9d1c00d4fa71f18343b0d4122a4d2df64dd6fee")
 	account1Addr                                   = crypto.PubkeyToAddress(account1Key.PublicKey) //0x703c4b2bD70c169f5717101CaeE543299Fc946C7
+	account1LeafKey                                = testhelpers.AddressToLeafKey(account1Addr)
 	account2Addr                                   = crypto.PubkeyToAddress(account2Key.PublicKey) //0x0D3ab14BBaD3D99F4203bd7a11aCB94882050E7e
+	account2LeafKey                                = testhelpers.AddressToLeafKey(account2Addr)
 	contractCode                                   = common.Hex2Bytes("608060405234801561001057600080fd5b50602060405190810160405280600160ff16815250600090600161003592919061003b565b506100a5565b826064810192821561006f579160200282015b8281111561006e578251829060ff1690559160200191906001019061004e565b5b50905061007c9190610080565b5090565b6100a291905b8082111561009e576000816000905550600101610086565b5090565b90565b610124806100b46000396000f3fe6080604052348015600f57600080fd5b5060043610604f576000357c01000000000000000000000000000000000000000000000000000000009004806360cd2685146054578063c16431b9146093575b600080fd5b607d60048036036020811015606857600080fd5b810190808035906020019092919050505060c8565b6040518082815260200191505060405180910390f35b60c66004803603604081101560a757600080fd5b81019080803590602001909291908035906020019092919050505060e0565b005b6000808260648110151560d757fe5b01549050919050565b8060008360648110151560ef57fe5b0181905550505056fea165627a7a7230582064e918c3140a117bf3aa65865a9b9e83fae21ad1720506e7933b2a9f54bb40260029")
 	contractAddr                                   common.Address
-	emptyAccountDiffEventualMap                    = make(map[common.Address]b.AccountDiff)
+	contractLeafKey                                common.Hash
-	emptyAccountDiffIncrementalMap                 = make(map[common.Address]b.AccountDiff)
+	emptyAccountDiffEventualMap                    = make(b.AccountDiffsMap)
+	emptyAccountDiffIncrementalMap                 = make(b.AccountDiffsMap)
 	block0Hash, block1Hash, block2Hash, block3Hash common.Hash
 	block0, block1, block2, block3                 *types.Block
 	builder                                        b.Builder
 	miningReward                                   = int64(2000000000000000000)
 	burnAddress                                    = common.HexToAddress("0x0")
+	burnLeafKey                                    = testhelpers.AddressToLeafKey(burnAddress)
 )
 
 func TestBuilder(t *testing.T) {
-	_, blockMap := makeChain(3, genesis)
+	_, blockMap, chain := makeChain(3, genesis)
+	contractLeafKey = testhelpers.AddressToLeafKey(contractAddr)
+	defer chain.Stop()
 	block0Hash = common.HexToHash("0xd1721cfd0b29c36fd7a68f25c128e86413fb666a6e1d68e89b875bd299262661")
 	block1Hash = common.HexToHash("0xbbe88de60ba33a3f18c0caa37d827bfb70252e19e40a07cd34041696c35ecb1a")
 	block2Hash = common.HexToHash("0xde75663f36a8497b4bdda2a4b52bd9540b705a2728c7391c59b8cb2cde5a2feb")
@@ -50,7 +59,7 @@ func TestBuilder(t *testing.T) {
 	block1 = blockMap[block1Hash]
 	block2 = blockMap[block2Hash]
 	block3 = blockMap[block3Hash]
-	builder = b.NewBuilder(testdb)
+	builder = b.NewBuilder(testdb, chain)
 
 	type arguments struct {
 		oldStateRoot common.Hash
@@ -113,15 +122,15 @@ func TestBuilder(t *testing.T) {
 			&b.StateDiff{
 				BlockNumber: block1.Number().Int64(),
 				BlockHash:   block1.Hash(),
-				CreatedAccounts: map[common.Address]b.AccountDiff{
+				CreatedAccounts: b.AccountDiffsMap{
-					account1Addr: {
+					account1LeafKey: {
 						Nonce:        b.DiffUint64{Value: &nonce0},
 						Balance:      b.DiffBigInt{Value: big.NewInt(balanceChange10000)},
 						CodeHash:     "0xc5d2460186f7233c927e7db2dcc703c0e500b653ca82273b7bfad8045d85a470",
 						ContractRoot: b.DiffString{Value: &originalContractRoot},
 						Storage:      map[string]b.DiffStorage{},
 					},
-					burnAddress: {
+					burnLeafKey: {
 						Nonce:        b.DiffUint64{Value: &nonce0},
 						Balance:      b.DiffBigInt{Value: big.NewInt(miningReward)},
 						CodeHash:     "0xc5d2460186f7233c927e7db2dcc703c0e500b653ca82273b7bfad8045d85a470",
@@ -130,8 +139,8 @@ func TestBuilder(t *testing.T) {
 					},
 				},
 				DeletedAccounts: emptyAccountDiffEventualMap,
-				UpdatedAccounts: map[common.Address]b.AccountDiff{
+				UpdatedAccounts: b.AccountDiffsMap{
-					testBankAddress: {
+					bankLeafKey: {
 						Nonce:        b.DiffUint64{Value: &nonce1},
 						Balance:      b.DiffBigInt{Value: big.NewInt(testBankFunds.Int64() - balanceChange10000)},
 						CodeHash:     "0xc5d2460186f7233c927e7db2dcc703c0e500b653ca82273b7bfad8045d85a470",
@@ -154,15 +163,15 @@ func TestBuilder(t *testing.T) {
 			&b.StateDiff{
 				BlockNumber: block2.Number().Int64(),
 				BlockHash:   block2.Hash(),
-				CreatedAccounts: map[common.Address]b.AccountDiff{
+				CreatedAccounts: b.AccountDiffsMap{
-					account2Addr: {
+					account2LeafKey: {
 						Nonce:        b.DiffUint64{Value: &nonce0},
 						Balance:      b.DiffBigInt{Value: big.NewInt(balanceChange1000)},
 						CodeHash:     "0xc5d2460186f7233c927e7db2dcc703c0e500b653ca82273b7bfad8045d85a470",
 						ContractRoot: b.DiffString{Value: &originalContractRoot},
 						Storage:      map[string]b.DiffStorage{},
 					},
-					contractAddr: {
+					contractLeafKey: {
 						Nonce:        b.DiffUint64{Value: &nonce1},
 						Balance:      b.DiffBigInt{Value: big.NewInt(0)},
 						CodeHash:     "0x753f98a8d4328b15636e46f66f2cb4bc860100aa17967cc145fcd17d1d4710ea",
@@ -175,22 +184,22 @@ func TestBuilder(t *testing.T) {
 					},
 				},
 				DeletedAccounts: emptyAccountDiffEventualMap,
-				UpdatedAccounts: map[common.Address]b.AccountDiff{
+				UpdatedAccounts: b.AccountDiffsMap{
-					testBankAddress: {
+					bankLeafKey: {
 						Nonce:        b.DiffUint64{Value: &nonce2},
 						Balance:      b.DiffBigInt{Value: big.NewInt(block1BankBalance - balanceChange1000)},
 						CodeHash:     "0xc5d2460186f7233c927e7db2dcc703c0e500b653ca82273b7bfad8045d85a470",
 						ContractRoot: b.DiffString{Value: &originalContractRoot},
 						Storage:      map[string]b.DiffStorage{},
 					},
-					account1Addr: {
+					account1LeafKey: {
 						Nonce:        b.DiffUint64{Value: &nonce2},
 						Balance:      b.DiffBigInt{Value: big.NewInt(block1Account1Balance - balanceChange1000 + balanceChange1000)},
 						CodeHash:     "0xc5d2460186f7233c927e7db2dcc703c0e500b653ca82273b7bfad8045d85a470",
 						ContractRoot: b.DiffString{Value: &originalContractRoot},
 						Storage:      map[string]b.DiffStorage{},
 					},
-					burnAddress: {
+					burnLeafKey: {
 						Nonce:        b.DiffUint64{Value: &nonce0},
 						Balance:      b.DiffBigInt{Value: big.NewInt(miningReward + miningReward)},
 						CodeHash:     "0xc5d2460186f7233c927e7db2dcc703c0e500b653ca82273b7bfad8045d85a470",
@@ -213,17 +222,17 @@ func TestBuilder(t *testing.T) {
 			&b.StateDiff{
 				BlockNumber:     block3.Number().Int64(),
 				BlockHash:       block3.Hash(),
-				CreatedAccounts: map[common.Address]b.AccountDiff{},
+				CreatedAccounts: b.AccountDiffsMap{},
 				DeletedAccounts: emptyAccountDiffEventualMap,
-				UpdatedAccounts: map[common.Address]b.AccountDiff{
+				UpdatedAccounts: b.AccountDiffsMap{
-					account2Addr: {
+					account2LeafKey: {
 						Nonce:        b.DiffUint64{Value: &nonce0},
 						Balance:      b.DiffBigInt{Value: big.NewInt(block2Account2Balance + miningReward)},
 						CodeHash:     "0xc5d2460186f7233c927e7db2dcc703c0e500b653ca82273b7bfad8045d85a470",
 						ContractRoot: b.DiffString{Value: &originalContractRoot},
 						Storage:      map[string]b.DiffStorage{},
 					},
-					contractAddr: {
+					contractLeafKey: {
 						Nonce:        b.DiffUint64{Value: &nonce1},
 						Balance:      b.DiffBigInt{Value: big.NewInt(0)},
 						CodeHash:     "0x753f98a8d4328b15636e46f66f2cb4bc860100aa17967cc145fcd17d1d4710ea",
@@ -234,7 +243,7 @@ func TestBuilder(t *testing.T) {
 								Value: &updatedStorageValue},
 						},
 					},
-					testBankAddress: {
+					bankLeafKey: {
 						Nonce:        b.DiffUint64{Value: &nonce3},
 						Balance:      b.DiffBigInt{Value: big.NewInt(99989000)},
 						CodeHash:     "0xc5d2460186f7233c927e7db2dcc703c0e500b653ca82273b7bfad8045d85a470",
@@ -252,7 +261,6 @@ func TestBuilder(t *testing.T) {
 		if err != nil {
 			t.Error(err)
 		}
-
 		fields := []string{"BlockNumber", "BlockHash", "DeletedAccounts", "UpdatedAccounts", "CreatedAccounts"}
 
 		for _, field := range fields {
@@ -287,8 +295,14 @@ func equals(actual, expected interface{}) (success bool) {
 // the returned hash chain is ordered head->parent. In addition, every 3rd block
 // contains a transaction and every 5th an uncle to allow testing correct block
 // reassembly.
-func makeChain(n int, parent *types.Block) ([]common.Hash, map[common.Hash]*types.Block) {
+func makeChain(n int, parent *types.Block) ([]common.Hash, map[common.Hash]*types.Block, *core.BlockChain) {
 	blocks, _ := core.GenerateChain(params.TestChainConfig, parent, ethash.NewFaker(), testdb, n, testChainGen)
+	headers := make([]*types.Header, len(blocks))
+	for i, block := range blocks {
+		headers[i] = block.Header()
+	}
+	chain, _ := core.NewBlockChain(testdb, nil, params.TestChainConfig, ethash.NewFaker(), vm.Config{}, nil)
+
 	hashes := make([]common.Hash, n+1)
 	hashes[len(hashes)-1] = parent.Hash()
 	blockm := make(map[common.Hash]*types.Block, n+1)
@@ -297,7 +311,7 @@ func makeChain(n int, parent *types.Block) ([]common.Hash, map[common.Hash]*type
 		hashes[len(hashes)-i-2] = b.Hash()
 		blockm[b.Hash()] = b
 	}
-	return hashes, blockm
+	return hashes, blockm, chain
 }
 
 func testChainGen(i int, block *core.BlockGen) {

statediff/builder/helpers.go

@@ -24,11 +24,10 @@ import (
 	"strings"
 
 	"github.com/ethereum/go-ethereum/common"
-	"github.com/ethereum/go-ethereum/core/state"
 	"github.com/ethereum/go-ethereum/trie"
 )
 
-func sortKeys(data map[common.Address]*state.Account) []string {
+func sortKeys(data AccountsMap) []string {
 	var keys []string
 	for key := range data {
 		keys = append(keys, key.Hex())

statediff/builder/struct.go

@@ -26,12 +26,13 @@ import (
 	"github.com/ethereum/go-ethereum/common"
 )
 
+type AccountDiffsMap map[common.Hash]AccountDiff
 type StateDiff struct {
-	BlockNumber     int64                          `json:"blockNumber"      gencodec:"required"`
+	BlockNumber     int64           `json:"blockNumber"      gencodec:"required"`
-	BlockHash       common.Hash                    `json:"blockHash"        gencodec:"required"`
+	BlockHash       common.Hash     `json:"blockHash"        gencodec:"required"`
-	CreatedAccounts map[common.Address]AccountDiff `json:"createdAccounts"  gencodec:"required"`
+	CreatedAccounts AccountDiffsMap `json:"createdAccounts"  gencodec:"required"`
-	DeletedAccounts map[common.Address]AccountDiff `json:"deletedAccounts"  gencodec:"required"`
+	DeletedAccounts AccountDiffsMap `json:"deletedAccounts"  gencodec:"required"`
-	UpdatedAccounts map[common.Address]AccountDiff `json:"updatedAccounts"  gencodec:"required"`
+	UpdatedAccounts AccountDiffsMap `json:"updatedAccounts"  gencodec:"required"`
 
 	encoded []byte
 	err     error

statediff/publisher/csv.go

@@ -15,7 +15,7 @@ var (
 	Headers = []string{
 		"blockNumber", "blockHash", "accountAction", "codeHash",
 		"nonceValue", "balanceValue", "contractRoot", "storageDiffPaths",
-		"accountAddress", "storageKey", "storageValue",
+		"accountLeafKey", "storageKey", "storageValue",
 	}
 
 	timeStampFormat      = "20060102150405.00000"
@@ -81,7 +81,7 @@ func accumulateAccountRows(sd builder.StateDiff) [][]string {
 	return accountRows
 }
 
-func formatAccountData(accountAddr common.Address, accountDiff builder.AccountDiff, sd builder.StateDiff, accountAction string) [][]string {
+func formatAccountData(accountAddr common.Hash, accountDiff builder.AccountDiff, sd builder.StateDiff, accountAction string) [][]string {
 	blockNumberString := strconv.FormatInt(sd.BlockNumber, 10)
 	blockHash := sd.BlockHash.String()
 	codeHash := accountDiff.CodeHash

statediff/publisher/publisher_test.go

@@ -35,7 +35,7 @@ var expectedCreatedAccountRow = []string{
 	strconv.FormatInt(testhelpers.NewBalanceValue, 10),
 	testhelpers.ContractRoot,
 	testhelpers.StoragePath,
-	testhelpers.ContractAddress,
+	testhelpers.ContractLeafKey.Hex(),
 	"0000000000000000000000000000000000000000000000000000000000000001",
 	testhelpers.StorageValue,
 }
@@ -49,7 +49,7 @@ var expectedCreatedAccountWithoutStorageUpdateRow = []string{
 	strconv.FormatInt(testhelpers.NewBalanceValue, 10),
 	testhelpers.ContractRoot,
 	"",
-	testhelpers.AnotherContractAddress,
+	testhelpers.AnotherContractLeafKey.Hex(),
 	"",
 	"",
 }
@@ -63,7 +63,7 @@ var expectedUpdatedAccountRow = []string{
 	strconv.FormatInt(testhelpers.NewBalanceValue, 10),
 	testhelpers.ContractRoot,
 	testhelpers.StoragePath,
-	testhelpers.ContractAddress,
+	testhelpers.ContractLeafKey.Hex(),
 	"0000000000000000000000000000000000000000000000000000000000000001",
 	testhelpers.StorageValue,
 }
@@ -77,7 +77,7 @@ var expectedDeletedAccountRow = []string{
 	strconv.FormatInt(testhelpers.NewBalanceValue, 10),
 	testhelpers.ContractRoot,
 	testhelpers.StoragePath,
-	testhelpers.ContractAddress,
+	testhelpers.ContractLeafKey.Hex(),
 	"0000000000000000000000000000000000000000000000000000000000000001",
 	testhelpers.StorageValue,
 }

statediff/service/service.go

@@ -19,6 +19,7 @@ import (
 type BlockChain interface {
 	SubscribeChainEvent(ch chan<- core.ChainEvent) event.Subscription
 	GetBlockByHash(hash common.Hash) *types.Block
+	AddToStateDiffProcessedCollection(hash common.Hash)
 }
 
 type StateDiffService struct {
@@ -28,7 +29,7 @@ type StateDiffService struct {
 }
 
 func NewStateDiffService(db ethdb.Database, blockChain *core.BlockChain, config statediff.Config) (*StateDiffService, error) {
-	builder := b.NewBuilder(db)
+	builder := b.NewBuilder(db, blockChain)
 	publisher, err := p.NewPublisher(config)
 	if err != nil {
 		return nil, err
@@ -94,6 +95,8 @@ HandleBlockChLoop:
 				log.Error("Error extracting statediff", "block number", currentBlock.Number(), "error", err)
 			} else {
 				log.Info("Statediff extracted", "block number", currentBlock.Number(), "location", stateDiffLocation)
+				sds.BlockChain.AddToStateDiffProcessedCollection(parentBlock.Root())
+				sds.BlockChain.AddToStateDiffProcessedCollection(currentBlock.Root())
 			}
 		case <-quitCh:
 			log.Debug("Quitting the statediff block channel")

statediff/testhelpers/mocks/blockchain.go

@@ -3,6 +3,8 @@ package mocks
 import (
 	"errors"
 
+	"time"
+
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/core"
 	"github.com/ethereum/go-ethereum/core/types"
@@ -16,6 +18,8 @@ type BlockChain struct {
 	ChainEvents          []core.ChainEvent
 }
 
+func (mc *BlockChain) AddToStateDiffProcessedCollection(hash common.Hash) {}
+
 func (mc *BlockChain) SetParentBlocksToReturn(blocks []*types.Block) {
 	mc.parentBlocksToReturn = blocks
 }
@@ -43,6 +47,7 @@ func (bc *BlockChain) SubscribeChainEvent(ch chan<- core.ChainEvent) event.Subsc
 	subscription := event.NewSubscription(func(quit <-chan struct{}) error {
 		for _, chainEvent := range bc.ChainEvents {
 			if eventCounter > 1 {
+				time.Sleep(250 * time.Millisecond)
 				return subErr
 			}
 			select {

statediff/testhelpers/test_data.go

@@ -5,9 +5,14 @@ import (
 	"math/rand"
 
 	"github.com/ethereum/go-ethereum/common"
+	"github.com/ethereum/go-ethereum/crypto"
 	"github.com/ethereum/go-ethereum/statediff/builder"
 )
 
+func AddressToLeafKey(address common.Address) common.Hash {
+	return common.BytesToHash(crypto.Keccak256(address[:]))
+}
+
 var (
 	BlockNumber     = rand.Int63()
 	BlockHash       = "0xfa40fbe2d98d98b3363a778d52f2bcd29d6790b9b3f3cab2b167fd12d3550f73"
@@ -24,18 +29,18 @@ var (
 	}}
 	emptyStorage           = map[string]builder.DiffStorage{}
 	address                = common.HexToAddress("0xaE9BEa628c4Ce503DcFD7E305CaB4e29E7476592")
+	ContractLeafKey        = AddressToLeafKey(address)
 	anotherAddress         = common.HexToAddress("0xaE9BEa628c4Ce503DcFD7E305CaB4e29E7476593")
-	ContractAddress        = address.String()
+	AnotherContractLeafKey = AddressToLeafKey(anotherAddress)
-	AnotherContractAddress = anotherAddress.String()
+	CreatedAccountDiffs    = builder.AccountDiffsMap{
-	CreatedAccountDiffs    = map[common.Address]builder.AccountDiff{
+		ContractLeafKey: {
-		address: {
 			Nonce:        builder.DiffUint64{Value: &NewNonceValue},
 			Balance:      builder.DiffBigInt{Value: big.NewInt(NewBalanceValue)},
 			ContractRoot: builder.DiffString{Value: &ContractRoot},
 			CodeHash:     CodeHash,
 			Storage:      storage,
 		},
-		anotherAddress: {
+		AnotherContractLeafKey: {
 			Nonce:        builder.DiffUint64{Value: &NewNonceValue},
 			Balance:      builder.DiffBigInt{Value: big.NewInt(NewBalanceValue)},
 			CodeHash:     CodeHash,
@@ -44,7 +49,7 @@ var (
 		},
 	}
 
-	UpdatedAccountDiffs = map[common.Address]builder.AccountDiff{address: {
+	UpdatedAccountDiffs = builder.AccountDiffsMap{ContractLeafKey: {
 		Nonce:        builder.DiffUint64{Value: &NewNonceValue},
 		Balance:      builder.DiffBigInt{Value: big.NewInt(NewBalanceValue)},
 		CodeHash:     CodeHash,
@@ -52,7 +57,7 @@ var (
 		Storage:      storage,
 	}}
 
-	DeletedAccountDiffs = map[common.Address]builder.AccountDiff{address: {
+	DeletedAccountDiffs = builder.AccountDiffsMap{ContractLeafKey: {
 		Nonce:        builder.DiffUint64{Value: &NewNonceValue},
 		Balance:      builder.DiffBigInt{Value: big.NewInt(NewBalanceValue)},
 		ContractRoot: builder.DiffString{Value: &ContractRoot},