1package arbo
  2
  3import (
  4	"encoding/hex"
  5	"fmt"
  6	"math/big"
  7	"testing"
  8
  9	"github.com/iden3/go-merkletree/db/memory"
 10	"github.com/stretchr/testify/assert"
 11	"github.com/stretchr/testify/require"
 12)
 13
 14func TestAddTestVectors(t *testing.T) {
 15	// Poseidon test vectors generated using https://github.com/iden3/circomlib smt.js
 16	testVectorsPoseidon := []string{
 17		"0000000000000000000000000000000000000000000000000000000000000000",
 18		"13578938674299138072471463694055224830892726234048532520316387704878000008795",
 19		"5412393676474193513566895793055462193090331607895808993925969873307089394741",
 20		"14204494359367183802864593755198662203838502594566452929175967972147978322084",
 21	}
 22	testAdd(t, HashFunctionPoseidon, testVectorsPoseidon)
 23
 24	testVectorsSha256 := []string{
 25		"0000000000000000000000000000000000000000000000000000000000000000",
 26		"46910109172468462938850740851377282682950237270676610513794735904325820156367",
 27		"59481735341404520835410489183267411392292882901306595567679529387376287440550",
 28		"20573794434149960984975763118181266662429997821552560184909083010514790081771",
 29	}
 30	testAdd(t, HashFunctionSha256, testVectorsSha256)
 31}
 32
 33func testAdd(t *testing.T, hashFunc HashFunction, testVectors []string) {
 34	tree, err := NewTree(memory.NewMemoryStorage(), 10, hashFunc)
 35	assert.Nil(t, err)
 36	defer tree.db.Close()
 37	assert.Equal(t, testVectors[0], hex.EncodeToString(tree.Root()))
 38
 39	err = tree.Add(
 40		BigIntToBytes(big.NewInt(1)),
 41		BigIntToBytes(big.NewInt(2)))
 42	assert.Nil(t, err)
 43	rootBI := BytesToBigInt(tree.Root())
 44	assert.Equal(t, testVectors[1], rootBI.String())
 45
 46	err = tree.Add(
 47		BigIntToBytes(big.NewInt(33)),
 48		BigIntToBytes(big.NewInt(44)))
 49	assert.Nil(t, err)
 50	rootBI = BytesToBigInt(tree.Root())
 51	assert.Equal(t, testVectors[2], rootBI.String())
 52
 53	err = tree.Add(
 54		BigIntToBytes(big.NewInt(1234)),
 55		BigIntToBytes(big.NewInt(9876)))
 56	assert.Nil(t, err)
 57	rootBI = BytesToBigInt(tree.Root())
 58	assert.Equal(t, testVectors[3], rootBI.String())
 59}
 60
 61func TestAdd1000(t *testing.T) {
 62	tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
 63	require.Nil(t, err)
 64
 65	defer tree.db.Close()
 66	for i := 0; i < 1000; i++ {
 67		k := BigIntToBytes(big.NewInt(int64(i)))
 68		v := BigIntToBytes(big.NewInt(0))
 69		if err := tree.Add(k, v); err != nil {
 70			t.Fatal(err)
 71		}
 72	}
 73
 74	rootBI := BytesToBigInt(tree.Root())
 75	assert.Equal(t,
 76		"296519252211642170490407814696803112091039265640052570497930797516015811235",
 77		rootBI.String())
 78}
 79
 80func TestAddDifferentOrder(t *testing.T) {
 81	tree1, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
 82	require.Nil(t, err)
 83
 84	defer tree1.db.Close()
 85	for i := 0; i < 16; i++ {
 86		k := SwapEndianness(big.NewInt(int64(i)).Bytes())
 87		v := SwapEndianness(big.NewInt(0).Bytes())
 88		if err := tree1.Add(k, v); err != nil {
 89			t.Fatal(err)
 90		}
 91	}
 92
 93	tree2, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
 94	require.Nil(t, err)
 95	defer tree2.db.Close()
 96	for i := 16 - 1; i >= 0; i-- {
 97		k := big.NewInt(int64(i)).Bytes()
 98		v := big.NewInt(0).Bytes()
 99		if err := tree2.Add(k, v); err != nil {
100			t.Fatal(err)
101		}
102	}
103
104	assert.Equal(t, hex.EncodeToString(tree1.Root()), hex.EncodeToString(tree2.Root()))
105	assert.Equal(t,
106		"3b89100bec24da9275c87bc188740389e1d5accfc7d88ba5688d7fa96a00d82f",
107		hex.EncodeToString(tree1.Root()))
108}
109
110func TestAddRepeatedIndex(t *testing.T) {
111	tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
112	require.Nil(t, err)
113	defer tree.db.Close()
114	k := big.NewInt(int64(3)).Bytes()
115	v := big.NewInt(int64(12)).Bytes()
116	if err := tree.Add(k, v); err != nil {
117		t.Fatal(err)
118	}
119	err = tree.Add(k, v)
120	assert.NotNil(t, err)
121	assert.Equal(t, fmt.Errorf("max virtual level 100"), err)
122}
123
124func TestAux(t *testing.T) {
125	tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
126	require.Nil(t, err)
127	defer tree.db.Close()
128	k := BigIntToBytes(big.NewInt(int64(1)))
129	v := BigIntToBytes(big.NewInt(int64(0)))
130	err = tree.Add(k, v)
131	assert.Nil(t, err)
132	k = BigIntToBytes(big.NewInt(int64(256)))
133	err = tree.Add(k, v)
134	assert.Nil(t, err)
135
136	k = BigIntToBytes(big.NewInt(int64(257)))
137	err = tree.Add(k, v)
138	assert.Nil(t, err)
139
140	k = BigIntToBytes(big.NewInt(int64(515)))
141	err = tree.Add(k, v)
142	assert.Nil(t, err)
143	k = BigIntToBytes(big.NewInt(int64(770)))
144	err = tree.Add(k, v)
145	assert.Nil(t, err)
146}
147
148func TestGenProof(t *testing.T) {
149	tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
150	require.Nil(t, err)
151
152	defer tree.db.Close()
153	for i := 0; i < 10; i++ {
154		k := BigIntToBytes(big.NewInt(int64(i)))
155		v := BigIntToBytes(big.NewInt(int64(i * 2)))
156		if err := tree.Add(k, v); err != nil {
157			t.Fatal(err)
158		}
159	}
160
161	k := BigIntToBytes(big.NewInt(int64(7)))
162	_, err = tree.GenProof(k)
163	assert.Nil(t, err)
164}
165
166func BenchmarkAdd(b *testing.B) {
167	// prepare inputs
168	var ks, vs [][]byte
169	for i := 0; i < 1000; i++ {
170		k := BigIntToBytes(big.NewInt(int64(i)))
171		v := BigIntToBytes(big.NewInt(int64(i)))
172		ks = append(ks, k)
173		vs = append(vs, v)
174	}
175
176	b.Run("Poseidon", func(b *testing.B) {
177		benchmarkAdd(b, HashFunctionPoseidon, ks, vs)
178	})
179	b.Run("Sha256", func(b *testing.B) {
180		benchmarkAdd(b, HashFunctionSha256, ks, vs)
181	})
182}
183
184func benchmarkAdd(b *testing.B, hashFunc HashFunction, ks, vs [][]byte) {
185	tree, err := NewTree(memory.NewMemoryStorage(), 140, hashFunc)
186	require.Nil(b, err)
187
188	defer tree.db.Close()
189	for i := 0; i < len(ks); i++ {
190		if err := tree.Add(ks[i], vs[i]); err != nil {
191			b.Fatal(err)
192		}
193	}
194}