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.

728 lines
20 KiB

  1. package arbo
  2. import (
  3. "crypto/rand"
  4. "encoding/hex"
  5. "fmt"
  6. "math/big"
  7. "runtime"
  8. "testing"
  9. "time"
  10. qt "github.com/frankban/quicktest"
  11. "github.com/iden3/go-merkletree/db/leveldb"
  12. "github.com/iden3/go-merkletree/db/memory"
  13. )
  14. var debug = true
  15. func debugTime(descr string, time1, time2 time.Duration) {
  16. if debug {
  17. fmt.Printf("%s was %f times faster than without AddBatch\n",
  18. descr, float64(time1)/float64(time2))
  19. }
  20. }
  21. func testInit(c *qt.C, n int) (*Tree, *Tree) {
  22. tree1, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
  23. c.Assert(err, qt.IsNil)
  24. defer tree1.db.Close()
  25. tree2, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
  26. c.Assert(err, qt.IsNil)
  27. defer tree2.db.Close()
  28. // add the initial leafs to fill a bit the trees before calling the
  29. // AddBatch method
  30. for i := 0; i < n; i++ {
  31. k := BigIntToBytes(big.NewInt(int64(i)))
  32. v := BigIntToBytes(big.NewInt(int64(i * 2)))
  33. if err := tree1.Add(k, v); err != nil {
  34. c.Fatal(err)
  35. }
  36. if err := tree2.Add(k, v); err != nil {
  37. c.Fatal(err)
  38. }
  39. }
  40. return tree1, tree2
  41. }
  42. func TestAddBatchCaseA(t *testing.T) {
  43. c := qt.New(t)
  44. nLeafs := 1024
  45. tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
  46. c.Assert(err, qt.IsNil)
  47. defer tree.db.Close()
  48. start := time.Now()
  49. for i := 0; i < nLeafs; i++ {
  50. k := BigIntToBytes(big.NewInt(int64(i)))
  51. v := BigIntToBytes(big.NewInt(int64(i * 2)))
  52. if err := tree.Add(k, v); err != nil {
  53. t.Fatal(err)
  54. }
  55. }
  56. time1 := time.Since(start)
  57. tree2, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
  58. c.Assert(err, qt.IsNil)
  59. defer tree2.db.Close()
  60. var keys, values [][]byte
  61. for i := 0; i < nLeafs; i++ {
  62. k := BigIntToBytes(big.NewInt(int64(i)))
  63. v := BigIntToBytes(big.NewInt(int64(i * 2)))
  64. keys = append(keys, k)
  65. values = append(values, v)
  66. }
  67. start = time.Now()
  68. indexes, err := tree2.AddBatch(keys, values)
  69. c.Assert(err, qt.IsNil)
  70. time2 := time.Since(start)
  71. debugTime("CASE A, AddBatch", time1, time2)
  72. c.Check(len(indexes), qt.Equals, 0)
  73. // check that both trees roots are equal
  74. c.Check(tree2.Root(), qt.DeepEquals, tree.Root())
  75. }
  76. func TestAddBatchCaseANotPowerOf2(t *testing.T) {
  77. c := qt.New(t)
  78. nLeafs := 1027
  79. tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
  80. c.Assert(err, qt.IsNil)
  81. defer tree.db.Close()
  82. for i := 0; i < nLeafs; i++ {
  83. k := BigIntToBytes(big.NewInt(int64(i)))
  84. v := BigIntToBytes(big.NewInt(int64(i * 2)))
  85. if err := tree.Add(k, v); err != nil {
  86. t.Fatal(err)
  87. }
  88. }
  89. tree2, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
  90. c.Assert(err, qt.IsNil)
  91. defer tree2.db.Close()
  92. var keys, values [][]byte
  93. for i := 0; i < nLeafs; i++ {
  94. k := BigIntToBytes(big.NewInt(int64(i)))
  95. v := BigIntToBytes(big.NewInt(int64(i * 2)))
  96. keys = append(keys, k)
  97. values = append(values, v)
  98. }
  99. indexes, err := tree2.AddBatch(keys, values)
  100. c.Assert(err, qt.IsNil)
  101. c.Check(len(indexes), qt.Equals, 0)
  102. // check that both trees roots are equal
  103. c.Check(tree2.Root(), qt.DeepEquals, tree.Root())
  104. }
  105. func randomBytes(n int) []byte {
  106. b := make([]byte, n)
  107. _, err := rand.Read(b)
  108. if err != nil {
  109. panic(err)
  110. }
  111. return b
  112. }
  113. func TestBuildTreeBottomUpSingleThread(t *testing.T) {
  114. c := qt.New(t)
  115. tree1, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionBlake2b)
  116. c.Assert(err, qt.IsNil)
  117. defer tree1.db.Close()
  118. tree2, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionBlake2b)
  119. c.Assert(err, qt.IsNil)
  120. defer tree2.db.Close()
  121. testvectorKeys := []string{
  122. "1c7c2265e368314ca58ed2e1f33a326f1220e234a566d55c3605439dbe411642",
  123. "2c9f0a578afff5bfa4e0992a43066460faaab9e8e500db0b16647c701cdb16bf",
  124. "1c45cb31f2fa39ec7b9ebf0fad40e0b8296016b5ce8844ae06ff77226379d9a5",
  125. "d8af98bbbb585129798ae54d5eabbc9d0561d583faf1663b3a3724d15bda4ec7",
  126. }
  127. var keys, values [][]byte
  128. for i := 0; i < len(testvectorKeys); i++ {
  129. key, err := hex.DecodeString(testvectorKeys[i])
  130. c.Assert(err, qt.IsNil)
  131. keys = append(keys, key)
  132. values = append(values, []byte{0})
  133. }
  134. for i := 0; i < len(keys); i++ {
  135. if err := tree1.Add(keys[i], values[i]); err != nil {
  136. t.Fatal(err)
  137. }
  138. }
  139. kvs, err := tree2.keysValuesToKvs(keys, values)
  140. c.Assert(err, qt.IsNil)
  141. sortKvs(kvs)
  142. tree2.tx, err = tree2.db.NewTx()
  143. c.Assert(err, qt.IsNil)
  144. // indexes, err := tree2.buildTreeBottomUpSingleThread(kvs)
  145. indexes, err := tree2.buildTreeBottomUp(4, kvs)
  146. c.Assert(err, qt.IsNil)
  147. // tree1.PrintGraphviz(nil)
  148. // tree2.PrintGraphviz(nil)
  149. c.Check(len(indexes), qt.Equals, 0)
  150. // check that both trees roots are equal
  151. c.Check(tree2.Root(), qt.DeepEquals, tree1.Root())
  152. // 15b6a23945ae6c81342b7eb14e70fff50812dc8791cb15ec791eb08f91784139
  153. }
  154. func TestAddBatchCaseATestVector(t *testing.T) {
  155. c := qt.New(t)
  156. tree1, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionBlake2b)
  157. c.Assert(err, qt.IsNil)
  158. defer tree1.db.Close()
  159. tree2, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionBlake2b)
  160. c.Assert(err, qt.IsNil)
  161. defer tree2.db.Close()
  162. // leafs in 2nd level subtrees: [ 6, 0, 1, 1]
  163. testvectorKeys := []string{
  164. "1c7c2265e368314ca58ed2e1f33a326f1220e234a566d55c3605439dbe411642",
  165. "2c9f0a578afff5bfa4e0992a43066460faaab9e8e500db0b16647c701cdb16bf",
  166. "1c45cb31f2fa39ec7b9ebf0fad40e0b8296016b5ce8844ae06ff77226379d9a5",
  167. "d8af98bbbb585129798ae54d5eabbc9d0561d583faf1663b3a3724d15bda4ec7",
  168. }
  169. var keys, values [][]byte
  170. for i := 0; i < len(testvectorKeys); i++ {
  171. key, err := hex.DecodeString(testvectorKeys[i])
  172. c.Assert(err, qt.IsNil)
  173. keys = append(keys, key)
  174. values = append(values, []byte{0})
  175. }
  176. for i := 0; i < len(keys); i++ {
  177. if err := tree1.Add(keys[i], values[i]); err != nil {
  178. t.Fatal(err)
  179. }
  180. }
  181. indexes, err := tree2.AddBatch(keys, values)
  182. c.Assert(err, qt.IsNil)
  183. // tree1.PrintGraphviz(nil)
  184. // tree2.PrintGraphviz(nil)
  185. c.Check(len(indexes), qt.Equals, 0)
  186. // check that both trees roots are equal
  187. // fmt.Println(hex.EncodeToString(tree1.Root()))
  188. c.Check(tree2.Root(), qt.DeepEquals, tree1.Root())
  189. }
  190. func TestAddBatchCaseARandomKeys(t *testing.T) {
  191. c := qt.New(t)
  192. nLeafs := 8
  193. tree1, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionBlake2b)
  194. c.Assert(err, qt.IsNil)
  195. defer tree1.db.Close()
  196. tree2, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionBlake2b)
  197. c.Assert(err, qt.IsNil)
  198. defer tree2.db.Close()
  199. var keys, values [][]byte
  200. for i := 0; i < nLeafs; i++ {
  201. keys = append(keys, randomBytes(32))
  202. // values = append(values, randomBytes(32))
  203. values = append(values, []byte{0})
  204. // fmt.Println("K", hex.EncodeToString(keys[i]))
  205. }
  206. // TMP:
  207. keys[0], _ = hex.DecodeString("1c7c2265e368314ca58ed2e1f33a326f1220e234a566d55c3605439dbe411642")
  208. keys[1], _ = hex.DecodeString("2c9f0a578afff5bfa4e0992a43066460faaab9e8e500db0b16647c701cdb16bf")
  209. keys[2], _ = hex.DecodeString("9cb87ec67e875c61390edcd1ab517f443591047709a4d4e45b0f9ed980857b8e")
  210. keys[3], _ = hex.DecodeString("9b4e9e92e974a589f426ceeb4cb291dc24893513fecf8e8460992dcf52621d4d")
  211. keys[4], _ = hex.DecodeString("1c45cb31f2fa39ec7b9ebf0fad40e0b8296016b5ce8844ae06ff77226379d9a5")
  212. keys[5], _ = hex.DecodeString("d8af98bbbb585129798ae54d5eabbc9d0561d583faf1663b3a3724d15bda4ec7")
  213. keys[6], _ = hex.DecodeString("3cd55dbfb8f975f20a0925dfbdabe79fa2d51dd0268afbb8ba6b01de9dfcdd3c")
  214. keys[7], _ = hex.DecodeString("5d0a9d6d9f197c091bf054fac9cb60e11ec723d6610ed8578e617b4d46cb43d5")
  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. // tree1.PrintGraphviz(nil)
  223. // tree2.PrintGraphviz(nil)
  224. c.Check(len(indexes), qt.Equals, 0)
  225. // check that both trees roots are equal
  226. c.Check(tree2.Root(), qt.DeepEquals, tree1.Root())
  227. }
  228. func TestAddBatchCaseB(t *testing.T) {
  229. c := qt.New(t)
  230. nLeafs := 1024
  231. initialNLeafs := 99 // TMP TODO use const minLeafsThreshold-1 once ready
  232. tree1, tree2 := testInit(c, initialNLeafs)
  233. start := time.Now()
  234. for i := initialNLeafs; i < nLeafs; i++ {
  235. k := BigIntToBytes(big.NewInt(int64(i)))
  236. v := BigIntToBytes(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(big.NewInt(int64(i)))
  246. v := BigIntToBytes(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. debugTime("CASE B, AddBatch", time1, time2)
  255. c.Check(len(indexes), qt.Equals, 0)
  256. // check that both trees roots are equal
  257. c.Check(tree2.Root(), qt.DeepEquals, tree1.Root())
  258. }
  259. func TestAddBatchCaseBRepeatedLeafs(t *testing.T) {
  260. c := qt.New(t)
  261. nLeafs := 1024
  262. initialNLeafs := 99 // TMP TODO use const minLeafsThreshold-1 once ready
  263. tree1, tree2 := testInit(c, initialNLeafs)
  264. for i := initialNLeafs; i < nLeafs; i++ {
  265. k := BigIntToBytes(big.NewInt(int64(i)))
  266. v := BigIntToBytes(big.NewInt(int64(i * 2)))
  267. if err := tree1.Add(k, v); err != nil {
  268. t.Fatal(err)
  269. }
  270. }
  271. // prepare the key-values to be added, including already existing keys
  272. var keys, values [][]byte
  273. for i := 0; i < nLeafs; i++ {
  274. k := BigIntToBytes(big.NewInt(int64(i)))
  275. v := BigIntToBytes(big.NewInt(int64(i * 2)))
  276. keys = append(keys, k)
  277. values = append(values, v)
  278. }
  279. indexes, err := tree2.AddBatch(keys, values)
  280. c.Assert(err, qt.IsNil)
  281. c.Check(len(indexes), qt.Equals, initialNLeafs)
  282. // check that both trees roots are equal
  283. c.Check(tree2.Root(), qt.DeepEquals, tree1.Root())
  284. }
  285. func TestCombineInKVSet(t *testing.T) {
  286. c := qt.New(t)
  287. var a, b, expected []kv
  288. for i := 0; i < 10; i++ {
  289. k := BigIntToBytes(big.NewInt(int64(i)))
  290. kv := kv{k: k}
  291. if i < 7 {
  292. a = append(a, kv)
  293. }
  294. if i >= 4 {
  295. b = append(b, kv)
  296. }
  297. expected = append(expected, kv)
  298. }
  299. r, invalids := combineInKVSet(a, b)
  300. for i := 0; i < len(r); i++ {
  301. c.Assert(r[i].k, qt.DeepEquals, expected[i].k)
  302. }
  303. c.Assert(len(invalids), qt.Equals, 7-4)
  304. }
  305. func TestGetKeysAtLevel(t *testing.T) {
  306. c := qt.New(t)
  307. tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
  308. c.Assert(err, qt.IsNil)
  309. defer tree.db.Close()
  310. for i := 0; i < 32; i++ {
  311. k := BigIntToBytes(big.NewInt(int64(i)))
  312. v := BigIntToBytes(big.NewInt(int64(i * 2)))
  313. if err := tree.Add(k, v); err != nil {
  314. t.Fatal(err)
  315. }
  316. }
  317. keys, err := tree.getKeysAtLevel(2)
  318. c.Assert(err, qt.IsNil)
  319. expected := []string{
  320. "a5d5f14fce7348e40751496cf25d107d91b0bd043435b9577d778a01f8aa6111",
  321. "e9e8dd9b28a7f81d1ff34cb5cefc0146dd848b31031a427b79bdadb62e7f6910",
  322. }
  323. for i := 0; i < len(keys); i++ {
  324. c.Assert(hex.EncodeToString(keys[i]), qt.Equals, expected[i])
  325. }
  326. keys, err = tree.getKeysAtLevel(3)
  327. c.Assert(err, qt.IsNil)
  328. expected = []string{
  329. "9f12c13e52bca96ad4882a26558e48ab67ddd63e062b839207e893d961390f01",
  330. "16d246dd6826ec7346c7328f11c4261facf82d4689f33263ff6e207956a77f21",
  331. "4a22cc901c6337daa17a431fa20170684b710e5f551509511492ec24e81a8f2f",
  332. "470d61abcbd154977bffc9a9ec5a8daff0caabcf2a25e8441f604c79daa0f82d",
  333. }
  334. for i := 0; i < len(keys); i++ {
  335. c.Assert(hex.EncodeToString(keys[i]), qt.Equals, expected[i])
  336. }
  337. keys, err = tree.getKeysAtLevel(4)
  338. c.Assert(err, qt.IsNil)
  339. expected = []string{
  340. "7a5d1c81f7b96318012de3417e53d4f13df5b1337718651cd29d0cb0a66edd20",
  341. "3408213e4e844bdf3355eb8781c74e31626812898c2dbe141ed6d2c92256fc1c",
  342. "dfd8a4d0b6954a3e9f3892e655b58d456eeedf9367f27dfdd9bc2dd6a5577312",
  343. "9e99fbec06fb2a6725997c12c4995f62725eb4cce4808523a5a5e80cca64b007",
  344. "0befa1e070231dbf4e8ff841c05878cdec823e0c09594c24910a248b3ff5a628",
  345. "b7131b0a15c772a57005a4dc5d0d6dd4b3414f5d9ee7408ce5e86c5ab3520e04",
  346. "6d1abe0364077846a56bab1deb1a04883eb796b74fe531a7676a9a370f83ab21",
  347. "4270116394bede69cf9cd72069eca018238080380bef5de75be8dcbbe968e105",
  348. }
  349. for i := 0; i < len(keys); i++ {
  350. c.Assert(hex.EncodeToString(keys[i]), qt.Equals, expected[i])
  351. }
  352. }
  353. func TestSplitInBuckets(t *testing.T) {
  354. c := qt.New(t)
  355. nLeafs := 16
  356. kvs := make([]kv, nLeafs)
  357. for i := 0; i < nLeafs; i++ {
  358. k := BigIntToBytes(big.NewInt(int64(i)))
  359. v := BigIntToBytes(big.NewInt(int64(i * 2)))
  360. keyPath := make([]byte, 32)
  361. copy(keyPath[:], k)
  362. kvs[i].pos = i
  363. kvs[i].keyPath = k
  364. kvs[i].k = k
  365. kvs[i].v = v
  366. }
  367. // check keyToBucket results for 4 buckets & 8 keys
  368. c.Assert(keyToBucket(kvs[0].k, 4), qt.Equals, 0)
  369. c.Assert(keyToBucket(kvs[1].k, 4), qt.Equals, 2)
  370. c.Assert(keyToBucket(kvs[2].k, 4), qt.Equals, 1)
  371. c.Assert(keyToBucket(kvs[3].k, 4), qt.Equals, 3)
  372. c.Assert(keyToBucket(kvs[4].k, 4), qt.Equals, 0)
  373. c.Assert(keyToBucket(kvs[5].k, 4), qt.Equals, 2)
  374. c.Assert(keyToBucket(kvs[6].k, 4), qt.Equals, 1)
  375. c.Assert(keyToBucket(kvs[7].k, 4), qt.Equals, 3)
  376. // check keyToBucket results for 8 buckets & 8 keys
  377. c.Assert(keyToBucket(kvs[0].k, 8), qt.Equals, 0)
  378. c.Assert(keyToBucket(kvs[1].k, 8), qt.Equals, 4)
  379. c.Assert(keyToBucket(kvs[2].k, 8), qt.Equals, 2)
  380. c.Assert(keyToBucket(kvs[3].k, 8), qt.Equals, 6)
  381. c.Assert(keyToBucket(kvs[4].k, 8), qt.Equals, 1)
  382. c.Assert(keyToBucket(kvs[5].k, 8), qt.Equals, 5)
  383. c.Assert(keyToBucket(kvs[6].k, 8), qt.Equals, 3)
  384. c.Assert(keyToBucket(kvs[7].k, 8), qt.Equals, 7)
  385. buckets := splitInBuckets(kvs, 4)
  386. expected := [][]string{
  387. {
  388. "00000000", // bucket 0
  389. "08000000",
  390. "04000000",
  391. "0c000000",
  392. },
  393. {
  394. "02000000", // bucket 1
  395. "0a000000",
  396. "06000000",
  397. "0e000000",
  398. },
  399. {
  400. "01000000", // bucket 2
  401. "09000000",
  402. "05000000",
  403. "0d000000",
  404. },
  405. {
  406. "03000000", // bucket 3
  407. "0b000000",
  408. "07000000",
  409. "0f000000",
  410. },
  411. }
  412. for i := 0; i < len(buckets); i++ {
  413. sortKvs(buckets[i])
  414. c.Assert(len(buckets[i]), qt.Equals, len(expected[i]))
  415. for j := 0; j < len(buckets[i]); j++ {
  416. c.Check(hex.EncodeToString(buckets[i][j].k[:4]), qt.Equals, expected[i][j])
  417. }
  418. }
  419. }
  420. func TestAddBatchCaseC(t *testing.T) {
  421. c := qt.New(t)
  422. nLeafs := 1024
  423. initialNLeafs := 101 // TMP TODO use const minLeafsThreshold+1 once ready
  424. tree1, tree2 := testInit(c, initialNLeafs)
  425. start := time.Now()
  426. for i := initialNLeafs; i < nLeafs; i++ {
  427. k := BigIntToBytes(big.NewInt(int64(i)))
  428. v := BigIntToBytes(big.NewInt(int64(i * 2)))
  429. if err := tree1.Add(k, v); err != nil {
  430. t.Fatal(err)
  431. }
  432. }
  433. time1 := time.Since(start)
  434. // prepare the key-values to be added
  435. var keys, values [][]byte
  436. for i := initialNLeafs; i < nLeafs; i++ {
  437. k := BigIntToBytes(big.NewInt(int64(i)))
  438. v := BigIntToBytes(big.NewInt(int64(i * 2)))
  439. keys = append(keys, k)
  440. values = append(values, v)
  441. }
  442. start = time.Now()
  443. indexes, err := tree2.AddBatch(keys, values)
  444. c.Assert(err, qt.IsNil)
  445. time2 := time.Since(start)
  446. debugTime("CASE C, AddBatch", time1, time2)
  447. c.Check(len(indexes), qt.Equals, 0)
  448. // check that both trees roots are equal
  449. c.Check(tree2.Root(), qt.DeepEquals, tree1.Root())
  450. }
  451. func TestAddBatchCaseD(t *testing.T) {
  452. c := qt.New(t)
  453. nLeafs := 4096
  454. initialNLeafs := 900
  455. tree1, tree2 := testInit(c, initialNLeafs)
  456. start := time.Now()
  457. for i := initialNLeafs; i < nLeafs; i++ {
  458. k := BigIntToBytes(big.NewInt(int64(i)))
  459. v := BigIntToBytes(big.NewInt(int64(i * 2)))
  460. if err := tree1.Add(k, v); err != nil {
  461. t.Fatal(err)
  462. }
  463. }
  464. time1 := time.Since(start)
  465. // prepare the key-values to be added
  466. var keys, values [][]byte
  467. for i := initialNLeafs; i < nLeafs; i++ {
  468. k := BigIntToBytes(big.NewInt(int64(i)))
  469. v := BigIntToBytes(big.NewInt(int64(i * 2)))
  470. keys = append(keys, k)
  471. values = append(values, v)
  472. }
  473. start = time.Now()
  474. indexes, err := tree2.AddBatch(keys, values)
  475. c.Assert(err, qt.IsNil)
  476. time2 := time.Since(start)
  477. debugTime("CASE D, AddBatch", time1, time2)
  478. c.Check(len(indexes), qt.Equals, 0)
  479. // check that both trees roots are equal
  480. c.Check(tree2.Root(), qt.DeepEquals, tree1.Root())
  481. }
  482. func TestAddBatchCaseE(t *testing.T) {
  483. c := qt.New(t)
  484. nLeafs := 4096
  485. initialNLeafs := 900
  486. tree1, _ := testInit(c, initialNLeafs)
  487. start := time.Now()
  488. for i := initialNLeafs; i < nLeafs; i++ {
  489. k := BigIntToBytes(big.NewInt(int64(i)))
  490. v := BigIntToBytes(big.NewInt(int64(i * 2)))
  491. if err := tree1.Add(k, v); err != nil {
  492. t.Fatal(err)
  493. }
  494. }
  495. time1 := time.Since(start)
  496. tree2, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
  497. c.Assert(err, qt.IsNil)
  498. defer tree2.db.Close()
  499. var keys, values [][]byte
  500. // add the initial leafs to fill a bit the tree before calling the
  501. // AddBatch method
  502. for i := 0; i < initialNLeafs; i++ {
  503. k := BigIntToBytes(big.NewInt(int64(i)))
  504. v := BigIntToBytes(big.NewInt(int64(i * 2)))
  505. // use only the keys of one bucket, store the not used ones for
  506. // later
  507. if i%4 != 0 {
  508. keys = append(keys, k)
  509. values = append(values, v)
  510. continue
  511. }
  512. if err := tree2.Add(k, v); err != nil {
  513. t.Fatal(err)
  514. }
  515. }
  516. for i := initialNLeafs; i < nLeafs; i++ {
  517. k := BigIntToBytes(big.NewInt(int64(i)))
  518. v := BigIntToBytes(big.NewInt(int64(i * 2)))
  519. keys = append(keys, k)
  520. values = append(values, v)
  521. }
  522. start = time.Now()
  523. indexes, err := tree2.AddBatch(keys, values)
  524. c.Assert(err, qt.IsNil)
  525. time2 := time.Since(start)
  526. debugTime("CASE E, AddBatch", time1, time2)
  527. c.Check(len(indexes), qt.Equals, 0)
  528. // check that both trees roots are equal
  529. c.Check(tree2.Root(), qt.DeepEquals, tree1.Root())
  530. }
  531. func TestFlp2(t *testing.T) {
  532. c := qt.New(t)
  533. c.Assert(flp2(31), qt.Equals, 16)
  534. c.Assert(flp2(32), qt.Equals, 32)
  535. c.Assert(flp2(33), qt.Equals, 32)
  536. c.Assert(flp2(63), qt.Equals, 32)
  537. c.Assert(flp2(64), qt.Equals, 64)
  538. c.Assert(flp2(9000), qt.Equals, 8192)
  539. }
  540. func TestAddBatchBench(t *testing.T) {
  541. nLeafs := 50_000
  542. fmt.Printf("TestAddBatchBench\n nCPU: %d, nLeafs: %d, hash: Blake2b, db: leveldb\n",
  543. runtime.NumCPU(), nLeafs)
  544. // prepare inputs
  545. var ks, vs [][]byte
  546. for i := 0; i < nLeafs; i++ {
  547. k := randomBytes(32)
  548. v := randomBytes(32)
  549. ks = append(ks, k)
  550. vs = append(vs, v)
  551. }
  552. benchAdd(t, ks, vs)
  553. benchAddBatch(t, ks, vs)
  554. }
  555. func benchAdd(t *testing.T, ks, vs [][]byte) {
  556. c := qt.New(t)
  557. dbDir := t.TempDir()
  558. storage, err := leveldb.NewLevelDbStorage(dbDir, false)
  559. c.Assert(err, qt.IsNil)
  560. tree, err := NewTree(storage, 140, HashFunctionBlake2b)
  561. c.Assert(err, qt.IsNil)
  562. start := time.Now()
  563. for i := 0; i < len(ks); i++ {
  564. err = tree.Add(ks[i], vs[i])
  565. c.Assert(err, qt.IsNil)
  566. }
  567. printRes(" Add loop", time.Since(start))
  568. }
  569. func benchAddBatch(t *testing.T, ks, vs [][]byte) {
  570. c := qt.New(t)
  571. dbDir := t.TempDir()
  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. invalids, err := tree.AddBatch(ks, vs)
  578. printRes(" AddBatch", time.Since(start))
  579. c.Assert(err, qt.IsNil)
  580. c.Assert(len(invalids), qt.Equals, 0)
  581. }
  582. func printRes(name string, duration time.Duration) {
  583. if debug {
  584. fmt.Printf(" %s: %s \n", name, duration)
  585. }
  586. }
  587. // func printLeafs(name string, t *Tree) {
  588. // w := bytes.NewBufferString("")
  589. //
  590. // err := t.Iterate(func(k, v []byte) {
  591. // if v[0] != PrefixValueLeaf {
  592. // return
  593. // }
  594. // leafK, _ := readLeafValue(v)
  595. // fmt.Fprintf(w, hex.EncodeToString(leafK[:4])+"\n")
  596. // })
  597. // if err != nil {
  598. // panic(err)
  599. // }
  600. // err = ioutil.WriteFile(name, w.Bytes(), 0644)
  601. // if err != nil {
  602. // panic(err)
  603. // }
  604. //
  605. // }
  606. // func TestComputeCosts(t *testing.T) {
  607. // fmt.Println(computeSimpleAddCost(10))
  608. // fmt.Println(computeBottomUpAddCost(10))
  609. //
  610. // fmt.Println(computeSimpleAddCost(1024))
  611. // fmt.Println(computeBottomUpAddCost(1024))
  612. // }
  613. // TODO test tree with nLeafs > minLeafsThreshold, but that at level L, there is
  614. // less keys than nBuckets (so CASE C could be applied if first few leafs are
  615. // added to balance the tree)
  616. // TODO test adding batch with repeated keys in the batch
  617. // TODO test adding batch with multiple invalid keys