refactor trie root hash computation

common/interface.go

@@ -51,7 +51,7 @@ type Trie interface {
 	GetSerializedNodes([]byte, uint64) ([][]byte, uint64, error)
 	GetSerializedNode([]byte) ([]byte, error)
 	GetAllLeavesOnChannel(allLeavesChan *TrieIteratorChannels, ctx context.Context, rootHash []byte, keyBuilder KeyBuilder, trieLeafParser TrieLeafParser) error
-	GetProof(key []byte) ([][]byte, []byte, error)
+	GetProof(key []byte, rootHash []byte) ([][]byte, []byte, error)
 	VerifyProof(rootHash []byte, key []byte, proof [][]byte) (bool, error)
 	GetStorageManager() StorageManager
 	IsMigratedToLatestVersion() (bool, error)

integrationTests/state/stateTrieSync/stateTrieSync_test.go

@@ -513,18 +513,16 @@ func testSyncMissingSnapshotNodes(t *testing.T, version int) {
 }
 
 func GetAllHashes(t *testing.T, tr common.Trie, rootHash []byte) [][]byte {
-	iterator, err := trie.NewDFSIterator(tr)
+	iterator, err := trie.NewDFSIterator(tr, rootHash)
 	require.Nil(t, err)
 
 	hashes := make([][]byte, 0)
-	hash, _ := iterator.GetHash()
-	hashes = append(hashes, hash)
+	hashes = append(hashes, iterator.GetHash())
 	for iterator.HasNext() {
 		err = iterator.Next()
 		require.Nil(t, err)
 
-		h, _ := iterator.GetHash()
-		hashes = append(hashes, h)
+		hashes = append(hashes, iterator.GetHash())
 	}
 
 	return hashes
@@ -550,7 +548,6 @@ func copyPartialState(t *testing.T, sourceNode, destinationNode *integrationTest
 		err = destStorage.Put(hash, val)
 		assert.Nil(t, err)
 	}
-
 }
 
 func getDataTriesHashes(t *testing.T, tr common.Trie, dataTriesRootHashes [][]byte) [][]byte {
@@ -560,7 +557,6 @@ func getDataTriesHashes(t *testing.T, tr common.Trie, dataTriesRootHashes [][]by
 		assert.Nil(t, err)
 
 		dtHashes := GetAllHashes(t, dt, rh)
-		assert.Nil(t, err)
 
 		hashes = append(hashes, dtHashes...)
 	}

node/node.go

@@ -1438,7 +1438,7 @@ func (n *Node) getProof(rootHash []byte, key []byte) (*common.GetProofResponse,
 		return nil, err
 	}
 
-	computedProof, value, err := tr.GetProof(key)
+	computedProof, value, err := tr.GetProof(key, rootHash)
 	if err != nil {
 		return nil, err
 	}

node/node_test.go

@@ -4004,7 +4004,7 @@ func TestNode_GetProofShouldWork(t *testing.T) {
 	stateComponents.AccountsAPI = &stateMock.AccountsStub{
 		GetTrieCalled: func(_ []byte) (common.Trie, error) {
 			return &trieMock.TrieStub{
-				GetProofCalled: func(key []byte) ([][]byte, []byte, error) {
+				GetProofCalled: func(key []byte, _ []byte) ([][]byte, []byte, error) {
 					assert.Equal(t, trieKey, hex.EncodeToString(key))
 					return proof, value, nil
 				},
@@ -4053,7 +4053,7 @@ func TestNode_getProofErrWhenComputingProof(t *testing.T) {
 	stateComponents.AccountsAPI = &stateMock.AccountsStub{
 		GetTrieCalled: func(_ []byte) (common.Trie, error) {
 			return &trieMock.TrieStub{
-				GetProofCalled: func(_ []byte) ([][]byte, []byte, error) {
+				GetProofCalled: func(_ []byte, _ []byte) ([][]byte, []byte, error) {
 					return nil, nil, expectedErr
 				},
 			}, nil
@@ -4129,7 +4129,7 @@ func TestNode_GetProofDataTrieShouldWork(t *testing.T) {
 	stateComponents.AccountsAPI = &stateMock.AccountsStub{
 		GetTrieCalled: func(_ []byte) (common.Trie, error) {
 			return &trieMock.TrieStub{
-				GetProofCalled: func(key []byte) ([][]byte, []byte, error) {
+				GetProofCalled: func(key []byte, _ []byte) ([][]byte, []byte, error) {
 					if hex.EncodeToString(key) == mainTrieKey {
 						return mainTrieProof, mainTrieValue, nil
 					}

testscommon/trie/trieStub.go

@@ -25,7 +25,7 @@ type TrieStub struct {
 	AppendToOldHashesCalled         func([][]byte)
 	GetSerializedNodesCalled        func([]byte, uint64) ([][]byte, uint64, error)
 	GetAllLeavesOnChannelCalled     func(leavesChannels *common.TrieIteratorChannels, ctx context.Context, rootHash []byte, keyBuilder common.KeyBuilder, trieLeafParser common.TrieLeafParser) error
-	GetProofCalled                  func(key []byte) ([][]byte, []byte, error)
+	GetProofCalled                  func(key []byte, rootHash []byte) ([][]byte, []byte, error)
 	VerifyProofCalled               func(rootHash []byte, key []byte, proof [][]byte) (bool, error)
 	GetStorageManagerCalled         func() common.StorageManager
 	GetSerializedNodeCalled         func(bytes []byte) ([]byte, error)
@@ -45,9 +45,9 @@ func (ts *TrieStub) GetStorageManager() common.StorageManager {
 }
 
 // GetProof -
-func (ts *TrieStub) GetProof(key []byte) ([][]byte, []byte, error) {
+func (ts *TrieStub) GetProof(key []byte, rootHash []byte) ([][]byte, []byte, error) {
 	if ts.GetProofCalled != nil {
-		return ts.GetProofCalled(key)
+		return ts.GetProofCalled(key, rootHash)
 	}
 
 	return nil, nil, nil

trie/baseIterator.go

@@ -12,11 +12,16 @@ type baseIterator struct {
 }
 
 // newBaseIterator creates a new instance of trie iterator
-func newBaseIterator(trie common.Trie) (*baseIterator, error) {
+func newBaseIterator(trie common.Trie, rootHash []byte) (*baseIterator, error) {
 	if check.IfNil(trie) {
 		return nil, ErrNilTrie
 	}
 
+	trie, err := trie.Recreate(rootHash)
+	if err != nil {
+		return nil, err
+	}
+
 	pmt, ok := trie.(*patriciaMerkleTrie)
 	if !ok {
 		return nil, ErrWrongTypeAssertion
@@ -56,20 +61,10 @@ func (it *baseIterator) next() ([]node, error) {
 
 // MarshalizedNode marshalizes the current node, and then returns the serialized node
 func (it *baseIterator) MarshalizedNode() ([]byte, error) {
-	err := it.currentNode.setHash()
-	if err != nil {
-		return nil, err
-	}
-
 	return it.currentNode.getEncodedNode()
 }
 
 // GetHash returns the current node hash
-func (it *baseIterator) GetHash() ([]byte, error) {
-	err := it.currentNode.setHash()
-	if err != nil {
-		return nil, err
-	}
-
-	return it.currentNode.getHash(), nil
+func (it *baseIterator) GetHash() []byte {
+	return it.currentNode.getHash()
 }

trie/baseIterator_test.go

@@ -14,8 +14,9 @@ func TestNewBaseIterator(t *testing.T) {
 	t.Parallel()
 
 	tr := initTrie()
+	rootHash, _ := tr.RootHash()
 
-	it, err := trie.NewBaseIterator(tr)
+	it, err := trie.NewBaseIterator(tr, rootHash)
 	assert.Nil(t, err)
 	assert.NotNil(t, it)
 }
@@ -25,7 +26,7 @@ func TestNewBaseIteratorNilTrieShouldErr(t *testing.T) {
 
 	var tr common.Trie
 
-	it, err := trie.NewBaseIterator(tr)
+	it, err := trie.NewBaseIterator(tr, nil)
 	assert.Nil(t, it)
 	assert.Equal(t, trie.ErrNilTrie, err)
 }
@@ -35,21 +36,24 @@ func TestBaseIterator_HasNext(t *testing.T) {
 
 	tr := emptyTrie()
 	_ = tr.Update([]byte("dog"), []byte("dog"))
-	trie.ExecuteUpdatesFromBatch(tr)
-	it, _ := trie.NewBaseIterator(tr)
+	_ = tr.Commit()
+	rootHash, _ := tr.RootHash()
+	it, _ := trie.NewBaseIterator(tr, rootHash)
 	assert.False(t, it.HasNext())
 
 	_ = tr.Update([]byte("doe"), []byte("doe"))
-	trie.ExecuteUpdatesFromBatch(tr)
-	it, _ = trie.NewBaseIterator(tr)
+	_ = tr.Commit()
+	rootHash, _ = tr.RootHash()
+	it, _ = trie.NewBaseIterator(tr, rootHash)
 	assert.True(t, it.HasNext())
 }
 
 func TestBaseIterator_GetMarshalizedNode(t *testing.T) {
 	t.Parallel()
 
 	tr := initTrie()
-	it, _ := trie.NewBaseIterator(tr)
+	rootHash, _ := tr.RootHash()
+	it, _ := trie.NewBaseIterator(tr, rootHash)
 
 	encNode, err := it.MarshalizedNode()
 	assert.Nil(t, err)
@@ -64,11 +68,11 @@ func TestBaseIterator_GetHash(t *testing.T) {
 	t.Parallel()
 
 	tr := initTrie()
+	_ = tr.Commit()
 	rootHash, _ := tr.RootHash()
-	it, _ := trie.NewBaseIterator(tr)
+	it, _ := trie.NewBaseIterator(tr, rootHash)
 
-	hash, err := it.GetHash()
-	assert.Nil(t, err)
+	hash := it.GetHash()
 	assert.Equal(t, rootHash, hash)
 }
 
@@ -80,7 +84,7 @@ func TestIterator_Search(t *testing.T) {
 	_ = tr.Update([]byte("dog"), []byte("puppy"))
 	_ = tr.Update([]byte("ddog"), []byte("cat"))
 	_ = tr.Update([]byte("ddoge"), []byte("foo"))
-	trie.ExecuteUpdatesFromBatch(tr)
+	_ = tr.Commit()
 
 	expectedHashes := []string{
 		"ecc2304769996585131ad6276c1422265813a2b79d60392130c4baa19a9b4e06",
@@ -109,9 +113,9 @@ func TestIterator_Search(t *testing.T) {
 			expectedHashes[8],
 		}
 
-		it, _ := trie.NewDFSIterator(tr)
-		nodeHash, err := it.GetHash()
-		require.Nil(t, err)
+		rootHash, _ := tr.RootHash()
+		it, _ := trie.NewDFSIterator(tr, rootHash)
+		nodeHash := it.GetHash()
 
 		nodesHashes := make([]string, 0)
 		nodesHashes = append(nodesHashes, hex.EncodeToString(nodeHash))
@@ -120,8 +124,7 @@ func TestIterator_Search(t *testing.T) {
 			err := it.Next()
 			require.Nil(t, err)
 
-			nodeHash, err := it.GetHash()
-			require.Nil(t, err)
+			nodeHash := it.GetHash()
 
 			nodesHashes = append(nodesHashes, hex.EncodeToString(nodeHash))
 		}

trie/baseNode.go

@@ -2,7 +2,7 @@ package trie
 
 import (
 	"sync"
-	
+
 	"github.com/multiversx/mx-chain-core-go/hashing"
 	"github.com/multiversx/mx-chain-core-go/marshal"
 )
@@ -16,20 +16,30 @@ type baseNode struct {
 }
 
 func (bn *baseNode) getHash() []byte {
-	//TODO add mutex protection for all methods
+	bn.mutex.RLock()
+	defer bn.mutex.RUnlock()
 
 	return bn.hash
 }
 
 func (bn *baseNode) setGivenHash(hash []byte) {
+	bn.mutex.Lock()
+	defer bn.mutex.Unlock()
+
 	bn.hash = hash
 }
 
 func (bn *baseNode) isDirty() bool {
+	bn.mutex.RLock()
+	defer bn.mutex.RUnlock()
+
 	return bn.dirty
 }
 
 func (bn *baseNode) setDirty(dirty bool) {
+	bn.mutex.Lock()
+	defer bn.mutex.Unlock()
+
 	bn.dirty = dirty
 }