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.

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