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}
 93
 94func TestVirtualTreeAddBatch(t *testing.T) {
 95	c := qt.New(t)
 96
 97	nLeafs := 2000
 98	maxLevels := 100
 99
100	keys := make([][]byte, nLeafs)
101	values := make([][]byte, nLeafs)
102	for i := 0; i < nLeafs; i++ {
103		keys[i] = randomBytes(32)
104		values[i] = randomBytes(32)
105	}
106
107	// normal tree, to have an expected root value
108	tree, err := NewTree(memory.NewMemoryStorage(), maxLevels, HashFunctionBlake2b)
109	c.Assert(err, qt.IsNil)
110	for i := 0; i < len(keys); i++ {
111		err := tree.Add(keys[i], values[i])
112		c.Assert(err, qt.IsNil)
113	}
114
115	// virtual tree
116	vTree := newVT(maxLevels, HashFunctionBlake2b)
117
118	c.Assert(vTree.root, qt.IsNil)
119
120	err = vTree.addBatch(keys, values)
121	c.Assert(err, qt.IsNil)
122
123	// compute hashes, and check Root
124	_, err = vTree.computeHashes()
125	c.Assert(err, qt.IsNil)
126	c.Assert(vTree.root.h, qt.DeepEquals, tree.root)
127}