1package arbo
 2
 3import (
 4	"encoding/hex"
 5	"math/big"
 6	"testing"
 7
 8	qt "github.com/frankban/quicktest"
 9	"github.com/iden3/go-merkletree/db/memory"
10)
11
12func TestVirtualTreeTestVectors(t *testing.T) {
13	c := qt.New(t)
14
15	keys := [][]byte{
16		BigIntToBytes(big.NewInt(1)),
17		BigIntToBytes(big.NewInt(33)),
18		BigIntToBytes(big.NewInt(1234)),
19		BigIntToBytes(big.NewInt(123456789)),
20	}
21	values := [][]byte{
22		BigIntToBytes(big.NewInt(2)),
23		BigIntToBytes(big.NewInt(44)),
24		BigIntToBytes(big.NewInt(9876)),
25		BigIntToBytes(big.NewInt(987654321)),
26	}
27
28	// check the root for different batches of leafs
29	testVirtualTree(c, 10, keys[:1], values[:1])
30	testVirtualTree(c, 10, keys[:2], values[:2])
31	testVirtualTree(c, 10, keys[:3], values[:3])
32	testVirtualTree(c, 10, keys[:4], values[:4])
33}
34
35func TestVirtualTreeRandomKeys(t *testing.T) {
36	c := qt.New(t)
37
38	// test with hardcoded values
39	keys := make([][]byte, 8)
40	values := make([][]byte, 8)
41	keys[0], _ = hex.DecodeString("1c7c2265e368314ca58ed2e1f33a326f1220e234a566d55c3605439dbe411642")
42	keys[1], _ = hex.DecodeString("2c9f0a578afff5bfa4e0992a43066460faaab9e8e500db0b16647c701cdb16bf")
43	keys[2], _ = hex.DecodeString("9cb87ec67e875c61390edcd1ab517f443591047709a4d4e45b0f9ed980857b8e")
44	keys[3], _ = hex.DecodeString("9b4e9e92e974a589f426ceeb4cb291dc24893513fecf8e8460992dcf52621d4d")
45	keys[4], _ = hex.DecodeString("1c45cb31f2fa39ec7b9ebf0fad40e0b8296016b5ce8844ae06ff77226379d9a5")
46	keys[5], _ = hex.DecodeString("d8af98bbbb585129798ae54d5eabbc9d0561d583faf1663b3a3724d15bda4ec7")
47	keys[6], _ = hex.DecodeString("3cd55dbfb8f975f20a0925dfbdabe79fa2d51dd0268afbb8ba6b01de9dfcdd3c")
48	keys[7], _ = hex.DecodeString("5d0a9d6d9f197c091bf054fac9cb60e11ec723d6610ed8578e617b4d46cb43d5")
49
50	// check the root for different batches of leafs
51	testVirtualTree(c, 10, keys[:1], values[:1])
52	testVirtualTree(c, 10, keys, values)
53
54	// test with random values
55	nLeafs := 1024
56
57	keys = make([][]byte, nLeafs)
58	values = make([][]byte, nLeafs)
59	for i := 0; i < nLeafs; i++ {
60		keys[i] = randomBytes(32)
61		values[i] = []byte{0}
62	}
63
64	testVirtualTree(c, 100, keys, values)
65}
66
67func testVirtualTree(c *qt.C, maxLevels int, keys, values [][]byte) {
68	c.Assert(len(keys), qt.Equals, len(values))
69
70	// normal tree, to have an expected root value
71	tree, err := NewTree(memory.NewMemoryStorage(), maxLevels, HashFunctionSha256)
72	c.Assert(err, qt.IsNil)
73	for i := 0; i < len(keys); i++ {
74		err := tree.Add(keys[i], values[i])
75		c.Assert(err, qt.IsNil)
76	}
77
78	// virtual tree
79	vTree := newVT(maxLevels, HashFunctionSha256)
80
81	c.Assert(vTree.root, qt.IsNil)
82
83	for i := 0; i < len(keys); i++ {
84		err := vTree.add(0, keys[i], values[i])
85		c.Assert(err, qt.IsNil)
86	}
87
88	// compute hashes, and check Root
89	_, err = vTree.computeHashes()
90	c.Assert(err, qt.IsNil)
91	c.Assert(vTree.root.h, qt.DeepEquals, tree.root)
92}