From 7c19e86698ee3d7a7387d9c6e02aa9b115560bb3 Mon Sep 17 00:00:00 2001
From: nabarun <nabarun@deepstacksoft.com>
Date: Thu, 12 May 2022 15:51:46 +0530
Subject: [PATCH] Call method for computing watchedAddressesLeafKeys

---
 pkg/builder.go |  4 ++--
 pkg/service.go | 20 ++++++++++++++++++++
 2 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/pkg/builder.go b/pkg/builder.go
index dfdf434..9890796 100644
--- a/pkg/builder.go
+++ b/pkg/builder.go
@@ -299,7 +299,7 @@ func (sdb *builder) buildStateDiff(args []iterPair, params sd.Params, output sdt
 	if params.IntermediateStateNodes {
 		diffAccountsAtB, diffPathsAtB, err = sdb.createdAndUpdatedStateWithIntermediateNodes(args[0], output)
 	} else {
-		diffAccountsAtB, diffPathsAtB, err = sdb.createdAndUpdatedState(args[0], params.WatchedAddressesLeafKeys)
+		diffAccountsAtB, diffPathsAtB, err = sdb.createdAndUpdatedState(args[0], params.WatchedAddressesLeafKeys())
 	}
 
 	if err != nil {
@@ -308,7 +308,7 @@ func (sdb *builder) buildStateDiff(args []iterPair, params sd.Params, output sdt
 
 	// collect a slice of all the nodes that existed at a path in A that doesn't exist in B
 	// a map of their leafkey to all the accounts that were touched and exist at A
-	diffAccountsAtA, err := sdb.deletedOrUpdatedState(args[1], diffAccountsAtB, diffPathsAtB, params.WatchedAddressesLeafKeys, params.IntermediateStorageNodes, output)
+	diffAccountsAtA, err := sdb.deletedOrUpdatedState(args[1], diffAccountsAtB, diffPathsAtB, params.WatchedAddressesLeafKeys(), params.IntermediateStorageNodes, output)
 	if err != nil {
 		return fmt.Errorf("error collecting deletedOrUpdatedNodes: %v", err)
 	}
diff --git a/pkg/service.go b/pkg/service.go
index 061e8b5..8d35d01 100644
--- a/pkg/service.go
+++ b/pkg/service.go
@@ -194,6 +194,10 @@ func (sds *Service) StateDiffAt(blockNumber uint64, params sd.Params) (*sd.Paylo
 		return nil, err
 	}
 	logrus.Infof("sending state diff at block %d", blockNumber)
+
+	// compute leaf keys of watched addresses in the params
+	params.ComputeWatchedAddressesLeafKeys()
+
 	if blockNumber == 0 {
 		return sds.processStateDiff(currentBlock, common.Hash{}, params)
 	}
@@ -212,6 +216,10 @@ func (sds *Service) StateDiffFor(blockHash common.Hash, params sd.Params) (*sd.P
 		return nil, err
 	}
 	logrus.Infof("sending state diff at block %s", blockHash.Hex())
+
+	// compute leaf keys of watched addresses in the params
+	params.ComputeWatchedAddressesLeafKeys()
+
 	if currentBlock.NumberU64() == 0 {
 		return sds.processStateDiff(currentBlock, common.Hash{}, params)
 	}
@@ -281,6 +289,10 @@ func (sds *Service) StateTrieAt(blockNumber uint64, params sd.Params) (*sd.Paylo
 		return nil, err
 	}
 	logrus.Infof("sending state trie at block %d", blockNumber)
+
+	// compute leaf keys of watched addresses in the params
+	params.ComputeWatchedAddressesLeafKeys()
+
 	return sds.processStateTrie(currentBlock, params)
 }
 
@@ -320,6 +332,10 @@ func (sds *Service) WriteStateDiffAt(blockNumber uint64, params sd.Params) error
 	if err != nil {
 		return err
 	}
+
+	// compute leaf keys of watched addresses in the params
+	params.ComputeWatchedAddressesLeafKeys()
+
 	parentRoot := common.Hash{}
 	if blockNumber != 0 {
 		parentBlock, err := sds.lvlDBReader.GetBlockByHash(currentBlock.ParentHash())
@@ -341,6 +357,10 @@ func (sds *Service) WriteStateDiffFor(blockHash common.Hash, params sd.Params) e
 	if err != nil {
 		return err
 	}
+
+	// compute leaf keys of watched addresses in the params
+	params.ComputeWatchedAddressesLeafKeys()
+
 	parentRoot := common.Hash{}
 	if currentBlock.NumberU64() != 0 {
 		parentBlock, err := sds.lvlDBReader.GetBlockByHash(currentBlock.ParentHash())