diff --git a/statediff/builder.go b/statediff/builder.go
index 10aea79c1..54bb720a8 100644
--- a/statediff/builder.go
+++ b/statediff/builder.go
@@ -98,10 +98,9 @@ func (sdb *builder) buildStateDiffWithIntermediateStateNodes(args Args, intermed
 	}
 
 	return StateDiff{
-		BlockNumber:       args.BlockNumber,
-		BlockHash:         args.BlockHash,
-		LeafNodes:         append(append(updatedAccounts, createdAccounts...), deletedAccounts...),
-		IntermediateNodes: append(createdOrUpdatedIntermediateNodes, deletedIntermediateNodes...),
+		BlockNumber: args.BlockNumber,
+		BlockHash:   args.BlockHash,
+		Nodes:       append(append(append(append(updatedAccounts, createdAccounts...), deletedAccounts...), createdOrUpdatedIntermediateNodes...), deletedIntermediateNodes...),
 	}, nil
 }
 
@@ -148,7 +147,7 @@ func (sdb *builder) buildStateDiffWithoutIntermediateStateNodes(args Args, param
 	return StateDiff{
 		BlockNumber: args.BlockNumber,
 		BlockHash:   args.BlockHash,
-		LeafNodes:   append(append(updatedAccounts, createdAccounts...), deletedAccounts...),
+		Nodes:       append(append(updatedAccounts, createdAccounts...), deletedAccounts...),
 	}, nil
 }
 
diff --git a/statediff/builder_test.go b/statediff/builder_test.go
index 82a9e8b74..67ec3b219 100644
--- a/statediff/builder_test.go
+++ b/statediff/builder_test.go
@@ -34,7 +34,7 @@ import (
 // TODO: add test that filters on address
 var (
 	contractLeafKey                []byte
-	emptyAccounts                  = make([]statediff.StateNode, 0)
+	emptyDiffs                     = make([]statediff.StateNode, 0)
 	emptyStorage                   = make([]statediff.StorageNode, 0)
 	block0, block1, block2, block3 *types.Block
 	builder                        statediff.Builder
@@ -266,60 +266,49 @@ var (
 	})
 )
 
-type arguments struct {
-	oldStateRoot common.Hash
-	newStateRoot common.Hash
-	blockNumber  *big.Int
-	blockHash    common.Hash
-}
-
 func TestBuilder(t *testing.T) {
-	blockHashes, blockMap, chain := testhelpers.MakeChain(3, testhelpers.Genesis)
+	BlockHashes, blockMap, chain := testhelpers.MakeChain(3, testhelpers.Genesis)
 	contractLeafKey = testhelpers.AddressToLeafKey(testhelpers.ContractAddr)
 	defer chain.Stop()
-	block0 = blockMap[blockHashes[3]]
-	block1 = blockMap[blockHashes[2]]
-	block2 = blockMap[blockHashes[1]]
-	block3 = blockMap[blockHashes[0]]
-	config := statediff.Config{
-		IntermediateNodes: false,
-	}
-	builder = statediff.NewBuilder(chain, config)
+	block0 = blockMap[BlockHashes[3]]
+	block1 = blockMap[BlockHashes[2]]
+	block2 = blockMap[BlockHashes[1]]
+	block3 = blockMap[BlockHashes[0]]
+	params := statediff.Params{}
+	builder = statediff.NewBuilder(chain.StateCache())
 
 	var tests = []struct {
 		name              string
-		startingArguments arguments
+		startingArguments statediff.Args
 		expected          *statediff.StateDiff
 	}{
 		{
 			"testEmptyDiff",
-			arguments{
-				oldStateRoot: block0.Root(),
-				newStateRoot: block0.Root(),
-				blockNumber:  block0.Number(),
-				blockHash:    block0.Hash(),
-			},
-			&statediff.StateDiff{
+			statediff.Args{
+				OldStateRoot: block0.Root(),
+				NewStateRoot: block0.Root(),
 				BlockNumber:  block0.Number(),
 				BlockHash:    block0.Hash(),
-				CreatedNodes: emptyAccounts,
-				DeletedNodes: emptyAccounts,
-				UpdatedNodes: emptyAccounts,
+			},
+			&statediff.StateDiff{
+				BlockNumber: block0.Number(),
+				BlockHash:   block0.Hash(),
+				Nodes:       emptyDiffs,
 			},
 		},
 		{
 			"testBlock0",
 			//10000 transferred from testBankAddress to account1Addr
-			arguments{
-				oldStateRoot: testhelpers.NullHash,
-				newStateRoot: block0.Root(),
-				blockNumber:  block0.Number(),
-				blockHash:    block0.Hash(),
+			statediff.Args{
+				OldStateRoot: testhelpers.NullHash,
+				NewStateRoot: block0.Root(),
+				BlockNumber:  block0.Number(),
+				BlockHash:    block0.Hash(),
 			},
 			&statediff.StateDiff{
 				BlockNumber: block0.Number(),
 				BlockHash:   block0.Hash(),
-				CreatedNodes: []statediff.StateNode{
+				Nodes: []statediff.StateNode{
 					{
 						Path:         []byte{},
 						NodeType:     statediff.Leaf,
@@ -328,23 +317,21 @@ func TestBuilder(t *testing.T) {
 						StorageDiffs: emptyStorage,
 					},
 				},
-				DeletedNodes: emptyAccounts,
-				UpdatedNodes: emptyAccounts,
 			},
 		},
 		{
 			"testBlock1",
 			//10000 transferred from testBankAddress to account1Addr
-			arguments{
-				oldStateRoot: block0.Root(),
-				newStateRoot: block1.Root(),
-				blockNumber:  block1.Number(),
-				blockHash:    block1.Hash(),
+			statediff.Args{
+				OldStateRoot: block0.Root(),
+				NewStateRoot: block1.Root(),
+				BlockNumber:  block1.Number(),
+				BlockHash:    block1.Hash(),
 			},
 			&statediff.StateDiff{
 				BlockNumber: block1.Number(),
 				BlockHash:   block1.Hash(),
-				CreatedNodes: []statediff.StateNode{
+				Nodes: []statediff.StateNode{
 					{
 						Path:         []byte{'\x00'},
 						NodeType:     statediff.Leaf,
@@ -367,16 +354,6 @@ func TestBuilder(t *testing.T) {
 						StorageDiffs: emptyStorage,
 					},
 				},
-				DeletedNodes: []statediff.StateNode{ // This leaf appears to be deleted since it is turned into a branch node and the account is moved to \x00
-					{ // It would instead show up in the UpdateAccounts as new branch node IF intermediate node diffing was turned on (as it is in the test below)
-						Path:         []byte{},
-						NodeType:     statediff.Removed,
-						LeafKey:      testhelpers.BankLeafKey,
-						NodeValue:    []byte{},
-						StorageDiffs: emptyStorage,
-					},
-				},
-				UpdatedNodes: emptyAccounts,
 			},
 		},
 		{
@@ -384,16 +361,37 @@ func TestBuilder(t *testing.T) {
 			// 1000 transferred from testBankAddress to account1Addr
 			// 1000 transferred from account1Addr to account2Addr
 			// account1addr creates a new contract
-			arguments{
-				oldStateRoot: block1.Root(),
-				newStateRoot: block2.Root(),
-				blockNumber:  block2.Number(),
-				blockHash:    block2.Hash(),
+			statediff.Args{
+				OldStateRoot: block1.Root(),
+				NewStateRoot: block2.Root(),
+				BlockNumber:  block2.Number(),
+				BlockHash:    block2.Hash(),
 			},
 			&statediff.StateDiff{
 				BlockNumber: block2.Number(),
 				BlockHash:   block2.Hash(),
-				CreatedNodes: []statediff.StateNode{
+				Nodes: []statediff.StateNode{
+					{
+						Path:         []byte{'\x00'},
+						NodeType:     statediff.Leaf,
+						LeafKey:      testhelpers.BankLeafKey,
+						NodeValue:    bankAccountAtBlock2LeafNode,
+						StorageDiffs: emptyStorage,
+					},
+					{
+						Path:         []byte{'\x05'},
+						NodeType:     statediff.Leaf,
+						LeafKey:      minerLeafKey,
+						NodeValue:    minerAccountAtBlock2LeafNode,
+						StorageDiffs: emptyStorage,
+					},
+					{
+						Path:         []byte{'\x0e'},
+						NodeType:     statediff.Leaf,
+						LeafKey:      testhelpers.Account1LeafKey,
+						NodeValue:    account1AtBlock2LeafNode,
+						StorageDiffs: emptyStorage,
+					},
 					{
 						Path:      []byte{'\x06'},
 						NodeType:  statediff.Leaf,
@@ -416,48 +414,22 @@ func TestBuilder(t *testing.T) {
 						StorageDiffs: emptyStorage,
 					},
 				},
-				DeletedNodes: emptyAccounts,
-				UpdatedNodes: []statediff.StateNode{
-					{
-						Path:         []byte{'\x00'},
-						NodeType:     statediff.Leaf,
-						LeafKey:      testhelpers.BankLeafKey,
-						NodeValue:    bankAccountAtBlock2LeafNode,
-						StorageDiffs: emptyStorage,
-					},
-					{
-						Path:         []byte{'\x05'},
-						NodeType:     statediff.Leaf,
-						LeafKey:      minerLeafKey,
-						NodeValue:    minerAccountAtBlock2LeafNode,
-						StorageDiffs: emptyStorage,
-					},
-					{
-						Path:         []byte{'\x0e'},
-						NodeType:     statediff.Leaf,
-						LeafKey:      testhelpers.Account1LeafKey,
-						NodeValue:    account1AtBlock2LeafNode,
-						StorageDiffs: emptyStorage,
-					},
-				},
 			},
 		},
 		{
 			"testBlock3",
 			//the contract's storage is changed
 			//and the block is mined by account 2
-			arguments{
-				oldStateRoot: block2.Root(),
-				newStateRoot: block3.Root(),
-				blockNumber:  block3.Number(),
-				blockHash:    block3.Hash(),
-			},
-			&statediff.StateDiff{
+			statediff.Args{
+				OldStateRoot: block2.Root(),
+				NewStateRoot: block3.Root(),
 				BlockNumber:  block3.Number(),
 				BlockHash:    block3.Hash(),
-				CreatedNodes: []statediff.StateNode{},
-				DeletedNodes: emptyAccounts,
-				UpdatedNodes: []statediff.StateNode{
+			},
+			&statediff.StateDiff{
+				BlockNumber: block3.Number(),
+				BlockHash:   block3.Hash(),
+				Nodes: []statediff.StateNode{
 					{
 						Path:         []byte{'\x00'},
 						NodeType:     statediff.Leaf,
@@ -498,8 +470,7 @@ func TestBuilder(t *testing.T) {
 	}
 
 	for _, test := range tests {
-		arguments := test.startingArguments
-		diff, err := builder.BuildStateDiff(arguments.oldStateRoot, arguments.newStateRoot, arguments.blockNumber, arguments.blockHash)
+		diff, err := builder.BuildStateDiff(test.startingArguments, params)
 		if err != nil {
 			t.Error(err)
 		}
@@ -521,52 +492,51 @@ func TestBuilder(t *testing.T) {
 }
 
 func TestBuilderWithIntermediateNodes(t *testing.T) {
-	blockHashes, blockMap, chain := testhelpers.MakeChain(3, testhelpers.Genesis)
+	BlockHashes, blockMap, chain := testhelpers.MakeChain(3, testhelpers.Genesis)
 	contractLeafKey = testhelpers.AddressToLeafKey(testhelpers.ContractAddr)
 	defer chain.Stop()
-	block0 = blockMap[blockHashes[3]]
-	block1 = blockMap[blockHashes[2]]
-	block2 = blockMap[blockHashes[1]]
-	block3 = blockMap[blockHashes[0]]
-	config := statediff.Config{
-		IntermediateNodes: true,
+	block0 = blockMap[BlockHashes[3]]
+	block1 = blockMap[BlockHashes[2]]
+	block2 = blockMap[BlockHashes[1]]
+	block3 = blockMap[BlockHashes[0]]
+	params := statediff.Params{
+		IntermediateStateNodes:   true,
+		IntermediateStorageNodes: true,
 	}
-	builder = statediff.NewBuilder(chain, config)
+	builder = statediff.NewBuilder(chain.StateCache())
 
 	var tests = []struct {
 		name              string
-		startingArguments arguments
+		startingArguments statediff.Args
 		expected          *statediff.StateDiff
 	}{
 		{
 			"testEmptyDiff",
-			arguments{
-				oldStateRoot: block0.Root(),
-				newStateRoot: block0.Root(),
-				blockNumber:  block0.Number(),
-				blockHash:    block0.Hash(),
-			},
-			&statediff.StateDiff{
+			statediff.Args{
+				OldStateRoot: block0.Root(),
+				NewStateRoot: block0.Root(),
 				BlockNumber:  block0.Number(),
 				BlockHash:    block0.Hash(),
-				CreatedNodes: emptyAccounts,
-				DeletedNodes: emptyAccounts,
-				UpdatedNodes: emptyAccounts,
+			},
+			&statediff.StateDiff{
+				BlockNumber: block0.Number(),
+				BlockHash:   block0.Hash(),
+				Nodes:       emptyDiffs,
 			},
 		},
 		{
 			"testBlock0",
 			//10000 transferred from testBankAddress to account1Addr
-			arguments{
-				oldStateRoot: testhelpers.NullHash,
-				newStateRoot: block0.Root(),
-				blockNumber:  block0.Number(),
-				blockHash:    block0.Hash(),
+			statediff.Args{
+				OldStateRoot: testhelpers.NullHash,
+				NewStateRoot: block0.Root(),
+				BlockNumber:  block0.Number(),
+				BlockHash:    block0.Hash(),
 			},
 			&statediff.StateDiff{
 				BlockNumber: block0.Number(),
 				BlockHash:   block0.Hash(),
-				CreatedNodes: []statediff.StateNode{
+				Nodes: []statediff.StateNode{
 					{
 						Path:         []byte{},
 						NodeType:     statediff.Leaf,
@@ -575,23 +545,27 @@ func TestBuilderWithIntermediateNodes(t *testing.T) {
 						StorageDiffs: emptyStorage,
 					},
 				},
-				DeletedNodes: emptyAccounts,
-				UpdatedNodes: emptyAccounts,
 			},
 		},
 		{
 			"testBlock1",
 			//10000 transferred from testBankAddress to account1Addr
-			arguments{
-				oldStateRoot: block0.Root(),
-				newStateRoot: block1.Root(),
-				blockNumber:  block1.Number(),
-				blockHash:    block1.Hash(),
+			statediff.Args{
+				OldStateRoot: block0.Root(),
+				NewStateRoot: block1.Root(),
+				BlockNumber:  block1.Number(),
+				BlockHash:    block1.Hash(),
 			},
 			&statediff.StateDiff{
 				BlockNumber: block1.Number(),
 				BlockHash:   block1.Hash(),
-				CreatedNodes: []statediff.StateNode{
+				Nodes: []statediff.StateNode{
+					{
+						Path:         []byte{},
+						NodeType:     statediff.Branch,
+						NodeValue:    block1BranchNode,
+						StorageDiffs: emptyStorage,
+					},
 					{
 						Path:         []byte{'\x00'},
 						NodeType:     statediff.Leaf,
@@ -614,15 +588,6 @@ func TestBuilderWithIntermediateNodes(t *testing.T) {
 						StorageDiffs: emptyStorage,
 					},
 				},
-				DeletedNodes: emptyAccounts,
-				UpdatedNodes: []statediff.StateNode{
-					{
-						Path:         []byte{},
-						NodeType:     statediff.Branch,
-						NodeValue:    block1BranchNode,
-						StorageDiffs: emptyStorage,
-					},
-				},
 			},
 		},
 		{
@@ -630,40 +595,16 @@ func TestBuilderWithIntermediateNodes(t *testing.T) {
 			// 1000 transferred from testBankAddress to account1Addr
 			// 1000 transferred from account1Addr to account2Addr
 			// account1addr creates a new contract
-			arguments{
-				oldStateRoot: block1.Root(),
-				newStateRoot: block2.Root(),
-				blockNumber:  block2.Number(),
-				blockHash:    block2.Hash(),
+			statediff.Args{
+				OldStateRoot: block1.Root(),
+				NewStateRoot: block2.Root(),
+				BlockNumber:  block2.Number(),
+				BlockHash:    block2.Hash(),
 			},
 			&statediff.StateDiff{
 				BlockNumber: block2.Number(),
 				BlockHash:   block2.Hash(),
-				CreatedNodes: []statediff.StateNode{
-					{
-						Path:      []byte{'\x06'},
-						NodeType:  statediff.Leaf,
-						LeafKey:   contractLeafKey,
-						NodeValue: contractAccountAtBlock2LeafNode,
-						StorageDiffs: []statediff.StorageNode{
-							{
-								Path:      []byte{},
-								NodeType:  statediff.Leaf,
-								LeafKey:   originalStorageKey,
-								NodeValue: originalStorageLeafNode,
-							},
-						},
-					},
-					{
-						Path:         []byte{'\x0c'},
-						NodeType:     statediff.Leaf,
-						LeafKey:      testhelpers.Account2LeafKey,
-						NodeValue:    account2AtBlock2LeafNode,
-						StorageDiffs: emptyStorage,
-					},
-				},
-				DeletedNodes: emptyAccounts,
-				UpdatedNodes: []statediff.StateNode{
+				Nodes: []statediff.StateNode{
 					{
 						Path:         []byte{},
 						NodeType:     statediff.Branch,
@@ -691,6 +632,27 @@ func TestBuilderWithIntermediateNodes(t *testing.T) {
 						NodeValue:    account1AtBlock2LeafNode,
 						StorageDiffs: emptyStorage,
 					},
+					{
+						Path:      []byte{'\x06'},
+						NodeType:  statediff.Leaf,
+						LeafKey:   contractLeafKey,
+						NodeValue: contractAccountAtBlock2LeafNode,
+						StorageDiffs: []statediff.StorageNode{
+							{
+								Path:      []byte{},
+								NodeType:  statediff.Leaf,
+								LeafKey:   originalStorageKey,
+								NodeValue: originalStorageLeafNode,
+							},
+						},
+					},
+					{
+						Path:         []byte{'\x0c'},
+						NodeType:     statediff.Leaf,
+						LeafKey:      testhelpers.Account2LeafKey,
+						NodeValue:    account2AtBlock2LeafNode,
+						StorageDiffs: emptyStorage,
+					},
 				},
 			},
 		},
@@ -698,18 +660,16 @@ func TestBuilderWithIntermediateNodes(t *testing.T) {
 			"testBlock3",
 			//the contract's storage is changed
 			//and the block is mined by account 2
-			arguments{
-				oldStateRoot: block2.Root(),
-				newStateRoot: block3.Root(),
-				blockNumber:  block3.Number(),
-				blockHash:    block3.Hash(),
-			},
-			&statediff.StateDiff{
+			statediff.Args{
+				OldStateRoot: block2.Root(),
+				NewStateRoot: block3.Root(),
 				BlockNumber:  block3.Number(),
 				BlockHash:    block3.Hash(),
-				CreatedNodes: []statediff.StateNode{},
-				DeletedNodes: emptyAccounts,
-				UpdatedNodes: []statediff.StateNode{
+			},
+			&statediff.StateDiff{
+				BlockNumber: block3.Number(),
+				BlockHash:   block3.Hash(),
+				Nodes: []statediff.StateNode{
 					{
 						Path:         []byte{},
 						NodeType:     statediff.Branch,
@@ -761,8 +721,7 @@ func TestBuilderWithIntermediateNodes(t *testing.T) {
 	}
 
 	for _, test := range tests {
-		arguments := test.startingArguments
-		diff, err := builder.BuildStateDiff(arguments.oldStateRoot, arguments.newStateRoot, arguments.blockNumber, arguments.blockHash)
+		diff, err := builder.BuildStateDiff(test.startingArguments, params)
 		if err != nil {
 			t.Error(err)
 		}
@@ -784,70 +743,65 @@ func TestBuilderWithIntermediateNodes(t *testing.T) {
 }
 
 func TestBuilderWithWatchedAddressList(t *testing.T) {
-	blockHashes, blockMap, chain := testhelpers.MakeChain(3, testhelpers.Genesis)
+	BlockHashes, blockMap, chain := testhelpers.MakeChain(3, testhelpers.Genesis)
 	contractLeafKey = testhelpers.AddressToLeafKey(testhelpers.ContractAddr)
 	defer chain.Stop()
-	block0 = blockMap[blockHashes[3]]
-	block1 = blockMap[blockHashes[2]]
-	block2 = blockMap[blockHashes[1]]
-	block3 = blockMap[blockHashes[0]]
-	config := statediff.Config{
-		IntermediateNodes: false,
-		WatchedAddresses:  []string{testhelpers.Account1Addr.Hex(), testhelpers.ContractAddr.Hex()},
+	block0 = blockMap[BlockHashes[3]]
+	block1 = blockMap[BlockHashes[2]]
+	block2 = blockMap[BlockHashes[1]]
+	block3 = blockMap[BlockHashes[0]]
+	params := statediff.Params{
+		WatchedAddresses: []string{testhelpers.Account1Addr.Hex(), testhelpers.ContractAddr.Hex()},
 	}
-	builder = statediff.NewBuilder(chain, config)
+	builder = statediff.NewBuilder(chain.StateCache())
 
 	var tests = []struct {
 		name              string
-		startingArguments arguments
+		startingArguments statediff.Args
 		expected          *statediff.StateDiff
 	}{
 		{
 			"testEmptyDiff",
-			arguments{
-				oldStateRoot: block0.Root(),
-				newStateRoot: block0.Root(),
-				blockNumber:  block0.Number(),
-				blockHash:    block0.Hash(),
-			},
-			&statediff.StateDiff{
+			statediff.Args{
+				OldStateRoot: block0.Root(),
+				NewStateRoot: block0.Root(),
 				BlockNumber:  block0.Number(),
 				BlockHash:    block0.Hash(),
-				CreatedNodes: emptyAccounts,
-				DeletedNodes: emptyAccounts,
-				UpdatedNodes: emptyAccounts,
+			},
+			&statediff.StateDiff{
+				BlockNumber: block0.Number(),
+				BlockHash:   block0.Hash(),
+				Nodes:       emptyDiffs,
 			},
 		},
 		{
 			"testBlock0",
 			//10000 transferred from testBankAddress to account1Addr
-			arguments{
-				oldStateRoot: testhelpers.NullHash,
-				newStateRoot: block0.Root(),
-				blockNumber:  block0.Number(),
-				blockHash:    block0.Hash(),
-			},
-			&statediff.StateDiff{
+			statediff.Args{
+				OldStateRoot: testhelpers.NullHash,
+				NewStateRoot: block0.Root(),
 				BlockNumber:  block0.Number(),
 				BlockHash:    block0.Hash(),
-				CreatedNodes: emptyAccounts,
-				DeletedNodes: emptyAccounts,
-				UpdatedNodes: emptyAccounts,
+			},
+			&statediff.StateDiff{
+				BlockNumber: block0.Number(),
+				BlockHash:   block0.Hash(),
+				Nodes:       emptyDiffs,
 			},
 		},
 		{
 			"testBlock1",
 			//10000 transferred from testBankAddress to account1Addr
-			arguments{
-				oldStateRoot: block0.Root(),
-				newStateRoot: block1.Root(),
-				blockNumber:  block1.Number(),
-				blockHash:    block1.Hash(),
+			statediff.Args{
+				OldStateRoot: block0.Root(),
+				NewStateRoot: block1.Root(),
+				BlockNumber:  block1.Number(),
+				BlockHash:    block1.Hash(),
 			},
 			&statediff.StateDiff{
 				BlockNumber: block1.Number(),
 				BlockHash:   block1.Hash(),
-				CreatedNodes: []statediff.StateNode{
+				Nodes: []statediff.StateNode{
 					{
 						Path:         []byte{'\x0e'},
 						NodeType:     statediff.Leaf,
@@ -856,24 +810,22 @@ func TestBuilderWithWatchedAddressList(t *testing.T) {
 						StorageDiffs: emptyStorage,
 					},
 				},
-				DeletedNodes: emptyAccounts,
-				UpdatedNodes: []statediff.StateNode{},
 			},
 		},
 		{
 			"testBlock2",
 			//1000 transferred from testBankAddress to account1Addr
 			//1000 transferred from account1Addr to account2Addr
-			arguments{
-				oldStateRoot: block1.Root(),
-				newStateRoot: block2.Root(),
-				blockNumber:  block2.Number(),
-				blockHash:    block2.Hash(),
+			statediff.Args{
+				OldStateRoot: block1.Root(),
+				NewStateRoot: block2.Root(),
+				BlockNumber:  block2.Number(),
+				BlockHash:    block2.Hash(),
 			},
 			&statediff.StateDiff{
 				BlockNumber: block2.Number(),
 				BlockHash:   block2.Hash(),
-				CreatedNodes: []statediff.StateNode{
+				Nodes: []statediff.StateNode{
 					{
 						Path:      []byte{'\x06'},
 						NodeType:  statediff.Leaf,
@@ -888,9 +840,6 @@ func TestBuilderWithWatchedAddressList(t *testing.T) {
 							},
 						},
 					},
-				},
-				DeletedNodes: emptyAccounts,
-				UpdatedNodes: []statediff.StateNode{
 					{
 						Path:         []byte{'\x0e'},
 						NodeType:     statediff.Leaf,
@@ -905,18 +854,16 @@ func TestBuilderWithWatchedAddressList(t *testing.T) {
 			"testBlock3",
 			//the contract's storage is changed
 			//and the block is mined by account 2
-			arguments{
-				oldStateRoot: block2.Root(),
-				newStateRoot: block3.Root(),
-				blockNumber:  block3.Number(),
-				blockHash:    block3.Hash(),
-			},
-			&statediff.StateDiff{
+			statediff.Args{
+				OldStateRoot: block2.Root(),
+				NewStateRoot: block3.Root(),
 				BlockNumber:  block3.Number(),
 				BlockHash:    block3.Hash(),
-				CreatedNodes: []statediff.StateNode{},
-				DeletedNodes: emptyAccounts,
-				UpdatedNodes: []statediff.StateNode{
+			},
+			&statediff.StateDiff{
+				BlockNumber: block3.Number(),
+				BlockHash:   block3.Hash(),
+				Nodes: []statediff.StateNode{
 					{
 						Path:      []byte{'\x06'},
 						NodeType:  statediff.Leaf,
@@ -943,8 +890,7 @@ func TestBuilderWithWatchedAddressList(t *testing.T) {
 	}
 
 	for _, test := range tests {
-		arguments := test.startingArguments
-		diff, err := builder.BuildStateDiff(arguments.oldStateRoot, arguments.newStateRoot, arguments.blockNumber, arguments.blockHash)
+		diff, err := builder.BuildStateDiff(test.startingArguments, params)
 		if err != nil {
 			t.Error(err)
 		}
diff --git a/statediff/service.go b/statediff/service.go
index b8f234008..505e6f8ce 100644
--- a/statediff/service.go
+++ b/statediff/service.go
@@ -140,6 +140,8 @@ func (sds *Service) Loop(chainEventCh chan core.ChainEvent) {
 			sds.streamStateDiff(currentBlock, parentBlock.Root())
 		case err := <-errCh:
 			log.Warn("Error from chain event subscription", "error", err)
+			sds.close()
+			return
 		case <-sds.QuitChan:
 			log.Info("Quitting the statediffing process")
 			sds.close()
diff --git a/statediff/service_test.go b/statediff/service_test.go
index 56cb58187..1c4475521 100644
--- a/statediff/service_test.go
+++ b/statediff/service_test.go
@@ -72,23 +72,30 @@ var (
 	event1 = core.ChainEvent{Block: testBlock1}
 	event2 = core.ChainEvent{Block: testBlock2}
 	event3 = core.ChainEvent{Block: testBlock3}
+
+	defaultParams = statediff.Params{
+		IncludeBlock:    true,
+		IncludeReceipts: true,
+		IncludeTD:       true,
+	}
 )
 
 func testErrorInChainEventLoop(t *testing.T) {
 	//the first chain event causes and error (in blockchain mock)
 	builder := mocks.Builder{}
 	blockChain := mocks.BlockChain{}
+	serviceQuit := make(chan bool)
 	service := statediff.Service{
-		Mutex:         sync.Mutex{},
-		Builder:       &builder,
-		BlockChain:    &blockChain,
-		QuitChan:      make(chan bool),
-		Subscriptions: make(map[rpc.ID]statediff.Subscription),
-		StreamBlock:   true,
+		Mutex:             sync.Mutex{},
+		Builder:           &builder,
+		BlockChain:        &blockChain,
+		QuitChan:          serviceQuit,
+		Subscriptions:     make(map[common.Hash]map[rpc.ID]statediff.Subscription),
+		SubscriptionTypes: make(map[common.Hash]statediff.Params),
 	}
 	payloadChan := make(chan statediff.Payload, 2)
 	quitChan := make(chan bool)
-	service.Subscribe(rpc.NewID(), payloadChan, quitChan)
+	service.Subscribe(rpc.NewID(), payloadChan, quitChan, defaultParams)
 	testRoot2 = common.HexToHash("0xTestRoot2")
 	blockMapping := make(map[common.Hash]*types.Block)
 	blockMapping[parentBlock1.Hash()] = parentBlock1
@@ -99,7 +106,7 @@ func testErrorInChainEventLoop(t *testing.T) {
 	blockChain.SetReceiptsForHash(testBlock2.Hash(), testReceipts2)
 
 	payloads := make([]statediff.Payload, 0, 2)
-	wg := sync.WaitGroup{}
+	wg := new(sync.WaitGroup)
 	go func() {
 		wg.Add(1)
 		for i := 0; i < 2; i++ {
@@ -111,7 +118,6 @@ func testErrorInChainEventLoop(t *testing.T) {
 		}
 		wg.Done()
 	}()
-
 	service.Loop(eventsChannel)
 	wg.Wait()
 	if len(payloads) != 2 {
@@ -135,23 +141,27 @@ func testErrorInChainEventLoop(t *testing.T) {
 		}
 	}
 
-	if !reflect.DeepEqual(builder.BlockHash, testBlock2.Hash()) {
+	if !reflect.DeepEqual(builder.Params, defaultParams) {
 		t.Error("Test failure:", t.Name())
-		t.Logf("Actual blockhash does not equal expected.\nactual:%+v\nexpected: %+v", builder.BlockHash, testBlock2.Hash())
+		t.Logf("Actual params does not equal expected.\nactual:%+v\nexpected: %+v", builder.Params, defaultParams)
 	}
-	if !bytes.Equal(builder.OldStateRoot.Bytes(), parentBlock2.Root().Bytes()) {
+	if !bytes.Equal(builder.Args.BlockHash.Bytes(), testBlock2.Hash().Bytes()) {
 		t.Error("Test failure:", t.Name())
-		t.Logf("Actual root does not equal expected.\nactual:%+v\nexpected: %+v", builder.OldStateRoot, parentBlock2.Root())
+		t.Logf("Actual blockhash does not equal expected.\nactual:%x\nexpected: %x", builder.Args.BlockHash.Bytes(), testBlock2.Hash().Bytes())
 	}
-	if !bytes.Equal(builder.NewStateRoot.Bytes(), testBlock2.Root().Bytes()) {
+	if !bytes.Equal(builder.Args.OldStateRoot.Bytes(), parentBlock2.Root().Bytes()) {
 		t.Error("Test failure:", t.Name())
-		t.Logf("Actual root does not equal expected.\nactual:%+v\nexpected: %+v", builder.NewStateRoot, testBlock2.Root())
+		t.Logf("Actual root does not equal expected.\nactual:%x\nexpected: %x", builder.Args.OldStateRoot.Bytes(), parentBlock2.Root().Bytes())
+	}
+	if !bytes.Equal(builder.Args.NewStateRoot.Bytes(), testBlock2.Root().Bytes()) {
+		t.Error("Test failure:", t.Name())
+		t.Logf("Actual root does not equal expected.\nactual:%x\nexpected: %x", builder.Args.NewStateRoot.Bytes(), testBlock2.Root().Bytes())
 	}
 	//look up the parent block from its hash
 	expectedHashes := []common.Hash{testBlock1.ParentHash(), testBlock2.ParentHash()}
 	if !reflect.DeepEqual(blockChain.HashesLookedUp, expectedHashes) {
 		t.Error("Test failure:", t.Name())
-		t.Logf("Actual parent hash does not equal expected.\nactual:%+v\nexpected: %+v", blockChain.HashesLookedUp, expectedHashes)
+		t.Logf("Actual looked up parent hashes does not equal expected.\nactual:%+v\nexpected: %+v", blockChain.HashesLookedUp, expectedHashes)
 	}
 }
 
@@ -160,14 +170,15 @@ func testErrorInBlockLoop(t *testing.T) {
 	builder := mocks.Builder{}
 	blockChain := mocks.BlockChain{}
 	service := statediff.Service{
-		Builder:       &builder,
-		BlockChain:    &blockChain,
-		QuitChan:      make(chan bool),
-		Subscriptions: make(map[rpc.ID]statediff.Subscription),
+		Builder:           &builder,
+		BlockChain:        &blockChain,
+		QuitChan:          make(chan bool),
+		Subscriptions:     make(map[common.Hash]map[rpc.ID]statediff.Subscription),
+		SubscriptionTypes: make(map[common.Hash]statediff.Params),
 	}
 	payloadChan := make(chan statediff.Payload)
 	quitChan := make(chan bool)
-	service.Subscribe(rpc.NewID(), payloadChan, quitChan)
+	service.Subscribe(rpc.NewID(), payloadChan, quitChan, defaultParams)
 	blockMapping := make(map[common.Hash]*types.Block)
 	blockMapping[parentBlock1.Hash()] = parentBlock1
 	blockChain.SetBlocksForHashes(blockMapping)
@@ -180,18 +191,21 @@ func testErrorInBlockLoop(t *testing.T) {
 		}
 	}()
 	service.Loop(eventsChannel)
-
-	if !bytes.Equal(builder.BlockHash.Bytes(), testBlock1.Hash().Bytes()) {
+	if !reflect.DeepEqual(builder.Params, defaultParams) {
 		t.Error("Test failure:", t.Name())
-		t.Logf("Actual does not equal expected.\nactual:%+v\nexpected: %+v", builder.BlockHash, testBlock1.Hash())
+		t.Logf("Actual params does not equal expected.\nactual:%+v\nexpected: %+v", builder.Params, defaultParams)
 	}
-	if !bytes.Equal(builder.OldStateRoot.Bytes(), parentBlock1.Root().Bytes()) {
+	if !bytes.Equal(builder.Args.BlockHash.Bytes(), testBlock1.Hash().Bytes()) {
 		t.Error("Test failure:", t.Name())
-		t.Logf("Actual does not equal expected.\nactual:%+v\nexpected: %+v", builder.OldStateRoot, parentBlock1.Root())
+		t.Logf("Actual blockhash does not equal expected.\nactual:%+v\nexpected: %x", builder.Args.BlockHash.Bytes(), testBlock1.Hash().Bytes())
 	}
-	if !bytes.Equal(builder.NewStateRoot.Bytes(), testBlock1.Root().Bytes()) {
+	if !bytes.Equal(builder.Args.OldStateRoot.Bytes(), parentBlock1.Root().Bytes()) {
 		t.Error("Test failure:", t.Name())
-		t.Logf("Actual does not equal expected.\nactual:%+v\nexpected: %+v", builder.NewStateRoot, testBlock1.Root())
+		t.Logf("Actual old state root does not equal expected.\nactual:%+v\nexpected: %x", builder.Args.OldStateRoot.Bytes(), parentBlock1.Root().Bytes())
+	}
+	if !bytes.Equal(builder.Args.NewStateRoot.Bytes(), testBlock1.Root().Bytes()) {
+		t.Error("Test failure:", t.Name())
+		t.Logf("Actual new state root does not equal expected.\nactual:%+v\nexpected: %x", builder.Args.NewStateRoot.Bytes(), testBlock1.Root().Bytes())
 	}
 }
 
@@ -234,14 +248,14 @@ func testErrorInStateDiffAt(t *testing.T) {
 	blockChain.SetBlockForNumber(testBlock1, testBlock1.NumberU64())
 	blockChain.SetReceiptsForHash(testBlock1.Hash(), testReceipts1)
 	service := statediff.Service{
-		Mutex:         sync.Mutex{},
-		Builder:       &builder,
-		BlockChain:    &blockChain,
-		QuitChan:      make(chan bool),
-		Subscriptions: make(map[rpc.ID]statediff.Subscription),
-		StreamBlock:   true,
+		Mutex:             sync.Mutex{},
+		Builder:           &builder,
+		BlockChain:        &blockChain,
+		QuitChan:          make(chan bool),
+		Subscriptions:     make(map[common.Hash]map[rpc.ID]statediff.Subscription),
+		SubscriptionTypes: make(map[common.Hash]statediff.Params),
 	}
-	stateDiffPayload, err := service.StateDiffAt(testBlock1.NumberU64())
+	stateDiffPayload, err := service.StateDiffAt(testBlock1.NumberU64(), defaultParams)
 	if err != nil {
 		t.Error(err)
 	}
@@ -249,20 +263,24 @@ func testErrorInStateDiffAt(t *testing.T) {
 	if err != nil {
 		t.Error(err)
 	}
-	if !bytes.Equal(builder.BlockHash.Bytes(), testBlock1.Hash().Bytes()) {
+	if !reflect.DeepEqual(builder.Params, defaultParams) {
 		t.Error("Test failure:", t.Name())
-		t.Logf("Actual does not equal expected.\nactual:%+v\nexpected: %+v", builder.BlockHash, testBlock1.Hash())
+		t.Logf("Actual params does not equal expected.\nactual:%+v\nexpected: %+v", builder.Params, defaultParams)
 	}
-	if !bytes.Equal(builder.OldStateRoot.Bytes(), parentBlock1.Root().Bytes()) {
+	if !bytes.Equal(builder.Args.BlockHash.Bytes(), testBlock1.Hash().Bytes()) {
 		t.Error("Test failure:", t.Name())
-		t.Logf("Actual does not equal expected.\nactual:%+v\nexpected: %+v", builder.OldStateRoot, parentBlock1.Root())
+		t.Logf("Actual blockhash does not equal expected.\nactual:%+v\nexpected: %x", builder.Args.BlockHash.Bytes(), testBlock1.Hash().Bytes())
 	}
-	if !bytes.Equal(builder.NewStateRoot.Bytes(), testBlock1.Root().Bytes()) {
+	if !bytes.Equal(builder.Args.OldStateRoot.Bytes(), parentBlock1.Root().Bytes()) {
 		t.Error("Test failure:", t.Name())
-		t.Logf("Actual does not equal expected.\nactual:%+v\nexpected: %+v", builder.NewStateRoot, testBlock1.Root())
+		t.Logf("Actual old state root does not equal expected.\nactual:%+v\nexpected: %x", builder.Args.OldStateRoot.Bytes(), parentBlock1.Root().Bytes())
+	}
+	if !bytes.Equal(builder.Args.NewStateRoot.Bytes(), testBlock1.Root().Bytes()) {
+		t.Error("Test failure:", t.Name())
+		t.Logf("Actual new state root does not equal expected.\nactual:%+v\nexpected: %x", builder.Args.NewStateRoot.Bytes(), testBlock1.Root().Bytes())
 	}
 	if !bytes.Equal(expectedStateDiffPayloadRlp, stateDiffPayloadRlp) {
 		t.Error("Test failure:", t.Name())
-		t.Logf("Actual does not equal expected.\nactual:%+v\nexpected: %+v", expectedStateDiffPayload, stateDiffPayload)
+		t.Logf("Actual state diff payload does not equal expected.\nactual:%+v\nexpected: %+v", expectedStateDiffPayload, stateDiffPayload)
 	}
 }
diff --git a/statediff/testhelpers/mocks/api.go b/statediff/testhelpers/mocks/api.go
index 625f6acc4..8e9f840f9 100644
--- a/statediff/testhelpers/mocks/api.go
+++ b/statediff/testhelpers/mocks/api.go
@@ -22,6 +22,8 @@ import (
 	"fmt"
 	"sync"
 
+	"github.com/ethereum/go-ethereum/common"
+	"github.com/ethereum/go-ethereum/crypto"
 	"github.com/ethereum/go-ethereum/rlp"
 
 	"github.com/ethereum/go-ethereum/core"
@@ -35,15 +37,15 @@ import (
 // MockStateDiffService is a mock state diff service
 type MockStateDiffService struct {
 	sync.Mutex
-	Builder         statediff.Builder
-	BlockChain      *BlockChain
-	ReturnProtocol  []p2p.Protocol
-	ReturnAPIs      []rpc.API
-	BlockChan       chan *types.Block
-	ParentBlockChan chan *types.Block
-	QuitChan        chan bool
-	Subscriptions   map[rpc.ID]statediff.Subscription
-	streamBlock     bool
+	Builder           statediff.Builder
+	BlockChain        *BlockChain
+	ReturnProtocol    []p2p.Protocol
+	ReturnAPIs        []rpc.API
+	BlockChan         chan *types.Block
+	ParentBlockChan   chan *types.Block
+	QuitChan          chan bool
+	Subscriptions     map[common.Hash]map[rpc.ID]statediff.Subscription
+	SubscriptionTypes map[common.Hash]statediff.Params
 }
 
 // Protocols mock method
@@ -78,12 +80,7 @@ func (sds *MockStateDiffService) Loop(chan core.ChainEvent) {
 					"current block number", currentBlock.Number())
 				continue
 			}
-			payload, err := sds.processStateDiff(currentBlock, parentBlock)
-			if err != nil {
-				log.Error("Error building statediff", "block number", currentBlock.Number(), "error", err)
-				continue
-			}
-			sds.send(*payload)
+			sds.streamStateDiff(currentBlock, parentBlock.Root())
 		case <-sds.QuitChan:
 			log.Debug("Quitting the statediff block channel")
 			sds.close()
@@ -92,13 +89,46 @@ func (sds *MockStateDiffService) Loop(chan core.ChainEvent) {
 	}
 }
 
-// processStateDiff method builds the state diff payload from the current and parent block and streams it to listening subscriptions
-func (sds *MockStateDiffService) processStateDiff(currentBlock, parentBlock *types.Block) (*statediff.Payload, error) {
-	stateDiff, err := sds.Builder.BuildStateDiff(parentBlock.Root(), currentBlock.Root(), currentBlock.Number(), currentBlock.Hash())
+// streamStateDiff method builds the state diff payload for each subscription according to their subscription type and sends them the result
+func (sds *MockStateDiffService) streamStateDiff(currentBlock *types.Block, parentRoot common.Hash) {
+	sds.Lock()
+	for ty, subs := range sds.Subscriptions {
+		params, ok := sds.SubscriptionTypes[ty]
+		if !ok {
+			log.Error(fmt.Sprintf("subscriptions type %s do not have a parameter set associated with them", ty.Hex()))
+			sds.closeType(ty)
+			continue
+		}
+		// create payload for this subscription type
+		payload, err := sds.processStateDiff(currentBlock, parentRoot, params)
+		if err != nil {
+			log.Error(fmt.Sprintf("statediff processing error for subscriptions with parameters: %+v", params))
+			sds.closeType(ty)
+			continue
+		}
+		for id, sub := range subs {
+			select {
+			case sub.PayloadChan <- *payload:
+				log.Debug(fmt.Sprintf("sending statediff payload to subscription %s", id))
+			default:
+				log.Info(fmt.Sprintf("unable to send statediff payload to subscription %s; channel has no receiver", id))
+			}
+		}
+	}
+	sds.Unlock()
+}
+
+// processStateDiff method builds the state diff payload from the current block, parent state root, and provided params
+func (sds *MockStateDiffService) processStateDiff(currentBlock *types.Block, parentRoot common.Hash, params statediff.Params) (*statediff.Payload, error) {
+	stateDiff, err := sds.Builder.BuildStateDiff(statediff.Args{
+		NewStateRoot: currentBlock.Root(),
+		OldStateRoot: parentRoot,
+		BlockHash:    currentBlock.Hash(),
+		BlockNumber:  currentBlock.Number(),
+	}, params)
 	if err != nil {
 		return nil, err
 	}
-
 	stateDiffRlp, err := rlp.EncodeToBytes(stateDiff)
 	if err != nil {
 		return nil, err
@@ -106,13 +136,17 @@ func (sds *MockStateDiffService) processStateDiff(currentBlock, parentBlock *typ
 	payload := statediff.Payload{
 		StateDiffRlp: stateDiffRlp,
 	}
-	if sds.streamBlock {
-		rlpBuff := new(bytes.Buffer)
-		if err = currentBlock.EncodeRLP(rlpBuff); err != nil {
+	if params.IncludeBlock {
+		blockBuff := new(bytes.Buffer)
+		if err = currentBlock.EncodeRLP(blockBuff); err != nil {
 			return nil, err
 		}
-		payload.BlockRlp = rlpBuff.Bytes()
+		payload.BlockRlp = blockBuff.Bytes()
+	}
+	if params.IncludeTD {
 		payload.TotalDifficulty = sds.BlockChain.GetTdByHash(currentBlock.Hash())
+	}
+	if params.IncludeReceipts {
 		receiptBuff := new(bytes.Buffer)
 		receipts := sds.BlockChain.GetReceiptsByHash(currentBlock.Hash())
 		if err = rlp.Encode(receiptBuff, receipts); err != nil {
@@ -123,53 +157,68 @@ func (sds *MockStateDiffService) processStateDiff(currentBlock, parentBlock *typ
 	return &payload, nil
 }
 
-// Subscribe mock method
-func (sds *MockStateDiffService) Subscribe(id rpc.ID, sub chan<- statediff.Payload, quitChan chan<- bool) {
-	log.Info("Subscribing to the mock statediff service")
+// Subscribe is used by the API to subscribe to the service loop
+func (sds *MockStateDiffService) Subscribe(id rpc.ID, sub chan<- statediff.Payload, quitChan chan<- bool, params statediff.Params) {
+	// Subscription type is defined as the hash of the rlp-serialized subscription params
+	by, err := rlp.EncodeToBytes(params)
+	if err != nil {
+		return
+	}
+	subscriptionType := crypto.Keccak256Hash(by)
+	// Add subscriber
 	sds.Lock()
-	sds.Subscriptions[id] = statediff.Subscription{
+	if sds.Subscriptions[subscriptionType] == nil {
+		sds.Subscriptions[subscriptionType] = make(map[rpc.ID]statediff.Subscription)
+	}
+	sds.Subscriptions[subscriptionType][id] = statediff.Subscription{
 		PayloadChan: sub,
 		QuitChan:    quitChan,
 	}
+	sds.SubscriptionTypes[subscriptionType] = params
 	sds.Unlock()
 }
 
-// Unsubscribe mock method
+// Unsubscribe is used to unsubscribe from the service loop
 func (sds *MockStateDiffService) Unsubscribe(id rpc.ID) error {
-	log.Info("Unsubscribing from the mock statediff service")
 	sds.Lock()
-	_, ok := sds.Subscriptions[id]
-	if !ok {
-		return fmt.Errorf("cannot unsubscribe; subscription for id %s does not exist", id)
+	for ty := range sds.Subscriptions {
+		delete(sds.Subscriptions[ty], id)
+		if len(sds.Subscriptions[ty]) == 0 {
+			// If we removed the last subscription of this type, remove the subscription type outright
+			delete(sds.Subscriptions, ty)
+			delete(sds.SubscriptionTypes, ty)
+		}
 	}
-	delete(sds.Subscriptions, id)
 	sds.Unlock()
 	return nil
 }
 
-func (sds *MockStateDiffService) send(payload statediff.Payload) {
-	sds.Lock()
-	for id, sub := range sds.Subscriptions {
-		select {
-		case sub.PayloadChan <- payload:
-			log.Info("sending state diff payload to subscription %s", id)
-		default:
-			log.Info("unable to send payload to subscription %s; channel has no receiver", id)
-		}
+// StateDiffAt mock method
+func (sds *MockStateDiffService) StateDiffAt(blockNumber uint64, params statediff.Params) (*statediff.Payload, error) {
+	currentBlock := sds.BlockChain.GetBlockByNumber(blockNumber)
+	log.Info(fmt.Sprintf("sending state diff at %d", blockNumber))
+	if blockNumber == 0 {
+		return sds.processStateDiff(currentBlock, common.Hash{}, params)
 	}
-	sds.Unlock()
+	parentBlock := sds.BlockChain.GetBlockByHash(currentBlock.ParentHash())
+	return sds.processStateDiff(currentBlock, parentBlock.Root(), params)
 }
 
+// close is used to close all listening subscriptions
 func (sds *MockStateDiffService) close() {
 	sds.Lock()
-	for id, sub := range sds.Subscriptions {
-		select {
-		case sub.QuitChan <- true:
-			delete(sds.Subscriptions, id)
-			log.Info("closing subscription %s", id)
-		default:
-			log.Info("unable to close subscription %s; channel has no receiver", id)
+	for ty, subs := range sds.Subscriptions {
+		for id, sub := range subs {
+			select {
+			case sub.QuitChan <- true:
+				log.Info(fmt.Sprintf("closing subscription %s", id))
+			default:
+				log.Info(fmt.Sprintf("unable to close subscription %s; channel has no receiver", id))
+			}
+			delete(sds.Subscriptions[ty], id)
 		}
+		delete(sds.Subscriptions, ty)
+		delete(sds.SubscriptionTypes, ty)
 	}
 	sds.Unlock()
 }
@@ -193,10 +242,22 @@ func (sds *MockStateDiffService) Stop() error {
 	return nil
 }
 
-// StateDiffAt mock method
-func (sds *MockStateDiffService) StateDiffAt(blockNumber uint64) (*statediff.Payload, error) {
-	currentBlock := sds.BlockChain.GetBlockByNumber(blockNumber)
-	parentBlock := sds.BlockChain.GetBlockByHash(currentBlock.ParentHash())
-	log.Info(fmt.Sprintf("sending state diff at %d", blockNumber))
-	return sds.processStateDiff(currentBlock, parentBlock)
+// closeType is used to close all subscriptions of given type
+// closeType needs to be called with subscription access locked
+func (sds *MockStateDiffService) closeType(subType common.Hash) {
+	subs := sds.Subscriptions[subType]
+	for id, sub := range subs {
+		sendNonBlockingQuit(id, sub)
+	}
+	delete(sds.Subscriptions, subType)
+	delete(sds.SubscriptionTypes, subType)
+}
+
+func sendNonBlockingQuit(id rpc.ID, sub statediff.Subscription) {
+	select {
+	case sub.QuitChan <- true:
+		log.Info(fmt.Sprintf("closing subscription %s", id))
+	default:
+		log.Info("unable to close subscription %s; channel has no receiver", id)
+	}
 }
diff --git a/statediff/testhelpers/mocks/builder.go b/statediff/testhelpers/mocks/builder.go
index 034af0415..75af7cf55 100644
--- a/statediff/testhelpers/mocks/builder.go
+++ b/statediff/testhelpers/mocks/builder.go
@@ -17,28 +17,21 @@
 package mocks
 
 import (
-	"math/big"
-
-	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/statediff"
 )
 
 // Builder is a mock state diff builder
 type Builder struct {
-	OldStateRoot common.Hash
-	NewStateRoot common.Hash
-	BlockNumber  *big.Int
-	BlockHash    common.Hash
+	Args         statediff.Args
+	Params       statediff.Params
 	stateDiff    statediff.StateDiff
 	builderError error
 }
 
 // BuildStateDiff mock method
-func (builder *Builder) BuildStateDiff(oldStateRoot, newStateRoot common.Hash, blockNumber *big.Int, blockHash common.Hash) (statediff.StateDiff, error) {
-	builder.OldStateRoot = oldStateRoot
-	builder.NewStateRoot = newStateRoot
-	builder.BlockNumber = blockNumber
-	builder.BlockHash = blockHash
+func (builder *Builder) BuildStateDiff(args statediff.Args, params statediff.Params) (statediff.StateDiff, error) {
+	builder.Args = args
+	builder.Params = params
 
 	return builder.stateDiff, builder.builderError
 }
diff --git a/statediff/types.go b/statediff/types.go
index faa372715..538691caa 100644
--- a/statediff/types.go
+++ b/statediff/types.go
@@ -81,10 +81,9 @@ func (sd *Payload) Encode() ([]byte, error) {
 
 // StateDiff is the final output structure from the builder
 type StateDiff struct {
-	BlockNumber       *big.Int    `json:"blockNumber"     gencodec:"required"`
-	BlockHash         common.Hash `json:"blockHash"       gencodec:"required"`
-	LeafNodes         []StateNode `json:"leafNodes" gencodec:"required"`
-	IntermediateNodes []StateNode `json:"intermediateNodes" gencodec:"required"`
+	BlockNumber *big.Int    `json:"blockNumber"     gencodec:"required"`
+	BlockHash   common.Hash `json:"blockHash"       gencodec:"required"`
+	Nodes       []StateNode `json:"Nodes" gencodec:"required"`
 
 	encoded []byte
 	err     error