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.

675 lines
20 KiB

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. "go.vocdoni.io/dvote/db"
  9. "go.vocdoni.io/dvote/db/badgerdb"
  10. )
  11. func checkRootBIString(c *qt.C, tree *Tree, expected string) {
  12. root, err := tree.Root()
  13. c.Assert(err, qt.IsNil)
  14. rootBI := BytesToBigInt(root)
  15. c.Check(rootBI.String(), qt.Equals, expected)
  16. }
  17. func TestDBTx(t *testing.T) {
  18. c := qt.New(t)
  19. database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
  20. c.Assert(err, qt.IsNil)
  21. wTx := database.WriteTx()
  22. _, err = wTx.Get([]byte("a"))
  23. c.Assert(err, qt.Equals, db.ErrKeyNotFound)
  24. err = wTx.Set([]byte("a"), []byte("b"))
  25. c.Assert(err, qt.IsNil)
  26. v, err := wTx.Get([]byte("a"))
  27. c.Assert(err, qt.IsNil)
  28. c.Assert(v, qt.DeepEquals, []byte("b"))
  29. }
  30. func TestAddTestVectors(t *testing.T) {
  31. c := qt.New(t)
  32. // Poseidon test vectors generated using https://github.com/iden3/circomlib smt.js
  33. testVectorsPoseidon := []string{
  34. "0000000000000000000000000000000000000000000000000000000000000000",
  35. "13578938674299138072471463694055224830892726234048532520316387704878000008795",
  36. "5412393676474193513566895793055462193090331607895808993925969873307089394741",
  37. "14204494359367183802864593755198662203838502594566452929175967972147978322084",
  38. }
  39. testAdd(c, HashFunctionPoseidon, testVectorsPoseidon)
  40. testVectorsSha256 := []string{
  41. "0000000000000000000000000000000000000000000000000000000000000000",
  42. "46910109172468462938850740851377282682950237270676610513794735904325820156367",
  43. "59481735341404520835410489183267411392292882901306595567679529387376287440550",
  44. "20573794434149960984975763118181266662429997821552560184909083010514790081771",
  45. }
  46. testAdd(c, HashFunctionSha256, testVectorsSha256)
  47. }
  48. func testAdd(c *qt.C, hashFunc HashFunction, testVectors []string) {
  49. database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
  50. c.Assert(err, qt.IsNil)
  51. tree, err := NewTree(database, 10, hashFunc)
  52. c.Assert(err, qt.IsNil)
  53. defer tree.db.Close() //nolint:errcheck
  54. root, err := tree.Root()
  55. c.Assert(err, qt.IsNil)
  56. c.Check(hex.EncodeToString(root), qt.Equals, testVectors[0])
  57. bLen := hashFunc.Len()
  58. err = tree.Add(
  59. BigIntToBytes(bLen, big.NewInt(1)),
  60. BigIntToBytes(bLen, big.NewInt(2)))
  61. c.Assert(err, qt.IsNil)
  62. checkRootBIString(c, tree, testVectors[1])
  63. err = tree.Add(
  64. BigIntToBytes(bLen, big.NewInt(33)),
  65. BigIntToBytes(bLen, big.NewInt(44)))
  66. c.Assert(err, qt.IsNil)
  67. checkRootBIString(c, tree, testVectors[2])
  68. err = tree.Add(
  69. BigIntToBytes(bLen, big.NewInt(1234)),
  70. BigIntToBytes(bLen, big.NewInt(9876)))
  71. c.Assert(err, qt.IsNil)
  72. checkRootBIString(c, tree, testVectors[3])
  73. }
  74. func TestAddBatch(t *testing.T) {
  75. c := qt.New(t)
  76. database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
  77. c.Assert(err, qt.IsNil)
  78. tree, err := NewTree(database, 100, HashFunctionPoseidon)
  79. c.Assert(err, qt.IsNil)
  80. defer tree.db.Close() //nolint:errcheck
  81. bLen := tree.HashFunction().Len()
  82. for i := 0; i < 1000; i++ {
  83. k := BigIntToBytes(bLen, big.NewInt(int64(i)))
  84. v := BigIntToBytes(bLen, big.NewInt(0))
  85. if err := tree.Add(k, v); err != nil {
  86. t.Fatal(err)
  87. }
  88. }
  89. checkRootBIString(c, tree,
  90. "296519252211642170490407814696803112091039265640052570497930797516015811235")
  91. database, err = badgerdb.New(badgerdb.Options{Path: c.TempDir()})
  92. c.Assert(err, qt.IsNil)
  93. tree2, err := NewTree(database, 100, HashFunctionPoseidon)
  94. c.Assert(err, qt.IsNil)
  95. defer tree2.db.Close() //nolint:errcheck
  96. var keys, values [][]byte
  97. for i := 0; i < 1000; i++ {
  98. k := BigIntToBytes(bLen, big.NewInt(int64(i)))
  99. v := BigIntToBytes(bLen, big.NewInt(0))
  100. keys = append(keys, k)
  101. values = append(values, v)
  102. }
  103. indexes, err := tree2.AddBatch(keys, values)
  104. c.Assert(err, qt.IsNil)
  105. c.Check(len(indexes), qt.Equals, 0)
  106. checkRootBIString(c, tree2,
  107. "296519252211642170490407814696803112091039265640052570497930797516015811235")
  108. }
  109. func TestAddDifferentOrder(t *testing.T) {
  110. c := qt.New(t)
  111. database1, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
  112. c.Assert(err, qt.IsNil)
  113. tree1, err := NewTree(database1, 100, HashFunctionPoseidon)
  114. c.Assert(err, qt.IsNil)
  115. defer tree1.db.Close() //nolint:errcheck
  116. bLen := tree1.HashFunction().Len()
  117. for i := 0; i < 16; i++ {
  118. k := BigIntToBytes(bLen, big.NewInt(int64(i)))
  119. v := BigIntToBytes(bLen, big.NewInt(0))
  120. if err := tree1.Add(k, v); err != nil {
  121. t.Fatal(err)
  122. }
  123. }
  124. database2, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
  125. c.Assert(err, qt.IsNil)
  126. tree2, err := NewTree(database2, 100, HashFunctionPoseidon)
  127. c.Assert(err, qt.IsNil)
  128. defer tree2.db.Close() //nolint:errcheck
  129. for i := 16 - 1; i >= 0; i-- {
  130. k := BigIntToBytes(bLen, big.NewInt(int64(i)))
  131. v := BigIntToBytes(bLen, big.NewInt(0))
  132. if err := tree2.Add(k, v); err != nil {
  133. t.Fatal(err)
  134. }
  135. }
  136. root1, err := tree1.Root()
  137. c.Assert(err, qt.IsNil)
  138. root2, err := tree2.Root()
  139. c.Assert(err, qt.IsNil)
  140. c.Check(hex.EncodeToString(root2), qt.Equals, hex.EncodeToString(root1))
  141. c.Check(hex.EncodeToString(root1), qt.Equals,
  142. "3b89100bec24da9275c87bc188740389e1d5accfc7d88ba5688d7fa96a00d82f")
  143. }
  144. func TestAddRepeatedIndex(t *testing.T) {
  145. c := qt.New(t)
  146. database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
  147. c.Assert(err, qt.IsNil)
  148. tree, err := NewTree(database, 100, HashFunctionPoseidon)
  149. c.Assert(err, qt.IsNil)
  150. defer tree.db.Close() //nolint:errcheck
  151. bLen := tree.HashFunction().Len()
  152. k := BigIntToBytes(bLen, big.NewInt(int64(3)))
  153. v := BigIntToBytes(bLen, big.NewInt(int64(12)))
  154. err = tree.Add(k, v)
  155. c.Assert(err, qt.IsNil)
  156. err = tree.Add(k, v) // repeating same key-value
  157. c.Check(err, qt.Equals, ErrKeyAlreadyExists)
  158. }
  159. func TestUpdate(t *testing.T) {
  160. c := qt.New(t)
  161. database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
  162. c.Assert(err, qt.IsNil)
  163. tree, err := NewTree(database, 100, HashFunctionPoseidon)
  164. c.Assert(err, qt.IsNil)
  165. defer tree.db.Close() //nolint:errcheck
  166. bLen := tree.HashFunction().Len()
  167. k := BigIntToBytes(bLen, big.NewInt(int64(20)))
  168. v := BigIntToBytes(bLen, big.NewInt(int64(12)))
  169. if err := tree.Add(k, v); err != nil {
  170. t.Fatal(err)
  171. }
  172. v = BigIntToBytes(bLen, big.NewInt(int64(11)))
  173. err = tree.Update(k, v)
  174. c.Assert(err, qt.IsNil)
  175. gettedKey, gettedValue, err := tree.Get(k)
  176. c.Assert(err, qt.IsNil)
  177. c.Check(gettedKey, qt.DeepEquals, k)
  178. c.Check(gettedValue, qt.DeepEquals, v)
  179. // add more leafs to the tree to do another test
  180. for i := 0; i < 16; i++ {
  181. k := BigIntToBytes(bLen, big.NewInt(int64(i)))
  182. v := BigIntToBytes(bLen, big.NewInt(int64(i*2)))
  183. if err := tree.Add(k, v); err != nil {
  184. t.Fatal(err)
  185. }
  186. }
  187. k = BigIntToBytes(bLen, big.NewInt(int64(3)))
  188. v = BigIntToBytes(bLen, big.NewInt(int64(11)))
  189. // check that before the Update, value for 3 is !=11
  190. gettedKey, gettedValue, err = tree.Get(k)
  191. c.Assert(err, qt.IsNil)
  192. c.Check(gettedKey, qt.DeepEquals, k)
  193. c.Check(gettedValue, qt.Not(qt.DeepEquals), v)
  194. c.Check(gettedValue, qt.DeepEquals, BigIntToBytes(bLen, big.NewInt(6)))
  195. err = tree.Update(k, v)
  196. c.Assert(err, qt.IsNil)
  197. // check that after Update, the value for 3 is ==11
  198. gettedKey, gettedValue, err = tree.Get(k)
  199. c.Assert(err, qt.IsNil)
  200. c.Check(gettedKey, qt.DeepEquals, k)
  201. c.Check(gettedValue, qt.DeepEquals, v)
  202. c.Check(gettedValue, qt.DeepEquals, BigIntToBytes(bLen, big.NewInt(11)))
  203. }
  204. func TestAux(t *testing.T) { // TODO split in proper tests
  205. c := qt.New(t)
  206. database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
  207. c.Assert(err, qt.IsNil)
  208. tree, err := NewTree(database, 100, HashFunctionPoseidon)
  209. c.Assert(err, qt.IsNil)
  210. defer tree.db.Close() //nolint:errcheck
  211. bLen := tree.HashFunction().Len()
  212. k := BigIntToBytes(bLen, big.NewInt(int64(1)))
  213. v := BigIntToBytes(bLen, big.NewInt(int64(0)))
  214. err = tree.Add(k, v)
  215. c.Assert(err, qt.IsNil)
  216. k = BigIntToBytes(bLen, big.NewInt(int64(256)))
  217. err = tree.Add(k, v)
  218. c.Assert(err, qt.IsNil)
  219. k = BigIntToBytes(bLen, big.NewInt(int64(257)))
  220. err = tree.Add(k, v)
  221. c.Assert(err, qt.IsNil)
  222. k = BigIntToBytes(bLen, big.NewInt(int64(515)))
  223. err = tree.Add(k, v)
  224. c.Assert(err, qt.IsNil)
  225. k = BigIntToBytes(bLen, big.NewInt(int64(770)))
  226. err = tree.Add(k, v)
  227. c.Assert(err, qt.IsNil)
  228. k = BigIntToBytes(bLen, big.NewInt(int64(388)))
  229. err = tree.Add(k, v)
  230. c.Assert(err, qt.IsNil)
  231. k = BigIntToBytes(bLen, big.NewInt(int64(900)))
  232. err = tree.Add(k, v)
  233. c.Assert(err, qt.IsNil)
  234. //
  235. // err = tree.PrintGraphviz(nil)
  236. // c.Assert(err, qt.IsNil)
  237. }
  238. func TestGet(t *testing.T) {
  239. c := qt.New(t)
  240. database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
  241. c.Assert(err, qt.IsNil)
  242. tree, err := NewTree(database, 100, HashFunctionPoseidon)
  243. c.Assert(err, qt.IsNil)
  244. defer tree.db.Close() //nolint:errcheck
  245. bLen := tree.HashFunction().Len()
  246. for i := 0; i < 10; i++ {
  247. k := BigIntToBytes(bLen, big.NewInt(int64(i)))
  248. v := BigIntToBytes(bLen, big.NewInt(int64(i*2)))
  249. if err := tree.Add(k, v); err != nil {
  250. t.Fatal(err)
  251. }
  252. }
  253. k := BigIntToBytes(bLen, big.NewInt(int64(7)))
  254. gettedKey, gettedValue, err := tree.Get(k)
  255. c.Assert(err, qt.IsNil)
  256. c.Check(gettedKey, qt.DeepEquals, k)
  257. c.Check(gettedValue, qt.DeepEquals, BigIntToBytes(bLen, big.NewInt(int64(7*2))))
  258. }
  259. func TestGenProofAndVerify(t *testing.T) {
  260. c := qt.New(t)
  261. database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
  262. c.Assert(err, qt.IsNil)
  263. tree, err := NewTree(database, 100, HashFunctionPoseidon)
  264. c.Assert(err, qt.IsNil)
  265. defer tree.db.Close() //nolint:errcheck
  266. bLen := tree.HashFunction().Len() - 1
  267. for i := 0; i < 10; i++ {
  268. k := BigIntToBytes(bLen, big.NewInt(int64(i)))
  269. v := BigIntToBytes(bLen, big.NewInt(int64(i*2)))
  270. if err := tree.Add(k, v); err != nil {
  271. t.Fatal(err)
  272. }
  273. }
  274. k := BigIntToBytes(bLen, big.NewInt(int64(7)))
  275. v := BigIntToBytes(bLen, big.NewInt(int64(14)))
  276. kAux, proofV, siblings, existence, err := tree.GenProof(k)
  277. c.Assert(err, qt.IsNil)
  278. c.Assert(proofV, qt.DeepEquals, v)
  279. c.Assert(k, qt.DeepEquals, kAux)
  280. c.Assert(existence, qt.IsTrue)
  281. root, err := tree.Root()
  282. c.Assert(err, qt.IsNil)
  283. verif, err := CheckProof(tree.hashFunction, k, v, root, siblings)
  284. c.Assert(err, qt.IsNil)
  285. c.Check(verif, qt.IsTrue)
  286. }
  287. func TestDumpAndImportDump(t *testing.T) {
  288. c := qt.New(t)
  289. database1, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
  290. c.Assert(err, qt.IsNil)
  291. tree1, err := NewTree(database1, 100, HashFunctionPoseidon)
  292. c.Assert(err, qt.IsNil)
  293. defer tree1.db.Close() //nolint:errcheck
  294. bLen := tree1.HashFunction().Len()
  295. for i := 0; i < 16; i++ {
  296. k := BigIntToBytes(bLen, big.NewInt(int64(i)))
  297. v := BigIntToBytes(bLen, big.NewInt(int64(i*2)))
  298. if err := tree1.Add(k, v); err != nil {
  299. t.Fatal(err)
  300. }
  301. }
  302. e, err := tree1.Dump(nil)
  303. c.Assert(err, qt.IsNil)
  304. database2, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
  305. c.Assert(err, qt.IsNil)
  306. tree2, err := NewTree(database2, 100, HashFunctionPoseidon)
  307. c.Assert(err, qt.IsNil)
  308. defer tree2.db.Close() //nolint:errcheck
  309. err = tree2.ImportDump(e)
  310. c.Assert(err, qt.IsNil)
  311. root1, err := tree1.Root()
  312. c.Assert(err, qt.IsNil)
  313. root2, err := tree2.Root()
  314. c.Assert(err, qt.IsNil)
  315. c.Check(root2, qt.DeepEquals, root1)
  316. c.Check(hex.EncodeToString(root2), qt.Equals,
  317. "0d93aaa3362b2f999f15e15728f123087c2eee716f01c01f56e23aae07f09f08")
  318. }
  319. func TestRWMutex(t *testing.T) {
  320. c := qt.New(t)
  321. database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
  322. c.Assert(err, qt.IsNil)
  323. tree, err := NewTree(database, 100, HashFunctionPoseidon)
  324. c.Assert(err, qt.IsNil)
  325. defer tree.db.Close() //nolint:errcheck
  326. bLen := tree.HashFunction().Len()
  327. var keys, values [][]byte
  328. for i := 0; i < 1000; i++ {
  329. k := BigIntToBytes(bLen, big.NewInt(int64(i)))
  330. v := BigIntToBytes(bLen, big.NewInt(0))
  331. keys = append(keys, k)
  332. values = append(values, v)
  333. }
  334. go func() {
  335. _, err = tree.AddBatch(keys, values)
  336. if err != nil {
  337. panic(err)
  338. }
  339. }()
  340. time.Sleep(500 * time.Millisecond)
  341. k := BigIntToBytes(bLen, big.NewInt(int64(99999)))
  342. v := BigIntToBytes(bLen, big.NewInt(int64(99999)))
  343. if err := tree.Add(k, v); err != nil {
  344. t.Fatal(err)
  345. }
  346. }
  347. // TODO UPDATE
  348. // func TestSetGetNLeafs(t *testing.T) {
  349. // c := qt.New(t)
  350. // database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
  351. // c.Assert(err, qt.IsNil)
  352. // tree, err := NewTree(database, 100, HashFunctionPoseidon)
  353. // c.Assert(err, qt.IsNil)
  354. //
  355. // // 0
  356. // tree.dbBatch = tree.db.NewBatch()
  357. //
  358. // err = tree.setNLeafs(0)
  359. // c.Assert(err, qt.IsNil)
  360. //
  361. // err = tree.dbBatch.Write()
  362. // c.Assert(err, qt.IsNil)
  363. //
  364. // n, err := tree.GetNLeafs()
  365. // c.Assert(err, qt.IsNil)
  366. // c.Assert(n, qt.Equals, 0)
  367. //
  368. // // 1024
  369. // tree.dbBatch = tree.db.NewBatch()
  370. //
  371. // err = tree.setNLeafs(1024)
  372. // c.Assert(err, qt.IsNil)
  373. //
  374. // err = tree.dbBatch.Write()
  375. // c.Assert(err, qt.IsNil)
  376. //
  377. // n, err = tree.GetNLeafs()
  378. // c.Assert(err, qt.IsNil)
  379. // c.Assert(n, qt.Equals, 1024)
  380. //
  381. // // 2**64 -1
  382. // tree.dbBatch = tree.db.NewBatch()
  383. //
  384. // maxUint := ^uint(0)
  385. // maxInt := int(maxUint >> 1)
  386. //
  387. // err = tree.setNLeafs(maxInt)
  388. // c.Assert(err, qt.IsNil)
  389. //
  390. // err = tree.dbBatch.Write()
  391. // c.Assert(err, qt.IsNil)
  392. //
  393. // n, err = tree.GetNLeafs()
  394. // c.Assert(err, qt.IsNil)
  395. // c.Assert(n, qt.Equals, maxInt)
  396. // }
  397. func TestAddBatchFullyUsed(t *testing.T) {
  398. c := qt.New(t)
  399. database1, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
  400. c.Assert(err, qt.IsNil)
  401. tree1, err := NewTree(database1, 4, HashFunctionPoseidon)
  402. c.Assert(err, qt.IsNil)
  403. database2, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
  404. c.Assert(err, qt.IsNil)
  405. tree2, err := NewTree(database2, 4, HashFunctionPoseidon)
  406. c.Assert(err, qt.IsNil)
  407. var keys, values [][]byte
  408. for i := 0; i < 16; i++ {
  409. k := BigIntToBytes(32, big.NewInt(int64(i)))
  410. v := k
  411. keys = append(keys, k)
  412. values = append(values, v)
  413. // add one by one expecting no error
  414. err := tree1.Add(k, v)
  415. c.Assert(err, qt.IsNil)
  416. }
  417. invalids, err := tree2.AddBatch(keys, values)
  418. c.Assert(err, qt.IsNil)
  419. c.Assert(0, qt.Equals, len(invalids))
  420. root1, err := tree1.Root()
  421. c.Assert(err, qt.IsNil)
  422. root2, err := tree2.Root()
  423. c.Assert(err, qt.IsNil)
  424. c.Assert(root1, qt.DeepEquals, root2)
  425. // get all key-values and check that are equal between both trees
  426. for i := 0; i < 16; i++ {
  427. auxK1, auxV1, err := tree1.Get(BigIntToBytes(32, big.NewInt(int64(i))))
  428. c.Assert(err, qt.IsNil)
  429. auxK2, auxV2, err := tree2.Get(BigIntToBytes(32, big.NewInt(int64(i))))
  430. c.Assert(err, qt.IsNil)
  431. c.Assert(auxK1, qt.DeepEquals, auxK2)
  432. c.Assert(auxV1, qt.DeepEquals, auxV2)
  433. }
  434. // try adding one more key to both trees (through Add & AddBatch) and
  435. // expect not being added due the tree is already full
  436. k := BigIntToBytes(32, big.NewInt(int64(16)))
  437. v := k
  438. err = tree1.Add(k, v)
  439. c.Assert(err, qt.Equals, ErrMaxVirtualLevel)
  440. invalids, err = tree2.AddBatch([][]byte{k}, [][]byte{v})
  441. c.Assert(err, qt.IsNil)
  442. c.Assert(1, qt.Equals, len(invalids))
  443. }
  444. func TestSetRoot(t *testing.T) {
  445. c := qt.New(t)
  446. database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
  447. c.Assert(err, qt.IsNil)
  448. tree, err := NewTree(database, 100, HashFunctionPoseidon)
  449. c.Assert(err, qt.IsNil)
  450. expectedRoot := "13742386369878513332697380582061714160370929283209286127733983161245560237407"
  451. // fill the tree
  452. bLen := tree.HashFunction().Len()
  453. var keys, values [][]byte
  454. for i := 0; i < 1000; i++ {
  455. k := BigIntToBytes(bLen, big.NewInt(int64(i)))
  456. v := BigIntToBytes(bLen, big.NewInt(int64(i)))
  457. keys = append(keys, k)
  458. values = append(values, v)
  459. }
  460. indexes, err := tree.AddBatch(keys, values)
  461. c.Assert(err, qt.IsNil)
  462. c.Check(len(indexes), qt.Equals, 0)
  463. checkRootBIString(c, tree,
  464. expectedRoot)
  465. // add one more k-v
  466. k := BigIntToBytes(bLen, big.NewInt(1000))
  467. v := BigIntToBytes(bLen, big.NewInt(1000))
  468. err = tree.Add(k, v)
  469. c.Assert(err, qt.IsNil)
  470. checkRootBIString(c, tree,
  471. "10747149055773881257049574592162159501044114324358186833013814735296193179713")
  472. // do a SetRoot, and expect the same root than the original tree
  473. pastRootBI, ok := new(big.Int).SetString(expectedRoot, 10)
  474. c.Assert(ok, qt.IsTrue)
  475. pastRoot := BigIntToBytes(32, pastRootBI)
  476. err = tree.SetRoot(pastRoot)
  477. c.Assert(err, qt.IsNil)
  478. checkRootBIString(c, tree, expectedRoot)
  479. // check that the tree can be updated
  480. err = tree.Add([]byte("test"), []byte("test"))
  481. c.Assert(err, qt.IsNil)
  482. err = tree.Update([]byte("test"), []byte("test"))
  483. c.Assert(err, qt.IsNil)
  484. // check that the k-v '1000' does not exist in the new tree
  485. _, _, err = tree.Get(k)
  486. c.Assert(err, qt.Equals, ErrKeyNotFound)
  487. // check that can be set an empty root
  488. err = tree.SetRoot(tree.emptyHash)
  489. c.Assert(err, qt.IsNil)
  490. }
  491. func TestSnapshot(t *testing.T) {
  492. c := qt.New(t)
  493. database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
  494. c.Assert(err, qt.IsNil)
  495. tree, err := NewTree(database, 100, HashFunctionPoseidon)
  496. c.Assert(err, qt.IsNil)
  497. // fill the tree
  498. bLen := tree.HashFunction().Len()
  499. var keys, values [][]byte
  500. for i := 0; i < 1000; i++ {
  501. k := BigIntToBytes(bLen, big.NewInt(int64(i)))
  502. v := BigIntToBytes(bLen, big.NewInt(int64(i)))
  503. keys = append(keys, k)
  504. values = append(values, v)
  505. }
  506. indexes, err := tree.AddBatch(keys, values)
  507. c.Assert(err, qt.IsNil)
  508. c.Check(len(indexes), qt.Equals, 0)
  509. checkRootBIString(c, tree,
  510. "13742386369878513332697380582061714160370929283209286127733983161245560237407")
  511. // do a snapshot, and expect the same root than the original tree
  512. snapshotTree, err := tree.Snapshot(nil)
  513. c.Assert(err, qt.IsNil)
  514. checkRootBIString(c, snapshotTree,
  515. "13742386369878513332697380582061714160370929283209286127733983161245560237407")
  516. // check that the snapshotTree can not be updated
  517. _, err = snapshotTree.AddBatch(keys, values)
  518. c.Assert(err, qt.Equals, ErrSnapshotNotEditable)
  519. err = snapshotTree.Add([]byte("test"), []byte("test"))
  520. c.Assert(err, qt.Equals, ErrSnapshotNotEditable)
  521. err = snapshotTree.Update([]byte("test"), []byte("test"))
  522. c.Assert(err, qt.Equals, ErrSnapshotNotEditable)
  523. err = snapshotTree.ImportDump(nil)
  524. c.Assert(err, qt.Equals, ErrSnapshotNotEditable)
  525. // update the original tree by adding a new key-value, and check that
  526. // snapshotTree still has the old root, but the original tree has a new
  527. // root
  528. err = tree.Add([]byte("test"), []byte("test"))
  529. c.Assert(err, qt.IsNil)
  530. checkRootBIString(c, snapshotTree,
  531. "13742386369878513332697380582061714160370929283209286127733983161245560237407")
  532. checkRootBIString(c, tree,
  533. "1025190963769001718196479367844646783678188389989148142691917685159698888868")
  534. }
  535. func TestGetFromSnapshotExpectArboErrKeyNotFound(t *testing.T) {
  536. c := qt.New(t)
  537. database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
  538. c.Assert(err, qt.IsNil)
  539. tree, err := NewTree(database, 100, HashFunctionPoseidon)
  540. c.Assert(err, qt.IsNil)
  541. defer tree.db.Close() //nolint:errcheck
  542. bLen := tree.HashFunction().Len()
  543. k := BigIntToBytes(bLen, big.NewInt(int64(3)))
  544. root, err := tree.Root()
  545. c.Assert(err, qt.IsNil)
  546. tree, err = tree.Snapshot(root)
  547. c.Assert(err, qt.IsNil)
  548. _, _, err = tree.Get(k)
  549. c.Assert(err, qt.Equals, ErrKeyNotFound) // and not equal to db.ErrKeyNotFound
  550. }
  551. func BenchmarkAdd(b *testing.B) {
  552. bLen := 32 // for both Poseidon & Sha256
  553. // prepare inputs
  554. var ks, vs [][]byte
  555. for i := 0; i < 1000; i++ {
  556. k := BigIntToBytes(bLen, big.NewInt(int64(i)))
  557. v := BigIntToBytes(bLen, big.NewInt(int64(i)))
  558. ks = append(ks, k)
  559. vs = append(vs, v)
  560. }
  561. b.Run("Poseidon", func(b *testing.B) {
  562. benchmarkAdd(b, HashFunctionPoseidon, ks, vs)
  563. })
  564. b.Run("Sha256", func(b *testing.B) {
  565. benchmarkAdd(b, HashFunctionSha256, ks, vs)
  566. })
  567. }
  568. func benchmarkAdd(b *testing.B, hashFunc HashFunction, ks, vs [][]byte) {
  569. c := qt.New(b)
  570. database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
  571. c.Assert(err, qt.IsNil)
  572. tree, err := NewTree(database, 140, hashFunc)
  573. c.Assert(err, qt.IsNil)
  574. defer tree.db.Close() //nolint:errcheck
  575. for i := 0; i < len(ks); i++ {
  576. if err := tree.Add(ks[i], vs[i]); err != nil {
  577. b.Fatal(err)
  578. }
  579. }
  580. }