refactor: only support tracked snapshots

tracking incremental state diffs is not feasible like this. because updates are processed after the
trie has been traversed, the iterators' state doesn't fully capture the progress of the diff. since
currently snapshots are the only use case for tracking, let's just support that.

builder.go

@@ -161,35 +161,23 @@ func (sdb *builder) WriteStateDiff(
 }
 
 // WriteStateDiff writes a statediff object to output sinks
-func (sdb *builder) WriteStateDiffTracked(
+func (sdb *builder) WriteStateSnapshot(
-	args Args, params Params,
+	stateRoot common.Hash, params Params,
 	nodeSink sdtypes.StateNodeSink,
 	ipldSink sdtypes.IPLDSink,
 	tracker tracker.IteratorTracker,
 ) error {
 	defer metrics.UpdateDuration(time.Now(), metrics.IndexerMetrics.WriteStateDiffTimer)
 	// Load tries for old and new states
-	triea, err := sdb.stateCache.OpenTrie(args.OldStateRoot)
+	tree, err := sdb.stateCache.OpenTrie(stateRoot)
-	if err != nil {
-		return fmt.Errorf("error opening old state trie: %w", err)
-	}
-	trieb, err := sdb.stateCache.OpenTrie(args.NewStateRoot)
 	if err != nil {
 		return fmt.Errorf("error opening new state trie: %w", err)
 	}
 
-	var subiters, bases []trie.NodeIterator
+	subiters, _, err := tracker.Restore(tree.NodeIterator)
-	// Constructor for difference iterator at a specific (recovered) path
-	makeIterator := func(key []byte) trie.NodeIterator {
-		a := triea.NodeIterator(key)
-		b := trieb.NodeIterator(key)
-		return utils.NewSymmetricDifferenceIterator(a, b)
-	}
-	subiters, bases, err = tracker.Restore(makeIterator)
 	if err != nil {
 		return fmt.Errorf("error restoring iterators: %w", err)
 	}
-
 	if len(subiters) != 0 {
 		// Completed iterators are not saved by the tracker, so restoring fewer than configured is ok,
 		// but having too many is a problem.
@@ -198,21 +186,21 @@ func (sdb *builder) WriteStateDiffTracked(
 				sdb.subtrieWorkers, len(subiters))
 		}
 	} else {
-		subiters = iterutils.SubtrieIterators(makeIterator, uint(sdb.subtrieWorkers))
+		subiters = iterutils.SubtrieIterators(tree.NodeIterator, uint(sdb.subtrieWorkers))
 		for i := range subiters {
 			subiters[i] = tracker.Tracked(subiters[i])
 		}
 	}
-	logger := log.New("hash", args.BlockHash, "number", args.BlockNumber)
 	// errgroup will cancel if any group fails
 	g, ctx := errgroup.WithContext(context.Background())
 	for i := range subiters {
 		func(subdiv uint) {
 			g.Go(func() error {
+				symdiff := utils.AlwaysBState()
 				return sdb.processAccounts(ctx,
-					subiters[subdiv], &bases[subdiv].(*utils.SymmDiffIterator).SymmDiffState,
+					subiters[subdiv], &symdiff,
 					params.watchedAddressesLeafPaths,
-					nodeSink, ipldSink, logger,
+					nodeSink, ipldSink, log.DefaultLogger,
 				)
 			})
 		}(uint(i))
@@ -299,8 +287,6 @@ func (sdb *builder) processAccounts(
 		// New inner trie nodes will be written to blockstore only.
 		// Reminder: this includes leaf nodes, since the geth iterator.Leaf() actually
 		// signifies a "value" node.
-
-		// TODO: A zero hash indicates what?
 		if it.Hash() == zeroHash {
 			continue
 		}
@@ -350,8 +336,6 @@ func (sdb *builder) processAccounts(
 			return err
 		}
 	}
-
-	metrics.IndexerMetrics.DifferenceIteratorCounter.Inc(int64(symdiff.Count()))
 	return it.Error()
 }
 
@@ -457,7 +441,8 @@ func (sdb *builder) processStorageCreations(
 	return it.Error()
 }
 
-// processStorageUpdates builds the storage diff node objects for all nodes that exist in a different state at B than A
+// processStorageUpdates builds the storage diff node objects for all nodes that exist in a
+// different state at B than A
 func (sdb *builder) processStorageUpdates(
 	oldroot common.Hash, newroot common.Hash,
 	storageSink sdtypes.StorageNodeSink,

builder_snapshot_test.go

@@ -0,0 +1,538 @@
+package statediff_test
+
+import (
+	"testing"
+
+	statediff "github.com/cerc-io/plugeth-statediff"
+	"github.com/cerc-io/plugeth-statediff/indexer/ipld"
+	"github.com/cerc-io/plugeth-statediff/test_helpers"
+	sdtypes "github.com/cerc-io/plugeth-statediff/types"
+	"github.com/ethereum/go-ethereum/common"
+	"github.com/ethereum/go-ethereum/crypto"
+)
+
+func TestBuilderSnapshot(t *testing.T) {
+	blocks, chain := test_helpers.MakeChain(3, test_helpers.Genesis, test_helpers.TestChainGen)
+	contractLeafKey = test_helpers.AddressToLeafKey(test_helpers.ContractAddr)
+	defer chain.Stop()
+	block0 = test_helpers.Genesis
+	block1 = blocks[0]
+	block2 = blocks[1]
+	block3 = blocks[2]
+	params := statediff.Params{}
+
+	tests := []test_helpers.SnapshotTestCase{
+		{
+			"testEmptyDiff",
+			common.Hash{},
+			&sdtypes.StateObject{
+				Nodes: emptyDiffs,
+			},
+		},
+		{
+			"testBlock0",
+			//10000 transferred from testBankAddress to account1Addr
+			block0.Root(),
+			&sdtypes.StateObject{
+				Nodes: []sdtypes.StateLeafNode{
+					{
+						Removed: false,
+						AccountWrapper: sdtypes.AccountWrapper{
+							Account: bankAccountAtBlock0,
+							LeafKey: test_helpers.BankLeafKey,
+							CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, crypto.Keccak256(bankAccountAtBlock0LeafNode)).String()},
+						StorageDiff: emptyStorage,
+					},
+				},
+				IPLDs: []sdtypes.IPLD{
+					{
+						CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, crypto.Keccak256(bankAccountAtBlock0LeafNode)).String(),
+						Content: bankAccountAtBlock0LeafNode,
+					},
+				},
+			},
+		},
+		{
+			"testBlock1",
+			//10000 transferred from testBankAddress to account1Addr
+			block1.Root(),
+			&sdtypes.StateObject{
+				Nodes: []sdtypes.StateLeafNode{
+					{
+						Removed: false,
+						AccountWrapper: sdtypes.AccountWrapper{
+							Account: bankAccountAtBlock1,
+							LeafKey: test_helpers.BankLeafKey,
+							CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, crypto.Keccak256(bankAccountAtBlock1LeafNode)).String()},
+						StorageDiff: emptyStorage,
+					},
+					{
+						Removed: false,
+						AccountWrapper: sdtypes.AccountWrapper{
+							Account: minerAccountAtBlock1,
+							LeafKey: minerLeafKey,
+							CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, crypto.Keccak256(minerAccountAtBlock1LeafNode)).String()},
+						StorageDiff: emptyStorage,
+					},
+					{
+						Removed: false,
+						AccountWrapper: sdtypes.AccountWrapper{
+							Account: account1AtBlock1,
+							LeafKey: test_helpers.Account1LeafKey,
+							CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, crypto.Keccak256(account1AtBlock1LeafNode)).String()},
+						StorageDiff: emptyStorage,
+					},
+				},
+				IPLDs: []sdtypes.IPLD{
+					{
+						CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, crypto.Keccak256(block1BranchRootNode)).String(),
+						Content: block1BranchRootNode,
+					},
+					{
+						CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, crypto.Keccak256(bankAccountAtBlock1LeafNode)).String(),
+						Content: bankAccountAtBlock1LeafNode,
+					},
+					{
+						CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, crypto.Keccak256(minerAccountAtBlock1LeafNode)).String(),
+						Content: minerAccountAtBlock1LeafNode,
+					},
+					{
+						CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, crypto.Keccak256(account1AtBlock1LeafNode)).String(),
+						Content: account1AtBlock1LeafNode,
+					},
+				},
+			},
+		},
+		{
+			"testBlock2",
+			//1000 transferred from testBankAddress to account1Addr
+			//1000 transferred from account1Addr to account2Addr
+			block2.Root(),
+			&sdtypes.StateObject{
+				Nodes: []sdtypes.StateLeafNode{
+					{
+						Removed: false,
+						AccountWrapper: sdtypes.AccountWrapper{
+							Account: bankAccountAtBlock2,
+							LeafKey: test_helpers.BankLeafKey,
+							CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, crypto.Keccak256(bankAccountAtBlock2LeafNode)).String()},
+						StorageDiff: emptyStorage,
+					},
+					{
+						Removed: false,
+						AccountWrapper: sdtypes.AccountWrapper{
+							Account: minerAccountAtBlock2,
+							LeafKey: minerLeafKey,
+							CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, crypto.Keccak256(minerAccountAtBlock2LeafNode)).String()},
+						StorageDiff: emptyStorage,
+					},
+					{
+						Removed: false,
+						AccountWrapper: sdtypes.AccountWrapper{
+							Account: account1AtBlock2,
+							LeafKey: test_helpers.Account1LeafKey,
+							CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, crypto.Keccak256(account1AtBlock2LeafNode)).String()},
+						StorageDiff: emptyStorage,
+					},
+					{
+						Removed: false,
+						AccountWrapper: sdtypes.AccountWrapper{
+							Account: contractAccountAtBlock2,
+							LeafKey: contractLeafKey,
+							CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, crypto.Keccak256(contractAccountAtBlock2LeafNode)).String()},
+						StorageDiff: []sdtypes.StorageLeafNode{
+							{
+								Removed: false,
+								Value:   slot0StorageValue,
+								LeafKey: slot0StorageKey.Bytes(),
+								CID:     ipld.Keccak256ToCid(ipld.MEthStorageTrie, crypto.Keccak256(slot0StorageLeafNode)).String(),
+							},
+							{
+								Removed: false,
+								Value:   slot1StorageValue,
+								LeafKey: slot1StorageKey.Bytes(),
+								CID:     ipld.Keccak256ToCid(ipld.MEthStorageTrie, crypto.Keccak256(slot1StorageLeafNode)).String(),
+							},
+						},
+					},
+					{
+						Removed: false,
+						AccountWrapper: sdtypes.AccountWrapper{
+							Account: account2AtBlock2,
+							LeafKey: test_helpers.Account2LeafKey,
+							CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, crypto.Keccak256(account2AtBlock2LeafNode)).String()},
+						StorageDiff: emptyStorage,
+					},
+				},
+				IPLDs: []sdtypes.IPLD{
+					{
+						CID:     ipld.Keccak256ToCid(ipld.RawBinary, test_helpers.CodeHash.Bytes()).String(),
+						Content: test_helpers.ByteCodeAfterDeployment,
+					},
+					{
+						CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, crypto.Keccak256(block2BranchRootNode)).String(),
+						Content: block2BranchRootNode,
+					},
+					{
+						CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, crypto.Keccak256(bankAccountAtBlock2LeafNode)).String(),
+						Content: bankAccountAtBlock2LeafNode,
+					},
+					{
+						CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, crypto.Keccak256(minerAccountAtBlock2LeafNode)).String(),
+						Content: minerAccountAtBlock2LeafNode,
+					},
+					{
+						CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, crypto.Keccak256(account1AtBlock2LeafNode)).String(),
+						Content: account1AtBlock2LeafNode,
+					},
+					{
+						CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, crypto.Keccak256(contractAccountAtBlock2LeafNode)).String(),
+						Content: contractAccountAtBlock2LeafNode,
+					},
+					{
+						CID:     ipld.Keccak256ToCid(ipld.MEthStorageTrie, crypto.Keccak256(block2StorageBranchRootNode)).String(),
+						Content: block2StorageBranchRootNode,
+					},
+					{
+						CID:     ipld.Keccak256ToCid(ipld.MEthStorageTrie, crypto.Keccak256(slot0StorageLeafNode)).String(),
+						Content: slot0StorageLeafNode,
+					},
+					{
+						CID:     ipld.Keccak256ToCid(ipld.MEthStorageTrie, crypto.Keccak256(slot1StorageLeafNode)).String(),
+						Content: slot1StorageLeafNode,
+					},
+					{
+						CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, crypto.Keccak256(account2AtBlock2LeafNode)).String(),
+						Content: account2AtBlock2LeafNode,
+					},
+				},
+			},
+		},
+		{
+			"testBlock3",
+			//the contract's storage is changed
+			//and the block is mined by account 2
+			block3.Root(),
+			&sdtypes.StateObject{
+				Nodes: []sdtypes.StateLeafNode{
+					{
+						Removed: false,
+						AccountWrapper: sdtypes.AccountWrapper{
+							Account: minerAccountAtBlock2,
+							LeafKey: minerLeafKey,
+							CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, crypto.Keccak256(minerAccountAtBlock2LeafNode)).String()},
+						StorageDiff: emptyStorage,
+					},
+					{
+						Removed: false,
+						AccountWrapper: sdtypes.AccountWrapper{
+							Account: account1AtBlock2,
+							LeafKey: test_helpers.Account1LeafKey,
+							CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, crypto.Keccak256(account1AtBlock2LeafNode)).String()},
+						StorageDiff: emptyStorage,
+					},
+					{
+						Removed: false,
+						AccountWrapper: sdtypes.AccountWrapper{
+							Account: bankAccountAtBlock3,
+							LeafKey: test_helpers.BankLeafKey,
+							CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, crypto.Keccak256(bankAccountAtBlock3LeafNode)).String()},
+						StorageDiff: emptyStorage,
+					},
+					{
+						Removed: false,
+						AccountWrapper: sdtypes.AccountWrapper{
+							Account: contractAccountAtBlock3,
+							LeafKey: contractLeafKey,
+							CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, crypto.Keccak256(contractAccountAtBlock3LeafNode)).String()},
+						StorageDiff: []sdtypes.StorageLeafNode{
+
+							{
+								Removed: false,
+								Value:   slot0StorageValue,
+								LeafKey: slot0StorageKey.Bytes(),
+								CID:     ipld.Keccak256ToCid(ipld.MEthStorageTrie, crypto.Keccak256(slot0StorageLeafNode)).String(),
+							},
+							{
+								Removed: false,
+								Value:   slot1StorageValue,
+								LeafKey: slot1StorageKey.Bytes(),
+								CID:     ipld.Keccak256ToCid(ipld.MEthStorageTrie, crypto.Keccak256(slot1StorageLeafNode)).String(),
+							},
+
+							{
+								Removed: false,
+								Value:   slot3StorageValue,
+								LeafKey: slot3StorageKey.Bytes(),
+								CID:     ipld.Keccak256ToCid(ipld.MEthStorageTrie, crypto.Keccak256(slot3StorageLeafNode)).String(),
+							},
+						},
+					},
+					{
+						Removed: false,
+						AccountWrapper: sdtypes.AccountWrapper{
+							Account: account2AtBlock3,
+							LeafKey: test_helpers.Account2LeafKey,
+							CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, crypto.Keccak256(account2AtBlock3LeafNode)).String()},
+						StorageDiff: emptyStorage,
+					},
+				},
+				IPLDs: []sdtypes.IPLD{
+					{
+						CID:     ipld.Keccak256ToCid(ipld.RawBinary, test_helpers.CodeHash.Bytes()).String(),
+						Content: test_helpers.ByteCodeAfterDeployment,
+					},
+					{
+						CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, crypto.Keccak256(minerAccountAtBlock2LeafNode)).String(),
+						Content: minerAccountAtBlock2LeafNode,
+					},
+					{
+						CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, crypto.Keccak256(account1AtBlock2LeafNode)).String(),
+						Content: account1AtBlock2LeafNode,
+					},
+					{
+						CID:     ipld.Keccak256ToCid(ipld.MEthStorageTrie, crypto.Keccak256(slot0StorageLeafNode)).String(),
+						Content: slot0StorageLeafNode,
+					},
+					{
+						CID:     ipld.Keccak256ToCid(ipld.MEthStorageTrie, crypto.Keccak256(slot1StorageLeafNode)).String(),
+						Content: slot1StorageLeafNode,
+					},
+					{
+						CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, crypto.Keccak256(block3BranchRootNode)).String(),
+						Content: block3BranchRootNode,
+					},
+					{
+						CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, crypto.Keccak256(bankAccountAtBlock3LeafNode)).String(),
+						Content: bankAccountAtBlock3LeafNode,
+					},
+					{
+						CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, crypto.Keccak256(contractAccountAtBlock3LeafNode)).String(),
+						Content: contractAccountAtBlock3LeafNode,
+					},
+					{
+						CID:     ipld.Keccak256ToCid(ipld.MEthStorageTrie, crypto.Keccak256(block3StorageBranchRootNode)).String(),
+						Content: block3StorageBranchRootNode,
+					},
+					{
+						CID:     ipld.Keccak256ToCid(ipld.MEthStorageTrie, crypto.Keccak256(slot3StorageLeafNode)).String(),
+						Content: slot3StorageLeafNode,
+					},
+					{
+						CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, crypto.Keccak256(account2AtBlock3LeafNode)).String(),
+						Content: account2AtBlock3LeafNode,
+					},
+				},
+			},
+		},
+	}
+
+	for _, test := range tests {
+		test_helpers.RunStateSnapshot(t, chain.StateCache(), test, params)
+	}
+}
+
+func TestBuilderSnapshotWithWatchedAddressList(t *testing.T) {
+	blocks, chain := test_helpers.MakeChain(3, test_helpers.Genesis, test_helpers.TestChainGen)
+	contractLeafKey = test_helpers.AddressToLeafKey(test_helpers.ContractAddr)
+	defer chain.Stop()
+	block0 = test_helpers.Genesis
+	block1 = blocks[0]
+	block2 = blocks[1]
+	block3 = blocks[2]
+	params := statediff.Params{
+		WatchedAddresses: []common.Address{test_helpers.Account1Addr, test_helpers.ContractAddr},
+	}
+	params.ComputeWatchedAddressesLeafPaths()
+
+	var tests = []test_helpers.SnapshotTestCase{
+		{
+			"testBlock0",
+			//10000 transferred from testBankAddress to account1Addr
+			block0.Root(),
+			&sdtypes.StateObject{
+				Nodes: emptyDiffs,
+			},
+		},
+		{
+			"testBlock1",
+			//10000 transferred from testBankAddress to account1Addr
+			block1.Root(),
+			&sdtypes.StateObject{
+				Nodes: []sdtypes.StateLeafNode{
+					{
+						Removed: false,
+						AccountWrapper: sdtypes.AccountWrapper{
+							Account: account1AtBlock1,
+							LeafKey: test_helpers.Account1LeafKey,
+							CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, crypto.Keccak256(account1AtBlock1LeafNode)).String()},
+						StorageDiff: emptyStorage,
+					},
+				},
+				IPLDs: []sdtypes.IPLD{
+					{
+						CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, crypto.Keccak256(block1BranchRootNode)).String(),
+						Content: block1BranchRootNode,
+					},
+					{
+						CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, crypto.Keccak256(account1AtBlock1LeafNode)).String(),
+						Content: account1AtBlock1LeafNode,
+					},
+				},
+			},
+		},
+		{
+			"testBlock2",
+			//1000 transferred from testBankAddress to account1Addr
+			//1000 transferred from account1Addr to account2Addr
+			block2.Root(),
+			&sdtypes.StateObject{
+				Nodes: []sdtypes.StateLeafNode{
+					{
+						Removed: false,
+						AccountWrapper: sdtypes.AccountWrapper{
+							Account: contractAccountAtBlock2,
+							LeafKey: contractLeafKey,
+							CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, crypto.Keccak256(contractAccountAtBlock2LeafNode)).String(),
+						},
+						StorageDiff: []sdtypes.StorageLeafNode{
+							{
+								Removed: false,
+								Value:   slot0StorageValue,
+								LeafKey: slot0StorageKey.Bytes(),
+								CID:     ipld.Keccak256ToCid(ipld.MEthStorageTrie, crypto.Keccak256(slot0StorageLeafNode)).String(),
+							},
+							{
+								Removed: false,
+								Value:   slot1StorageValue,
+								LeafKey: slot1StorageKey.Bytes(),
+								CID:     ipld.Keccak256ToCid(ipld.MEthStorageTrie, crypto.Keccak256(slot1StorageLeafNode)).String(),
+							},
+						},
+					},
+					{
+						Removed: false,
+						AccountWrapper: sdtypes.AccountWrapper{
+							Account: account1AtBlock2,
+							LeafKey: test_helpers.Account1LeafKey,
+							CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, crypto.Keccak256(account1AtBlock2LeafNode)).String()},
+						StorageDiff: emptyStorage,
+					},
+				},
+				IPLDs: []sdtypes.IPLD{
+					{
+						CID:     ipld.Keccak256ToCid(ipld.RawBinary, test_helpers.CodeHash.Bytes()).String(),
+						Content: test_helpers.ByteCodeAfterDeployment,
+					},
+					{
+						CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, crypto.Keccak256(block2BranchRootNode)).String(),
+						Content: block2BranchRootNode,
+					},
+					{
+						CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, crypto.Keccak256(contractAccountAtBlock2LeafNode)).String(),
+						Content: contractAccountAtBlock2LeafNode,
+					},
+					{
+						CID:     ipld.Keccak256ToCid(ipld.MEthStorageTrie, crypto.Keccak256(block2StorageBranchRootNode)).String(),
+						Content: block2StorageBranchRootNode,
+					},
+					{
+						CID:     ipld.Keccak256ToCid(ipld.MEthStorageTrie, crypto.Keccak256(slot0StorageLeafNode)).String(),
+						Content: slot0StorageLeafNode,
+					},
+					{
+						CID:     ipld.Keccak256ToCid(ipld.MEthStorageTrie, crypto.Keccak256(slot1StorageLeafNode)).String(),
+						Content: slot1StorageLeafNode,
+					},
+					{
+						CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, crypto.Keccak256(account1AtBlock2LeafNode)).String(),
+						Content: account1AtBlock2LeafNode,
+					},
+				},
+			},
+		},
+		{
+			"testBlock3",
+			//the contract's storage is changed
+			//and the block is mined by account 2
+			block3.Root(),
+			&sdtypes.StateObject{
+				Nodes: []sdtypes.StateLeafNode{
+					{
+						Removed: false,
+						AccountWrapper: sdtypes.AccountWrapper{
+							Account: account1AtBlock2,
+							LeafKey: test_helpers.Account1LeafKey,
+							CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, crypto.Keccak256(account1AtBlock2LeafNode)).String()},
+						StorageDiff: emptyStorage,
+					},
+					{
+						Removed: false,
+						AccountWrapper: sdtypes.AccountWrapper{
+							Account: contractAccountAtBlock3,
+							LeafKey: contractLeafKey,
+							CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, crypto.Keccak256(contractAccountAtBlock3LeafNode)).String()},
+						StorageDiff: []sdtypes.StorageLeafNode{
+							{
+								Removed: false,
+								Value:   slot0StorageValue,
+								LeafKey: slot0StorageKey.Bytes(),
+								CID:     ipld.Keccak256ToCid(ipld.MEthStorageTrie, crypto.Keccak256(slot0StorageLeafNode)).String(),
+							},
+							{
+								Removed: false,
+								Value:   slot1StorageValue,
+								LeafKey: slot1StorageKey.Bytes(),
+								CID:     ipld.Keccak256ToCid(ipld.MEthStorageTrie, crypto.Keccak256(slot1StorageLeafNode)).String(),
+							},
+							{
+								Removed: false,
+								Value:   slot3StorageValue,
+								LeafKey: slot3StorageKey.Bytes(),
+								CID:     ipld.Keccak256ToCid(ipld.MEthStorageTrie, crypto.Keccak256(slot3StorageLeafNode)).String(),
+							},
+						},
+					},
+				},
+				IPLDs: []sdtypes.IPLD{
+					{
+						CID:     ipld.Keccak256ToCid(ipld.RawBinary, test_helpers.CodeHash.Bytes()).String(),
+						Content: test_helpers.ByteCodeAfterDeployment,
+					},
+					{
+						CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, crypto.Keccak256(account1AtBlock2LeafNode)).String(),
+						Content: account1AtBlock2LeafNode,
+					},
+					{
+						CID:     ipld.Keccak256ToCid(ipld.MEthStorageTrie, crypto.Keccak256(slot0StorageLeafNode)).String(),
+						Content: slot0StorageLeafNode,
+					},
+					{
+						CID:     ipld.Keccak256ToCid(ipld.MEthStorageTrie, crypto.Keccak256(slot1StorageLeafNode)).String(),
+						Content: slot1StorageLeafNode,
+					},
+					{
+						CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, crypto.Keccak256(block3BranchRootNode)).String(),
+						Content: block3BranchRootNode,
+					},
+					{
+						CID:     ipld.Keccak256ToCid(ipld.MEthStateTrie, crypto.Keccak256(contractAccountAtBlock3LeafNode)).String(),
+						Content: contractAccountAtBlock3LeafNode,
+					},
+					{
+						CID:     ipld.Keccak256ToCid(ipld.MEthStorageTrie, crypto.Keccak256(block3StorageBranchRootNode)).String(),
+						Content: block3StorageBranchRootNode,
+					},
+					{
+						CID:     ipld.Keccak256ToCid(ipld.MEthStorageTrie, crypto.Keccak256(slot3StorageLeafNode)).String(),
+						Content: slot3StorageLeafNode,
+					},
+				},
+			},
+		},
+	}
+
+	for _, test := range tests {
+		test_helpers.RunStateSnapshot(t, chain.StateCache(), test, params)
+	}
+}

builder_test.go

@@ -503,7 +503,7 @@ func TestBuilder(t *testing.T) {
 	block3 = blocks[2]
 	params := statediff.Params{}
 
-	var tests = []test_helpers.TestCase{
+	var tests = []test_helpers.DiffTestCase{
 		{
 			"testEmptyDiff",
 			statediff.Args{
@@ -795,7 +795,7 @@ func TestBuilder(t *testing.T) {
 		},
 	}
 
-	test_helpers.RunBuilderTests(t, chain.StateCache(), tests, params, []uint{1, 8, 32})
+	test_helpers.RunBuildStateDiff(t, chain.StateCache(), tests, params)
 	test_helpers.CheckedRoots{
 		block0: bankAccountAtBlock0LeafNode,
 		block1: block1BranchRootNode,
@@ -817,7 +817,7 @@ func TestBuilderWithWatchedAddressList(t *testing.T) {
 	}
 	params.ComputeWatchedAddressesLeafPaths()
 
-	var tests = []test_helpers.TestCase{
+	var tests = []test_helpers.DiffTestCase{
 		{
 			"testEmptyDiff",
 			statediff.Args{
@@ -1009,7 +1009,7 @@ func TestBuilderWithWatchedAddressList(t *testing.T) {
 		},
 	}
 
-	test_helpers.RunBuilderTests(t, chain.StateCache(), tests, params, []uint{1, 8, 32})
+	test_helpers.RunBuildStateDiff(t, chain.StateCache(), tests, params)
 	test_helpers.CheckedRoots{
 		block0: bankAccountAtBlock0LeafNode,
 		block1: block1BranchRootNode,
@@ -1028,7 +1028,7 @@ func TestBuilderWithRemovedAccountAndStorage(t *testing.T) {
 	block6 = blocks[5]
 	params := statediff.Params{}
 
-	var tests = []test_helpers.TestCase{
+	var tests = []test_helpers.DiffTestCase{
 		// blocks 0-3 are the same as in TestBuilderWithIntermediateNodes
 		{
 			"testBlock4",
@@ -1260,7 +1260,7 @@ func TestBuilderWithRemovedAccountAndStorage(t *testing.T) {
 		},
 	}
 
-	test_helpers.RunBuilderTests(t, chain.StateCache(), tests, params, []uint{1, 8, 32})
+	test_helpers.RunBuildStateDiff(t, chain.StateCache(), tests, params)
 	test_helpers.CheckedRoots{
 		block4: block4BranchRootNode,
 		block5: block5BranchRootNode,
@@ -1281,7 +1281,7 @@ func TestBuilderWithRemovedNonWatchedAccount(t *testing.T) {
 	}
 	params.ComputeWatchedAddressesLeafPaths()
 
-	var tests = []test_helpers.TestCase{
+	var tests = []test_helpers.DiffTestCase{
 		{
 			"testBlock4",
 			statediff.Args{
@@ -1395,7 +1395,7 @@ func TestBuilderWithRemovedNonWatchedAccount(t *testing.T) {
 		},
 	}
 
-	test_helpers.RunBuilderTests(t, chain.StateCache(), tests, params, []uint{1, 8, 32})
+	test_helpers.RunBuildStateDiff(t, chain.StateCache(), tests, params)
 	test_helpers.CheckedRoots{
 		block4: block4BranchRootNode,
 		block5: block5BranchRootNode,
@@ -1416,7 +1416,7 @@ func TestBuilderWithRemovedWatchedAccount(t *testing.T) {
 	}
 	params.ComputeWatchedAddressesLeafPaths()
 
-	var tests = []test_helpers.TestCase{
+	var tests = []test_helpers.DiffTestCase{
 		{
 			"testBlock4",
 			statediff.Args{
@@ -1599,7 +1599,7 @@ func TestBuilderWithRemovedWatchedAccount(t *testing.T) {
 		},
 	}
 
-	test_helpers.RunBuilderTests(t, chain.StateCache(), tests, params, []uint{1, 8, 32})
+	test_helpers.RunBuildStateDiff(t, chain.StateCache(), tests, params)
 	test_helpers.CheckedRoots{
 		block4: block4BranchRootNode,
 		block5: block5BranchRootNode,
@@ -1700,7 +1700,7 @@ func TestBuilderWithMovedAccount(t *testing.T) {
 	block2 = blocks[1]
 	params := statediff.Params{}
 
-	var tests = []test_helpers.TestCase{
+	var tests = []test_helpers.DiffTestCase{
 		{
 			"testBlock1",
 			statediff.Args{
@@ -1827,7 +1827,7 @@ func TestBuilderWithMovedAccount(t *testing.T) {
 		},
 	}
 
-	test_helpers.RunBuilderTests(t, chain.StateCache(), tests, params, []uint{1, 8, 32})
+	test_helpers.RunBuildStateDiff(t, chain.StateCache(), tests, params)
 	test_helpers.CheckedRoots{
 		block1: block01BranchRootNode,
 		block2: bankAccountAtBlock02LeafNode,
@@ -2088,7 +2088,7 @@ func TestBuilderWithInternalizedLeafNode(t *testing.T) {
 	block3 = blocks[2]
 	params := statediff.Params{}
 
-	var tests = []test_helpers.TestCase{
+	var tests = []test_helpers.DiffTestCase{
 		{
 			"testEmptyDiff",
 			statediff.Args{
@@ -2354,7 +2354,7 @@ func TestBuilderWithInternalizedLeafNode(t *testing.T) {
 		},
 	}
 
-	test_helpers.RunBuilderTests(t, chain.StateCache(), tests, params, []uint{1, 8, 32})
+	test_helpers.RunBuildStateDiff(t, chain.StateCache(), tests, params)
 	test_helpers.CheckedRoots{
 		block1: block1bBranchRootNode,
 		block2: block2bBranchRootNode,
@@ -2377,7 +2377,7 @@ func TestBuilderWithInternalizedLeafNodeAndWatchedAddress(t *testing.T) {
 	}
 	params.ComputeWatchedAddressesLeafPaths()
 
-	var tests = []test_helpers.TestCase{
+	var tests = []test_helpers.DiffTestCase{
 		{
 			"testEmptyDiff",
 			statediff.Args{
@@ -2556,7 +2556,7 @@ func TestBuilderWithInternalizedLeafNodeAndWatchedAddress(t *testing.T) {
 		},
 	}
 
-	test_helpers.RunBuilderTests(t, chain.StateCache(), tests, params, []uint{1, 8, 32})
+	test_helpers.RunBuildStateDiff(t, chain.StateCache(), tests, params)
 	test_helpers.CheckedRoots{
 		block1: block1bBranchRootNode,
 		block2: block2bBranchRootNode,

mainnet_tests/builder_test.go

@@ -444,7 +444,7 @@ func TestBuilderOnMainnetBlocks(t *testing.T) {
 	}
 	params := statediff.Params{}
 
-	var tests = []test_helpers.TestCase{
+	var tests = []test_helpers.DiffTestCase{
 		// note that block0 (genesis) has over 1000 nodes due to the pre-allocation for the crowd-sale
 		// it is not feasible to write a unit test of that size at this time
 		{
@@ -624,7 +624,7 @@ func TestBuilderOnMainnetBlocks(t *testing.T) {
 		},
 	}
 
-	test_helpers.RunBuilderTests(t, chain.StateCache(), tests, params, []uint{1, 8, 32})
+	test_helpers.RunBuildStateDiff(t, chain.StateCache(), tests, params)
 	test_helpers.CheckedRoots{
 		block1: block1RootBranchNode,
 		block2: block2RootBranchNode,

test_helpers/builder.go

@@ -4,9 +4,13 @@ import (
 	"bytes"
 	"fmt"
 	"math/big"
+	"math/rand"
+	"path/filepath"
 	"sort"
+	"sync"
 	"testing"
 
+	"github.com/cerc-io/eth-iterator-utils/tracker"
 	statediff "github.com/cerc-io/plugeth-statediff"
 	"github.com/cerc-io/plugeth-statediff/adapt"
 	sdtypes "github.com/cerc-io/plugeth-statediff/types"
@@ -17,12 +21,20 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
-type TestCase struct {
+var subtrieCounts = []uint{1, 8, 32}
+
+type DiffTestCase struct {
 	Name     string
 	Args     statediff.Args
 	Expected *sdtypes.StateObject
 }
 
+type SnapshotTestCase struct {
+	Name      string
+	StateRoot common.Hash
+	Expected  *sdtypes.StateObject
+}
+
 type CheckedRoots map[*types.Block][]byte
 
 // Replicates the statediff object, but indexes nodes by CID
@@ -33,12 +45,11 @@ type normalizedStateDiff struct {
 	IPLDs       map[string]sdtypes.IPLD
 }
 
-func RunBuilderTests(
+func RunBuildStateDiff(
 	t *testing.T,
 	sdb state.Database,
-	tests []TestCase,
+	tests []DiffTestCase,
 	params statediff.Params,
-	subtrieCounts []uint,
 ) {
 	builder := statediff.NewBuilder(adapt.GethStateView(sdb))
 	for _, test := range tests {
@@ -58,6 +69,81 @@ func RunBuilderTests(
 	}
 }
 
+func RunStateSnapshot(
+	t *testing.T,
+	sdb state.Database,
+	test SnapshotTestCase,
+	params statediff.Params,
+) {
+	builder := statediff.NewBuilder(adapt.GethStateView(sdb))
+
+	for _, subtries := range subtrieCounts {
+		// Skip the recovery test for empty diffs
+		doRecovery := len(test.Expected.Nodes) != 0
+
+		t.Run(fmt.Sprintf("%s with %d subtries", test.Name, subtries), func(t *testing.T) {
+			builder.SetSubtrieWorkers(subtries)
+			var stateNodes []sdtypes.StateLeafNode
+			var iplds []sdtypes.IPLD
+			interrupt := randomInterrupt(len(test.Expected.IPLDs))
+			stateAppender := failingSyncedAppender(&stateNodes, -1)
+			ipldAppender := failingSyncedAppender(&iplds, interrupt)
+			recoveryFile := filepath.Join(t.TempDir(), "recovery.txt")
+			build := func() error {
+				tr := tracker.New(recoveryFile, subtries)
+				defer tr.CloseAndSave()
+				return builder.WriteStateSnapshot(
+					test.StateRoot, params, stateAppender, ipldAppender, tr,
+				)
+			}
+			if doRecovery {
+				// First attempt fails, second succeeds
+				if build() == nil {
+					t.Fatal("expected an error")
+				}
+			}
+			// Ensure we don't exceed the expected number of nodes. If we do, it implies the
+			// failed attempt got further than intended, and we have duplicates.
+			// ipldAppender = failingSyncedAppender(&iplds, len(test.Expected.IPLDs))
+			ipldAppender = failingSyncedAppender(&iplds, -1)
+			if err := build(); err != nil {
+				t.Fatal(err)
+			}
+			diff := sdtypes.StateObject{
+				Nodes: stateNodes,
+				IPLDs: iplds,
+			}
+			require.Equal(t,
+				normalize(test.Expected),
+				normalize(&diff),
+			)
+		})
+	}
+
+}
+
+// an appender which fails on a configured trigger
+func failingSyncedAppender[T any](to *[]T, failAt int) func(T) error {
+	var mtx sync.Mutex
+	return func(item T) error {
+		mtx.Lock()
+		defer mtx.Unlock()
+		if len(*to) == failAt {
+			return fmt.Errorf("failing at %d items", failAt)
+		}
+		*to = append(*to, item)
+		return nil
+	}
+}
+
+// function to pick random int between N/4 and 3N/4
+func randomInterrupt(N int) int {
+	if N < 2 {
+		return 0
+	}
+	return rand.Intn(N/2) + N/4
+}
+
 func (roots CheckedRoots) Check(t *testing.T) {
 	// Let's also confirm that our root state nodes form the state root hash in the headers
 	for block, node := range roots {

utils/iterator.go

@@ -184,3 +184,9 @@ func compareNodes(a, b trie.NodeIterator) int {
 	}
 	return 0
 }
+
+// AlwaysBState returns a dummy SymmDiffState that indicates all elements are from B, and have no
+// common paths with A. This is equivalent to a diff against an empty A.
+func AlwaysBState() SymmDiffState {
+	return SymmDiffState{yieldFromA: false, eqPathIndex: -2}
+}