From cd5810d6cf00b1bc17a981e68dda54e1f57d8ffd Mon Sep 17 00:00:00 2001
From: Roman Khimov <roman@nspcc.ru>
Date: Thu, 27 Jan 2022 14:25:11 +0300
Subject: [PATCH] mpt: simplify makeStorageKey() interface

---
 pkg/core/mpt/batch_test.go  |  2 +-
 pkg/core/mpt/billet.go      |  4 ++--
 pkg/core/mpt/billet_test.go |  2 +-
 pkg/core/mpt/proof.go       |  2 +-
 pkg/core/mpt/trie.go        | 10 +++++-----
 pkg/core/mpt/trie_test.go   |  2 +-
 6 files changed, 11 insertions(+), 11 deletions(-)

diff --git a/pkg/core/mpt/batch_test.go b/pkg/core/mpt/batch_test.go
index 6d763b61e..3df6f7bc0 100644
--- a/pkg/core/mpt/batch_test.go
+++ b/pkg/core/mpt/batch_test.go
@@ -274,7 +274,7 @@ func TestTrie_PutBatchHash(t *testing.T) {
 		}
 		tr1.Collapse(1)
 		tr2.Collapse(1)
-		key := makeStorageKey(tr1.root.(*BranchNode).Children[2].Hash().BytesBE())
+		key := makeStorageKey(tr1.root.(*BranchNode).Children[2].Hash())
 		require.NoError(t, tr1.Store.Delete(key))
 		require.NoError(t, tr2.Store.Delete(key))
 		testIncompletePut(t, ps, 1, tr1, tr2)
diff --git a/pkg/core/mpt/billet.go b/pkg/core/mpt/billet.go
index 39c4770fe..1840e04a5 100644
--- a/pkg/core/mpt/billet.go
+++ b/pkg/core/mpt/billet.go
@@ -177,7 +177,7 @@ func (b *Billet) putIntoHash(curr *HashNode, path []byte, val Node) (Node, error
 }
 
 func (b *Billet) incrementRefAndStore(h util.Uint256, bs []byte) {
-	key := makeStorageKey(h.BytesBE())
+	key := makeStorageKey(h)
 	if b.refcountEnabled {
 		var (
 			err  error
@@ -325,7 +325,7 @@ func (b *Billet) tryCollapseBranch(curr *BranchNode) Node {
 
 // GetFromStore returns MPT node from the storage.
 func (b *Billet) GetFromStore(h util.Uint256) (Node, error) {
-	data, err := b.Store.Get(makeStorageKey(h.BytesBE()))
+	data, err := b.Store.Get(makeStorageKey(h))
 	if err != nil {
 		return nil, err
 	}
diff --git a/pkg/core/mpt/billet_test.go b/pkg/core/mpt/billet_test.go
index 4893d0d41..be24f9a83 100644
--- a/pkg/core/mpt/billet_test.go
+++ b/pkg/core/mpt/billet_test.go
@@ -16,7 +16,7 @@ func TestBillet_RestoreHashNode(t *testing.T) {
 		_ = expectedRoot.Hash()
 		_ = tr.root.Hash()
 		require.Equal(t, expectedRoot, tr.root)
-		expectedBytes, err := tr.Store.Get(makeStorageKey(expectedNode.Hash().BytesBE()))
+		expectedBytes, err := tr.Store.Get(makeStorageKey(expectedNode.Hash()))
 		if expectedRefCount != 0 {
 			require.NoError(t, err)
 			require.Equal(t, expectedRefCount, binary.LittleEndian.Uint32(expectedBytes[len(expectedBytes)-4:]))
diff --git a/pkg/core/mpt/proof.go b/pkg/core/mpt/proof.go
index cc8b5876d..a769550ca 100644
--- a/pkg/core/mpt/proof.go
+++ b/pkg/core/mpt/proof.go
@@ -70,7 +70,7 @@ func VerifyProof(rh util.Uint256, key []byte, proofs [][]byte) ([]byte, bool) {
 	for i := range proofs {
 		h := hash.DoubleSha256(proofs[i])
 		// no errors in Put to memory store
-		_ = tr.Store.Put(makeStorageKey(h[:]), proofs[i])
+		_ = tr.Store.Put(makeStorageKey(h), proofs[i])
 	}
 	_, leaf, _, err := tr.getWithPath(tr.root, path, true)
 	if err != nil {
diff --git a/pkg/core/mpt/trie.go b/pkg/core/mpt/trie.go
index da466096d..7bc1c603e 100644
--- a/pkg/core/mpt/trie.go
+++ b/pkg/core/mpt/trie.go
@@ -372,8 +372,8 @@ func (t *Trie) StateRoot() util.Uint256 {
 	return t.root.Hash()
 }
 
-func makeStorageKey(mptKey []byte) []byte {
-	return append([]byte{byte(storage.DataMPT)}, mptKey...)
+func makeStorageKey(mptKey util.Uint256) []byte {
+	return append([]byte{byte(storage.DataMPT)}, mptKey[:]...)
 }
 
 // Flush puts every node in the trie except Hash ones to the storage.
@@ -392,7 +392,7 @@ func (t *Trie) Flush() {
 					delete(t.refcount, h)
 				}
 			} else if node.refcount > 0 {
-				_ = t.Store.Put(makeStorageKey(h.BytesBE()), node.bytes)
+				_ = t.Store.Put(makeStorageKey(h), node.bytes)
 			}
 			node.refcount = 0
 		} else {
@@ -407,7 +407,7 @@ func (t *Trie) updateRefCount(h util.Uint256) int32 {
 		panic("`updateRefCount` is called, but GC is disabled")
 	}
 	var data []byte
-	key := makeStorageKey(h.BytesBE())
+	key := makeStorageKey(h)
 	node := t.refcount[h]
 	cnt := node.initial
 	if cnt == 0 {
@@ -466,7 +466,7 @@ func (t *Trie) removeRef(h util.Uint256, bs []byte) {
 }
 
 func (t *Trie) getFromStore(h util.Uint256) (Node, error) {
-	data, err := t.Store.Get(makeStorageKey(h.BytesBE()))
+	data, err := t.Store.Get(makeStorageKey(h))
 	if err != nil {
 		return nil, err
 	}
diff --git a/pkg/core/mpt/trie_test.go b/pkg/core/mpt/trie_test.go
index 81d7c9fe3..465d7964b 100644
--- a/pkg/core/mpt/trie_test.go
+++ b/pkg/core/mpt/trie_test.go
@@ -251,7 +251,7 @@ func (tr *Trie) putToStore(n Node) {
 		}
 		tr.updateRefCount(n.Hash())
 	} else {
-		_ = tr.Store.Put(makeStorageKey(n.Hash().BytesBE()), n.Bytes())
+		_ = tr.Store.Put(makeStorageKey(n.Hash()), n.Bytes())
 	}
 }