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.

464 lines
12 KiB

  1. package arbo
  2. import (
  3. "encoding/hex"
  4. "fmt"
  5. "math/big"
  6. "testing"
  7. "time"
  8. qt "github.com/frankban/quicktest"
  9. "github.com/iden3/go-merkletree/db/memory"
  10. )
  11. func testInit(c *qt.C, n int) (*Tree, *Tree) {
  12. tree1, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
  13. c.Assert(err, qt.IsNil)
  14. defer tree1.db.Close()
  15. tree2, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
  16. c.Assert(err, qt.IsNil)
  17. defer tree2.db.Close()
  18. // add the initial leafs to fill a bit the trees before calling the
  19. // AddBatch method
  20. for i := 0; i < n; i++ {
  21. k := BigIntToBytes(big.NewInt(int64(i)))
  22. v := BigIntToBytes(big.NewInt(int64(i * 2)))
  23. if err := tree1.Add(k, v); err != nil {
  24. c.Fatal(err)
  25. }
  26. if err := tree2.Add(k, v); err != nil {
  27. c.Fatal(err)
  28. }
  29. }
  30. return tree1, tree2
  31. }
  32. func ratio(t1, t2 time.Duration) float64 {
  33. a := float64(t1)
  34. b := float64(t2)
  35. return (a / b)
  36. }
  37. func TestAddBatchCaseA(t *testing.T) {
  38. c := qt.New(t)
  39. nLeafs := 1024
  40. tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
  41. c.Assert(err, qt.IsNil)
  42. defer tree.db.Close()
  43. start := time.Now()
  44. for i := 0; i < nLeafs; i++ {
  45. k := BigIntToBytes(big.NewInt(int64(i)))
  46. v := BigIntToBytes(big.NewInt(int64(i * 2)))
  47. if err := tree.Add(k, v); err != nil {
  48. t.Fatal(err)
  49. }
  50. }
  51. time1 := time.Since(start)
  52. tree2, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
  53. c.Assert(err, qt.IsNil)
  54. defer tree2.db.Close()
  55. var keys, values [][]byte
  56. for i := 0; i < nLeafs; i++ {
  57. k := BigIntToBytes(big.NewInt(int64(i)))
  58. v := BigIntToBytes(big.NewInt(int64(i * 2)))
  59. keys = append(keys, k)
  60. values = append(values, v)
  61. }
  62. start = time.Now()
  63. indexes, err := tree2.AddBatchOpt(keys, values)
  64. c.Assert(err, qt.IsNil)
  65. time2 := time.Since(start)
  66. fmt.Printf("CASE A, AddBatch was %f times faster than without AddBatch\n",
  67. ratio(time1, time2))
  68. c.Check(len(indexes), qt.Equals, 0)
  69. // check that both trees roots are equal
  70. c.Check(tree2.Root(), qt.DeepEquals, tree.Root())
  71. }
  72. func TestAddBatchCaseANotPowerOf2(t *testing.T) {
  73. c := qt.New(t)
  74. nLeafs := 1027
  75. tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
  76. c.Assert(err, qt.IsNil)
  77. defer tree.db.Close()
  78. for i := 0; i < nLeafs; i++ {
  79. k := BigIntToBytes(big.NewInt(int64(i)))
  80. v := BigIntToBytes(big.NewInt(int64(i * 2)))
  81. if err := tree.Add(k, v); err != nil {
  82. t.Fatal(err)
  83. }
  84. }
  85. tree2, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
  86. c.Assert(err, qt.IsNil)
  87. defer tree2.db.Close()
  88. var keys, values [][]byte
  89. for i := 0; i < nLeafs; i++ {
  90. k := BigIntToBytes(big.NewInt(int64(i)))
  91. v := BigIntToBytes(big.NewInt(int64(i * 2)))
  92. keys = append(keys, k)
  93. values = append(values, v)
  94. }
  95. indexes, err := tree2.AddBatchOpt(keys, values)
  96. c.Assert(err, qt.IsNil)
  97. c.Check(len(indexes), qt.Equals, 0)
  98. // check that both trees roots are equal
  99. c.Check(tree2.Root(), qt.DeepEquals, tree.Root())
  100. }
  101. func TestAddBatchCaseB(t *testing.T) {
  102. c := qt.New(t)
  103. nLeafs := 1024
  104. initialNLeafs := 99 // TMP TODO use const minLeafsThreshold-1 once ready
  105. tree1, tree2 := testInit(c, initialNLeafs)
  106. start := time.Now()
  107. for i := initialNLeafs; i < nLeafs; i++ {
  108. k := BigIntToBytes(big.NewInt(int64(i)))
  109. v := BigIntToBytes(big.NewInt(int64(i * 2)))
  110. if err := tree1.Add(k, v); err != nil {
  111. t.Fatal(err)
  112. }
  113. }
  114. time1 := time.Since(start)
  115. // prepare the key-values to be added
  116. var keys, values [][]byte
  117. for i := initialNLeafs; i < nLeafs; i++ {
  118. k := BigIntToBytes(big.NewInt(int64(i)))
  119. v := BigIntToBytes(big.NewInt(int64(i * 2)))
  120. keys = append(keys, k)
  121. values = append(values, v)
  122. }
  123. start = time.Now()
  124. indexes, err := tree2.AddBatchOpt(keys, values)
  125. c.Assert(err, qt.IsNil)
  126. time2 := time.Since(start)
  127. fmt.Printf("CASE B, AddBatch was %f times faster than without AddBatch\n",
  128. ratio(time1, time2))
  129. c.Check(len(indexes), qt.Equals, 0)
  130. // check that both trees roots are equal
  131. c.Check(tree2.Root(), qt.DeepEquals, tree1.Root())
  132. }
  133. func TestGetKeysAtLevel(t *testing.T) {
  134. c := qt.New(t)
  135. tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
  136. c.Assert(err, qt.IsNil)
  137. defer tree.db.Close()
  138. for i := 0; i < 32; i++ {
  139. k := BigIntToBytes(big.NewInt(int64(i)))
  140. v := BigIntToBytes(big.NewInt(int64(i * 2)))
  141. if err := tree.Add(k, v); err != nil {
  142. t.Fatal(err)
  143. }
  144. }
  145. keys, err := tree.getKeysAtLevel(2)
  146. c.Assert(err, qt.IsNil)
  147. expected := []string{
  148. "a5d5f14fce7348e40751496cf25d107d91b0bd043435b9577d778a01f8aa6111",
  149. "e9e8dd9b28a7f81d1ff34cb5cefc0146dd848b31031a427b79bdadb62e7f6910",
  150. }
  151. for i := 0; i < len(keys); i++ {
  152. c.Assert(hex.EncodeToString(keys[i]), qt.Equals, expected[i])
  153. }
  154. keys, err = tree.getKeysAtLevel(3)
  155. c.Assert(err, qt.IsNil)
  156. expected = []string{
  157. "9f12c13e52bca96ad4882a26558e48ab67ddd63e062b839207e893d961390f01",
  158. "16d246dd6826ec7346c7328f11c4261facf82d4689f33263ff6e207956a77f21",
  159. "4a22cc901c6337daa17a431fa20170684b710e5f551509511492ec24e81a8f2f",
  160. "470d61abcbd154977bffc9a9ec5a8daff0caabcf2a25e8441f604c79daa0f82d",
  161. }
  162. for i := 0; i < len(keys); i++ {
  163. c.Assert(hex.EncodeToString(keys[i]), qt.Equals, expected[i])
  164. }
  165. keys, err = tree.getKeysAtLevel(4)
  166. c.Assert(err, qt.IsNil)
  167. expected = []string{
  168. "7a5d1c81f7b96318012de3417e53d4f13df5b1337718651cd29d0cb0a66edd20",
  169. "3408213e4e844bdf3355eb8781c74e31626812898c2dbe141ed6d2c92256fc1c",
  170. "dfd8a4d0b6954a3e9f3892e655b58d456eeedf9367f27dfdd9bc2dd6a5577312",
  171. "9e99fbec06fb2a6725997c12c4995f62725eb4cce4808523a5a5e80cca64b007",
  172. "0befa1e070231dbf4e8ff841c05878cdec823e0c09594c24910a248b3ff5a628",
  173. "b7131b0a15c772a57005a4dc5d0d6dd4b3414f5d9ee7408ce5e86c5ab3520e04",
  174. "6d1abe0364077846a56bab1deb1a04883eb796b74fe531a7676a9a370f83ab21",
  175. "4270116394bede69cf9cd72069eca018238080380bef5de75be8dcbbe968e105",
  176. }
  177. for i := 0; i < len(keys); i++ {
  178. c.Assert(hex.EncodeToString(keys[i]), qt.Equals, expected[i])
  179. }
  180. }
  181. func TestSplitInBuckets(t *testing.T) {
  182. c := qt.New(t)
  183. nLeafs := 16
  184. kvs := make([]kv, nLeafs)
  185. for i := 0; i < nLeafs; i++ {
  186. k := BigIntToBytes(big.NewInt(int64(i)))
  187. v := BigIntToBytes(big.NewInt(int64(i * 2)))
  188. keyPath := make([]byte, 32)
  189. copy(keyPath[:], k)
  190. kvs[i].pos = i
  191. kvs[i].keyPath = k
  192. kvs[i].k = k
  193. kvs[i].v = v
  194. }
  195. // check keyToBucket results for 4 buckets & 8 keys
  196. c.Assert(keyToBucket(kvs[0].k, 4), qt.Equals, 0)
  197. c.Assert(keyToBucket(kvs[1].k, 4), qt.Equals, 2)
  198. c.Assert(keyToBucket(kvs[2].k, 4), qt.Equals, 1)
  199. c.Assert(keyToBucket(kvs[3].k, 4), qt.Equals, 3)
  200. c.Assert(keyToBucket(kvs[4].k, 4), qt.Equals, 0)
  201. c.Assert(keyToBucket(kvs[5].k, 4), qt.Equals, 2)
  202. c.Assert(keyToBucket(kvs[6].k, 4), qt.Equals, 1)
  203. c.Assert(keyToBucket(kvs[7].k, 4), qt.Equals, 3)
  204. // check keyToBucket results for 8 buckets & 8 keys
  205. c.Assert(keyToBucket(kvs[0].k, 8), qt.Equals, 0)
  206. c.Assert(keyToBucket(kvs[1].k, 8), qt.Equals, 4)
  207. c.Assert(keyToBucket(kvs[2].k, 8), qt.Equals, 2)
  208. c.Assert(keyToBucket(kvs[3].k, 8), qt.Equals, 6)
  209. c.Assert(keyToBucket(kvs[4].k, 8), qt.Equals, 1)
  210. c.Assert(keyToBucket(kvs[5].k, 8), qt.Equals, 5)
  211. c.Assert(keyToBucket(kvs[6].k, 8), qt.Equals, 3)
  212. c.Assert(keyToBucket(kvs[7].k, 8), qt.Equals, 7)
  213. buckets := splitInBuckets(kvs, 4)
  214. expected := [][]string{
  215. {
  216. "00000000", // bucket 0
  217. "08000000",
  218. "04000000",
  219. "0c000000",
  220. },
  221. {
  222. "02000000", // bucket 1
  223. "0a000000",
  224. "06000000",
  225. "0e000000",
  226. },
  227. {
  228. "01000000", // bucket 2
  229. "09000000",
  230. "05000000",
  231. "0d000000",
  232. },
  233. {
  234. "03000000", // bucket 3
  235. "0b000000",
  236. "07000000",
  237. "0f000000",
  238. },
  239. }
  240. for i := 0; i < len(buckets); i++ {
  241. sortKvs(buckets[i])
  242. c.Assert(len(buckets[i]), qt.Equals, len(expected[i]))
  243. for j := 0; j < len(buckets[i]); j++ {
  244. c.Check(hex.EncodeToString(buckets[i][j].k[:4]), qt.Equals, expected[i][j])
  245. }
  246. }
  247. }
  248. func TestAddBatchCaseC(t *testing.T) {
  249. c := qt.New(t)
  250. nLeafs := 1024
  251. initialNLeafs := 101 // TMP TODO use const minLeafsThreshold+1 once ready
  252. tree1, tree2 := testInit(c, initialNLeafs)
  253. start := time.Now()
  254. for i := initialNLeafs; i < nLeafs; i++ {
  255. k := BigIntToBytes(big.NewInt(int64(i)))
  256. v := BigIntToBytes(big.NewInt(int64(i * 2)))
  257. if err := tree1.Add(k, v); err != nil {
  258. t.Fatal(err)
  259. }
  260. }
  261. time1 := time.Since(start)
  262. // prepare the key-values to be added
  263. var keys, values [][]byte
  264. for i := initialNLeafs; i < nLeafs; i++ {
  265. k := BigIntToBytes(big.NewInt(int64(i)))
  266. v := BigIntToBytes(big.NewInt(int64(i * 2)))
  267. keys = append(keys, k)
  268. values = append(values, v)
  269. }
  270. start = time.Now()
  271. indexes, err := tree2.AddBatchOpt(keys, values)
  272. c.Assert(err, qt.IsNil)
  273. time2 := time.Since(start)
  274. fmt.Printf("CASE C, AddBatch was %f times faster than without AddBatch\n",
  275. ratio(time1, time2))
  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 TestAddBatchCaseD(t *testing.T) {
  281. c := qt.New(t)
  282. nLeafs := 4096
  283. initialNLeafs := 900
  284. tree1, tree2 := testInit(c, initialNLeafs)
  285. start := time.Now()
  286. for i := initialNLeafs; i < nLeafs; i++ {
  287. k := BigIntToBytes(big.NewInt(int64(i)))
  288. v := BigIntToBytes(big.NewInt(int64(i * 2)))
  289. if err := tree1.Add(k, v); err != nil {
  290. t.Fatal(err)
  291. }
  292. }
  293. time1 := time.Since(start)
  294. // prepare the key-values to be added
  295. var keys, values [][]byte
  296. for i := initialNLeafs; i < nLeafs; i++ {
  297. k := BigIntToBytes(big.NewInt(int64(i)))
  298. v := BigIntToBytes(big.NewInt(int64(i * 2)))
  299. keys = append(keys, k)
  300. values = append(values, v)
  301. }
  302. start = time.Now()
  303. indexes, err := tree2.AddBatchOpt(keys, values)
  304. c.Assert(err, qt.IsNil)
  305. time2 := time.Since(start)
  306. fmt.Printf("CASE D, AddBatch was %f times faster than without AddBatch\n",
  307. ratio(time1, time2))
  308. c.Check(len(indexes), qt.Equals, 0)
  309. // check that both trees roots are equal
  310. c.Check(tree2.Root(), qt.DeepEquals, tree1.Root())
  311. }
  312. func TestAddBatchCaseE(t *testing.T) {
  313. c := qt.New(t)
  314. nLeafs := 4096
  315. initialNLeafs := 900
  316. tree1, _ := testInit(c, initialNLeafs)
  317. start := time.Now()
  318. for i := initialNLeafs; i < nLeafs; i++ {
  319. k := BigIntToBytes(big.NewInt(int64(i)))
  320. v := BigIntToBytes(big.NewInt(int64(i * 2)))
  321. if err := tree1.Add(k, v); err != nil {
  322. t.Fatal(err)
  323. }
  324. }
  325. time1 := time.Since(start)
  326. tree2, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
  327. c.Assert(err, qt.IsNil)
  328. defer tree2.db.Close()
  329. var keys, values [][]byte
  330. // add the initial leafs to fill a bit the tree before calling the
  331. // AddBatch method
  332. for i := 0; i < initialNLeafs; i++ {
  333. k := BigIntToBytes(big.NewInt(int64(i)))
  334. v := BigIntToBytes(big.NewInt(int64(i * 2)))
  335. // use only the keys of one bucket, store the not used ones for
  336. // later
  337. if i%4 != 0 {
  338. keys = append(keys, k)
  339. values = append(values, v)
  340. continue
  341. }
  342. if err := tree2.Add(k, v); err != nil {
  343. t.Fatal(err)
  344. }
  345. }
  346. for i := initialNLeafs; i < nLeafs; i++ {
  347. k := BigIntToBytes(big.NewInt(int64(i)))
  348. v := BigIntToBytes(big.NewInt(int64(i * 2)))
  349. keys = append(keys, k)
  350. values = append(values, v)
  351. }
  352. start = time.Now()
  353. indexes, err := tree2.AddBatchOpt(keys, values)
  354. c.Assert(err, qt.IsNil)
  355. time2 := time.Since(start)
  356. fmt.Printf("CASE E, AddBatch was %f times faster than without AddBatch\n",
  357. ratio(time1, time2))
  358. c.Check(len(indexes), qt.Equals, 0)
  359. // check that both trees roots are equal
  360. c.Check(tree2.Root(), qt.DeepEquals, tree1.Root())
  361. }
  362. func TestHighestPowerOfTwo(t *testing.T) {
  363. c := qt.New(t)
  364. c.Assert(highestPowerOfTwo(31), qt.Equals, 16)
  365. c.Assert(highestPowerOfTwo(32), qt.Equals, 32)
  366. c.Assert(highestPowerOfTwo(33), qt.Equals, 32)
  367. c.Assert(highestPowerOfTwo(63), qt.Equals, 32)
  368. c.Assert(highestPowerOfTwo(64), qt.Equals, 64)
  369. }
  370. // func printLeafs(name string, t *Tree) {
  371. // w := bytes.NewBufferString("")
  372. //
  373. // err := t.Iterate(func(k, v []byte) {
  374. // if v[0] != PrefixValueLeaf {
  375. // return
  376. // }
  377. // leafK, _ := readLeafValue(v)
  378. // fmt.Fprintf(w, hex.EncodeToString(leafK[:4])+"\n")
  379. // })
  380. // if err != nil {
  381. // panic(err)
  382. // }
  383. // err = ioutil.WriteFile(name, w.Bytes(), 0644)
  384. // if err != nil {
  385. // panic(err)
  386. // }
  387. //
  388. // }
  389. // func TestComputeCosts(t *testing.T) {
  390. // fmt.Println(computeSimpleAddCost(10))
  391. // fmt.Println(computeBottomUpAddCost(10))
  392. //
  393. // fmt.Println(computeSimpleAddCost(1024))
  394. // fmt.Println(computeBottomUpAddCost(1024))
  395. // }
  396. // TODO test tree with nLeafs > minLeafsThreshold, but that at level L, there is
  397. // less keys than nBuckets (so CASE C could be applied if first few leafs are
  398. // added to balance the tree)
  399. // TODO test adding batch with repeated keys in the batch
  400. // TODO test adding batch with multiple invalid keys