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.

800 lines
20 KiB

  1. package arbo
  2. import (
  3. "crypto/rand"
  4. "encoding/hex"
  5. "fmt"
  6. "math/big"
  7. "os"
  8. "runtime"
  9. "sort"
  10. "testing"
  11. "time"
  12. qt "github.com/frankban/quicktest"
  13. "github.com/iden3/go-merkletree/db/leveldb"
  14. "github.com/iden3/go-merkletree/db/memory"
  15. )
  16. var debug = true
  17. func printTestContext(prefix string, nLeafs int, hashName, dbName string) {
  18. if debug {
  19. fmt.Printf("%snCPU: %d, nLeafs: %d, hash: %s, db: %s\n",
  20. prefix, runtime.NumCPU(), nLeafs, hashName, dbName)
  21. }
  22. }
  23. func printRes(name string, duration time.Duration) {
  24. if debug {
  25. fmt.Printf("%s: %s \n", name, duration)
  26. }
  27. }
  28. func debugTime(descr string, time1, time2 time.Duration) {
  29. if debug {
  30. fmt.Printf("%s was %.02fx times faster than without AddBatch\n",
  31. descr, float64(time1)/float64(time2))
  32. }
  33. }
  34. func testInit(c *qt.C, n int) (*Tree, *Tree) {
  35. tree1, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
  36. c.Assert(err, qt.IsNil)
  37. defer tree1.db.Close()
  38. tree2, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
  39. c.Assert(err, qt.IsNil)
  40. defer tree2.db.Close()
  41. bLen := HashFunctionPoseidon.Len()
  42. // add the initial leafs to fill a bit the trees before calling the
  43. // AddBatch method
  44. for i := 0; i < n; i++ {
  45. k := BigIntToBytes(bLen, big.NewInt(int64(i)))
  46. v := BigIntToBytes(bLen, big.NewInt(int64(i*2)))
  47. if err := tree1.Add(k, v); err != nil {
  48. c.Fatal(err)
  49. }
  50. if err := tree2.Add(k, v); err != nil {
  51. c.Fatal(err)
  52. }
  53. }
  54. return tree1, tree2
  55. }
  56. func TestAddBatchTreeEmpty(t *testing.T) {
  57. c := qt.New(t)
  58. nLeafs := 1024
  59. tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
  60. c.Assert(err, qt.IsNil)
  61. defer tree.db.Close()
  62. bLen := tree.HashFunction().Len()
  63. start := time.Now()
  64. for i := 0; i < nLeafs; i++ {
  65. k := BigIntToBytes(bLen, big.NewInt(int64(i)))
  66. v := BigIntToBytes(bLen, big.NewInt(int64(i*2)))
  67. if err := tree.Add(k, v); err != nil {
  68. t.Fatal(err)
  69. }
  70. }
  71. time1 := time.Since(start)
  72. tree2, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
  73. c.Assert(err, qt.IsNil)
  74. defer tree2.db.Close()
  75. tree2.dbgInit()
  76. var keys, values [][]byte
  77. for i := 0; i < nLeafs; i++ {
  78. k := BigIntToBytes(bLen, big.NewInt(int64(i)))
  79. v := BigIntToBytes(bLen, big.NewInt(int64(i*2)))
  80. keys = append(keys, k)
  81. values = append(values, v)
  82. }
  83. start = time.Now()
  84. indexes, err := tree2.AddBatch(keys, values)
  85. c.Assert(err, qt.IsNil)
  86. time2 := time.Since(start)
  87. if debug {
  88. debugTime("Case tree empty, AddBatch", time1, time2)
  89. printTestContext(" ", nLeafs, "Poseidon", "memory")
  90. tree2.dbg.print(" ")
  91. }
  92. c.Check(len(indexes), qt.Equals, 0)
  93. // check that both trees roots are equal
  94. c.Check(tree2.Root(), qt.DeepEquals, tree.Root())
  95. }
  96. func TestAddBatchTreeEmptyNotPowerOf2(t *testing.T) {
  97. c := qt.New(t)
  98. nLeafs := 1027
  99. tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
  100. c.Assert(err, qt.IsNil)
  101. defer tree.db.Close()
  102. bLen := tree.HashFunction().Len()
  103. for i := 0; i < nLeafs; i++ {
  104. k := BigIntToBytes(bLen, big.NewInt(int64(i)))
  105. v := BigIntToBytes(bLen, big.NewInt(int64(i*2)))
  106. if err := tree.Add(k, v); err != nil {
  107. t.Fatal(err)
  108. }
  109. }
  110. tree2, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
  111. c.Assert(err, qt.IsNil)
  112. defer tree2.db.Close()
  113. var keys, values [][]byte
  114. for i := 0; i < nLeafs; i++ {
  115. k := BigIntToBytes(bLen, big.NewInt(int64(i)))
  116. v := BigIntToBytes(bLen, big.NewInt(int64(i*2)))
  117. keys = append(keys, k)
  118. values = append(values, v)
  119. }
  120. indexes, err := tree2.AddBatch(keys, values)
  121. c.Assert(err, qt.IsNil)
  122. c.Check(len(indexes), qt.Equals, 0)
  123. // check that both trees roots are equal
  124. c.Check(tree2.Root(), qt.DeepEquals, tree.Root())
  125. }
  126. func randomBytes(n int) []byte {
  127. b := make([]byte, n)
  128. _, err := rand.Read(b)
  129. if err != nil {
  130. panic(err)
  131. }
  132. return b
  133. }
  134. func TestAddBatchTreeEmptyTestVector(t *testing.T) {
  135. c := qt.New(t)
  136. tree1, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionBlake2b)
  137. c.Assert(err, qt.IsNil)
  138. defer tree1.db.Close()
  139. tree2, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionBlake2b)
  140. c.Assert(err, qt.IsNil)
  141. defer tree2.db.Close()
  142. // leafs in 2nd level subtrees: [ 6, 0, 1, 1]
  143. testvectorKeys := []string{
  144. "1c7c2265e368314ca58ed2e1f33a326f1220e234a566d55c3605439dbe411642",
  145. "2c9f0a578afff5bfa4e0992a43066460faaab9e8e500db0b16647c701cdb16bf",
  146. "1c45cb31f2fa39ec7b9ebf0fad40e0b8296016b5ce8844ae06ff77226379d9a5",
  147. "d8af98bbbb585129798ae54d5eabbc9d0561d583faf1663b3a3724d15bda4ec7",
  148. }
  149. var keys, values [][]byte
  150. for i := 0; i < len(testvectorKeys); i++ {
  151. key, err := hex.DecodeString(testvectorKeys[i])
  152. c.Assert(err, qt.IsNil)
  153. keys = append(keys, key)
  154. values = append(values, []byte{0})
  155. }
  156. for i := 0; i < len(keys); i++ {
  157. if err := tree1.Add(keys[i], values[i]); err != nil {
  158. t.Fatal(err)
  159. }
  160. }
  161. indexes, err := tree2.AddBatch(keys, values)
  162. c.Assert(err, qt.IsNil)
  163. c.Check(len(indexes), qt.Equals, 0)
  164. // check that both trees roots are equal
  165. c.Check(tree2.Root(), qt.DeepEquals, tree1.Root())
  166. // 2nd test vectors
  167. tree1, err = NewTree(memory.NewMemoryStorage(), 100, HashFunctionBlake2b)
  168. c.Assert(err, qt.IsNil)
  169. defer tree1.db.Close()
  170. tree2, err = NewTree(memory.NewMemoryStorage(), 100, HashFunctionBlake2b)
  171. c.Assert(err, qt.IsNil)
  172. defer tree2.db.Close()
  173. testvectorKeys = []string{
  174. "1c7c2265e368314ca58ed2e1f33a326f1220e234a566d55c3605439dbe411642",
  175. "2c9f0a578afff5bfa4e0992a43066460faaab9e8e500db0b16647c701cdb16bf",
  176. "9cb87ec67e875c61390edcd1ab517f443591047709a4d4e45b0f9ed980857b8e",
  177. "9b4e9e92e974a589f426ceeb4cb291dc24893513fecf8e8460992dcf52621d4d",
  178. "1c45cb31f2fa39ec7b9ebf0fad40e0b8296016b5ce8844ae06ff77226379d9a5",
  179. "d8af98bbbb585129798ae54d5eabbc9d0561d583faf1663b3a3724d15bda4ec7",
  180. "3cd55dbfb8f975f20a0925dfbdabe79fa2d51dd0268afbb8ba6b01de9dfcdd3c",
  181. "5d0a9d6d9f197c091bf054fac9cb60e11ec723d6610ed8578e617b4d46cb43d5",
  182. }
  183. keys = [][]byte{}
  184. values = [][]byte{}
  185. for i := 0; i < len(testvectorKeys); i++ {
  186. key, err := hex.DecodeString(testvectorKeys[i])
  187. c.Assert(err, qt.IsNil)
  188. keys = append(keys, key)
  189. values = append(values, []byte{0})
  190. }
  191. for i := 0; i < len(keys); i++ {
  192. if err := tree1.Add(keys[i], values[i]); err != nil {
  193. t.Fatal(err)
  194. }
  195. }
  196. indexes, err = tree2.AddBatch(keys, values)
  197. c.Assert(err, qt.IsNil)
  198. c.Check(len(indexes), qt.Equals, 0)
  199. // check that both trees roots are equal
  200. c.Check(tree2.Root(), qt.DeepEquals, tree1.Root())
  201. }
  202. func TestAddBatchTreeEmptyRandomKeys(t *testing.T) {
  203. c := qt.New(t)
  204. nLeafs := 8
  205. tree1, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionBlake2b)
  206. c.Assert(err, qt.IsNil)
  207. defer tree1.db.Close()
  208. tree2, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionBlake2b)
  209. c.Assert(err, qt.IsNil)
  210. defer tree2.db.Close()
  211. var keys, values [][]byte
  212. for i := 0; i < nLeafs; i++ {
  213. keys = append(keys, randomBytes(32))
  214. values = append(values, randomBytes(32))
  215. }
  216. for i := 0; i < len(keys); i++ {
  217. if err := tree1.Add(keys[i], values[i]); err != nil {
  218. t.Fatal(err)
  219. }
  220. }
  221. indexes, err := tree2.AddBatch(keys, values)
  222. c.Assert(err, qt.IsNil)
  223. c.Check(len(indexes), qt.Equals, 0)
  224. // check that both trees roots are equal
  225. c.Check(tree2.Root(), qt.DeepEquals, tree1.Root())
  226. }
  227. func TestAddBatchTreeNotEmptyFewLeafs(t *testing.T) {
  228. c := qt.New(t)
  229. nLeafs := 1024
  230. initialNLeafs := 99
  231. tree1, tree2 := testInit(c, initialNLeafs)
  232. tree2.dbgInit()
  233. bLen := tree1.HashFunction().Len()
  234. start := time.Now()
  235. for i := initialNLeafs; i < nLeafs; i++ {
  236. k := BigIntToBytes(bLen, big.NewInt(int64(i)))
  237. v := BigIntToBytes(bLen, big.NewInt(int64(i*2)))
  238. if err := tree1.Add(k, v); err != nil {
  239. t.Fatal(err)
  240. }
  241. }
  242. time1 := time.Since(start)
  243. // prepare the key-values to be added
  244. var keys, values [][]byte
  245. for i := initialNLeafs; i < nLeafs; i++ {
  246. k := BigIntToBytes(bLen, big.NewInt(int64(i)))
  247. v := BigIntToBytes(bLen, big.NewInt(int64(i*2)))
  248. keys = append(keys, k)
  249. values = append(values, v)
  250. }
  251. start = time.Now()
  252. indexes, err := tree2.AddBatch(keys, values)
  253. c.Assert(err, qt.IsNil)
  254. time2 := time.Since(start)
  255. if debug {
  256. debugTime("Case tree not empty w/ few leafs, AddBatch", time1, time2)
  257. printTestContext(" ", nLeafs, "Poseidon", "memory")
  258. tree2.dbg.print(" ")
  259. }
  260. c.Check(len(indexes), qt.Equals, 0)
  261. // check that both trees roots are equal
  262. c.Check(tree2.Root(), qt.DeepEquals, tree1.Root())
  263. }
  264. func TestAddBatchTreeNotEmptyEnoughLeafs(t *testing.T) {
  265. c := qt.New(t)
  266. nLeafs := 1024
  267. initialNLeafs := 500
  268. tree1, tree2 := testInit(c, initialNLeafs)
  269. tree2.dbgInit()
  270. bLen := tree1.HashFunction().Len()
  271. start := time.Now()
  272. for i := initialNLeafs; i < nLeafs; i++ {
  273. k := BigIntToBytes(bLen, big.NewInt(int64(i)))
  274. v := BigIntToBytes(bLen, big.NewInt(int64(i*2)))
  275. if err := tree1.Add(k, v); err != nil {
  276. t.Fatal(err)
  277. }
  278. }
  279. time1 := time.Since(start)
  280. // prepare the key-values to be added
  281. var keys, values [][]byte
  282. for i := initialNLeafs; i < nLeafs; i++ {
  283. k := BigIntToBytes(bLen, big.NewInt(int64(i)))
  284. v := BigIntToBytes(bLen, big.NewInt(int64(i*2)))
  285. keys = append(keys, k)
  286. values = append(values, v)
  287. }
  288. start = time.Now()
  289. indexes, err := tree2.AddBatch(keys, values)
  290. c.Assert(err, qt.IsNil)
  291. time2 := time.Since(start)
  292. if debug {
  293. debugTime("Case tree not empty w/ enough leafs, AddBatch", time1, time2)
  294. printTestContext(" ", nLeafs, "Poseidon", "memory")
  295. tree2.dbg.print(" ")
  296. }
  297. c.Check(len(indexes), qt.Equals, 0)
  298. // check that both trees roots are equal
  299. c.Check(tree2.Root(), qt.DeepEquals, tree1.Root())
  300. }
  301. func TestAddBatchTreeEmptyRepeatedLeafs(t *testing.T) {
  302. c := qt.New(t)
  303. nLeafs := 1024
  304. nRepeatedKeys := 99
  305. tree1, tree2 := testInit(c, 0)
  306. bLen := tree1.HashFunction().Len()
  307. // prepare the key-values to be added
  308. var keys, values [][]byte
  309. for i := 0; i < nLeafs; i++ {
  310. k := BigIntToBytes(bLen, big.NewInt(int64(i)))
  311. v := BigIntToBytes(bLen, big.NewInt(int64(i*2)))
  312. keys = append(keys, k)
  313. values = append(values, v)
  314. }
  315. // add repeated key-values
  316. for i := 0; i < nRepeatedKeys; i++ {
  317. k := BigIntToBytes(bLen, big.NewInt(int64(i)))
  318. v := BigIntToBytes(bLen, big.NewInt(int64(i*2)))
  319. keys = append(keys, k)
  320. values = append(values, v)
  321. }
  322. // add the non-repeated key-values in tree1 with .Add loop
  323. for i := 0; i < nLeafs; i++ {
  324. if err := tree1.Add(keys[i], values[i]); err != nil {
  325. t.Fatal(err)
  326. }
  327. }
  328. indexes, err := tree2.AddBatch(keys, values)
  329. c.Assert(err, qt.IsNil)
  330. c.Check(len(indexes), qt.Equals, nRepeatedKeys)
  331. // check that both trees roots are equal
  332. c.Check(tree2.Root(), qt.DeepEquals, tree1.Root())
  333. }
  334. func TestAddBatchTreeNotEmptyFewLeafsRepeatedLeafs(t *testing.T) {
  335. c := qt.New(t)
  336. nLeafs := 1024
  337. initialNLeafs := 99
  338. tree1, tree2 := testInit(c, initialNLeafs)
  339. bLen := tree1.HashFunction().Len()
  340. // prepare the key-values to be added
  341. var keys, values [][]byte
  342. for i := 0; i < nLeafs; i++ {
  343. k := BigIntToBytes(bLen, big.NewInt(int64(i)))
  344. v := BigIntToBytes(bLen, big.NewInt(int64(i*2)))
  345. keys = append(keys, k)
  346. values = append(values, v)
  347. }
  348. // add the keys that will be existing when AddBatch is called
  349. for i := initialNLeafs; i < nLeafs; i++ {
  350. if err := tree1.Add(keys[i], values[i]); err != nil {
  351. t.Fatal(err)
  352. }
  353. }
  354. indexes, err := tree2.AddBatch(keys, values)
  355. c.Assert(err, qt.IsNil)
  356. c.Check(len(indexes), qt.Equals, initialNLeafs)
  357. // check that both trees roots are equal
  358. c.Check(tree2.Root(), qt.DeepEquals, tree1.Root())
  359. }
  360. func TestSplitInBuckets(t *testing.T) {
  361. c := qt.New(t)
  362. bLen := HashFunctionPoseidon.Len()
  363. nLeafs := 16
  364. kvs := make([]kv, nLeafs)
  365. for i := 0; i < nLeafs; i++ {
  366. k := BigIntToBytes(bLen, big.NewInt(int64(i)))
  367. v := BigIntToBytes(bLen, big.NewInt(int64(i*2)))
  368. keyPath := make([]byte, 32)
  369. copy(keyPath[:], k)
  370. kvs[i].pos = i
  371. kvs[i].keyPath = k
  372. kvs[i].k = k
  373. kvs[i].v = v
  374. }
  375. // check keyToBucket results for 4 buckets & 8 keys
  376. c.Assert(keyToBucket(kvs[0].k, 4), qt.Equals, 0)
  377. c.Assert(keyToBucket(kvs[1].k, 4), qt.Equals, 2)
  378. c.Assert(keyToBucket(kvs[2].k, 4), qt.Equals, 1)
  379. c.Assert(keyToBucket(kvs[3].k, 4), qt.Equals, 3)
  380. c.Assert(keyToBucket(kvs[4].k, 4), qt.Equals, 0)
  381. c.Assert(keyToBucket(kvs[5].k, 4), qt.Equals, 2)
  382. c.Assert(keyToBucket(kvs[6].k, 4), qt.Equals, 1)
  383. c.Assert(keyToBucket(kvs[7].k, 4), qt.Equals, 3)
  384. // check keyToBucket results for 8 buckets & 8 keys
  385. c.Assert(keyToBucket(kvs[0].k, 8), qt.Equals, 0)
  386. c.Assert(keyToBucket(kvs[1].k, 8), qt.Equals, 4)
  387. c.Assert(keyToBucket(kvs[2].k, 8), qt.Equals, 2)
  388. c.Assert(keyToBucket(kvs[3].k, 8), qt.Equals, 6)
  389. c.Assert(keyToBucket(kvs[4].k, 8), qt.Equals, 1)
  390. c.Assert(keyToBucket(kvs[5].k, 8), qt.Equals, 5)
  391. c.Assert(keyToBucket(kvs[6].k, 8), qt.Equals, 3)
  392. c.Assert(keyToBucket(kvs[7].k, 8), qt.Equals, 7)
  393. buckets := splitInBuckets(kvs, 4)
  394. expected := [][]string{
  395. {
  396. "00000000", // bucket 0
  397. "08000000",
  398. "04000000",
  399. "0c000000",
  400. },
  401. {
  402. "02000000", // bucket 1
  403. "0a000000",
  404. "06000000",
  405. "0e000000",
  406. },
  407. {
  408. "01000000", // bucket 2
  409. "09000000",
  410. "05000000",
  411. "0d000000",
  412. },
  413. {
  414. "03000000", // bucket 3
  415. "0b000000",
  416. "07000000",
  417. "0f000000",
  418. },
  419. }
  420. for i := 0; i < len(buckets); i++ {
  421. sortKvs(buckets[i])
  422. c.Assert(len(buckets[i]), qt.Equals, len(expected[i]))
  423. for j := 0; j < len(buckets[i]); j++ {
  424. c.Check(hex.EncodeToString(buckets[i][j].k[:4]),
  425. qt.Equals, expected[i][j])
  426. }
  427. }
  428. }
  429. // compareBytes compares byte slices where the bytes are compared from left to
  430. // right and each byte is compared by bit from right to left
  431. func compareBytes(a, b []byte) bool {
  432. // WIP
  433. for i := 0; i < len(a); i++ {
  434. for j := 0; j < 8; j++ {
  435. aBit := a[i] & (1 << j)
  436. bBit := b[i] & (1 << j)
  437. if aBit > bBit {
  438. return false
  439. } else if aBit < bBit {
  440. return true
  441. }
  442. }
  443. }
  444. return false
  445. }
  446. // sortKvs sorts the kv by path
  447. func sortKvs(kvs []kv) {
  448. sort.Slice(kvs, func(i, j int) bool {
  449. return compareBytes(kvs[i].keyPath, kvs[j].keyPath)
  450. })
  451. }
  452. func TestAddBatchTreeNotEmpty(t *testing.T) {
  453. c := qt.New(t)
  454. nLeafs := 4096
  455. initialNLeafs := 900
  456. tree1, tree2 := testInit(c, initialNLeafs)
  457. tree2.dbgInit()
  458. bLen := tree1.HashFunction().Len()
  459. start := time.Now()
  460. for i := initialNLeafs; i < nLeafs; i++ {
  461. k := BigIntToBytes(bLen, big.NewInt(int64(i)))
  462. v := BigIntToBytes(bLen, big.NewInt(int64(i*2)))
  463. if err := tree1.Add(k, v); err != nil {
  464. t.Fatal(err)
  465. }
  466. }
  467. time1 := time.Since(start)
  468. // prepare the key-values to be added
  469. var keys, values [][]byte
  470. for i := initialNLeafs; i < nLeafs; i++ {
  471. k := BigIntToBytes(bLen, big.NewInt(int64(i)))
  472. v := BigIntToBytes(bLen, big.NewInt(int64(i*2)))
  473. keys = append(keys, k)
  474. values = append(values, v)
  475. }
  476. start = time.Now()
  477. indexes, err := tree2.AddBatch(keys, values)
  478. c.Assert(err, qt.IsNil)
  479. time2 := time.Since(start)
  480. if debug {
  481. debugTime("Case tree not empty, AddBatch", time1, time2)
  482. printTestContext(" ", nLeafs, "Poseidon", "memory")
  483. tree2.dbg.print(" ")
  484. }
  485. c.Check(len(indexes), qt.Equals, 0)
  486. // check that both trees roots are equal
  487. c.Check(tree2.Root(), qt.DeepEquals, tree1.Root())
  488. }
  489. func TestAddBatchNotEmptyUnbalanced(t *testing.T) {
  490. c := qt.New(t)
  491. nLeafs := 4096
  492. initialNLeafs := 900
  493. tree1, _ := testInit(c, initialNLeafs)
  494. bLen := tree1.HashFunction().Len()
  495. start := time.Now()
  496. for i := initialNLeafs; i < nLeafs; i++ {
  497. k := BigIntToBytes(bLen, big.NewInt(int64(i)))
  498. v := BigIntToBytes(bLen, big.NewInt(int64(i*2)))
  499. if err := tree1.Add(k, v); err != nil {
  500. t.Fatal(err)
  501. }
  502. }
  503. time1 := time.Since(start)
  504. tree2, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
  505. c.Assert(err, qt.IsNil)
  506. defer tree2.db.Close()
  507. tree2.dbgInit()
  508. var keys, values [][]byte
  509. // add the initial leafs to fill a bit the tree before calling the
  510. // AddBatch method
  511. for i := 0; i < initialNLeafs; i++ {
  512. k := BigIntToBytes(bLen, big.NewInt(int64(i)))
  513. v := BigIntToBytes(bLen, big.NewInt(int64(i*2)))
  514. // use only the keys of one bucket, store the not used ones for
  515. // later
  516. if i%4 != 0 {
  517. keys = append(keys, k)
  518. values = append(values, v)
  519. continue
  520. }
  521. if err := tree2.Add(k, v); err != nil {
  522. t.Fatal(err)
  523. }
  524. }
  525. for i := initialNLeafs; i < nLeafs; i++ {
  526. k := BigIntToBytes(bLen, big.NewInt(int64(i)))
  527. v := BigIntToBytes(bLen, big.NewInt(int64(i*2)))
  528. keys = append(keys, k)
  529. values = append(values, v)
  530. }
  531. start = time.Now()
  532. indexes, err := tree2.AddBatch(keys, values)
  533. c.Assert(err, qt.IsNil)
  534. time2 := time.Since(start)
  535. if debug {
  536. debugTime("Case tree not empty & unbalanced, AddBatch", time1, time2)
  537. printTestContext(" ", nLeafs, "Poseidon", "memory")
  538. tree2.dbg.print(" ")
  539. }
  540. c.Check(len(indexes), qt.Equals, 0)
  541. // check that both trees roots are equal
  542. c.Check(tree2.Root(), qt.DeepEquals, tree1.Root())
  543. }
  544. func TestFlp2(t *testing.T) {
  545. c := qt.New(t)
  546. c.Assert(flp2(31), qt.Equals, 16)
  547. c.Assert(flp2(32), qt.Equals, 32)
  548. c.Assert(flp2(33), qt.Equals, 32)
  549. c.Assert(flp2(63), qt.Equals, 32)
  550. c.Assert(flp2(64), qt.Equals, 64)
  551. c.Assert(flp2(9000), qt.Equals, 8192)
  552. }
  553. func TestAddBatchBench(t *testing.T) {
  554. nLeafs := 50_000
  555. printTestContext("TestAddBatchBench: ", nLeafs, "Blake2b", "leveldb")
  556. // prepare inputs
  557. var ks, vs [][]byte
  558. for i := 0; i < nLeafs; i++ {
  559. k := randomBytes(32)
  560. v := randomBytes(32)
  561. ks = append(ks, k)
  562. vs = append(vs, v)
  563. }
  564. benchAdd(t, ks, vs)
  565. benchAddBatch(t, ks, vs)
  566. }
  567. func benchAdd(t *testing.T, ks, vs [][]byte) {
  568. c := qt.New(t)
  569. dbDir := t.TempDir()
  570. defer os.RemoveAll(dbDir) //nolint:errcheck
  571. // storage, err := pebble.NewPebbleStorage(dbDir, false)
  572. storage, err := leveldb.NewLevelDbStorage(dbDir, false)
  573. c.Assert(err, qt.IsNil)
  574. tree, err := NewTree(storage, 140, HashFunctionBlake2b)
  575. c.Assert(err, qt.IsNil)
  576. start := time.Now()
  577. for i := 0; i < len(ks); i++ {
  578. err = tree.Add(ks[i], vs[i])
  579. c.Assert(err, qt.IsNil)
  580. }
  581. if debug {
  582. printRes(" Add loop", time.Since(start))
  583. tree.dbg.print(" ")
  584. }
  585. }
  586. func benchAddBatch(t *testing.T, ks, vs [][]byte) {
  587. c := qt.New(t)
  588. dbDir := t.TempDir()
  589. defer os.RemoveAll(dbDir) //nolint:errcheck
  590. // storage, err := pebble.NewPebbleStorage(dbDir, false)
  591. storage, err := leveldb.NewLevelDbStorage(dbDir, false)
  592. c.Assert(err, qt.IsNil)
  593. tree, err := NewTree(storage, 140, HashFunctionBlake2b)
  594. c.Assert(err, qt.IsNil)
  595. tree.dbgInit()
  596. start := time.Now()
  597. invalids, err := tree.AddBatch(ks, vs)
  598. if debug {
  599. printRes(" AddBatch", time.Since(start))
  600. tree.dbg.print(" ")
  601. }
  602. c.Assert(err, qt.IsNil)
  603. c.Assert(len(invalids), qt.Equals, 0)
  604. }
  605. func TestDbgStats(t *testing.T) {
  606. c := qt.New(t)
  607. nLeafs := 10_000
  608. // prepare inputs
  609. var ks, vs [][]byte
  610. for i := 0; i < nLeafs; i++ {
  611. k := randomBytes(32)
  612. v := randomBytes(32)
  613. ks = append(ks, k)
  614. vs = append(vs, v)
  615. }
  616. // 1
  617. tree1, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionBlake2b)
  618. c.Assert(err, qt.IsNil)
  619. defer tree1.db.Close()
  620. tree1.dbgInit()
  621. for i := 0; i < len(ks); i++ {
  622. err = tree1.Add(ks[i], vs[i])
  623. c.Assert(err, qt.IsNil)
  624. }
  625. // 2
  626. tree2, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionBlake2b)
  627. c.Assert(err, qt.IsNil)
  628. defer tree2.db.Close()
  629. tree2.dbgInit()
  630. invalids, err := tree2.AddBatch(ks, vs)
  631. c.Assert(err, qt.IsNil)
  632. c.Assert(len(invalids), qt.Equals, 0)
  633. // 3
  634. tree3, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionBlake2b)
  635. c.Assert(err, qt.IsNil)
  636. defer tree3.db.Close()
  637. tree3.dbgInit()
  638. // add few key-values
  639. // invalids, err = tree3.AddBatch(ks[:], vs[:])
  640. invalids, err = tree3.AddBatch(ks[:1000], vs[:1000])
  641. c.Assert(err, qt.IsNil)
  642. c.Assert(len(invalids), qt.Equals, 0)
  643. // add the rest of key-values
  644. invalids, err = tree3.AddBatch(ks[1000:], vs[1000:])
  645. c.Assert(err, qt.IsNil)
  646. c.Assert(len(invalids), qt.Equals, 0)
  647. c.Check(tree2.Root(), qt.DeepEquals, tree1.Root())
  648. c.Check(tree3.Root(), qt.DeepEquals, tree1.Root())
  649. if debug {
  650. fmt.Println("TestDbgStats")
  651. tree1.dbg.print(" add in loop in emptyTree ")
  652. tree2.dbg.print(" addbatch caseEmptyTree ")
  653. tree3.dbg.print(" addbatch caseNotEmptyTree ")
  654. }
  655. }
  656. func TestLoadVT(t *testing.T) {
  657. c := qt.New(t)
  658. nLeafs := 1024
  659. tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
  660. c.Assert(err, qt.IsNil)
  661. defer tree.db.Close()
  662. var keys, values [][]byte
  663. for i := 0; i < nLeafs; i++ {
  664. k := randomBytes(31)
  665. v := randomBytes(31)
  666. keys = append(keys, k)
  667. values = append(values, v)
  668. }
  669. indexes, err := tree.AddBatch(keys, values)
  670. c.Assert(err, qt.IsNil)
  671. c.Check(len(indexes), qt.Equals, 0)
  672. vt, err := tree.loadVT()
  673. c.Assert(err, qt.IsNil)
  674. _, err = vt.computeHashes()
  675. c.Assert(err, qt.IsNil)
  676. // check that tree & vt roots are equal
  677. c.Check(tree.Root(), qt.DeepEquals, vt.root.h)
  678. }
  679. // TODO test adding batch with multiple invalid keys