Added merkletree package (0xPolygonHermez#805)

* Added merkletree package

* fix test

* remove uneeded txBundleID

* remove unneeded nested type

Makefile

@@ -240,7 +240,7 @@ generate-mocks: ## Generates mocks for the tests, using mockery tool
 
 .PHONY: generate-code-from-proto
 generate-code-from-proto: ## Generates code from proto files
-	cd proto/src/proto/mt/v1 && protoc --proto_path=. --go_out=../../../../../state/tree/pb --go-grpc_out=../../../../../state/tree/pb --go_opt=paths=source_relative --go-grpc_opt=paths=source_relative mt.proto
+	cd proto/src/proto/mt/v1 && protoc --proto_path=. --proto_path=../../../../include --go_out=../../../../../merkletree/pb --go-grpc_out=../../../../../merkletree/pb --go_opt=paths=source_relative --go-grpc_opt=paths=source_relative mt.proto
 	cd proto/src/proto/zkprover/v1 && protoc --proto_path=. --go_out=../../../../../proverclient/pb --go-grpc_out=../../../../../proverclient/pb --go_opt=paths=source_relative --go-grpc_opt=paths=source_relative zk-prover.proto
 	cd proto/src/proto/zkprover/v1 && protoc --proto_path=. --go_out=../../../../../proverservice/pb --go-grpc_out=../../../../../proverservice/pb --go-grpc_opt=paths=source_relative --go_opt=paths=source_relative zk-prover.proto
 	cd proto/src/proto/executor/v1 && protoc --proto_path=. --go_out=../../../../../statev2/runtime/executor/pb --go-grpc_out=../../../../../statev2/runtime/executor/pb --go-grpc_opt=paths=source_relative --go_opt=paths=source_relative executor.proto

merkletree/key.go

@@ -0,0 +1,109 @@
+package merkletree
+
+import (
+	"math/big"
+
+	"github.com/ethereum/go-ethereum/common"
+	poseidon "github.com/iden3/go-iden3-crypto/goldenposeidon"
+)
+
+// Key stores key of the leaf
+type Key [32]byte
+
+const (
+	HashPoseidonAllZeroes = "0xc71603f33a1144ca7953db0ab48808f4c4055e3364a246c33c18a9786cb0b359"
+)
+
+// keyEthAddr is the common code for all the keys related to ethereum addresses.
+func keyEthAddr(ethAddr common.Address, leafType leafType, key1Capacity [4]uint64) ([]byte, error) {
+	ethAddrBI := new(big.Int).SetBytes(ethAddr.Bytes())
+	ethAddrArr := scalar2fea(ethAddrBI)
+
+	key1 := [8]uint64{
+		ethAddrArr[0],
+		ethAddrArr[1],
+		ethAddrArr[2],
+		ethAddrArr[3],
+		ethAddrArr[4],
+		0,
+		uint64(leafType),
+		0,
+	}
+
+	result, err := poseidon.Hash(key1, key1Capacity)
+	if err != nil {
+		return nil, err
+	}
+
+	return h4ToFilledByteSlice(result[:]), nil
+}
+
+func defaultCapIn() ([4]uint64, error) {
+	capIn, err := stringToh4(HashPoseidonAllZeroes)
+	if err != nil {
+		return [4]uint64{}, err
+	}
+
+	return [4]uint64{capIn[0], capIn[1], capIn[2], capIn[3]}, nil
+}
+
+// KeyEthAddrBalance returns the key of balance leaf:
+// hk0: H([0, 0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0])
+// key: H([ethAddr[0:4], ethAddr[4:8], ethAddr[8:12], ethAddr[12:16], ethAddr[16:20], 0, 0, 0], [hk0[0], hk0[1], hk0[2], hk0[3]])
+func KeyEthAddrBalance(ethAddr common.Address) ([]byte, error) {
+	capIn, err := defaultCapIn()
+	if err != nil {
+		return nil, err
+	}
+
+	return keyEthAddr(ethAddr, leafTypeBalance, capIn)
+}
+
+// KeyEthAddrNonce returns the key of nonce leaf:
+// hk0: H([0, 0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0])
+// key: H([ethAddr[0:4], ethAddr[4:8], ethAddr[8:12], ethAddr[12:16], ethAddr[16:20], 0, 1, 0], [hk0[0], hk0[1], hk0[2], hk0[3]]
+func KeyEthAddrNonce(ethAddr common.Address) ([]byte, error) {
+	capIn, err := defaultCapIn()
+	if err != nil {
+		return nil, err
+	}
+
+	return keyEthAddr(ethAddr, leafTypeNonce, capIn)
+}
+
+// KeyContractCode returns the key of contract code leaf:
+// hk0: H([0, 0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0])
+// key: H([ethAddr[0:4], ethAddr[4:8], ethAddr[8:12], ethAddr[12:16], ethAddr[16:20], 0, 2, 0], [hk0[0], hk0[1], hk0[2], hk0[3]]
+func KeyContractCode(ethAddr common.Address) ([]byte, error) {
+	capIn, err := defaultCapIn()
+	if err != nil {
+		return nil, err
+	}
+
+	return keyEthAddr(ethAddr, leafTypeCode, capIn)
+}
+
+// KeyContractStorage returns the key of contract storage position leaf:
+// hk0: H([stoPos[0:4], stoPos[4:8], stoPos[8:12], stoPos[12:16], stoPos[16:20], stoPos[20:24], stoPos[24:28], stoPos[28:32], [0, 0, 0, 0])
+// key: H([ethAddr[0:4], ethAddr[4:8], ethAddr[8:12], ethAddr[12:16], ethAddr[16:20], 0, 3, 0], [hk0[0], hk0[1], hk0[2], hk0[3])
+func KeyContractStorage(ethAddr common.Address, storagePos []byte) ([]byte, error) {
+	storageBI := new(big.Int).SetBytes(storagePos)
+
+	storageArr := scalar2fea(storageBI)
+
+	hk0, err := poseidon.Hash([8]uint64{
+		storageArr[0],
+		storageArr[1],
+		storageArr[2],
+		storageArr[3],
+		storageArr[4],
+		storageArr[5],
+		storageArr[6],
+		storageArr[7],
+	}, [4]uint64{})
+	if err != nil {
+		return nil, err
+	}
+
+	return keyEthAddr(ethAddr, leafTypeStorage, hk0)
+}

merkletree/key_test.go

@@ -0,0 +1,100 @@
+package merkletree
+
+import (
+	"encoding/hex"
+	"encoding/json"
+	"fmt"
+	"math/big"
+	"os"
+	"path"
+	"runtime"
+	"testing"
+
+	"github.com/ethereum/go-ethereum/common"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+type testVectorKey struct {
+	EthAddr         string `json:"ethAddr"`
+	StoragePosition string `json:"storagePosition"`
+	ExpectedKey     string `json:"expectedKey"`
+}
+
+func init() {
+	// Change dir to project root
+	// This is important because we have relative paths to files containing test vectors
+	_, filename, _, _ := runtime.Caller(0)
+	dir := path.Join(path.Dir(filename), "../")
+	err := os.Chdir(dir)
+	if err != nil {
+		panic(err)
+	}
+}
+
+func Test_CommonKeys(t *testing.T) {
+	tcs := []struct {
+		description    string
+		testVectorFile string
+		keyFunc        func(common.Address) ([]byte, error)
+	}{
+		{
+			description:    "keyEthAddressBalance",
+			testVectorFile: "test/vectors/src/merkle-tree/smt-key-eth-balance.json",
+			keyFunc:        KeyEthAddrBalance,
+		},
+		{
+			description:    "keyEthAddressNonce",
+			testVectorFile: "test/vectors/src/merkle-tree/smt-key-eth-nonce.json",
+			keyFunc:        KeyEthAddrNonce,
+		},
+		{
+			description:    "keyContractCode",
+			testVectorFile: "test/vectors/src/merkle-tree/smt-key-contract-code.json",
+			keyFunc:        KeyContractCode,
+		},
+	}
+	for _, tc := range tcs {
+		tc := tc
+
+		data, err := os.ReadFile(tc.testVectorFile)
+		require.NoError(t, err)
+
+		var testVectors []testVectorKey
+		err = json.Unmarshal(data, &testVectors)
+		require.NoError(t, err)
+
+		for ti, testVector := range testVectors {
+			t.Run(fmt.Sprintf("%s, test vector %d", tc.description, ti), func(t *testing.T) {
+				key, err := tc.keyFunc(common.HexToAddress(testVector.EthAddr))
+				require.NoError(t, err)
+				require.Equal(t, len(key), maxBigIntLen)
+
+				expected, _ := new(big.Int).SetString(testVector.ExpectedKey, 10)
+				assert.Equal(t, hex.EncodeToString(expected.Bytes()), hex.EncodeToString(key))
+			})
+		}
+	}
+}
+
+func Test_KeyContractStorage(t *testing.T) {
+	data, err := os.ReadFile("test/vectors/src/merkle-tree/smt-key-contract-storage.json")
+	require.NoError(t, err)
+
+	var testVectors []testVectorKey
+	err = json.Unmarshal(data, &testVectors)
+	require.NoError(t, err)
+
+	for ti, testVector := range testVectors {
+		t.Run(fmt.Sprintf("Test vector %d", ti), func(t *testing.T) {
+			storagePosition, ok := new(big.Int).SetString(testVector.StoragePosition, 10)
+			require.True(t, ok)
+			key, err := KeyContractStorage(common.HexToAddress(testVector.EthAddr), storagePosition.Bytes())
+			require.NoError(t, err)
+			require.Equal(t, len(key), maxBigIntLen)
+
+			expected, _ := new(big.Int).SetString(testVector.ExpectedKey, 10)
+			assert.Equal(t, hex.EncodeToString(expected.Bytes()), hex.EncodeToString(key))
+		})
+	}
+}

merkletree/leaf.go

@@ -0,0 +1,15 @@
+package merkletree
+
+// leafType specifies type of the leaf
+type leafType uint8
+
+const (
+	// leafTypeBalance specifies that leaf stores Balance
+	leafTypeBalance leafType = 0
+	// leafTypeNonce specifies that leaf stores Nonce
+	leafTypeNonce leafType = 1
+	// leafTypeCode specifies that leaf stores Code
+	leafTypeCode leafType = 2
+	// leafTypeStorage specifies that leaf stores Storage Value
+	leafTypeStorage leafType = 3
+)