Generate generates deduplicated encoded proof nodes for N full keys

qdm12 · qdm12 · commit a475bfa5a784 · 2022-07-04T14:40:01.000Z
- Load trie only once from database
- Move StorageState `GenerateTrieProof` to `lib/trie/proof` package
- Update tests
diff --git a/dot/state/storage.go b/dot/state/storage.go
@@ -303,30 +303,6 @@ func (s *StorageState) LoadCodeHash(hash *common.Hash) (common.Hash, error) {
 
 // GenerateTrieProof returns the proofs related to the keys on the state root trie
 func (s *StorageState) GenerateTrieProof(stateRoot common.Hash, keys [][]byte) (
-	proofs [][]byte, err error) {
-	proofHashes := make(map[common.Hash]struct{}) // to avoid duplicates
-	for _, key := range keys {
-		encodedProofNodes, err := proof.Generate(stateRoot[:], key, s.db)
-		if err != nil {
-			return nil, err
-		}
-
-		for _, encodedProofNode := range encodedProofNodes {
-			proofHash, err := common.Blake2bHash(encodedProofNode)
-			if err != nil {
-				return nil, fmt.Errorf("blake2b hash: %w", err)
-			}
-
-			_, exists := proofHashes[proofHash]
-			if exists {
-				// this encoded node was encountered previously, skip it.
-				continue
-			}
-
-			proofHashes[proofHash] = struct{}{}
-			proofs = append(proofs, encodedProofNode)
-		}
-	}
-
-	return proofs, nil
+	encodedProofNodes [][]byte, err error) {
+	return proof.Generate(stateRoot[:], keys, s.db)
 }
diff --git a/lib/runtime/wasmer/imports_test.go b/lib/runtime/wasmer/imports_test.go
@@ -1797,12 +1797,8 @@ func Test_ext_trie_blake2_256_verify_proof_version_1(t *testing.T) {
 	root := hash.ToBytes()
 	otherRoot := otherHash.ToBytes()
 
-	var allProofs [][]byte
-	for _, key := range keys {
-		singleProof, err := proof.Generate(root, key, memdb)
-		require.NoError(t, err)
-		allProofs = append(allProofs, singleProof...)
-	}
+	allProofs, err := proof.Generate(root, keys, memdb)
+	require.NoError(t, err)
 
 	testcases := map[string]struct {
 		root, key, value []byte
diff --git a/lib/trie/proof/generate.go b/lib/trie/proof/generate.go
@@ -24,24 +24,43 @@ type Database interface {
 	Get(key []byte) (value []byte, err error)
 }
 
-// Generate returns the encoded proof nodes for the trie
-// corresponding to the root hash given, and for the (Little Endian)
-// full key given. The database given is used to load the trie
-// using the root hash given.
-func Generate(rootHash []byte, fullKey []byte, database Database) (
+// Generate generates and deduplicates the encoded proof nodes
+// for the trie corresponding to the root hash given, and for
+// the slice of (Little Endian) full keys given. The database given
+// is used to load the trie using the root hash given.
+func Generate(rootHash []byte, fullKeys [][]byte, database Database) (
 	encodedProofNodes [][]byte, err error) {
 	trie := trie.NewEmptyTrie()
 	if err := trie.Load(database, common.BytesToHash(rootHash)); err != nil {
 		return nil, fmt.Errorf("loading trie: %w", err)
 	}
-
 	rootNode := trie.RootNode()
-	fullKeyNibbles := codec.KeyLEToNibbles(fullKey)
-	encodedProofNodes, err = walk(rootNode, fullKeyNibbles)
-	if err != nil {
-		// Note we wrap the full key context here since walk is recursive and
-		// may not be aware of the initial full key.
-		return nil, fmt.Errorf("walking to node at key 0x%x: %w", fullKey, err)
+
+	hashesSeen := make(map[string]struct{})
+	for _, fullKey := range fullKeys {
+		fullKeyNibbles := codec.KeyLEToNibbles(fullKey)
+		newEncodedProofNodes, err := walk(rootNode, fullKeyNibbles)
+		if err != nil {
+			// Note we wrap the full key context here since walk is recursive and
+			// may not be aware of the initial full key.
+			return nil, fmt.Errorf("walking to node at key 0x%x: %w", fullKey, err)
+		}
+
+		for _, encodedProofNode := range newEncodedProofNodes {
+			digest, err := common.Blake2bHash(encodedProofNode)
+			if err != nil {
+				return nil, fmt.Errorf("blake2b hash: %w", err)
+			}
+			hashString := string(digest.ToBytes())
+
+			_, seen := hashesSeen[hashString]
+			if seen {
+				continue
+			}
+			hashesSeen[hashString] = struct{}{}
+
+			encodedProofNodes = append(encodedProofNodes, encodedProofNode)
+		}
 	}
 
 	return encodedProofNodes, nil
diff --git a/lib/trie/proof/generate_test.go b/lib/trie/proof/generate_test.go
@@ -27,7 +27,7 @@ func Test_Generate(t *testing.T) {
 
 	testCases := map[string]struct {
 		rootHash          []byte
-		fullKey           []byte // nibbles
+		fullKeysNibbles   [][]byte
 		databaseBuilder   func(ctrl *gomock.Controller) Database
 		encodedProofNodes [][]byte
 		errWrapped        error
@@ -48,7 +48,8 @@ func Test_Generate(t *testing.T) {
 				"test error",
 		},
 		"walk error": {
-			rootHash: someHash,
+			rootHash:        someHash,
+			fullKeysNibbles: [][]byte{{1}},
 			databaseBuilder: func(ctrl *gomock.Controller) Database {
 				mockDatabase := NewMockDatabase(ctrl)
 				encodedRoot := encodeNode(t, node.Node{
@@ -59,12 +60,12 @@ func Test_Generate(t *testing.T) {
 					Return(encodedRoot, nil)
 				return mockDatabase
 			},
-			fullKey:    []byte{1},
 			errWrapped: ErrKeyNotFound,
 			errMessage: "walking to node at key 0x01: key not found",
 		},
 		"leaf root": {
-			rootHash: someHash,
+			rootHash:        someHash,
+			fullKeysNibbles: [][]byte{{}},
 			databaseBuilder: func(ctrl *gomock.Controller) Database {
 				mockDatabase := NewMockDatabase(ctrl)
 				encodedRoot := encodeNode(t, node.Node{
@@ -83,7 +84,8 @@ func Test_Generate(t *testing.T) {
 			},
 		},
 		"branch root": {
-			rootHash: someHash,
+			rootHash:        someHash,
+			fullKeysNibbles: [][]byte{{}},
 			databaseBuilder: func(ctrl *gomock.Controller) Database {
 				mockDatabase := NewMockDatabase(ctrl)
 				encodedRoot := encodeNode(t, node.Node{
@@ -115,6 +117,108 @@ func Test_Generate(t *testing.T) {
 				}),
 			},
 		},
+		"target leaf of branch": {
+			rootHash: someHash,
+			fullKeysNibbles: [][]byte{
+				{1, 2, 3, 4},
+			},
+			databaseBuilder: func(ctrl *gomock.Controller) Database {
+				mockDatabase := NewMockDatabase(ctrl)
+				encodedRoot := encodeNode(t, node.Node{
+					Key:   []byte{1, 2},
+					Value: []byte{2},
+					Children: padRightChildren([]*node.Node{
+						nil, nil, nil,
+						{ // full key 1, 2, 3, 4
+							Key:   []byte{4},
+							Value: []byte{4},
+						},
+					}),
+				})
+				mockDatabase.EXPECT().Get(someHash).
+					Return(encodedRoot, nil)
+				return mockDatabase
+			},
+			encodedProofNodes: [][]byte{
+				encodeNode(t, node.Node{
+					Key:   []byte{1, 2},
+					Value: []byte{2},
+					Children: padRightChildren([]*node.Node{
+						nil, nil, nil,
+						{
+							Key:   []byte{4},
+							Value: []byte{4},
+						},
+					}),
+				}),
+				encodeNode(t, node.Node{
+					Key:   []byte{4},
+					Value: []byte{4},
+				}),
+			},
+		},
+		"deduplicate proof nodes": {
+			rootHash: someHash,
+			fullKeysNibbles: [][]byte{
+				{1, 2, 3, 4},
+				{1, 2, 4, 4},
+				{1, 2, 5, 4},
+			},
+			databaseBuilder: func(ctrl *gomock.Controller) Database {
+				mockDatabase := NewMockDatabase(ctrl)
+				encodedRoot := encodeNode(t, node.Node{
+					Key:   []byte{1, 2},
+					Value: []byte{2},
+					Children: padRightChildren([]*node.Node{
+						nil, nil, nil,
+						{ // full key 1, 2, 3, 4
+							Key:   []byte{4},
+							Value: []byte{4},
+						},
+						{ // full key 1, 2, 4, 4
+							Key:   []byte{4},
+							Value: []byte{4},
+						},
+						{ // full key 1, 2, 5, 4
+							Key:   []byte{4},
+							Value: []byte{5},
+						},
+					}),
+				})
+				mockDatabase.EXPECT().Get(someHash).
+					Return(encodedRoot, nil)
+				return mockDatabase
+			},
+			encodedProofNodes: [][]byte{
+				encodeNode(t, node.Node{
+					Key:   []byte{1, 2},
+					Value: []byte{2},
+					Children: padRightChildren([]*node.Node{
+						nil, nil, nil,
+						{ // full key 1, 2, 3, 4
+							Key:   []byte{4},
+							Value: []byte{4},
+						},
+						{ // full key 1, 2, 4, 4
+							Key:   []byte{4},
+							Value: []byte{4},
+						},
+						{ // full key 1, 2, 5, 4
+							Key:   []byte{4},
+							Value: []byte{5},
+						},
+					}),
+				}),
+				encodeNode(t, node.Node{
+					Key:   []byte{4},
+					Value: []byte{4},
+				}),
+				encodeNode(t, node.Node{
+					Key:   []byte{4},
+					Value: []byte{5},
+				}),
+			},
+		},
 	}
 
 	for name, testCase := range testCases {
@@ -124,9 +228,13 @@ func Test_Generate(t *testing.T) {
 			ctrl := gomock.NewController(t)
 
 			database := testCase.databaseBuilder(ctrl)
+			fullKeysLE := make([][]byte, len(testCase.fullKeysNibbles))
+			for i, fullKeyNibbles := range testCase.fullKeysNibbles {
+				fullKeysLE[i] = codec.NibblesToKeyLE(fullKeyNibbles)
+			}
 
 			encodedProofNodes, err := Generate(testCase.rootHash,
-				testCase.fullKey, database)
+				fullKeysLE, database)
 
 			assert.ErrorIs(t, err, testCase.errWrapped)
 			if testCase.errWrapped != nil {
diff --git a/lib/trie/proof/proof_test.go b/lib/trie/proof/proof_test.go
@@ -41,7 +41,8 @@ func Test_Generate_Verify(t *testing.T) {
 	require.NoError(t, err)
 
 	for i, key := range keys {
-		proof, err := Generate(rootHash.ToBytes(), []byte(key), database)
+		fullKeys := [][]byte{[]byte(key)}
+		proof, err := Generate(rootHash.ToBytes(), fullKeys, database)
 		require.NoError(t, err)
 
 		expectedValue := fmt.Sprintf("%x-%d", key, i)
