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.

1009 lines
27 KiB

  1. package arbo
  2. import (
  3. "crypto/rand"
  4. "encoding/hex"
  5. "fmt"
  6. "math/big"
  7. "runtime"
  8. "sort"
  9. "testing"
  10. "time"
  11. qt "github.com/frankban/quicktest"
  12. "go.vocdoni.io/dvote/db/badgerdb"
  13. )
  14. var debug = true
  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 %.02fx 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. database1, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
  34. c.Assert(err, qt.IsNil)
  35. tree1, err := NewTree(database1, 100, HashFunctionPoseidon)
  36. c.Assert(err, qt.IsNil)
  37. database2, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
  38. c.Assert(err, qt.IsNil)
  39. tree2, err := NewTree(database2, 100, HashFunctionPoseidon)
  40. c.Assert(err, qt.IsNil)
  41. bLen := HashFunctionPoseidon.Len()
  42. // add the initial leafs to fill a bit the trees before calling the
  43. // AddBatch method
  44. for i := 0; i < n; i++ {
  45. k := BigIntToBytes(bLen, big.NewInt(int64(i)))
  46. v := BigIntToBytes(bLen, big.NewInt(int64(i*2)))
  47. if err := tree1.Add(k, v); err != nil {
  48. c.Fatal(err)
  49. }
  50. if err := tree2.Add(k, v); err != nil {
  51. c.Fatal(err)
  52. }
  53. }
  54. return tree1, tree2
  55. }
  56. func TestAddBatchTreeEmpty(t *testing.T) {
  57. c := qt.New(t)
  58. nLeafs := 1024
  59. database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
  60. c.Assert(err, qt.IsNil)
  61. tree, err := NewTree(database, 100, HashFunctionPoseidon)
  62. c.Assert(err, qt.IsNil)
  63. defer tree.db.Close() //nolint:errcheck
  64. bLen := tree.HashFunction().Len()
  65. var keys, values [][]byte
  66. for i := 0; i < nLeafs; i++ {
  67. k := BigIntToBytes(bLen, big.NewInt(int64(i)))
  68. v := BigIntToBytes(bLen, big.NewInt(int64(i*2)))
  69. keys = append(keys, k)
  70. values = append(values, v)
  71. }
  72. start := time.Now()
  73. for i := 0; i < nLeafs; i++ {
  74. if err := tree.Add(keys[i], values[i]); err != nil {
  75. t.Fatal(err)
  76. }
  77. }
  78. time1 := time.Since(start)
  79. database2, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
  80. c.Assert(err, qt.IsNil)
  81. tree2, err := NewTree(database2, 100, HashFunctionPoseidon)
  82. c.Assert(err, qt.IsNil)
  83. defer tree2.db.Close() //nolint:errcheck
  84. tree2.dbgInit()
  85. start = time.Now()
  86. indexes, err := tree2.AddBatch(keys, values)
  87. c.Assert(err, qt.IsNil)
  88. time2 := time.Since(start)
  89. if debug {
  90. debugTime("Case tree empty, AddBatch", time1, time2)
  91. printTestContext(" ", nLeafs, "Poseidon", "memory")
  92. tree2.dbg.print(" ")
  93. }
  94. c.Check(len(indexes), qt.Equals, 0)
  95. // check that both trees roots are equal
  96. checkRoots(c, tree, tree2)
  97. }
  98. func TestAddBatchTreeEmptyNotPowerOf2(t *testing.T) {
  99. c := qt.New(t)
  100. nLeafs := 1027
  101. database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
  102. c.Assert(err, qt.IsNil)
  103. tree, err := NewTree(database, 100, HashFunctionPoseidon)
  104. c.Assert(err, qt.IsNil)
  105. defer tree.db.Close() //nolint:errcheck
  106. bLen := tree.HashFunction().Len()
  107. for i := 0; i < nLeafs; i++ {
  108. k := BigIntToBytes(bLen, big.NewInt(int64(i)))
  109. v := BigIntToBytes(bLen, big.NewInt(int64(i*2)))
  110. if err := tree.Add(k, v); err != nil {
  111. t.Fatal(err)
  112. }
  113. }
  114. database2, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
  115. c.Assert(err, qt.IsNil)
  116. tree2, err := NewTree(database2, 100, HashFunctionPoseidon)
  117. c.Assert(err, qt.IsNil)
  118. defer tree2.db.Close() //nolint:errcheck
  119. var keys, values [][]byte
  120. for i := 0; i < nLeafs; i++ {
  121. k := BigIntToBytes(bLen, big.NewInt(int64(i)))
  122. v := BigIntToBytes(bLen, big.NewInt(int64(i*2)))
  123. keys = append(keys, k)
  124. values = append(values, v)
  125. }
  126. indexes, err := tree2.AddBatch(keys, values)
  127. c.Assert(err, qt.IsNil)
  128. c.Check(len(indexes), qt.Equals, 0)
  129. // check that both trees roots are equal
  130. checkRoots(c, tree, tree2)
  131. }
  132. func randomBytes(n int) []byte {
  133. b := make([]byte, n)
  134. _, err := rand.Read(b)
  135. if err != nil {
  136. panic(err)
  137. }
  138. return b
  139. }
  140. func TestAddBatchTestVector1(t *testing.T) {
  141. c := qt.New(t)
  142. database1, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
  143. c.Assert(err, qt.IsNil)
  144. tree1, err := NewTree(database1, 100, HashFunctionBlake2b)
  145. c.Assert(err, qt.IsNil)
  146. defer tree1.db.Close() //nolint:errcheck
  147. database2, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
  148. c.Assert(err, qt.IsNil)
  149. tree2, err := NewTree(database2, 100, HashFunctionBlake2b)
  150. c.Assert(err, qt.IsNil)
  151. defer tree2.db.Close() //nolint:errcheck
  152. // leafs in 2nd level subtrees: [ 6, 0, 1, 1]
  153. testvectorKeys := []string{
  154. "1c7c2265e368314ca58ed2e1f33a326f1220e234a566d55c3605439dbe411642",
  155. "2c9f0a578afff5bfa4e0992a43066460faaab9e8e500db0b16647c701cdb16bf",
  156. "1c45cb31f2fa39ec7b9ebf0fad40e0b8296016b5ce8844ae06ff77226379d9a5",
  157. "d8af98bbbb585129798ae54d5eabbc9d0561d583faf1663b3a3724d15bda4ec7",
  158. }
  159. var keys, values [][]byte
  160. for i := 0; i < len(testvectorKeys); i++ {
  161. key, err := hex.DecodeString(testvectorKeys[i])
  162. c.Assert(err, qt.IsNil)
  163. keys = append(keys, key)
  164. values = append(values, []byte{0})
  165. }
  166. for i := 0; i < len(keys); i++ {
  167. if err := tree1.Add(keys[i], values[i]); err != nil {
  168. t.Fatal(err)
  169. }
  170. }
  171. indexes, err := tree2.AddBatch(keys, values)
  172. c.Assert(err, qt.IsNil)
  173. c.Check(len(indexes), qt.Equals, 0)
  174. // check that both trees roots are equal
  175. checkRoots(c, tree1, tree2)
  176. // 2nd test vectors
  177. database1, err = badgerdb.New(badgerdb.Options{Path: c.TempDir()})
  178. c.Assert(err, qt.IsNil)
  179. tree1, err = NewTree(database1, 100, HashFunctionBlake2b)
  180. c.Assert(err, qt.IsNil)
  181. defer tree1.db.Close() //nolint:errcheck
  182. database2, err = badgerdb.New(badgerdb.Options{Path: c.TempDir()})
  183. c.Assert(err, qt.IsNil)
  184. tree2, err = NewTree(database2, 100, HashFunctionBlake2b)
  185. c.Assert(err, qt.IsNil)
  186. defer tree2.db.Close() //nolint:errcheck
  187. testvectorKeys = []string{
  188. "1c7c2265e368314ca58ed2e1f33a326f1220e234a566d55c3605439dbe411642",
  189. "2c9f0a578afff5bfa4e0992a43066460faaab9e8e500db0b16647c701cdb16bf",
  190. "9cb87ec67e875c61390edcd1ab517f443591047709a4d4e45b0f9ed980857b8e",
  191. "9b4e9e92e974a589f426ceeb4cb291dc24893513fecf8e8460992dcf52621d4d",
  192. "1c45cb31f2fa39ec7b9ebf0fad40e0b8296016b5ce8844ae06ff77226379d9a5",
  193. "d8af98bbbb585129798ae54d5eabbc9d0561d583faf1663b3a3724d15bda4ec7",
  194. "3cd55dbfb8f975f20a0925dfbdabe79fa2d51dd0268afbb8ba6b01de9dfcdd3c",
  195. "5d0a9d6d9f197c091bf054fac9cb60e11ec723d6610ed8578e617b4d46cb43d5",
  196. }
  197. keys = [][]byte{}
  198. values = [][]byte{}
  199. for i := 0; i < len(testvectorKeys); i++ {
  200. key, err := hex.DecodeString(testvectorKeys[i])
  201. c.Assert(err, qt.IsNil)
  202. keys = append(keys, key)
  203. values = append(values, []byte{0})
  204. }
  205. for i := 0; i < len(keys); i++ {
  206. if err := tree1.Add(keys[i], values[i]); err != nil {
  207. t.Fatal(err)
  208. }
  209. }
  210. indexes, err = tree2.AddBatch(keys, values)
  211. c.Assert(err, qt.IsNil)
  212. c.Check(len(indexes), qt.Equals, 0)
  213. // check that both trees roots are equal
  214. checkRoots(c, tree1, tree2)
  215. }
  216. func TestAddBatchTestVector2(t *testing.T) {
  217. // test vector with unbalanced tree
  218. c := qt.New(t)
  219. database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
  220. c.Assert(err, qt.IsNil)
  221. tree1, err := NewTree(database, 100, HashFunctionPoseidon)
  222. c.Assert(err, qt.IsNil)
  223. defer tree1.db.Close() //nolint:errcheck
  224. database2, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
  225. c.Assert(err, qt.IsNil)
  226. tree2, err := NewTree(database2, 100, HashFunctionPoseidon)
  227. c.Assert(err, qt.IsNil)
  228. defer tree2.db.Close() //nolint:errcheck
  229. bLen := tree1.HashFunction().Len()
  230. var keys, values [][]byte
  231. // 1
  232. keys = append(keys, BigIntToBytes(bLen, big.NewInt(int64(1))))
  233. values = append(values, BigIntToBytes(bLen, big.NewInt(int64(1))))
  234. // 2
  235. keys = append(keys, BigIntToBytes(bLen, big.NewInt(int64(2))))
  236. values = append(values, BigIntToBytes(bLen, big.NewInt(int64(2))))
  237. // 3
  238. keys = append(keys, BigIntToBytes(bLen, big.NewInt(int64(3))))
  239. values = append(values, BigIntToBytes(bLen, big.NewInt(int64(3))))
  240. // 5
  241. keys = append(keys, BigIntToBytes(bLen, big.NewInt(int64(5))))
  242. values = append(values, BigIntToBytes(bLen, big.NewInt(int64(5))))
  243. for i := 0; i < len(keys); i++ {
  244. if err := tree1.Add(keys[i], values[i]); err != nil {
  245. t.Fatal(err)
  246. }
  247. }
  248. indexes, err := tree2.AddBatch(keys, values)
  249. c.Assert(err, qt.IsNil)
  250. c.Check(len(indexes), qt.Equals, 0)
  251. // check that both trees roots are equal
  252. checkRoots(c, tree1, tree2)
  253. }
  254. func TestAddBatchTestVector3(t *testing.T) {
  255. // test vector with unbalanced tree
  256. c := qt.New(t)
  257. database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
  258. c.Assert(err, qt.IsNil)
  259. tree1, err := NewTree(database, 100, HashFunctionPoseidon)
  260. c.Assert(err, qt.IsNil)
  261. defer tree1.db.Close() //nolint:errcheck
  262. database2, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
  263. c.Assert(err, qt.IsNil)
  264. tree2, err := NewTree(database2, 100, HashFunctionPoseidon)
  265. c.Assert(err, qt.IsNil)
  266. defer tree2.db.Close() //nolint:errcheck
  267. bLen := tree1.HashFunction().Len()
  268. var keys, values [][]byte
  269. // 0
  270. keys = append(keys, BigIntToBytes(bLen, big.NewInt(int64(0))))
  271. values = append(values, BigIntToBytes(bLen, big.NewInt(int64(0))))
  272. // 3
  273. keys = append(keys, BigIntToBytes(bLen, big.NewInt(int64(3))))
  274. values = append(values, BigIntToBytes(bLen, big.NewInt(int64(3))))
  275. // 7
  276. keys = append(keys, BigIntToBytes(bLen, big.NewInt(int64(7))))
  277. values = append(values, BigIntToBytes(bLen, big.NewInt(int64(7))))
  278. // 135
  279. keys = append(keys, BigIntToBytes(bLen, big.NewInt(int64(135))))
  280. values = append(values, BigIntToBytes(bLen, big.NewInt(int64(135))))
  281. for i := 0; i < len(keys); i++ {
  282. if err := tree1.Add(keys[i], values[i]); err != nil {
  283. t.Fatal(err)
  284. }
  285. }
  286. indexes, err := tree2.AddBatch(keys, values)
  287. c.Assert(err, qt.IsNil)
  288. c.Check(len(indexes), qt.Equals, 0)
  289. // check that both trees roots are equal
  290. checkRoots(c, tree1, tree2)
  291. //
  292. // tree1.PrintGraphvizFirstNLevels(nil, 100)
  293. // tree2.PrintGraphvizFirstNLevels(nil, 100)
  294. }
  295. func TestAddBatchTreeEmptyRandomKeys(t *testing.T) {
  296. c := qt.New(t)
  297. nLeafs := 8
  298. database1, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
  299. c.Assert(err, qt.IsNil)
  300. tree1, err := NewTree(database1, 100, HashFunctionBlake2b)
  301. c.Assert(err, qt.IsNil)
  302. defer tree1.db.Close() //nolint:errcheck
  303. database2, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
  304. c.Assert(err, qt.IsNil)
  305. tree2, err := NewTree(database2, 100, HashFunctionBlake2b)
  306. c.Assert(err, qt.IsNil)
  307. defer tree2.db.Close() //nolint:errcheck
  308. var keys, values [][]byte
  309. for i := 0; i < nLeafs; i++ {
  310. keys = append(keys, randomBytes(32))
  311. values = append(values, randomBytes(32))
  312. }
  313. for i := 0; i < len(keys); i++ {
  314. if err := tree1.Add(keys[i], values[i]); err != nil {
  315. t.Fatal(err)
  316. }
  317. }
  318. indexes, err := tree2.AddBatch(keys, values)
  319. c.Assert(err, qt.IsNil)
  320. c.Check(len(indexes), qt.Equals, 0)
  321. // check that both trees roots are equal
  322. checkRoots(c, tree1, tree2)
  323. }
  324. func TestAddBatchTreeNotEmptyFewLeafs(t *testing.T) {
  325. c := qt.New(t)
  326. nLeafs := 1024
  327. initialNLeafs := 99
  328. tree1, tree2 := testInit(c, initialNLeafs)
  329. tree2.dbgInit()
  330. bLen := tree1.HashFunction().Len()
  331. start := time.Now()
  332. for i := initialNLeafs; i < nLeafs; i++ {
  333. k := BigIntToBytes(bLen, big.NewInt(int64(i)))
  334. v := BigIntToBytes(bLen, big.NewInt(int64(i*2)))
  335. if err := tree1.Add(k, v); err != nil {
  336. t.Fatal(err)
  337. }
  338. }
  339. time1 := time.Since(start)
  340. // prepare the key-values to be added
  341. var keys, values [][]byte
  342. for i := initialNLeafs; i < nLeafs; i++ {
  343. k := BigIntToBytes(bLen, big.NewInt(int64(i)))
  344. v := BigIntToBytes(bLen, big.NewInt(int64(i*2)))
  345. keys = append(keys, k)
  346. values = append(values, v)
  347. }
  348. start = time.Now()
  349. indexes, err := tree2.AddBatch(keys, values)
  350. c.Assert(err, qt.IsNil)
  351. time2 := time.Since(start)
  352. if debug {
  353. debugTime("Case tree not empty w/ few leafs, AddBatch", time1, time2)
  354. printTestContext(" ", nLeafs, "Poseidon", "memory")
  355. tree2.dbg.print(" ")
  356. }
  357. c.Check(len(indexes), qt.Equals, 0)
  358. // check that both trees roots are equal
  359. checkRoots(c, tree1, tree2)
  360. }
  361. func TestAddBatchTreeNotEmptyEnoughLeafs(t *testing.T) {
  362. c := qt.New(t)
  363. nLeafs := 1024
  364. initialNLeafs := 500
  365. tree1, tree2 := testInit(c, initialNLeafs)
  366. tree2.dbgInit()
  367. bLen := tree1.HashFunction().Len()
  368. start := time.Now()
  369. for i := initialNLeafs; i < nLeafs; i++ {
  370. k := BigIntToBytes(bLen, big.NewInt(int64(i)))
  371. v := BigIntToBytes(bLen, big.NewInt(int64(i*2)))
  372. if err := tree1.Add(k, v); err != nil {
  373. t.Fatal(err)
  374. }
  375. }
  376. time1 := time.Since(start)
  377. // prepare the key-values to be added
  378. var keys, values [][]byte
  379. for i := initialNLeafs; i < nLeafs; i++ {
  380. k := BigIntToBytes(bLen, big.NewInt(int64(i)))
  381. v := BigIntToBytes(bLen, big.NewInt(int64(i*2)))
  382. keys = append(keys, k)
  383. values = append(values, v)
  384. }
  385. start = time.Now()
  386. indexes, err := tree2.AddBatch(keys, values)
  387. c.Assert(err, qt.IsNil)
  388. time2 := time.Since(start)
  389. if debug {
  390. debugTime("Case tree not empty w/ enough leafs, AddBatch", time1, time2)
  391. printTestContext(" ", nLeafs, "Poseidon", "memory")
  392. tree2.dbg.print(" ")
  393. }
  394. c.Check(len(indexes), qt.Equals, 0)
  395. // check that both trees roots are equal
  396. checkRoots(c, tree1, tree2)
  397. }
  398. func TestAddBatchTreeEmptyRepeatedLeafs(t *testing.T) {
  399. c := qt.New(t)
  400. nLeafs := 1024
  401. nRepeatedKeys := 99
  402. tree1, tree2 := testInit(c, 0)
  403. bLen := tree1.HashFunction().Len()
  404. // prepare the key-values to be added
  405. var keys, values [][]byte
  406. for i := 0; i < nLeafs; i++ {
  407. k := BigIntToBytes(bLen, big.NewInt(int64(i)))
  408. v := BigIntToBytes(bLen, big.NewInt(int64(i*2)))
  409. keys = append(keys, k)
  410. values = append(values, v)
  411. }
  412. // add repeated key-values
  413. for i := 0; i < nRepeatedKeys; i++ {
  414. k := BigIntToBytes(bLen, big.NewInt(int64(i)))
  415. v := BigIntToBytes(bLen, big.NewInt(int64(i*2)))
  416. keys = append(keys, k)
  417. values = append(values, v)
  418. }
  419. // add the non-repeated key-values in tree1 with .Add loop
  420. for i := 0; i < nLeafs; i++ {
  421. if err := tree1.Add(keys[i], values[i]); err != nil {
  422. t.Fatal(err)
  423. }
  424. }
  425. indexes, err := tree2.AddBatch(keys, values)
  426. c.Assert(err, qt.IsNil)
  427. c.Check(len(indexes), qt.Equals, nRepeatedKeys)
  428. // check that both trees roots are equal
  429. checkRoots(c, tree1, tree2)
  430. }
  431. func TestAddBatchTreeNotEmptyFewLeafsRepeatedLeafs(t *testing.T) {
  432. c := qt.New(t)
  433. nLeafs := 1024
  434. initialNLeafs := 99
  435. tree1, tree2 := testInit(c, initialNLeafs)
  436. bLen := tree1.HashFunction().Len()
  437. // prepare the key-values to be added
  438. var keys, values [][]byte
  439. for i := 0; i < nLeafs; i++ {
  440. k := BigIntToBytes(bLen, big.NewInt(int64(i)))
  441. v := BigIntToBytes(bLen, big.NewInt(int64(i*2)))
  442. keys = append(keys, k)
  443. values = append(values, v)
  444. }
  445. // add the keys that will be existing when AddBatch is called
  446. for i := initialNLeafs; i < nLeafs; i++ {
  447. if err := tree1.Add(keys[i], values[i]); err != nil {
  448. t.Fatal(err)
  449. }
  450. }
  451. indexes, err := tree2.AddBatch(keys, values)
  452. c.Assert(err, qt.IsNil)
  453. c.Check(len(indexes), qt.Equals, initialNLeafs)
  454. // check that both trees roots are equal
  455. checkRoots(c, tree1, tree2)
  456. }
  457. func TestSplitInBuckets(t *testing.T) {
  458. c := qt.New(t)
  459. bLen := HashFunctionPoseidon.Len()
  460. nLeafs := 16
  461. kvs := make([]kv, nLeafs)
  462. for i := 0; i < nLeafs; i++ {
  463. k := BigIntToBytes(bLen, big.NewInt(int64(i)))
  464. v := BigIntToBytes(bLen, big.NewInt(int64(i*2)))
  465. keyPath := make([]byte, 32)
  466. copy(keyPath[:], k)
  467. kvs[i].pos = i
  468. kvs[i].keyPath = k
  469. kvs[i].k = k
  470. kvs[i].v = v
  471. }
  472. // check keyToBucket results for 4 buckets & 8 keys
  473. c.Assert(keyToBucket(kvs[0].k, 4), qt.Equals, 0)
  474. c.Assert(keyToBucket(kvs[1].k, 4), qt.Equals, 2)
  475. c.Assert(keyToBucket(kvs[2].k, 4), qt.Equals, 1)
  476. c.Assert(keyToBucket(kvs[3].k, 4), qt.Equals, 3)
  477. c.Assert(keyToBucket(kvs[4].k, 4), qt.Equals, 0)
  478. c.Assert(keyToBucket(kvs[5].k, 4), qt.Equals, 2)
  479. c.Assert(keyToBucket(kvs[6].k, 4), qt.Equals, 1)
  480. c.Assert(keyToBucket(kvs[7].k, 4), qt.Equals, 3)
  481. // check keyToBucket results for 8 buckets & 8 keys
  482. c.Assert(keyToBucket(kvs[0].k, 8), qt.Equals, 0)
  483. c.Assert(keyToBucket(kvs[1].k, 8), qt.Equals, 4)
  484. c.Assert(keyToBucket(kvs[2].k, 8), qt.Equals, 2)
  485. c.Assert(keyToBucket(kvs[3].k, 8), qt.Equals, 6)
  486. c.Assert(keyToBucket(kvs[4].k, 8), qt.Equals, 1)
  487. c.Assert(keyToBucket(kvs[5].k, 8), qt.Equals, 5)
  488. c.Assert(keyToBucket(kvs[6].k, 8), qt.Equals, 3)
  489. c.Assert(keyToBucket(kvs[7].k, 8), qt.Equals, 7)
  490. buckets := splitInBuckets(kvs, 4)
  491. expected := [][]string{
  492. {
  493. "00000000", // bucket 0
  494. "08000000",
  495. "04000000",
  496. "0c000000",
  497. },
  498. {
  499. "02000000", // bucket 1
  500. "0a000000",
  501. "06000000",
  502. "0e000000",
  503. },
  504. {
  505. "01000000", // bucket 2
  506. "09000000",
  507. "05000000",
  508. "0d000000",
  509. },
  510. {
  511. "03000000", // bucket 3
  512. "0b000000",
  513. "07000000",
  514. "0f000000",
  515. },
  516. }
  517. for i := 0; i < len(buckets); i++ {
  518. sortKvs(buckets[i])
  519. c.Assert(len(buckets[i]), qt.Equals, len(expected[i]))
  520. for j := 0; j < len(buckets[i]); j++ {
  521. c.Check(hex.EncodeToString(buckets[i][j].k[:4]),
  522. qt.Equals, expected[i][j])
  523. }
  524. }
  525. }
  526. // compareBytes compares byte slices where the bytes are compared from left to
  527. // right and each byte is compared by bit from right to left
  528. func compareBytes(a, b []byte) bool {
  529. // WIP
  530. for i := 0; i < len(a); i++ {
  531. for j := 0; j < 8; j++ {
  532. aBit := a[i] & (1 << j)
  533. bBit := b[i] & (1 << j)
  534. if aBit > bBit {
  535. return false
  536. } else if aBit < bBit {
  537. return true
  538. }
  539. }
  540. }
  541. return false
  542. }
  543. // sortKvs sorts the kv by path
  544. func sortKvs(kvs []kv) {
  545. sort.Slice(kvs, func(i, j int) bool {
  546. return compareBytes(kvs[i].keyPath, kvs[j].keyPath)
  547. })
  548. }
  549. func TestAddBatchTreeNotEmpty(t *testing.T) {
  550. c := qt.New(t)
  551. nLeafs := 4096
  552. initialNLeafs := 900
  553. tree1, tree2 := testInit(c, initialNLeafs)
  554. tree2.dbgInit()
  555. bLen := tree1.HashFunction().Len()
  556. start := time.Now()
  557. for i := initialNLeafs; i < nLeafs; i++ {
  558. k := BigIntToBytes(bLen, big.NewInt(int64(i)))
  559. v := BigIntToBytes(bLen, big.NewInt(int64(i*2)))
  560. if err := tree1.Add(k, v); err != nil {
  561. t.Fatal(err)
  562. }
  563. }
  564. time1 := time.Since(start)
  565. // prepare the key-values to be added
  566. var keys, values [][]byte
  567. for i := initialNLeafs; i < nLeafs; i++ {
  568. k := BigIntToBytes(bLen, big.NewInt(int64(i)))
  569. v := BigIntToBytes(bLen, big.NewInt(int64(i*2)))
  570. keys = append(keys, k)
  571. values = append(values, v)
  572. }
  573. start = time.Now()
  574. indexes, err := tree2.AddBatch(keys, values)
  575. c.Assert(err, qt.IsNil)
  576. time2 := time.Since(start)
  577. if debug {
  578. debugTime("Case tree not empty, AddBatch", time1, time2)
  579. printTestContext(" ", nLeafs, "Poseidon", "memory")
  580. tree2.dbg.print(" ")
  581. }
  582. c.Check(len(indexes), qt.Equals, 0)
  583. // check that both trees roots are equal
  584. checkRoots(c, tree1, tree2)
  585. }
  586. func TestAddBatchNotEmptyUnbalanced(t *testing.T) {
  587. c := qt.New(t)
  588. nLeafs := 4096
  589. initialNLeafs := 900
  590. tree1, _ := testInit(c, initialNLeafs)
  591. bLen := tree1.HashFunction().Len()
  592. start := time.Now()
  593. for i := initialNLeafs; i < nLeafs; i++ {
  594. k := BigIntToBytes(bLen, big.NewInt(int64(i)))
  595. v := BigIntToBytes(bLen, big.NewInt(int64(i*2)))
  596. if err := tree1.Add(k, v); err != nil {
  597. t.Fatal(err)
  598. }
  599. }
  600. time1 := time.Since(start)
  601. database2, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
  602. c.Assert(err, qt.IsNil)
  603. tree2, err := NewTree(database2, 100, HashFunctionPoseidon)
  604. c.Assert(err, qt.IsNil)
  605. defer tree2.db.Close() //nolint:errcheck
  606. tree2.dbgInit()
  607. var keys, values [][]byte
  608. // add the initial leafs to fill a bit the tree before calling the
  609. // AddBatch method
  610. for i := 0; i < initialNLeafs; i++ {
  611. k := BigIntToBytes(bLen, big.NewInt(int64(i)))
  612. v := BigIntToBytes(bLen, big.NewInt(int64(i*2)))
  613. // use only the keys of one bucket, store the not used ones for
  614. // later
  615. if i%4 != 0 {
  616. keys = append(keys, k)
  617. values = append(values, v)
  618. continue
  619. }
  620. if err := tree2.Add(k, v); err != nil {
  621. t.Fatal(err)
  622. }
  623. }
  624. for i := initialNLeafs; i < nLeafs; i++ {
  625. k := BigIntToBytes(bLen, big.NewInt(int64(i)))
  626. v := BigIntToBytes(bLen, big.NewInt(int64(i*2)))
  627. keys = append(keys, k)
  628. values = append(values, v)
  629. }
  630. start = time.Now()
  631. indexes, err := tree2.AddBatch(keys, values)
  632. c.Assert(err, qt.IsNil)
  633. time2 := time.Since(start)
  634. if debug {
  635. debugTime("Case tree not empty & unbalanced, AddBatch", time1, time2)
  636. printTestContext(" ", nLeafs, "Poseidon", "memory")
  637. tree2.dbg.print(" ")
  638. }
  639. c.Check(len(indexes), qt.Equals, 0)
  640. // check that both trees roots are equal
  641. checkRoots(c, tree1, tree2)
  642. }
  643. func TestFlp2(t *testing.T) {
  644. c := qt.New(t)
  645. c.Assert(flp2(31), qt.Equals, 16)
  646. c.Assert(flp2(32), qt.Equals, 32)
  647. c.Assert(flp2(33), qt.Equals, 32)
  648. c.Assert(flp2(63), qt.Equals, 32)
  649. c.Assert(flp2(64), qt.Equals, 64)
  650. c.Assert(flp2(9000), qt.Equals, 8192)
  651. }
  652. func TestAddBatchBench(t *testing.T) {
  653. nLeafs := 50_000
  654. printTestContext("TestAddBatchBench: ", nLeafs, "Blake2b", "badgerdb")
  655. // prepare inputs
  656. var ks, vs [][]byte
  657. for i := 0; i < nLeafs; i++ {
  658. k := randomBytes(32)
  659. v := randomBytes(32)
  660. ks = append(ks, k)
  661. vs = append(vs, v)
  662. }
  663. benchAdd(t, ks, vs)
  664. benchAddBatch(t, ks, vs)
  665. }
  666. func benchAdd(t *testing.T, ks, vs [][]byte) {
  667. c := qt.New(t)
  668. database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
  669. c.Assert(err, qt.IsNil)
  670. tree, err := NewTree(database, 140, HashFunctionBlake2b)
  671. c.Assert(err, qt.IsNil)
  672. defer tree.db.Close() //nolint:errcheck
  673. start := time.Now()
  674. for i := 0; i < len(ks); i++ {
  675. err = tree.Add(ks[i], vs[i])
  676. c.Assert(err, qt.IsNil)
  677. }
  678. if debug {
  679. printRes(" Add loop", time.Since(start))
  680. tree.dbg.print(" ")
  681. }
  682. }
  683. func benchAddBatch(t *testing.T, ks, vs [][]byte) {
  684. c := qt.New(t)
  685. database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
  686. c.Assert(err, qt.IsNil)
  687. tree, err := NewTree(database, 140, HashFunctionBlake2b)
  688. c.Assert(err, qt.IsNil)
  689. defer tree.db.Close() //nolint:errcheck
  690. tree.dbgInit()
  691. start := time.Now()
  692. invalids, err := tree.AddBatch(ks, vs)
  693. if debug {
  694. printRes(" AddBatch", time.Since(start))
  695. tree.dbg.print(" ")
  696. }
  697. c.Assert(err, qt.IsNil)
  698. c.Assert(len(invalids), qt.Equals, 0)
  699. }
  700. func TestDbgStats(t *testing.T) {
  701. c := qt.New(t)
  702. nLeafs := 10_000
  703. // prepare inputs
  704. var ks, vs [][]byte
  705. for i := 0; i < nLeafs; i++ {
  706. k := randomBytes(32)
  707. v := randomBytes(32)
  708. ks = append(ks, k)
  709. vs = append(vs, v)
  710. }
  711. // 1
  712. database1, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
  713. c.Assert(err, qt.IsNil)
  714. tree1, err := NewTree(database1, 100, HashFunctionBlake2b)
  715. c.Assert(err, qt.IsNil)
  716. defer tree1.db.Close() //nolint:errcheck
  717. tree1.dbgInit()
  718. for i := 0; i < len(ks); i++ {
  719. err = tree1.Add(ks[i], vs[i])
  720. c.Assert(err, qt.IsNil)
  721. }
  722. // 2
  723. database2, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
  724. c.Assert(err, qt.IsNil)
  725. tree2, err := NewTree(database2, 100, HashFunctionBlake2b)
  726. c.Assert(err, qt.IsNil)
  727. defer tree2.db.Close() //nolint:errcheck
  728. tree2.dbgInit()
  729. invalids, err := tree2.AddBatch(ks, vs)
  730. c.Assert(err, qt.IsNil)
  731. c.Assert(len(invalids), qt.Equals, 0)
  732. // 3
  733. database3, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
  734. c.Assert(err, qt.IsNil)
  735. tree3, err := NewTree(database3, 100, HashFunctionBlake2b)
  736. c.Assert(err, qt.IsNil)
  737. defer tree3.db.Close() //nolint:errcheck
  738. tree3.dbgInit()
  739. // add few key-values
  740. // invalids, err = tree3.AddBatch(ks[:], vs[:])
  741. invalids, err = tree3.AddBatch(ks[:1000], vs[:1000])
  742. c.Assert(err, qt.IsNil)
  743. c.Assert(len(invalids), qt.Equals, 0)
  744. // add the rest of key-values
  745. invalids, err = tree3.AddBatch(ks[1000:], vs[1000:])
  746. c.Assert(err, qt.IsNil)
  747. c.Assert(len(invalids), qt.Equals, 0)
  748. checkRoots(c, tree1, tree2)
  749. checkRoots(c, tree1, tree3)
  750. if debug {
  751. fmt.Println("TestDbgStats")
  752. tree1.dbg.print(" add in loop in emptyTree ")
  753. tree2.dbg.print(" addbatch caseEmptyTree ")
  754. tree3.dbg.print(" addbatch caseNotEmptyTree ")
  755. }
  756. }
  757. func TestLoadVT(t *testing.T) {
  758. c := qt.New(t)
  759. nLeafs := 1024
  760. database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
  761. c.Assert(err, qt.IsNil)
  762. tree, err := NewTree(database, 100, HashFunctionPoseidon)
  763. c.Assert(err, qt.IsNil)
  764. defer tree.db.Close() //nolint:errcheck
  765. var keys, values [][]byte
  766. for i := 0; i < nLeafs; i++ {
  767. k := randomBytes(31)
  768. v := randomBytes(31)
  769. keys = append(keys, k)
  770. values = append(values, v)
  771. }
  772. indexes, err := tree.AddBatch(keys, values)
  773. c.Assert(err, qt.IsNil)
  774. c.Check(len(indexes), qt.Equals, 0)
  775. rTx := tree.db.ReadTx()
  776. defer rTx.Discard()
  777. vt, err := tree.loadVT(rTx)
  778. c.Assert(err, qt.IsNil)
  779. _, err = vt.computeHashes()
  780. c.Assert(err, qt.IsNil)
  781. // check that tree & vt roots are equal
  782. root, err := tree.Root()
  783. c.Assert(err, qt.IsNil)
  784. c.Check(root, qt.DeepEquals, vt.root.h)
  785. }
  786. // TestAddKeysWithEmptyValues calls AddBatch giving an array of empty values
  787. func TestAddKeysWithEmptyValues(t *testing.T) {
  788. c := qt.New(t)
  789. nLeafs := 1024
  790. database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
  791. c.Assert(err, qt.IsNil)
  792. tree, err := NewTree(database, 100, HashFunctionPoseidon)
  793. c.Assert(err, qt.IsNil)
  794. defer tree.db.Close() //nolint:errcheck
  795. bLen := tree.HashFunction().Len()
  796. var keys, values [][]byte
  797. for i := 0; i < nLeafs; i++ {
  798. k := BigIntToBytes(bLen, big.NewInt(int64(i)))
  799. v := []byte{}
  800. keys = append(keys, k)
  801. values = append(values, v)
  802. }
  803. for i := 0; i < nLeafs; i++ {
  804. if err := tree.Add(keys[i], values[i]); err != nil {
  805. t.Fatal(err)
  806. }
  807. }
  808. database2, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
  809. c.Assert(err, qt.IsNil)
  810. tree2, err := NewTree(database2, 100, HashFunctionPoseidon)
  811. c.Assert(err, qt.IsNil)
  812. defer tree2.db.Close() //nolint:errcheck
  813. tree2.dbgInit()
  814. indexes, err := tree2.AddBatch(keys, values)
  815. c.Assert(err, qt.IsNil)
  816. c.Check(len(indexes), qt.Equals, 0)
  817. // check that both trees roots are equal
  818. checkRoots(c, tree, tree2)
  819. // use tree3 to add nil value array
  820. database3, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
  821. c.Assert(err, qt.IsNil)
  822. tree3, err := NewTree(database3, 100, HashFunctionPoseidon)
  823. c.Assert(err, qt.IsNil)
  824. defer tree3.db.Close() //nolint:errcheck
  825. indexes, err = tree3.AddBatch(keys, nil)
  826. c.Assert(err, qt.IsNil)
  827. c.Check(len(indexes), qt.Equals, 0)
  828. checkRoots(c, tree, tree3)
  829. kAux, proofV, siblings, existence, err := tree2.GenProof(keys[9])
  830. c.Assert(err, qt.IsNil)
  831. c.Assert(proofV, qt.DeepEquals, values[9])
  832. c.Assert(keys[9], qt.DeepEquals, kAux)
  833. c.Assert(existence, qt.IsTrue)
  834. // check with empty array
  835. root, err := tree.Root()
  836. c.Assert(err, qt.IsNil)
  837. verif, err := CheckProof(tree.hashFunction, keys[9], []byte{}, root, siblings)
  838. c.Assert(err, qt.IsNil)
  839. c.Check(verif, qt.IsTrue)
  840. // check with array with only 1 zero
  841. verif, err = CheckProof(tree.hashFunction, keys[9], []byte{0}, root, siblings)
  842. c.Assert(err, qt.IsNil)
  843. c.Check(verif, qt.IsTrue)
  844. // check with array with 32 zeroes
  845. e32 := []byte{0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
  846. 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0}
  847. c.Assert(len(e32), qt.Equals, 32)
  848. verif, err = CheckProof(tree.hashFunction, keys[9], e32, root, siblings)
  849. c.Assert(err, qt.IsNil)
  850. c.Check(verif, qt.IsTrue)
  851. // check with array with value!=0 returns false at verification
  852. verif, err = CheckProof(tree.hashFunction, keys[9], []byte{0, 1}, root, siblings)
  853. c.Assert(err, qt.IsNil)
  854. c.Check(verif, qt.IsFalse)
  855. }
  856. // TODO test adding batch with multiple invalid keys
  857. // TODO for tests of AddBatch, if the root does not match the Add root, bulk
  858. // all the leafs of both trees into a log file to later be able to debug and
  859. // recreate the case