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.

127 lines
3.5 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. err = vTree.addBatch(keys, values)
  97. c.Assert(err, qt.IsNil)
  98. // compute hashes, and check Root
  99. _, err = vTree.computeHashes()
  100. c.Assert(err, qt.IsNil)
  101. c.Assert(vTree.root.h, qt.DeepEquals, tree.root)
  102. }