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.

425 lines
12 KiB

3 years ago
3 years ago
3 years ago
3 years ago
3 years ago
3 years ago
3 years ago
  1. package arbo
  2. import (
  3. "encoding/hex"
  4. "math/big"
  5. "testing"
  6. "time"
  7. qt "github.com/frankban/quicktest"
  8. "github.com/iden3/go-merkletree/db/memory"
  9. )
  10. func TestAddTestVectors(t *testing.T) {
  11. c := qt.New(t)
  12. // Poseidon test vectors generated using https://github.com/iden3/circomlib smt.js
  13. testVectorsPoseidon := []string{
  14. "0000000000000000000000000000000000000000000000000000000000000000",
  15. "13578938674299138072471463694055224830892726234048532520316387704878000008795",
  16. "5412393676474193513566895793055462193090331607895808993925969873307089394741",
  17. "14204494359367183802864593755198662203838502594566452929175967972147978322084",
  18. }
  19. testAdd(c, HashFunctionPoseidon, testVectorsPoseidon)
  20. testVectorsSha256 := []string{
  21. "0000000000000000000000000000000000000000000000000000000000000000",
  22. "46910109172468462938850740851377282682950237270676610513794735904325820156367",
  23. "59481735341404520835410489183267411392292882901306595567679529387376287440550",
  24. "20573794434149960984975763118181266662429997821552560184909083010514790081771",
  25. }
  26. testAdd(c, HashFunctionSha256, testVectorsSha256)
  27. }
  28. func testAdd(c *qt.C, hashFunc HashFunction, testVectors []string) {
  29. tree, err := NewTree(memory.NewMemoryStorage(), 10, hashFunc)
  30. c.Assert(err, qt.IsNil)
  31. defer tree.db.Close()
  32. c.Check(hex.EncodeToString(tree.Root()), qt.Equals, testVectors[0])
  33. bLen := hashFunc.Len()
  34. err = tree.Add(
  35. BigIntToBytes(bLen, big.NewInt(1)),
  36. BigIntToBytes(bLen, big.NewInt(2)))
  37. c.Assert(err, qt.IsNil)
  38. rootBI := BytesToBigInt(tree.Root())
  39. c.Check(rootBI.String(), qt.Equals, testVectors[1])
  40. err = tree.Add(
  41. BigIntToBytes(bLen, big.NewInt(33)),
  42. BigIntToBytes(bLen, big.NewInt(44)))
  43. c.Assert(err, qt.IsNil)
  44. rootBI = BytesToBigInt(tree.Root())
  45. c.Check(rootBI.String(), qt.Equals, testVectors[2])
  46. err = tree.Add(
  47. BigIntToBytes(bLen, big.NewInt(1234)),
  48. BigIntToBytes(bLen, big.NewInt(9876)))
  49. c.Assert(err, qt.IsNil)
  50. rootBI = BytesToBigInt(tree.Root())
  51. c.Check(rootBI.String(), qt.Equals, testVectors[3])
  52. }
  53. func TestAddBatch(t *testing.T) {
  54. c := qt.New(t)
  55. tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
  56. c.Assert(err, qt.IsNil)
  57. defer tree.db.Close()
  58. bLen := tree.HashFunction().Len()
  59. for i := 0; i < 1000; i++ {
  60. k := BigIntToBytes(bLen, big.NewInt(int64(i)))
  61. v := BigIntToBytes(bLen, big.NewInt(0))
  62. if err := tree.Add(k, v); err != nil {
  63. t.Fatal(err)
  64. }
  65. }
  66. rootBI := BytesToBigInt(tree.Root())
  67. c.Check(rootBI.String(), qt.Equals,
  68. "296519252211642170490407814696803112091039265640052570497930797516015811235")
  69. tree2, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
  70. c.Assert(err, qt.IsNil)
  71. defer tree2.db.Close()
  72. var keys, values [][]byte
  73. for i := 0; i < 1000; i++ {
  74. k := BigIntToBytes(bLen, big.NewInt(int64(i)))
  75. v := BigIntToBytes(bLen, big.NewInt(0))
  76. keys = append(keys, k)
  77. values = append(values, v)
  78. }
  79. indexes, err := tree2.AddBatch(keys, values)
  80. c.Assert(err, qt.IsNil)
  81. c.Check(len(indexes), qt.Equals, 0)
  82. rootBI = BytesToBigInt(tree2.Root())
  83. c.Check(rootBI.String(), qt.Equals,
  84. "296519252211642170490407814696803112091039265640052570497930797516015811235")
  85. }
  86. func TestAddDifferentOrder(t *testing.T) {
  87. c := qt.New(t)
  88. tree1, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
  89. c.Assert(err, qt.IsNil)
  90. defer tree1.db.Close()
  91. bLen := tree1.HashFunction().Len()
  92. for i := 0; i < 16; i++ {
  93. k := BigIntToBytes(bLen, big.NewInt(int64(i)))
  94. v := BigIntToBytes(bLen, big.NewInt(0))
  95. if err := tree1.Add(k, v); err != nil {
  96. t.Fatal(err)
  97. }
  98. }
  99. tree2, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
  100. c.Assert(err, qt.IsNil)
  101. defer tree2.db.Close()
  102. for i := 16 - 1; i >= 0; i-- {
  103. k := BigIntToBytes(bLen, big.NewInt(int64(i)))
  104. v := BigIntToBytes(bLen, big.NewInt(0))
  105. if err := tree2.Add(k, v); err != nil {
  106. t.Fatal(err)
  107. }
  108. }
  109. c.Check(hex.EncodeToString(tree2.Root()), qt.Equals, hex.EncodeToString(tree1.Root()))
  110. c.Check(hex.EncodeToString(tree1.Root()), qt.Equals,
  111. "3b89100bec24da9275c87bc188740389e1d5accfc7d88ba5688d7fa96a00d82f")
  112. }
  113. func TestAddRepeatedIndex(t *testing.T) {
  114. c := qt.New(t)
  115. tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
  116. c.Assert(err, qt.IsNil)
  117. defer tree.db.Close()
  118. bLen := tree.HashFunction().Len()
  119. k := BigIntToBytes(bLen, big.NewInt(int64(3)))
  120. v := BigIntToBytes(bLen, big.NewInt(int64(12)))
  121. err = tree.Add(k, v)
  122. c.Assert(err, qt.IsNil)
  123. err = tree.Add(k, v) // repeating same key-value
  124. c.Check(err, qt.Equals, ErrKeyAlreadyExists)
  125. }
  126. func TestUpdate(t *testing.T) {
  127. c := qt.New(t)
  128. tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
  129. c.Assert(err, qt.IsNil)
  130. defer tree.db.Close()
  131. bLen := tree.HashFunction().Len()
  132. k := BigIntToBytes(bLen, big.NewInt(int64(20)))
  133. v := BigIntToBytes(bLen, big.NewInt(int64(12)))
  134. if err := tree.Add(k, v); err != nil {
  135. t.Fatal(err)
  136. }
  137. v = BigIntToBytes(bLen, big.NewInt(int64(11)))
  138. err = tree.Update(k, v)
  139. c.Assert(err, qt.IsNil)
  140. gettedKey, gettedValue, err := tree.Get(k)
  141. c.Assert(err, qt.IsNil)
  142. c.Check(gettedKey, qt.DeepEquals, k)
  143. c.Check(gettedValue, qt.DeepEquals, v)
  144. // add more leafs to the tree to do another test
  145. for i := 0; i < 16; i++ {
  146. k := BigIntToBytes(bLen, big.NewInt(int64(i)))
  147. v := BigIntToBytes(bLen, big.NewInt(int64(i*2)))
  148. if err := tree.Add(k, v); err != nil {
  149. t.Fatal(err)
  150. }
  151. }
  152. k = BigIntToBytes(bLen, big.NewInt(int64(3)))
  153. v = BigIntToBytes(bLen, big.NewInt(int64(11)))
  154. // check that before the Update, value for 3 is !=11
  155. gettedKey, gettedValue, err = tree.Get(k)
  156. c.Assert(err, qt.IsNil)
  157. c.Check(gettedKey, qt.DeepEquals, k)
  158. c.Check(gettedValue, qt.Not(qt.DeepEquals), v)
  159. c.Check(gettedValue, qt.DeepEquals, BigIntToBytes(bLen, big.NewInt(6)))
  160. err = tree.Update(k, v)
  161. c.Assert(err, qt.IsNil)
  162. // check that after Update, the value for 3 is ==11
  163. gettedKey, gettedValue, err = tree.Get(k)
  164. c.Assert(err, qt.IsNil)
  165. c.Check(gettedKey, qt.DeepEquals, k)
  166. c.Check(gettedValue, qt.DeepEquals, v)
  167. c.Check(gettedValue, qt.DeepEquals, BigIntToBytes(bLen, big.NewInt(11)))
  168. }
  169. func TestAux(t *testing.T) { // TODO split in proper tests
  170. c := qt.New(t)
  171. tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
  172. c.Assert(err, qt.IsNil)
  173. defer tree.db.Close()
  174. bLen := tree.HashFunction().Len()
  175. k := BigIntToBytes(bLen, big.NewInt(int64(1)))
  176. v := BigIntToBytes(bLen, big.NewInt(int64(0)))
  177. err = tree.Add(k, v)
  178. c.Assert(err, qt.IsNil)
  179. k = BigIntToBytes(bLen, big.NewInt(int64(256)))
  180. err = tree.Add(k, v)
  181. c.Assert(err, qt.IsNil)
  182. k = BigIntToBytes(bLen, big.NewInt(int64(257)))
  183. err = tree.Add(k, v)
  184. c.Assert(err, qt.IsNil)
  185. k = BigIntToBytes(bLen, big.NewInt(int64(515)))
  186. err = tree.Add(k, v)
  187. c.Assert(err, qt.IsNil)
  188. k = BigIntToBytes(bLen, big.NewInt(int64(770)))
  189. err = tree.Add(k, v)
  190. c.Assert(err, qt.IsNil)
  191. k = BigIntToBytes(bLen, big.NewInt(int64(388)))
  192. err = tree.Add(k, v)
  193. c.Assert(err, qt.IsNil)
  194. k = BigIntToBytes(bLen, big.NewInt(int64(900)))
  195. err = tree.Add(k, v)
  196. c.Assert(err, qt.IsNil)
  197. //
  198. // err = tree.PrintGraphviz(nil)
  199. // c.Assert(err, qt.IsNil)
  200. }
  201. func TestGet(t *testing.T) {
  202. c := qt.New(t)
  203. tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
  204. c.Assert(err, qt.IsNil)
  205. defer tree.db.Close()
  206. bLen := tree.HashFunction().Len()
  207. for i := 0; i < 10; i++ {
  208. k := BigIntToBytes(bLen, big.NewInt(int64(i)))
  209. v := BigIntToBytes(bLen, big.NewInt(int64(i*2)))
  210. if err := tree.Add(k, v); err != nil {
  211. t.Fatal(err)
  212. }
  213. }
  214. k := BigIntToBytes(bLen, big.NewInt(int64(7)))
  215. gettedKey, gettedValue, err := tree.Get(k)
  216. c.Assert(err, qt.IsNil)
  217. c.Check(gettedKey, qt.DeepEquals, k)
  218. c.Check(gettedValue, qt.DeepEquals, BigIntToBytes(bLen, big.NewInt(int64(7*2))))
  219. }
  220. func TestGenProofAndVerify(t *testing.T) {
  221. c := qt.New(t)
  222. tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
  223. c.Assert(err, qt.IsNil)
  224. defer tree.db.Close()
  225. bLen := tree.HashFunction().Len()
  226. for i := 0; i < 10; i++ {
  227. k := BigIntToBytes(bLen, big.NewInt(int64(i)))
  228. v := BigIntToBytes(bLen, big.NewInt(int64(i*2)))
  229. if err := tree.Add(k, v); err != nil {
  230. t.Fatal(err)
  231. }
  232. }
  233. k := BigIntToBytes(bLen, big.NewInt(int64(7)))
  234. v := BigIntToBytes(bLen, big.NewInt(int64(14)))
  235. proofV, siblings, err := tree.GenProof(k)
  236. c.Assert(err, qt.IsNil)
  237. c.Assert(proofV, qt.DeepEquals, v)
  238. verif, err := CheckProof(tree.hashFunction, k, v, tree.Root(), siblings)
  239. c.Assert(err, qt.IsNil)
  240. c.Check(verif, qt.IsTrue)
  241. }
  242. func TestDumpAndImportDump(t *testing.T) {
  243. c := qt.New(t)
  244. tree1, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
  245. c.Assert(err, qt.IsNil)
  246. defer tree1.db.Close()
  247. bLen := tree1.HashFunction().Len()
  248. for i := 0; i < 16; i++ {
  249. k := BigIntToBytes(bLen, big.NewInt(int64(i)))
  250. v := BigIntToBytes(bLen, big.NewInt(int64(i*2)))
  251. if err := tree1.Add(k, v); err != nil {
  252. t.Fatal(err)
  253. }
  254. }
  255. e, err := tree1.Dump(nil)
  256. c.Assert(err, qt.IsNil)
  257. tree2, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
  258. c.Assert(err, qt.IsNil)
  259. defer tree2.db.Close()
  260. err = tree2.ImportDump(e)
  261. c.Assert(err, qt.IsNil)
  262. c.Check(tree2.Root(), qt.DeepEquals, tree1.Root())
  263. c.Check(hex.EncodeToString(tree2.Root()), qt.Equals,
  264. "0d93aaa3362b2f999f15e15728f123087c2eee716f01c01f56e23aae07f09f08")
  265. }
  266. func TestRWMutex(t *testing.T) {
  267. c := qt.New(t)
  268. tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
  269. c.Assert(err, qt.IsNil)
  270. defer tree.db.Close()
  271. bLen := tree.HashFunction().Len()
  272. var keys, values [][]byte
  273. for i := 0; i < 1000; i++ {
  274. k := BigIntToBytes(bLen, big.NewInt(int64(i)))
  275. v := BigIntToBytes(bLen, big.NewInt(0))
  276. keys = append(keys, k)
  277. values = append(values, v)
  278. }
  279. go func() {
  280. _, err = tree.AddBatch(keys, values)
  281. if err != nil {
  282. panic(err)
  283. }
  284. }()
  285. time.Sleep(500 * time.Millisecond)
  286. k := BigIntToBytes(bLen, big.NewInt(int64(99999)))
  287. v := BigIntToBytes(bLen, big.NewInt(int64(99999)))
  288. if err := tree.Add(k, v); err != nil {
  289. t.Fatal(err)
  290. }
  291. }
  292. func TestSetGetNLeafs(t *testing.T) {
  293. c := qt.New(t)
  294. tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
  295. c.Assert(err, qt.IsNil)
  296. // 0
  297. tree.tx, err = tree.db.NewTx()
  298. c.Assert(err, qt.IsNil)
  299. err = tree.setNLeafs(0)
  300. c.Assert(err, qt.IsNil)
  301. err = tree.tx.Commit()
  302. c.Assert(err, qt.IsNil)
  303. n, err := tree.GetNLeafs()
  304. c.Assert(err, qt.IsNil)
  305. c.Assert(n, qt.Equals, 0)
  306. // 1024
  307. tree.tx, err = tree.db.NewTx()
  308. c.Assert(err, qt.IsNil)
  309. err = tree.setNLeafs(1024)
  310. c.Assert(err, qt.IsNil)
  311. err = tree.tx.Commit()
  312. c.Assert(err, qt.IsNil)
  313. n, err = tree.GetNLeafs()
  314. c.Assert(err, qt.IsNil)
  315. c.Assert(n, qt.Equals, 1024)
  316. // 2**64 -1
  317. tree.tx, err = tree.db.NewTx()
  318. c.Assert(err, qt.IsNil)
  319. maxUint := ^uint(0)
  320. maxInt := int(maxUint >> 1)
  321. err = tree.setNLeafs(maxInt)
  322. c.Assert(err, qt.IsNil)
  323. err = tree.tx.Commit()
  324. c.Assert(err, qt.IsNil)
  325. n, err = tree.GetNLeafs()
  326. c.Assert(err, qt.IsNil)
  327. c.Assert(n, qt.Equals, maxInt)
  328. }
  329. func BenchmarkAdd(b *testing.B) {
  330. bLen := 32 // for both Poseidon & Sha256
  331. // prepare inputs
  332. var ks, vs [][]byte
  333. for i := 0; i < 1000; i++ {
  334. k := BigIntToBytes(bLen, big.NewInt(int64(i)))
  335. v := BigIntToBytes(bLen, big.NewInt(int64(i)))
  336. ks = append(ks, k)
  337. vs = append(vs, v)
  338. }
  339. b.Run("Poseidon", func(b *testing.B) {
  340. benchmarkAdd(b, HashFunctionPoseidon, ks, vs)
  341. })
  342. b.Run("Sha256", func(b *testing.B) {
  343. benchmarkAdd(b, HashFunctionSha256, ks, vs)
  344. })
  345. }
  346. func benchmarkAdd(b *testing.B, hashFunc HashFunction, ks, vs [][]byte) {
  347. c := qt.New(b)
  348. tree, err := NewTree(memory.NewMemoryStorage(), 140, hashFunc)
  349. c.Assert(err, qt.IsNil)
  350. defer tree.db.Close()
  351. for i := 0; i < len(ks); i++ {
  352. if err := tree.Add(ks[i], vs[i]); err != nil {
  353. b.Fatal(err)
  354. }
  355. }
  356. }