You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

289 lines
6.9 KiB

  1. package arbo
  2. import (
  3. "encoding/hex"
  4. "math/big"
  5. "testing"
  6. qt "github.com/frankban/quicktest"
  7. "github.com/iden3/go-merkletree/db/memory"
  8. )
  9. func TestVirtualTreeTestVectors(t *testing.T) {
  10. c := qt.New(t)
  11. keys := [][]byte{
  12. BigIntToBytes(big.NewInt(1)),
  13. BigIntToBytes(big.NewInt(33)),
  14. BigIntToBytes(big.NewInt(1234)),
  15. BigIntToBytes(big.NewInt(123456789)),
  16. }
  17. values := [][]byte{
  18. BigIntToBytes(big.NewInt(2)),
  19. BigIntToBytes(big.NewInt(44)),
  20. BigIntToBytes(big.NewInt(9876)),
  21. BigIntToBytes(big.NewInt(987654321)),
  22. }
  23. // check the root for different batches of leafs
  24. // testVirtualTree(c, 10, keys[:1], values[:1])
  25. // testVirtualTree(c, 10, keys[:2], values[:2])
  26. // testVirtualTree(c, 10, keys[:3], values[:3])
  27. testVirtualTree(c, 10, keys[:4], values[:4])
  28. }
  29. func TestVirtualTreeRandomKeys(t *testing.T) {
  30. c := qt.New(t)
  31. // test with hardcoded values
  32. keys := make([][]byte, 8)
  33. values := make([][]byte, 8)
  34. keys[0], _ = hex.DecodeString("1c7c2265e368314ca58ed2e1f33a326f1220e234a566d55c3605439dbe411642")
  35. keys[1], _ = hex.DecodeString("2c9f0a578afff5bfa4e0992a43066460faaab9e8e500db0b16647c701cdb16bf")
  36. keys[2], _ = hex.DecodeString("9cb87ec67e875c61390edcd1ab517f443591047709a4d4e45b0f9ed980857b8e")
  37. keys[3], _ = hex.DecodeString("9b4e9e92e974a589f426ceeb4cb291dc24893513fecf8e8460992dcf52621d4d")
  38. keys[4], _ = hex.DecodeString("1c45cb31f2fa39ec7b9ebf0fad40e0b8296016b5ce8844ae06ff77226379d9a5")
  39. keys[5], _ = hex.DecodeString("d8af98bbbb585129798ae54d5eabbc9d0561d583faf1663b3a3724d15bda4ec7")
  40. keys[6], _ = hex.DecodeString("3cd55dbfb8f975f20a0925dfbdabe79fa2d51dd0268afbb8ba6b01de9dfcdd3c")
  41. keys[7], _ = hex.DecodeString("5d0a9d6d9f197c091bf054fac9cb60e11ec723d6610ed8578e617b4d46cb43d5")
  42. // check the root for different batches of leafs
  43. testVirtualTree(c, 10, keys[:1], values[:1])
  44. testVirtualTree(c, 10, keys, values)
  45. // test with random values
  46. nLeafs := 1024
  47. keys = make([][]byte, nLeafs)
  48. values = make([][]byte, nLeafs)
  49. for i := 0; i < nLeafs; i++ {
  50. keys[i] = randomBytes(32)
  51. values[i] = []byte{0}
  52. }
  53. testVirtualTree(c, 100, keys, values)
  54. }
  55. func testVirtualTree(c *qt.C, maxLevels int, keys, values [][]byte) {
  56. c.Assert(len(keys), qt.Equals, len(values))
  57. // normal tree, to have an expected root value
  58. tree, err := NewTree(memory.NewMemoryStorage(), maxLevels, HashFunctionSha256)
  59. c.Assert(err, qt.IsNil)
  60. for i := 0; i < len(keys); i++ {
  61. err := tree.Add(keys[i], values[i])
  62. c.Assert(err, qt.IsNil)
  63. }
  64. // virtual tree
  65. vTree := newVT(maxLevels, HashFunctionSha256)
  66. c.Assert(vTree.root, qt.IsNil)
  67. for i := 0; i < len(keys); i++ {
  68. err := vTree.add(0, keys[i], values[i])
  69. c.Assert(err, qt.IsNil)
  70. }
  71. // compute hashes, and check Root
  72. _, err = vTree.computeHashes()
  73. c.Assert(err, qt.IsNil)
  74. c.Assert(vTree.root.h, qt.DeepEquals, tree.root)
  75. }
  76. func TestVirtualTreeAddBatch(t *testing.T) {
  77. c := qt.New(t)
  78. nLeafs := 2000
  79. maxLevels := 100
  80. keys := make([][]byte, nLeafs)
  81. values := make([][]byte, nLeafs)
  82. for i := 0; i < nLeafs; i++ {
  83. keys[i] = randomBytes(32)
  84. values[i] = randomBytes(32)
  85. }
  86. // normal tree, to have an expected root value
  87. tree, err := NewTree(memory.NewMemoryStorage(), maxLevels, HashFunctionBlake2b)
  88. c.Assert(err, qt.IsNil)
  89. for i := 0; i < len(keys); i++ {
  90. err := tree.Add(keys[i], values[i])
  91. c.Assert(err, qt.IsNil)
  92. }
  93. // virtual tree
  94. vTree := newVT(maxLevels, HashFunctionBlake2b)
  95. c.Assert(vTree.root, qt.IsNil)
  96. invalids, err := vTree.addBatch(keys, values)
  97. c.Assert(err, qt.IsNil)
  98. c.Assert(len(invalids), qt.Equals, 0)
  99. // compute hashes, and check Root
  100. _, err = vTree.computeHashes()
  101. c.Assert(err, qt.IsNil)
  102. c.Assert(vTree.root.h, qt.DeepEquals, tree.root)
  103. }
  104. func TestGetNodesAtLevel(t *testing.T) {
  105. c := qt.New(t)
  106. tree0 := vt{
  107. params: &params{
  108. maxLevels: 100,
  109. hashFunction: HashFunctionBlake2b,
  110. emptyHash: make([]byte, HashFunctionBlake2b.Len()),
  111. },
  112. root: nil,
  113. }
  114. tree1 := vt{
  115. params: &params{
  116. maxLevels: 100,
  117. hashFunction: HashFunctionBlake2b,
  118. emptyHash: make([]byte, HashFunctionBlake2b.Len()),
  119. },
  120. root: &node{
  121. l: &node{
  122. l: &node{
  123. k: []byte{0, 0, 0, 0},
  124. v: []byte{0, 0, 0, 0},
  125. },
  126. r: &node{
  127. k: []byte{0, 0, 0, 1},
  128. v: []byte{0, 0, 0, 1},
  129. },
  130. },
  131. r: &node{
  132. l: &node{
  133. k: []byte{0, 0, 0, 2},
  134. v: []byte{0, 0, 0, 2},
  135. },
  136. r: &node{
  137. k: []byte{0, 0, 0, 3},
  138. v: []byte{0, 0, 0, 3},
  139. },
  140. },
  141. },
  142. }
  143. // tree1.printGraphviz()
  144. tree2 := vt{
  145. params: &params{
  146. maxLevels: 100,
  147. hashFunction: HashFunctionBlake2b,
  148. emptyHash: make([]byte, HashFunctionBlake2b.Len()),
  149. },
  150. root: &node{
  151. l: nil,
  152. r: &node{
  153. l: &node{
  154. l: &node{
  155. l: &node{
  156. k: []byte{0, 0, 0, 0},
  157. v: []byte{0, 0, 0, 0},
  158. },
  159. r: &node{
  160. k: []byte{0, 0, 0, 1},
  161. v: []byte{0, 0, 0, 1},
  162. },
  163. },
  164. r: &node{
  165. k: []byte{0, 0, 0, 2},
  166. v: []byte{0, 0, 0, 2},
  167. },
  168. },
  169. r: &node{
  170. k: []byte{0, 0, 0, 3},
  171. v: []byte{0, 0, 0, 3},
  172. },
  173. },
  174. },
  175. }
  176. // tree2.printGraphviz()
  177. tree3 := vt{
  178. params: &params{
  179. maxLevels: 100,
  180. hashFunction: HashFunctionBlake2b,
  181. emptyHash: make([]byte, HashFunctionBlake2b.Len()),
  182. },
  183. root: &node{
  184. l: nil,
  185. r: &node{
  186. l: &node{
  187. l: &node{
  188. l: &node{
  189. k: []byte{0, 0, 0, 0},
  190. v: []byte{0, 0, 0, 0},
  191. },
  192. r: &node{
  193. k: []byte{0, 0, 0, 1},
  194. v: []byte{0, 0, 0, 1},
  195. },
  196. },
  197. r: &node{
  198. k: []byte{0, 0, 0, 2},
  199. v: []byte{0, 0, 0, 2},
  200. },
  201. },
  202. r: nil,
  203. },
  204. },
  205. }
  206. // tree3.printGraphviz()
  207. nodes0, err := tree0.getNodesAtLevel(2)
  208. c.Assert(err, qt.IsNil)
  209. c.Assert(len(nodes0), qt.DeepEquals, 4)
  210. c.Assert("0000", qt.DeepEquals, getNotNils(nodes0))
  211. nodes1, err := tree1.getNodesAtLevel(2)
  212. c.Assert(err, qt.IsNil)
  213. c.Assert(len(nodes1), qt.DeepEquals, 4)
  214. c.Assert("1111", qt.DeepEquals, getNotNils(nodes1))
  215. nodes1, err = tree1.getNodesAtLevel(3)
  216. c.Assert(err, qt.IsNil)
  217. c.Assert(len(nodes1), qt.DeepEquals, 8)
  218. c.Assert("00000000", qt.DeepEquals, getNotNils(nodes1))
  219. nodes2, err := tree2.getNodesAtLevel(2)
  220. c.Assert(err, qt.IsNil)
  221. c.Assert(len(nodes2), qt.DeepEquals, 4)
  222. c.Assert("0011", qt.DeepEquals, getNotNils(nodes2))
  223. nodes2, err = tree2.getNodesAtLevel(3)
  224. c.Assert(err, qt.IsNil)
  225. c.Assert(len(nodes2), qt.DeepEquals, 8)
  226. c.Assert("00001100", qt.DeepEquals, getNotNils(nodes2))
  227. nodes3, err := tree3.getNodesAtLevel(2)
  228. c.Assert(err, qt.IsNil)
  229. c.Assert(len(nodes3), qt.DeepEquals, 4)
  230. c.Assert("0010", qt.DeepEquals, getNotNils(nodes3))
  231. nodes3, err = tree3.getNodesAtLevel(3)
  232. c.Assert(err, qt.IsNil)
  233. c.Assert(len(nodes3), qt.DeepEquals, 8)
  234. c.Assert("00001100", qt.DeepEquals, getNotNils(nodes3))
  235. nodes3, err = tree3.getNodesAtLevel(4)
  236. c.Assert(err, qt.IsNil)
  237. c.Assert(len(nodes3), qt.DeepEquals, 16)
  238. c.Assert("0000000011000000", qt.DeepEquals, getNotNils(nodes3))
  239. }
  240. func getNotNils(nodes []*node) string {
  241. s := ""
  242. for i := 0; i < len(nodes); i++ {
  243. if nodes[i] == nil {
  244. s += "0"
  245. } else {
  246. s += "1"
  247. }
  248. }
  249. return s
  250. }