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.

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