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.

713 lines
20 KiB

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