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.

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