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.

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