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.

796 lines
20 KiB

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