Merge pull request #26 from vocdoni/feature/addbatch-bigtrees-indisk

Feature/addbatch bigtrees indisk

addbatch_test.go

@@ -1,6 +1,7 @@
 package arbo
 
 import (
+	"bytes"
 	"crypto/rand"
 	"encoding/hex"
 	"fmt"
@@ -11,7 +12,9 @@ import (
 	"time"
 
 	qt "github.com/frankban/quicktest"
+	"go.vocdoni.io/dvote/db"
 	"go.vocdoni.io/dvote/db/badgerdb"
+	"go.vocdoni.io/dvote/db/pebbledb"
 )
 
 var debug = true
@@ -37,14 +40,16 @@ func debugTime(descr string, time1, time2 time.Duration) {
 }
 
 func testInit(c *qt.C, n int) (*Tree, *Tree) {
-	database1, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
+	database1, err := badgerdb.New(db.Options{Path: c.TempDir()})
 	c.Assert(err, qt.IsNil)
-	tree1, err := NewTree(database1, 256, HashFunctionPoseidon)
+	tree1, err := NewTree(Config{Database: database1, MaxLevels: 256,
+		HashFunction: HashFunctionPoseidon})
 	c.Assert(err, qt.IsNil)
 
-	database2, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
+	database2, err := badgerdb.New(db.Options{Path: c.TempDir()})
 	c.Assert(err, qt.IsNil)
-	tree2, err := NewTree(database2, 256, HashFunctionPoseidon)
+	tree2, err := NewTree(Config{Database: database2, MaxLevels: 256,
+		HashFunction: HashFunctionPoseidon})
 	c.Assert(err, qt.IsNil)
 
 	bLen := HashFunctionPoseidon.Len()
@@ -68,9 +73,10 @@ func TestAddBatchTreeEmpty(t *testing.T) {
 
 	nLeafs := 1024
 
-	database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
+	database, err := badgerdb.New(db.Options{Path: c.TempDir()})
 	c.Assert(err, qt.IsNil)
-	tree, err := NewTree(database, 256, HashFunctionPoseidon)
+	tree, err := NewTree(Config{database, 256, DefaultThresholdNLeafs,
+		HashFunctionPoseidon})
 	c.Assert(err, qt.IsNil)
 	defer tree.db.Close() //nolint:errcheck
 
@@ -91,15 +97,16 @@ func TestAddBatchTreeEmpty(t *testing.T) {
 	}
 	time1 := time.Since(start)
 
-	database2, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
+	database2, err := badgerdb.New(db.Options{Path: c.TempDir()})
 	c.Assert(err, qt.IsNil)
-	tree2, err := NewTree(database2, 256, HashFunctionPoseidon)
+	tree2, err := NewTree(Config{database2, 256, DefaultThresholdNLeafs,
+		HashFunctionPoseidon})
 	c.Assert(err, qt.IsNil)
 	defer tree2.db.Close() //nolint:errcheck
 	tree2.dbgInit()
 
 	start = time.Now()
-	indexes, err := tree2.AddBatch(keys, values)
+	invalids, err := tree2.AddBatch(keys, values)
 	c.Assert(err, qt.IsNil)
 	time2 := time.Since(start)
 	if debug {
@@ -107,7 +114,7 @@ func TestAddBatchTreeEmpty(t *testing.T) {
 		printTestContext("	", nLeafs, "Poseidon", "memory")
 		tree2.dbg.print("	")
 	}
-	c.Check(len(indexes), qt.Equals, 0)
+	c.Check(len(invalids), qt.Equals, 0)
 
 	// check that both trees roots are equal
 	checkRoots(c, tree, tree2)
@@ -118,9 +125,10 @@ func TestAddBatchTreeEmptyNotPowerOf2(t *testing.T) {
 
 	nLeafs := 1027
 
-	database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
+	database, err := badgerdb.New(db.Options{Path: c.TempDir()})
 	c.Assert(err, qt.IsNil)
-	tree, err := NewTree(database, 256, HashFunctionPoseidon)
+	tree, err := NewTree(Config{database, 256, DefaultThresholdNLeafs,
+		HashFunctionPoseidon})
 	c.Assert(err, qt.IsNil)
 	defer tree.db.Close() //nolint:errcheck
 
@@ -133,9 +141,10 @@ func TestAddBatchTreeEmptyNotPowerOf2(t *testing.T) {
 		}
 	}
 
-	database2, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
+	database2, err := badgerdb.New(db.Options{Path: c.TempDir()})
 	c.Assert(err, qt.IsNil)
-	tree2, err := NewTree(database2, 256, HashFunctionPoseidon)
+	tree2, err := NewTree(Config{database2, 256, DefaultThresholdNLeafs,
+		HashFunctionPoseidon})
 	c.Assert(err, qt.IsNil)
 	defer tree2.db.Close() //nolint:errcheck
 
@@ -146,9 +155,9 @@ func TestAddBatchTreeEmptyNotPowerOf2(t *testing.T) {
 		keys = append(keys, k)
 		values = append(values, v)
 	}
-	indexes, err := tree2.AddBatch(keys, values)
+	invalids, err := tree2.AddBatch(keys, values)
 	c.Assert(err, qt.IsNil)
-	c.Check(len(indexes), qt.Equals, 0)
+	c.Check(len(invalids), qt.Equals, 0)
 
 	// check that both trees roots are equal
 	checkRoots(c, tree, tree2)
@@ -165,15 +174,17 @@ func randomBytes(n int) []byte {
 
 func TestAddBatchTestVector1(t *testing.T) {
 	c := qt.New(t)
-	database1, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
+	database1, err := badgerdb.New(db.Options{Path: c.TempDir()})
 	c.Assert(err, qt.IsNil)
-	tree1, err := NewTree(database1, 256, HashFunctionBlake2b)
+	tree1, err := NewTree(Config{database1, 256, DefaultThresholdNLeafs,
+		HashFunctionBlake2b})
 	c.Assert(err, qt.IsNil)
 	defer tree1.db.Close() //nolint:errcheck
 
-	database2, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
+	database2, err := badgerdb.New(db.Options{Path: c.TempDir()})
 	c.Assert(err, qt.IsNil)
-	tree2, err := NewTree(database2, 256, HashFunctionBlake2b)
+	tree2, err := NewTree(Config{database2, 256, DefaultThresholdNLeafs,
+		HashFunctionBlake2b})
 	c.Assert(err, qt.IsNil)
 	defer tree2.db.Close() //nolint:errcheck
 
@@ -198,22 +209,24 @@ func TestAddBatchTestVector1(t *testing.T) {
 		}
 	}
 
-	indexes, err := tree2.AddBatch(keys, values)
+	invalids, err := tree2.AddBatch(keys, values)
 	c.Assert(err, qt.IsNil)
-	c.Check(len(indexes), qt.Equals, 0)
+	c.Check(len(invalids), qt.Equals, 0)
 	// check that both trees roots are equal
 	checkRoots(c, tree1, tree2)
 
 	// 2nd test vectors
-	database1, err = badgerdb.New(badgerdb.Options{Path: c.TempDir()})
+	database1, err = badgerdb.New(db.Options{Path: c.TempDir()})
 	c.Assert(err, qt.IsNil)
-	tree1, err = NewTree(database1, 256, HashFunctionBlake2b)
+	tree1, err = NewTree(Config{database1, 256, DefaultThresholdNLeafs,
+		HashFunctionBlake2b})
 	c.Assert(err, qt.IsNil)
 	defer tree1.db.Close() //nolint:errcheck
 
-	database2, err = badgerdb.New(badgerdb.Options{Path: c.TempDir()})
+	database2, err = badgerdb.New(db.Options{Path: c.TempDir()})
 	c.Assert(err, qt.IsNil)
-	tree2, err = NewTree(database2, 256, HashFunctionBlake2b)
+	tree2, err = NewTree(Config{database2, 256, DefaultThresholdNLeafs,
+		HashFunctionBlake2b})
 	c.Assert(err, qt.IsNil)
 	defer tree2.db.Close() //nolint:errcheck
 
@@ -242,9 +255,9 @@ func TestAddBatchTestVector1(t *testing.T) {
 		}
 	}
 
-	indexes, err = tree2.AddBatch(keys, values)
+	invalids, err = tree2.AddBatch(keys, values)
 	c.Assert(err, qt.IsNil)
-	c.Check(len(indexes), qt.Equals, 0)
+	c.Check(len(invalids), qt.Equals, 0)
 	// check that both trees roots are equal
 	checkRoots(c, tree1, tree2)
 }
@@ -253,15 +266,17 @@ func TestAddBatchTestVector2(t *testing.T) {
 	// test vector with unbalanced tree
 	c := qt.New(t)
 
-	database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
+	database, err := badgerdb.New(db.Options{Path: c.TempDir()})
 	c.Assert(err, qt.IsNil)
-	tree1, err := NewTree(database, 256, HashFunctionPoseidon)
+	tree1, err := NewTree(Config{database, 256, DefaultThresholdNLeafs,
+		HashFunctionPoseidon})
 	c.Assert(err, qt.IsNil)
 	defer tree1.db.Close() //nolint:errcheck
 
-	database2, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
+	database2, err := badgerdb.New(db.Options{Path: c.TempDir()})
 	c.Assert(err, qt.IsNil)
-	tree2, err := NewTree(database2, 256, HashFunctionPoseidon)
+	tree2, err := NewTree(Config{database2, 256, DefaultThresholdNLeafs,
+		HashFunctionPoseidon})
 	c.Assert(err, qt.IsNil)
 	defer tree2.db.Close() //nolint:errcheck
 
@@ -286,9 +301,9 @@ func TestAddBatchTestVector2(t *testing.T) {
 		}
 	}
 
-	indexes, err := tree2.AddBatch(keys, values)
+	invalids, err := tree2.AddBatch(keys, values)
 	c.Assert(err, qt.IsNil)
-	c.Check(len(indexes), qt.Equals, 0)
+	c.Check(len(invalids), qt.Equals, 0)
 
 	// check that both trees roots are equal
 	checkRoots(c, tree1, tree2)
@@ -298,15 +313,17 @@ func TestAddBatchTestVector3(t *testing.T) {
 	// test vector with unbalanced tree
 	c := qt.New(t)
 
-	database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
+	database, err := badgerdb.New(db.Options{Path: c.TempDir()})
 	c.Assert(err, qt.IsNil)
-	tree1, err := NewTree(database, 256, HashFunctionPoseidon)
+	tree1, err := NewTree(Config{database, 256, DefaultThresholdNLeafs,
+		HashFunctionPoseidon})
 	c.Assert(err, qt.IsNil)
 	defer tree1.db.Close() //nolint:errcheck
 
-	database2, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
+	database2, err := badgerdb.New(db.Options{Path: c.TempDir()})
 	c.Assert(err, qt.IsNil)
-	tree2, err := NewTree(database2, 256, HashFunctionPoseidon)
+	tree2, err := NewTree(Config{database2, 256, DefaultThresholdNLeafs,
+		HashFunctionPoseidon})
 	c.Assert(err, qt.IsNil)
 	defer tree2.db.Close() //nolint:errcheck
 
@@ -331,9 +348,9 @@ func TestAddBatchTestVector3(t *testing.T) {
 		}
 	}
 
-	indexes, err := tree2.AddBatch(keys, values)
+	invalids, err := tree2.AddBatch(keys, values)
 	c.Assert(err, qt.IsNil)
-	c.Check(len(indexes), qt.Equals, 0)
+	c.Check(len(invalids), qt.Equals, 0)
 
 	// check that both trees roots are equal
 	checkRoots(c, tree1, tree2)
@@ -347,15 +364,17 @@ func TestAddBatchTreeEmptyRandomKeys(t *testing.T) {
 
 	nLeafs := 8
 
-	database1, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
+	database1, err := badgerdb.New(db.Options{Path: c.TempDir()})
 	c.Assert(err, qt.IsNil)
-	tree1, err := NewTree(database1, 256, HashFunctionBlake2b)
+	tree1, err := NewTree(Config{database1, 256, DefaultThresholdNLeafs,
+		HashFunctionBlake2b})
 	c.Assert(err, qt.IsNil)
 	defer tree1.db.Close() //nolint:errcheck
 
-	database2, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
+	database2, err := badgerdb.New(db.Options{Path: c.TempDir()})
 	c.Assert(err, qt.IsNil)
-	tree2, err := NewTree(database2, 256, HashFunctionBlake2b)
+	tree2, err := NewTree(Config{database2, 256, DefaultThresholdNLeafs,
+		HashFunctionBlake2b})
 	c.Assert(err, qt.IsNil)
 	defer tree2.db.Close() //nolint:errcheck
 
@@ -371,9 +390,9 @@ func TestAddBatchTreeEmptyRandomKeys(t *testing.T) {
 		}
 	}
 
-	indexes, err := tree2.AddBatch(keys, values)
+	invalids, err := tree2.AddBatch(keys, values)
 	c.Assert(err, qt.IsNil)
-	c.Check(len(indexes), qt.Equals, 0)
+	c.Check(len(invalids), qt.Equals, 0)
 	// check that both trees roots are equal
 	checkRoots(c, tree1, tree2)
 }
@@ -407,7 +426,7 @@ func TestAddBatchTreeNotEmptyFewLeafs(t *testing.T) {
 		values = append(values, v)
 	}
 	start = time.Now()
-	indexes, err := tree2.AddBatch(keys, values)
+	invalids, err := tree2.AddBatch(keys, values)
 	c.Assert(err, qt.IsNil)
 	time2 := time.Since(start)
 	if debug {
@@ -415,7 +434,7 @@ func TestAddBatchTreeNotEmptyFewLeafs(t *testing.T) {
 		printTestContext("	", nLeafs, "Poseidon", "memory")
 		tree2.dbg.print("	")
 	}
-	c.Check(len(indexes), qt.Equals, 0)
+	c.Check(len(invalids), qt.Equals, 0)
 
 	// check that both trees roots are equal
 	checkRoots(c, tree1, tree2)
@@ -450,7 +469,7 @@ func TestAddBatchTreeNotEmptyEnoughLeafs(t *testing.T) {
 		values = append(values, v)
 	}
 	start = time.Now()
-	indexes, err := tree2.AddBatch(keys, values)
+	invalids, err := tree2.AddBatch(keys, values)
 	c.Assert(err, qt.IsNil)
 	time2 := time.Since(start)
 	if debug {
@@ -458,7 +477,7 @@ func TestAddBatchTreeNotEmptyEnoughLeafs(t *testing.T) {
 		printTestContext("	", nLeafs, "Poseidon", "memory")
 		tree2.dbg.print("	")
 	}
-	c.Check(len(indexes), qt.Equals, 0)
+	c.Check(len(invalids), qt.Equals, 0)
 	// check that both trees roots are equal
 	checkRoots(c, tree1, tree2)
 }
@@ -495,9 +514,9 @@ func TestAddBatchTreeEmptyRepeatedLeafs(t *testing.T) {
 		}
 	}
 
-	indexes, err := tree2.AddBatch(keys, values)
+	invalids, err := tree2.AddBatch(keys, values)
 	c.Assert(err, qt.IsNil)
-	c.Check(len(indexes), qt.Equals, nRepeatedKeys)
+	c.Check(len(invalids), qt.Equals, nRepeatedKeys)
 	// check that both trees roots are equal
 	checkRoots(c, tree1, tree2)
 }
@@ -527,9 +546,9 @@ func TestAddBatchTreeNotEmptyFewLeafsRepeatedLeafs(t *testing.T) {
 		}
 	}
 
-	indexes, err := tree2.AddBatch(keys, values)
+	invalids, err := tree2.AddBatch(keys, values)
 	c.Assert(err, qt.IsNil)
-	c.Check(len(indexes), qt.Equals, initialNLeafs)
+	c.Assert(len(invalids), qt.Equals, initialNLeafs)
 	// check that both trees roots are equal
 	checkRoots(c, tree1, tree2)
 }
@@ -664,7 +683,7 @@ func TestAddBatchTreeNotEmpty(t *testing.T) {
 		values = append(values, v)
 	}
 	start = time.Now()
-	indexes, err := tree2.AddBatch(keys, values)
+	invalids, err := tree2.AddBatch(keys, values)
 	c.Assert(err, qt.IsNil)
 	time2 := time.Since(start)
 	if debug {
@@ -672,7 +691,7 @@ func TestAddBatchTreeNotEmpty(t *testing.T) {
 		printTestContext("	", nLeafs, "Poseidon", "memory")
 		tree2.dbg.print("	")
 	}
-	c.Check(len(indexes), qt.Equals, 0)
+	c.Check(len(invalids), qt.Equals, 0)
 
 	// check that both trees roots are equal
 	checkRoots(c, tree1, tree2)
@@ -697,9 +716,10 @@ func TestAddBatchNotEmptyUnbalanced(t *testing.T) {
 	}
 	time1 := time.Since(start)
 
-	database2, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
+	database2, err := badgerdb.New(db.Options{Path: c.TempDir()})
 	c.Assert(err, qt.IsNil)
-	tree2, err := NewTree(database2, 256, HashFunctionPoseidon)
+	tree2, err := NewTree(Config{database2, 256, DefaultThresholdNLeafs,
+		HashFunctionPoseidon})
 	c.Assert(err, qt.IsNil)
 	defer tree2.db.Close() //nolint:errcheck
 	tree2.dbgInit()
@@ -729,7 +749,7 @@ func TestAddBatchNotEmptyUnbalanced(t *testing.T) {
 		values = append(values, v)
 	}
 	start = time.Now()
-	indexes, err := tree2.AddBatch(keys, values)
+	invalids, err := tree2.AddBatch(keys, values)
 	c.Assert(err, qt.IsNil)
 	time2 := time.Since(start)
 	if debug {
@@ -737,7 +757,7 @@ func TestAddBatchNotEmptyUnbalanced(t *testing.T) {
 		printTestContext("	", nLeafs, "Poseidon", "memory")
 		tree2.dbg.print("	")
 	}
-	c.Check(len(indexes), qt.Equals, 0)
+	c.Check(len(invalids), qt.Equals, 0)
 
 	// check that both trees roots are equal
 	checkRoots(c, tree1, tree2)
@@ -774,9 +794,10 @@ func TestAddBatchBench(t *testing.T) {
 func benchAdd(t *testing.T, ks, vs [][]byte) {
 	c := qt.New(t)
 
-	database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
+	database, err := badgerdb.New(db.Options{Path: c.TempDir()})
 	c.Assert(err, qt.IsNil)
-	tree, err := NewTree(database, 256, HashFunctionBlake2b)
+	tree, err := NewTree(Config{database, 256, DefaultThresholdNLeafs,
+		HashFunctionBlake2b})
 	c.Assert(err, qt.IsNil)
 	defer tree.db.Close() //nolint:errcheck
 
@@ -794,9 +815,10 @@ func benchAdd(t *testing.T, ks, vs [][]byte) {
 func benchAddBatch(t *testing.T, ks, vs [][]byte) {
 	c := qt.New(t)
 
-	database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
+	database, err := badgerdb.New(db.Options{Path: c.TempDir()})
 	c.Assert(err, qt.IsNil)
-	tree, err := NewTree(database, 256, HashFunctionBlake2b)
+	tree, err := NewTree(Config{database, 256, DefaultThresholdNLeafs,
+		HashFunctionBlake2b})
 	c.Assert(err, qt.IsNil)
 	defer tree.db.Close() //nolint:errcheck
 
@@ -827,9 +849,10 @@ func TestDbgStats(t *testing.T) {
 	}
 
 	// 1
-	database1, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
+	database1, err := badgerdb.New(db.Options{Path: c.TempDir()})
 	c.Assert(err, qt.IsNil)
-	tree1, err := NewTree(database1, 256, HashFunctionBlake2b)
+	tree1, err := NewTree(Config{database1, 256, DefaultThresholdNLeafs,
+		HashFunctionBlake2b})
 	c.Assert(err, qt.IsNil)
 	defer tree1.db.Close() //nolint:errcheck
 
@@ -841,9 +864,10 @@ func TestDbgStats(t *testing.T) {
 	}
 
 	// 2
-	database2, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
+	database2, err := badgerdb.New(db.Options{Path: c.TempDir()})
 	c.Assert(err, qt.IsNil)
-	tree2, err := NewTree(database2, 256, HashFunctionBlake2b)
+	tree2, err := NewTree(Config{database2, 256, DefaultThresholdNLeafs,
+		HashFunctionBlake2b})
 	c.Assert(err, qt.IsNil)
 	defer tree2.db.Close() //nolint:errcheck
 
@@ -854,9 +878,10 @@ func TestDbgStats(t *testing.T) {
 	c.Assert(len(invalids), qt.Equals, 0)
 
 	// 3
-	database3, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
+	database3, err := badgerdb.New(db.Options{Path: c.TempDir()})
 	c.Assert(err, qt.IsNil)
-	tree3, err := NewTree(database3, 256, HashFunctionBlake2b)
+	tree3, err := NewTree(Config{database3, 256, DefaultThresholdNLeafs,
+		HashFunctionBlake2b})
 	c.Assert(err, qt.IsNil)
 	defer tree3.db.Close() //nolint:errcheck
 
@@ -889,9 +914,10 @@ func TestLoadVT(t *testing.T) {
 
 	nLeafs := 1024
 
-	database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
+	database, err := badgerdb.New(db.Options{Path: c.TempDir()})
 	c.Assert(err, qt.IsNil)
-	tree, err := NewTree(database, 256, HashFunctionPoseidon)
+	tree, err := NewTree(Config{database, 256, DefaultThresholdNLeafs,
+		HashFunctionPoseidon})
 	c.Assert(err, qt.IsNil)
 	defer tree.db.Close() //nolint:errcheck
 
@@ -902,9 +928,9 @@ func TestLoadVT(t *testing.T) {
 		keys = append(keys, k)
 		values = append(values, v)
 	}
-	indexes, err := tree.AddBatch(keys, values)
+	invalids, err := tree.AddBatch(keys, values)
 	c.Assert(err, qt.IsNil)
-	c.Check(len(indexes), qt.Equals, 0)
+	c.Check(len(invalids), qt.Equals, 0)
 
 	rTx := tree.db.ReadTx()
 	defer rTx.Discard()
@@ -925,9 +951,10 @@ func TestAddKeysWithEmptyValues(t *testing.T) {
 
 	nLeafs := 1024
 
-	database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
+	database, err := badgerdb.New(db.Options{Path: c.TempDir()})
 	c.Assert(err, qt.IsNil)
-	tree, err := NewTree(database, 256, HashFunctionPoseidon)
+	tree, err := NewTree(Config{database, 256, DefaultThresholdNLeafs,
+		HashFunctionPoseidon})
 	c.Assert(err, qt.IsNil)
 	defer tree.db.Close() //nolint:errcheck
 
@@ -946,29 +973,31 @@ func TestAddKeysWithEmptyValues(t *testing.T) {
 		}
 	}
 
-	database2, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
+	database2, err := badgerdb.New(db.Options{Path: c.TempDir()})
 	c.Assert(err, qt.IsNil)
-	tree2, err := NewTree(database2, 256, HashFunctionPoseidon)
+	tree2, err := NewTree(Config{database2, 256, DefaultThresholdNLeafs,
+		HashFunctionPoseidon})
 	c.Assert(err, qt.IsNil)
 	defer tree2.db.Close() //nolint:errcheck
 	tree2.dbgInit()
 
-	indexes, err := tree2.AddBatch(keys, values)
+	invalids, err := tree2.AddBatch(keys, values)
 	c.Assert(err, qt.IsNil)
-	c.Check(len(indexes), qt.Equals, 0)
+	c.Check(len(invalids), qt.Equals, 0)
 	// check that both trees roots are equal
 	checkRoots(c, tree, tree2)
 
 	// use tree3 to add nil value array
-	database3, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
+	database3, err := badgerdb.New(db.Options{Path: c.TempDir()})
 	c.Assert(err, qt.IsNil)
-	tree3, err := NewTree(database3, 256, HashFunctionPoseidon)
+	tree3, err := NewTree(Config{database3, 256, DefaultThresholdNLeafs,
+		HashFunctionPoseidon})
 	c.Assert(err, qt.IsNil)
 	defer tree3.db.Close() //nolint:errcheck
 
-	indexes, err = tree3.AddBatch(keys, nil)
+	invalids, err = tree3.AddBatch(keys, nil)
 	c.Assert(err, qt.IsNil)
-	c.Check(len(indexes), qt.Equals, 0)
+	c.Check(len(invalids), qt.Equals, 0)
 	checkRoots(c, tree, tree3)
 
 	kAux, proofV, siblings, existence, err := tree2.GenProof(keys[9])
@@ -1003,6 +1032,180 @@ func TestAddKeysWithEmptyValues(t *testing.T) {
 	c.Check(verif, qt.IsFalse)
 }
 
+func TestAddBatchThresholdInDisk(t *testing.T) {
+	c := qt.New(t)
+
+	// customize thresholdNLeafs for the test
+	testThresholdNLeafs := 1024
+
+	database1, err := badgerdb.New(db.Options{Path: c.TempDir()})
+	c.Assert(err, qt.IsNil)
+	tree1, err := NewTree(Config{database1, 256, testThresholdNLeafs,
+		HashFunctionBlake2b})
+	c.Assert(err, qt.IsNil)
+	defer tree1.db.Close() //nolint:errcheck
+
+	database2, err := pebbledb.New(db.Options{Path: c.TempDir()})
+	c.Assert(err, qt.IsNil)
+	tree2, err := NewTree(Config{database2, 256, testThresholdNLeafs,
+		HashFunctionBlake2b})
+	c.Assert(err, qt.IsNil)
+	defer tree2.db.Close() //nolint:errcheck
+
+	database3, err := pebbledb.New(db.Options{Path: c.TempDir()})
+	c.Assert(err, qt.IsNil)
+	tree3, err := NewTree(Config{database3, 256, testThresholdNLeafs,
+		HashFunctionBlake2b})
+	c.Assert(err, qt.IsNil)
+	defer tree3.db.Close() //nolint:errcheck
+
+	var keys, values [][]byte
+	for i := 0; i < 3*testThresholdNLeafs; i++ {
+		k := randomBytes(32)
+		v := randomBytes(32)
+		if err := tree1.Add(k, v); err != nil {
+			t.Fatal(err)
+		}
+		if i < testThresholdNLeafs+1 {
+			if err := tree2.Add(k, v); err != nil {
+				t.Fatal(err)
+			}
+		}
+		// store for later addition through AddBatch
+		keys = append(keys, k)
+		values = append(values, v)
+	}
+
+	invalids, err := tree2.AddBatch(keys[testThresholdNLeafs+1:], values[testThresholdNLeafs+1:])
+	c.Assert(err, qt.IsNil)
+	c.Check(len(invalids), qt.Equals, 0)
+	// check that both trees roots are equal
+	checkRoots(c, tree1, tree2)
+
+	// call directly the tree3.addBatchInDisk to ensure that is tested
+	wTx := tree3.db.WriteTx()
+	defer wTx.Discard()
+	invalids, err = tree3.addBatchInDisk(wTx, keys, values)
+	c.Assert(err, qt.IsNil)
+	err = wTx.Commit()
+	c.Assert(err, qt.IsNil)
+	c.Check(len(invalids), qt.Equals, 0)
+	// check that both trees roots are equal
+	checkRoots(c, tree1, tree3)
+
+	// now add one leaf more to the trees to ensure that the previous
+	// actions did not left the tree in an invalid state
+	k := randomBytes(32)
+	v := randomBytes(32)
+	err = tree1.Add(k, v)
+	c.Assert(err, qt.IsNil)
+	err = tree2.Add(k, v)
+	c.Assert(err, qt.IsNil)
+	err = tree3.Add(k, v)
+	c.Assert(err, qt.IsNil)
+}
+
+func initTestUpFromSubRoots(c *qt.C) (*Tree, *Tree) {
+	database1, err := badgerdb.New(db.Options{Path: c.TempDir()})
+	c.Assert(err, qt.IsNil)
+	tree1, err := NewTree(Config{database1, 256, DefaultThresholdNLeafs,
+		HashFunctionBlake2b})
+	c.Assert(err, qt.IsNil)
+
+	database2, err := badgerdb.New(db.Options{Path: c.TempDir()})
+	c.Assert(err, qt.IsNil)
+	tree2, err := NewTree(Config{database2, 256, DefaultThresholdNLeafs,
+		HashFunctionBlake2b})
+	c.Assert(err, qt.IsNil)
+	return tree1, tree2
+}
+
+func testUpFromSubRoots(c *qt.C, tree1, tree2 *Tree, preSubRoots [][]byte) {
+	// add the preSubRoots to the tree1
+	for i := 0; i < len(preSubRoots); i++ {
+		if bytes.Equal(preSubRoots[i], tree1.emptyHash) {
+			continue
+		}
+		err := tree1.Add(preSubRoots[i], nil)
+		c.Assert(err, qt.IsNil)
+	}
+	root1, err := tree1.Root()
+	c.Assert(err, qt.IsNil)
+
+	wTx := tree2.db.WriteTx()
+	subRoots := make([][]byte, len(preSubRoots))
+	for i := 0; i < len(preSubRoots); i++ {
+		if preSubRoots[i] == nil || bytes.Equal(preSubRoots[i], tree1.emptyHash) {
+			subRoots[i] = tree1.emptyHash
+			continue
+		}
+		leafKey, leafValue, err := tree2.newLeafValue(preSubRoots[i], nil)
+		c.Assert(err, qt.IsNil)
+		subRoots[i] = leafKey
+
+		err = wTx.Set(leafKey, leafValue)
+		c.Assert(err, qt.IsNil)
+	}
+	// first fill the leaf nodes
+	// then call upFromSubRoots
+	root2FromUp, err := tree2.upFromSubRoots(wTx, subRoots)
+	c.Assert(err, qt.IsNil)
+
+	err = tree2.SetRootWithTx(wTx, root2FromUp)
+	c.Assert(err, qt.IsNil)
+	err = wTx.Commit()
+	c.Assert(err, qt.IsNil)
+
+	root2, err := tree2.Root()
+	c.Assert(err, qt.IsNil)
+
+	c.Assert(root1, qt.DeepEquals, root2)
+}
+
+func testUpFromSubRootsWithEmpties(c *qt.C, preSubRoots [][]byte, indexEmpties []int) {
+	tree1, tree2 := initTestUpFromSubRoots(c)
+	defer tree1.db.Close() //nolint:errcheck
+	defer tree2.db.Close() //nolint:errcheck
+
+	testPreSubRoots := make([][]byte, len(preSubRoots))
+	copy(testPreSubRoots[:], preSubRoots[:])
+	for i := 0; i < len(indexEmpties); i++ {
+		testPreSubRoots[indexEmpties[i]] = tree1.emptyHash
+	}
+	testUpFromSubRoots(c, tree1, tree2, testPreSubRoots)
+}
+
+func TestUpFromSubRoots(t *testing.T) {
+	c := qt.New(t)
+
+	// prepare preSubRoots
+	preSubRoots := [][]byte{
+		BigIntToBytes(32, big.NewInt(4)),
+		BigIntToBytes(32, big.NewInt(2)),
+		BigIntToBytes(32, big.NewInt(1)),
+		BigIntToBytes(32, big.NewInt(3)),
+	}
+
+	// test using the full 4 leafs as subRoots
+	testUpFromSubRootsWithEmpties(c, preSubRoots, nil)
+	// 1st subRoot empty
+	testUpFromSubRootsWithEmpties(c, preSubRoots, []int{0})
+	// 2nd subRoot empty
+	testUpFromSubRootsWithEmpties(c, preSubRoots, []int{1})
+	// 3rd subRoot empty
+	testUpFromSubRootsWithEmpties(c, preSubRoots, []int{2})
+	// 4th subRoot empty
+	testUpFromSubRootsWithEmpties(c, preSubRoots, []int{3})
+
+	// other combinations of empty SubRoots
+	testUpFromSubRootsWithEmpties(c, preSubRoots, []int{0, 1, 2, 3})
+	testUpFromSubRootsWithEmpties(c, preSubRoots, []int{0, 1})
+	testUpFromSubRootsWithEmpties(c, preSubRoots, []int{1, 2})
+	testUpFromSubRootsWithEmpties(c, preSubRoots, []int{1, 3})
+	testUpFromSubRootsWithEmpties(c, preSubRoots, []int{2, 3})
+	testUpFromSubRootsWithEmpties(c, preSubRoots, []int{0, 2, 3})
+}
+
 // TODO test adding batch with multiple invalid keys
 // TODO for tests of AddBatch, if the root does not match the Add root, bulk
 // all the leafs of both trees into a log file to later be able to debug and

circomproofs_test.go

@@ -6,14 +6,16 @@ import (
 	"testing"
 
 	qt "github.com/frankban/quicktest"
+	"go.vocdoni.io/dvote/db"
 	"go.vocdoni.io/dvote/db/badgerdb"
 )
 
 func TestCircomVerifierProof(t *testing.T) {
 	c := qt.New(t)
-	database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
+	database, err := badgerdb.New(db.Options{Path: c.TempDir()})
 	c.Assert(err, qt.IsNil)
-	tree, err := NewTree(database, 4, HashFunctionPoseidon)
+	tree, err := NewTree(Config{Database: database, MaxLevels: 4,
+		HashFunction: HashFunctionPoseidon})
 	c.Assert(err, qt.IsNil)
 	defer tree.db.Close() //nolint:errcheck
 

go.mod

@@ -5,6 +5,6 @@ go 1.16
 require (
 	github.com/frankban/quicktest v1.13.0
 	github.com/iden3/go-iden3-crypto v0.0.6-0.20210308142348-8f85683b2cef
-	go.vocdoni.io/dvote v1.0.4-0.20210806163627-9494efbc5382
-	golang.org/x/crypto v0.0.0-20210711020723-a769d52b0f97
+	go.vocdoni.io/dvote v1.0.4-0.20211025120558-83c64f440044
+	golang.org/x/crypto v0.0.0-20210920023735-84f357641f63
 )

hash.go

@@ -77,7 +77,9 @@ func (f HashPoseidon) Len() int {
 	return 32 //nolint:gomnd
 }
 
-// Hash implements the hash method for the HashFunction HashPoseidon
+// Hash implements the hash method for the HashFunction HashPoseidon. It
+// expects the byte arrays to be little-endian representations of big.Int
+// values.
 func (f HashPoseidon) Hash(b ...[]byte) ([]byte, error) {
 	var toHash []*big.Int
 	for i := 0; i < len(b); i++ {

helpers_test.go

@@ -9,6 +9,7 @@ import (
 	"time"
 
 	qt "github.com/frankban/quicktest"
+	"go.vocdoni.io/dvote/db"
 	"go.vocdoni.io/dvote/db/badgerdb"
 )
 
@@ -87,14 +88,16 @@ func TestReadTreeDBG(t *testing.T) {
 
 	c := qt.New(t)
 
-	database1, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
+	database1, err := badgerdb.New(db.Options{Path: c.TempDir()})
 	c.Assert(err, qt.IsNil)
-	tree1, err := NewTree(database1, 100, HashFunctionBlake2b)
+	tree1, err := NewTree(Config{Database: database1, MaxLevels: 100,
+		HashFunction: HashFunctionBlake2b})
 	c.Assert(err, qt.IsNil)
 
-	database2, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
+	database2, err := badgerdb.New(db.Options{Path: c.TempDir()})
 	c.Assert(err, qt.IsNil)
-	tree2, err := NewTree(database2, 100, HashFunctionBlake2b)
+	tree2, err := NewTree(Config{Database: database2, MaxLevels: 100,
+		HashFunction: HashFunctionBlake2b})
 	c.Assert(err, qt.IsNil)
 
 	// tree1 is generated by a loop of .Add

testvectors/circom/go-data-generator/generator_test.go

@@ -8,14 +8,16 @@ import (
 
 	qt "github.com/frankban/quicktest"
 	"github.com/vocdoni/arbo"
+	"go.vocdoni.io/dvote/db"
 	"go.vocdoni.io/dvote/db/badgerdb"
 )
 
 func TestGenerator(t *testing.T) {
 	c := qt.New(t)
-	database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
+	database, err := badgerdb.New(db.Options{Path: c.TempDir()})
 	c.Assert(err, qt.IsNil)
-	tree, err := arbo.NewTree(database, 4, arbo.HashFunctionPoseidon)
+	tree, err := arbo.NewTree(arbo.Config{Database: database, MaxLevels: 4,
+		HashFunction: arbo.HashFunctionPoseidon})
 	c.Assert(err, qt.IsNil)
 
 	testVector := [][]int64{

tree.go

@@ -17,6 +17,7 @@ import (
 	"fmt"
 	"io"
 	"math"
+	"runtime"
 	"sync"
 
 	"go.vocdoni.io/dvote/db"
@@ -45,6 +46,12 @@ const (
 )
 
 var (
+	// DefaultThresholdNLeafs defines the threshold number of leafs in the
+	// tree that determines if AddBatch will work in memory or in disk.  It
+	// is defined when calling NewTree, and if set to 0 it will work always
+	// in disk.
+	DefaultThresholdNLeafs = 65536
+
 	dbKeyRoot   = []byte("root")
 	dbKeyNLeafs = []byte("nleafs")
 	emptyValue  = []byte{0}
@@ -78,9 +85,14 @@ var (
 type Tree struct {
 	sync.Mutex
 
-	db           db.Database
-	maxLevels    int
-	snapshotRoot []byte
+	db        db.Database
+	maxLevels int
+	// thresholdNLeafs defines the threshold number of leafs in the tree
+	// that determines if AddBatch will work in memory or in disk.  It is
+	// defined when calling NewTree, and if set to 0 it will work always in
+	// disk.
+	thresholdNLeafs int
+	snapshotRoot    []byte
 
 	hashFunction HashFunction
 	// TODO in the methods that use it, check if emptyHash param is len>0
@@ -90,13 +102,21 @@ type Tree struct {
 	dbg *dbgStats
 }
 
+// Config defines the configuration for calling NewTree & NewTreeWithTx methods
+type Config struct {
+	Database        db.Database
+	MaxLevels       int
+	ThresholdNLeafs int
+	HashFunction    HashFunction
+}
+
 // NewTree returns a new Tree, if there is a Tree still in the given database, it
 // will load it.
-func NewTree(database db.Database, maxLevels int, hash HashFunction) (*Tree, error) {
-	wTx := database.WriteTx()
+func NewTree(cfg Config) (*Tree, error) {
+	wTx := cfg.Database.WriteTx()
 	defer wTx.Discard()
 
-	t, err := NewTreeWithTx(wTx, database, maxLevels, hash)
+	t, err := NewTreeWithTx(wTx, cfg)
 	if err != nil {
 		return nil, err
 	}
@@ -110,9 +130,10 @@ func NewTree(database db.Database, maxLevels int, hash HashFunction) (*Tree, err
 // NewTreeWithTx returns a new Tree using the given db.WriteTx, which will not
 // be ccommited inside this method, if there is a Tree still in the given
 // database, it will load it.
-func NewTreeWithTx(wTx db.WriteTx, database db.Database,
-	maxLevels int, hash HashFunction) (*Tree, error) {
-	t := Tree{db: database, maxLevels: maxLevels, hashFunction: hash}
+func NewTreeWithTx(wTx db.WriteTx, cfg Config) (*Tree, error) {
+	// if thresholdNLeafs is set to 0, use the DefaultThresholdNLeafs
+	t := Tree{db: cfg.Database, maxLevels: cfg.MaxLevels,
+		thresholdNLeafs: cfg.ThresholdNLeafs, hashFunction: cfg.HashFunction}
 	t.emptyHash = make([]byte, t.hashFunction.Len()) // empty
 
 	_, err := wTx.Get(dbKeyRoot)
@@ -197,11 +218,6 @@ func (t *Tree) AddBatchWithTx(wTx db.WriteTx, keys, values [][]byte) ([]Invalid,
 		return nil, ErrSnapshotNotEditable
 	}
 
-	vt, err := t.loadVT(wTx)
-	if err != nil {
-		return nil, err
-	}
-
 	e := []byte{}
 	// equal the number of keys & values
 	if len(keys) > len(values) {
@@ -214,6 +230,223 @@ func (t *Tree) AddBatchWithTx(wTx db.WriteTx, keys, values [][]byte) ([]Invalid,
 		values = values[:len(keys)]
 	}
 
+	nLeafs, err := t.GetNLeafsWithTx(wTx)
+	if err != nil {
+		return nil, err
+	}
+	if nLeafs > t.thresholdNLeafs {
+		return t.addBatchInDisk(wTx, keys, values)
+	}
+	return t.addBatchInMemory(wTx, keys, values)
+}
+
+func (t *Tree) addBatchInDisk(wTx db.WriteTx, keys, values [][]byte) ([]Invalid, error) {
+	nCPU := flp2(runtime.NumCPU())
+	if nCPU == 1 || len(keys) < nCPU {
+		var invalids []Invalid
+		for i := 0; i < len(keys); i++ {
+			if err := t.addWithTx(wTx, keys[i], values[i]); err != nil {
+				invalids = append(invalids, Invalid{i, err})
+			}
+		}
+		return invalids, nil
+	}
+
+	kvs, invalids, err := keysValuesToKvs(t.maxLevels, keys, values)
+	if err != nil {
+		return nil, err
+	}
+
+	buckets := splitInBuckets(kvs, nCPU)
+
+	root, err := t.RootWithTx(wTx)
+	if err != nil {
+		return nil, err
+	}
+
+	l := int(math.Log2(float64(nCPU)))
+	subRoots, err := t.getSubRootsAtLevel(wTx, root, l+1)
+	if err != nil {
+		return nil, err
+	}
+	if len(subRoots) != nCPU {
+		// Already populated Tree but Unbalanced.
+
+		// add one key at each bucket, and then continue with the flow
+		for i := 0; i < len(buckets); i++ {
+			// add one leaf of the bucket, if there is an error when
+			// adding the k-v, try to add the next one of the bucket
+			// (until one is added)
+			inserted := -1
+			for j := 0; j < len(buckets[i]); j++ {
+				if newRoot, err := t.add(wTx, root, 0,
+					buckets[i][j].k, buckets[i][j].v); err == nil {
+					inserted = j
+					root = newRoot
+					break
+				}
+			}
+
+			// remove the inserted element from buckets[i]
+			if inserted != -1 {
+				buckets[i] = append(buckets[i][:inserted], buckets[i][inserted+1:]...)
+			}
+		}
+		subRoots, err = t.getSubRootsAtLevel(wTx, root, l+1)
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	if len(subRoots) != nCPU {
+		return nil, fmt.Errorf("This error should not be reached."+
+			" len(subRoots) != nCPU, len(subRoots)=%d, nCPU=%d."+
+			" Please report it in a new issue:"+
+			" https://github.com/vocdoni/arbo/issues/new", len(subRoots), nCPU)
+	}
+
+	invalidsInBucket := make([][]Invalid, nCPU)
+	txs := make([]db.WriteTx, nCPU)
+	for i := 0; i < nCPU; i++ {
+		txs[i] = t.db.WriteTx()
+		err := txs[i].Apply(wTx)
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	var wg sync.WaitGroup
+	wg.Add(nCPU)
+	for i := 0; i < nCPU; i++ {
+		go func(cpu int) {
+			// use different wTx for each cpu, after once all
+			// are done, iter over the cpuWTxs and copy their
+			// content into the main wTx
+			for j := 0; j < len(buckets[cpu]); j++ {
+				newSubRoot, err := t.add(txs[cpu], subRoots[cpu],
+					l, buckets[cpu][j].k, buckets[cpu][j].v)
+				if err != nil {
+					invalidsInBucket[cpu] = append(invalidsInBucket[cpu],
+						Invalid{buckets[cpu][j].pos, err})
+					continue
+				}
+				// if there has not been errors, set the new subRoots[cpu]
+				subRoots[cpu] = newSubRoot
+			}
+			wg.Done()
+		}(i)
+	}
+	wg.Wait()
+
+	for i := 0; i < nCPU; i++ {
+		if err := wTx.Apply(txs[i]); err != nil {
+			return nil, err
+		}
+		txs[i].Discard()
+	}
+
+	for i := 0; i < len(invalidsInBucket); i++ {
+		invalids = append(invalids, invalidsInBucket[i]...)
+	}
+
+	newRoot, err := t.upFromSubRoots(wTx, subRoots)
+	if err != nil {
+		return nil, err
+	}
+
+	// update dbKeyNLeafs
+	if err := t.SetRootWithTx(wTx, newRoot); err != nil {
+		return nil, err
+	}
+
+	// update nLeafs
+	if err := t.incNLeafs(wTx, len(keys)-len(invalids)); err != nil {
+		return nil, err
+	}
+
+	return invalids, nil
+}
+
+func (t *Tree) upFromSubRoots(wTx db.WriteTx, subRoots [][]byte) ([]byte, error) {
+	// is a method of Tree just to get access to t.hashFunction and
+	// t.emptyHash.
+
+	// go up from subRoots to up, storing nodes in the given WriteTx
+	// once up at the root, store it in the WriteTx using the dbKeyRoot
+	if len(subRoots) == 1 {
+		return subRoots[0], nil
+	}
+	// get the subRoots values to know the node types of each subRoot
+	nodeTypes := make([]byte, len(subRoots))
+	for i := 0; i < len(subRoots); i++ {
+		if bytes.Equal(subRoots[i], t.emptyHash) {
+			nodeTypes[i] = PrefixValueEmpty
+			continue
+		}
+		v, err := wTx.Get(subRoots[i])
+		if err != nil {
+			return nil, err
+		}
+		nodeTypes[i] = v[0]
+	}
+
+	var newSubRoots [][]byte
+	for i := 0; i < len(subRoots); i += 2 {
+		if (bytes.Equal(subRoots[i], t.emptyHash) && bytes.Equal(subRoots[i+1], t.emptyHash)) ||
+			(nodeTypes[i] == PrefixValueLeaf && bytes.Equal(subRoots[i+1], t.emptyHash)) {
+			// when both sub nodes are empty, the parent is also empty
+			// or
+			// when 1st sub node is a leaf but the 2nd is empty, the
+			// leaf is used as 'parent'
+
+			newSubRoots = append(newSubRoots, subRoots[i])
+			continue
+		}
+		if bytes.Equal(subRoots[i], t.emptyHash) && nodeTypes[i+1] == PrefixValueLeaf {
+			// when 2nd sub node is a leaf but the 1st is empty,
+			// the leaf is used as 'parent'
+			newSubRoots = append(newSubRoots, subRoots[i+1])
+			continue
+		}
+
+		k, v, err := t.newIntermediate(subRoots[i], subRoots[i+1])
+		if err != nil {
+			return nil, err
+		}
+		// store k-v to db
+		if err = wTx.Set(k, v); err != nil {
+			return nil, err
+		}
+		newSubRoots = append(newSubRoots, k)
+	}
+
+	return t.upFromSubRoots(wTx, newSubRoots)
+}
+
+func (t *Tree) getSubRootsAtLevel(rTx db.ReadTx, root []byte, l int) ([][]byte, error) {
+	// go at level l and return each node key, where each node key is the
+	// subRoot of the subTree that starts there
+
+	var subRoots [][]byte
+	err := t.iterWithStop(rTx, root, 0, func(currLvl int, k, v []byte) bool {
+		if currLvl == l && !bytes.Equal(k, t.emptyHash) {
+			subRoots = append(subRoots, k)
+		}
+		if currLvl >= l {
+			return true // to stop the iter from going down
+		}
+		return false
+	})
+
+	return subRoots, err
+}
+
+func (t *Tree) addBatchInMemory(wTx db.WriteTx, keys, values [][]byte) ([]Invalid, error) {
+	vt, err := t.loadVT(wTx)
+	if err != nil {
+		return nil, err
+	}
+
 	invalids, err := vt.addBatch(keys, values)
 	if err != nil {
 		return nil, err
@@ -298,7 +531,12 @@ func (t *Tree) AddWithTx(wTx db.WriteTx, k, v []byte) error {
 	if !t.editable() {
 		return ErrSnapshotNotEditable
 	}
+	return t.addWithTx(wTx, k, v)
+}
 
+// warning: addWithTx does not use the Tree mutex, the mutex is responsibility
+// of the methods calling this method, and same with t.editable().
+func (t *Tree) addWithTx(wTx db.WriteTx, k, v []byte) error {
 	root, err := t.RootWithTx(wTx)
 	if err != nil {
 		return err
@@ -935,7 +1173,8 @@ func (t *Tree) SetRootWithTx(wTx db.WriteTx, root []byte) error {
 	// check that the root exists in the db
 	if !bytes.Equal(root, t.emptyHash) {
 		if _, err := wTx.Get(root); err == ErrKeyNotFound {
-			return fmt.Errorf("can not SetRoot with root %x, as it does not exist in the db", root)
+			return fmt.Errorf("can not SetRoot with root %x, as it"+
+				" does not exist in the db", root)
 		} else if err != nil {
 			return err
 		}
@@ -960,8 +1199,8 @@ func (t *Tree) Snapshot(fromRoot []byte) (*Tree, error) {
 	if !bytes.Equal(fromRoot, t.emptyHash) {
 		if _, err := rTx.Get(fromRoot); err == ErrKeyNotFound {
 			return nil,
-				fmt.Errorf("can not do a Snapshot with root %x, as it does not exist in the db",
-					fromRoot)
+				fmt.Errorf("can not do a Snapshot with root %x,"+
+					" as it does not exist in the db", fromRoot)
 		} else if err != nil {
 			return nil, err
 		}

tree_test.go

@@ -11,6 +11,7 @@ import (
 	qt "github.com/frankban/quicktest"
 	"go.vocdoni.io/dvote/db"
 	"go.vocdoni.io/dvote/db/badgerdb"
+	"go.vocdoni.io/dvote/db/pebbledb"
 )
 
 func checkRootBIString(c *qt.C, tree *Tree, expected string) {
@@ -23,11 +24,20 @@ func checkRootBIString(c *qt.C, tree *Tree, expected string) {
 func TestDBTx(t *testing.T) {
 	c := qt.New(t)
 
-	database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
+	dbBadger, err := badgerdb.New(db.Options{Path: c.TempDir()})
 	c.Assert(err, qt.IsNil)
+	testDBTx(c, dbBadger)
+
+	dbPebble, err := pebbledb.New(db.Options{Path: c.TempDir()})
+	c.Assert(err, qt.IsNil)
+	testDBTx(c, dbPebble)
+}
+
+func testDBTx(c *qt.C, database db.Database) {
 	wTx := database.WriteTx()
+	defer wTx.Discard()
 
-	_, err = wTx.Get([]byte("a"))
+	_, err := wTx.Get([]byte("a"))
 	c.Assert(err, qt.Equals, db.ErrKeyNotFound)
 
 	err = wTx.Set([]byte("a"), []byte("b"))
@@ -60,9 +70,10 @@ func TestAddTestVectors(t *testing.T) {
 }
 
 func testAdd(c *qt.C, hashFunc HashFunction, testVectors []string) {
-	database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
+	database, err := badgerdb.New(db.Options{Path: c.TempDir()})
 	c.Assert(err, qt.IsNil)
-	tree, err := NewTree(database, 256, hashFunc)
+	tree, err := NewTree(Config{Database: database, MaxLevels: 256,
+		HashFunction: hashFunc})
 	c.Assert(err, qt.IsNil)
 	defer tree.db.Close() //nolint:errcheck
 
@@ -92,9 +103,10 @@ func testAdd(c *qt.C, hashFunc HashFunction, testVectors []string) {
 
 func TestAddBatch(t *testing.T) {
 	c := qt.New(t)
-	database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
+	database, err := badgerdb.New(db.Options{Path: c.TempDir()})
 	c.Assert(err, qt.IsNil)
-	tree, err := NewTree(database, 256, HashFunctionPoseidon)
+	tree, err := NewTree(Config{Database: database, MaxLevels: 256,
+		HashFunction: HashFunctionPoseidon})
 	c.Assert(err, qt.IsNil)
 	defer tree.db.Close() //nolint:errcheck
 
@@ -110,9 +122,10 @@ func TestAddBatch(t *testing.T) {
 	checkRootBIString(c, tree,
 		"296519252211642170490407814696803112091039265640052570497930797516015811235")
 
-	database, err = badgerdb.New(badgerdb.Options{Path: c.TempDir()})
+	database, err = badgerdb.New(db.Options{Path: c.TempDir()})
 	c.Assert(err, qt.IsNil)
-	tree2, err := NewTree(database, 256, HashFunctionPoseidon)
+	tree2, err := NewTree(Config{Database: database, MaxLevels: 256,
+		HashFunction: HashFunctionPoseidon})
 	c.Assert(err, qt.IsNil)
 	defer tree2.db.Close() //nolint:errcheck
 
@@ -133,9 +146,10 @@ func TestAddBatch(t *testing.T) {
 
 func TestAddDifferentOrder(t *testing.T) {
 	c := qt.New(t)
-	database1, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
+	database1, err := badgerdb.New(db.Options{Path: c.TempDir()})
 	c.Assert(err, qt.IsNil)
-	tree1, err := NewTree(database1, 256, HashFunctionPoseidon)
+	tree1, err := NewTree(Config{Database: database1, MaxLevels: 256,
+		HashFunction: HashFunctionPoseidon})
 	c.Assert(err, qt.IsNil)
 	defer tree1.db.Close() //nolint:errcheck
 
@@ -148,9 +162,10 @@ func TestAddDifferentOrder(t *testing.T) {
 		}
 	}
 
-	database2, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
+	database2, err := badgerdb.New(db.Options{Path: c.TempDir()})
 	c.Assert(err, qt.IsNil)
-	tree2, err := NewTree(database2, 256, HashFunctionPoseidon)
+	tree2, err := NewTree(Config{Database: database2, MaxLevels: 256,
+		HashFunction: HashFunctionPoseidon})
 	c.Assert(err, qt.IsNil)
 	defer tree2.db.Close() //nolint:errcheck
 
@@ -173,9 +188,10 @@ func TestAddDifferentOrder(t *testing.T) {
 
 func TestAddRepeatedIndex(t *testing.T) {
 	c := qt.New(t)
-	database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
+	database, err := badgerdb.New(db.Options{Path: c.TempDir()})
 	c.Assert(err, qt.IsNil)
-	tree, err := NewTree(database, 256, HashFunctionPoseidon)
+	tree, err := NewTree(Config{Database: database, MaxLevels: 256,
+		HashFunction: HashFunctionPoseidon})
 	c.Assert(err, qt.IsNil)
 	defer tree.db.Close() //nolint:errcheck
 
@@ -191,9 +207,10 @@ func TestAddRepeatedIndex(t *testing.T) {
 
 func TestUpdate(t *testing.T) {
 	c := qt.New(t)
-	database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
+	database, err := badgerdb.New(db.Options{Path: c.TempDir()})
 	c.Assert(err, qt.IsNil)
-	tree, err := NewTree(database, 256, HashFunctionPoseidon)
+	tree, err := NewTree(Config{Database: database, MaxLevels: 256,
+		HashFunction: HashFunctionPoseidon})
 	c.Assert(err, qt.IsNil)
 	defer tree.db.Close() //nolint:errcheck
 
@@ -244,9 +261,10 @@ func TestUpdate(t *testing.T) {
 
 func TestAux(t *testing.T) { // TODO split in proper tests
 	c := qt.New(t)
-	database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
+	database, err := badgerdb.New(db.Options{Path: c.TempDir()})
 	c.Assert(err, qt.IsNil)
-	tree, err := NewTree(database, 256, HashFunctionPoseidon)
+	tree, err := NewTree(Config{Database: database, MaxLevels: 256,
+		HashFunction: HashFunctionPoseidon})
 	c.Assert(err, qt.IsNil)
 	defer tree.db.Close() //nolint:errcheck
 
@@ -283,9 +301,10 @@ func TestAux(t *testing.T) { // TODO split in proper tests
 
 func TestGet(t *testing.T) {
 	c := qt.New(t)
-	database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
+	database, err := badgerdb.New(db.Options{Path: c.TempDir()})
 	c.Assert(err, qt.IsNil)
-	tree, err := NewTree(database, 256, HashFunctionPoseidon)
+	tree, err := NewTree(Config{Database: database, MaxLevels: 256,
+		HashFunction: HashFunctionPoseidon})
 	c.Assert(err, qt.IsNil)
 	defer tree.db.Close() //nolint:errcheck
 
@@ -407,9 +426,10 @@ func TestPackAndUnpackSiblings(t *testing.T) {
 
 func TestGenProofAndVerify(t *testing.T) {
 	c := qt.New(t)
-	database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
+	database, err := badgerdb.New(db.Options{Path: c.TempDir()})
 	c.Assert(err, qt.IsNil)
-	tree, err := NewTree(database, 256, HashFunctionPoseidon)
+	tree, err := NewTree(Config{Database: database, MaxLevels: 256,
+		HashFunction: HashFunctionPoseidon})
 	c.Assert(err, qt.IsNil)
 	defer tree.db.Close() //nolint:errcheck
 
@@ -439,9 +459,10 @@ func TestGenProofAndVerify(t *testing.T) {
 
 func TestDumpAndImportDump(t *testing.T) {
 	c := qt.New(t)
-	database1, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
+	database1, err := badgerdb.New(db.Options{Path: c.TempDir()})
 	c.Assert(err, qt.IsNil)
-	tree1, err := NewTree(database1, 256, HashFunctionPoseidon)
+	tree1, err := NewTree(Config{Database: database1, MaxLevels: 256,
+		HashFunction: HashFunctionPoseidon})
 	c.Assert(err, qt.IsNil)
 	defer tree1.db.Close() //nolint:errcheck
 
@@ -457,9 +478,10 @@ func TestDumpAndImportDump(t *testing.T) {
 	e, err := tree1.Dump(nil)
 	c.Assert(err, qt.IsNil)
 
-	database2, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
+	database2, err := badgerdb.New(db.Options{Path: c.TempDir()})
 	c.Assert(err, qt.IsNil)
-	tree2, err := NewTree(database2, 256, HashFunctionPoseidon)
+	tree2, err := NewTree(Config{Database: database2, MaxLevels: 256,
+		HashFunction: HashFunctionPoseidon})
 	c.Assert(err, qt.IsNil)
 	defer tree2.db.Close() //nolint:errcheck
 	err = tree2.ImportDump(e)
@@ -476,9 +498,10 @@ func TestDumpAndImportDump(t *testing.T) {
 
 func TestRWMutex(t *testing.T) {
 	c := qt.New(t)
-	database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
+	database, err := badgerdb.New(db.Options{Path: c.TempDir()})
 	c.Assert(err, qt.IsNil)
-	tree, err := NewTree(database, 256, HashFunctionPoseidon)
+	tree, err := NewTree(Config{Database: database, MaxLevels: 256,
+		HashFunction: HashFunctionPoseidon})
 	c.Assert(err, qt.IsNil)
 	defer tree.db.Close() //nolint:errcheck
 
@@ -508,7 +531,7 @@ func TestRWMutex(t *testing.T) {
 // TODO UPDATE
 // func TestSetGetNLeafs(t *testing.T) {
 //         c := qt.New(t)
-//         database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
+//         database, err := badgerdb.New(db.Options{Path: c.TempDir()})
 //         c.Assert(err, qt.IsNil)
 //         tree, err := NewTree(database, 100, HashFunctionPoseidon)
 //         c.Assert(err, qt.IsNil)
@@ -559,14 +582,16 @@ func TestRWMutex(t *testing.T) {
 func TestAddBatchFullyUsed(t *testing.T) {
 	c := qt.New(t)
 
-	database1, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
+	database1, err := badgerdb.New(db.Options{Path: c.TempDir()})
 	c.Assert(err, qt.IsNil)
-	tree1, err := NewTree(database1, 4, HashFunctionPoseidon)
+	tree1, err := NewTree(Config{Database: database1, MaxLevels: 4,
+		HashFunction: HashFunctionPoseidon})
 	c.Assert(err, qt.IsNil)
 
-	database2, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
+	database2, err := badgerdb.New(db.Options{Path: c.TempDir()})
 	c.Assert(err, qt.IsNil)
-	tree2, err := NewTree(database2, 4, HashFunctionPoseidon)
+	tree2, err := NewTree(Config{Database: database2, MaxLevels: 4,
+		HashFunction: HashFunctionPoseidon})
 	c.Assert(err, qt.IsNil)
 
 	var keys, values [][]byte
@@ -618,9 +643,10 @@ func TestAddBatchFullyUsed(t *testing.T) {
 
 func TestSetRoot(t *testing.T) {
 	c := qt.New(t)
-	database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
+	database, err := badgerdb.New(db.Options{Path: c.TempDir()})
 	c.Assert(err, qt.IsNil)
-	tree, err := NewTree(database, 256, HashFunctionPoseidon)
+	tree, err := NewTree(Config{Database: database, MaxLevels: 256,
+		HashFunction: HashFunctionPoseidon})
 	c.Assert(err, qt.IsNil)
 
 	expectedRoot := "13742386369878513332697380582061714160370929283209286127733983161245560237407"
@@ -674,9 +700,10 @@ func TestSetRoot(t *testing.T) {
 
 func TestSnapshot(t *testing.T) {
 	c := qt.New(t)
-	database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
+	database, err := badgerdb.New(db.Options{Path: c.TempDir()})
 	c.Assert(err, qt.IsNil)
-	tree, err := NewTree(database, 256, HashFunctionPoseidon)
+	tree, err := NewTree(Config{Database: database, MaxLevels: 256,
+		HashFunction: HashFunctionPoseidon})
 	c.Assert(err, qt.IsNil)
 
 	// fill the tree
@@ -724,9 +751,10 @@ func TestSnapshot(t *testing.T) {
 func TestGetFromSnapshotExpectArboErrKeyNotFound(t *testing.T) {
 	c := qt.New(t)
 
-	database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
+	database, err := badgerdb.New(db.Options{Path: c.TempDir()})
 	c.Assert(err, qt.IsNil)
-	tree, err := NewTree(database, 256, HashFunctionPoseidon)
+	tree, err := NewTree(Config{Database: database, MaxLevels: 256,
+		HashFunction: HashFunctionPoseidon})
 	c.Assert(err, qt.IsNil)
 	defer tree.db.Close() //nolint:errcheck
 
@@ -744,11 +772,12 @@ func TestGetFromSnapshotExpectArboErrKeyNotFound(t *testing.T) {
 
 func TestKeyLen(t *testing.T) {
 	c := qt.New(t)
-	database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
+	database, err := badgerdb.New(db.Options{Path: c.TempDir()})
 	c.Assert(err, qt.IsNil)
 	// maxLevels is 100, keyPath length = ceil(maxLevels/8) = 13
 	maxLevels := 100
-	tree, err := NewTree(database, maxLevels, HashFunctionBlake2b)
+	tree, err := NewTree(Config{Database: database, MaxLevels: maxLevels,
+		HashFunction: HashFunctionBlake2b})
 	c.Assert(err, qt.IsNil)
 
 	// expect no errors when adding a key of only 4 bytes (when the
@@ -778,9 +807,10 @@ func TestKeyLen(t *testing.T) {
 	// expect errors when adding a key bigger than maximum capacity of the
 	// tree: ceil(maxLevels/8)
 	maxLevels = 32
-	database, err = badgerdb.New(badgerdb.Options{Path: c.TempDir()})
+	database, err = badgerdb.New(db.Options{Path: c.TempDir()})
 	c.Assert(err, qt.IsNil)
-	tree, err = NewTree(database, maxLevels, HashFunctionBlake2b)
+	tree, err = NewTree(Config{Database: database, MaxLevels: maxLevels,
+		HashFunction: HashFunctionBlake2b})
 	c.Assert(err, qt.IsNil)
 
 	maxKeyLen := int(math.Ceil(float64(maxLevels) / float64(8))) //nolint:gomnd
@@ -848,9 +878,10 @@ func TestKeyLen(t *testing.T) {
 func TestKeyLenBiggerThan32(t *testing.T) {
 	c := qt.New(t)
 	maxLevels := 264
-	database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
+	database, err := badgerdb.New(db.Options{Path: c.TempDir()})
 	c.Assert(err, qt.IsNil)
-	tree, err := NewTree(database, maxLevels, HashFunctionBlake2b)
+	tree, err := NewTree(Config{Database: database, MaxLevels: maxLevels,
+		HashFunction: HashFunctionBlake2b})
 	c.Assert(err, qt.IsNil)
 
 	bLen := 33
@@ -890,9 +921,10 @@ func BenchmarkAdd(b *testing.B) {
 
 func benchmarkAdd(b *testing.B, hashFunc HashFunction, ks, vs [][]byte) {
 	c := qt.New(b)
-	database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
+	database, err := badgerdb.New(db.Options{Path: c.TempDir()})
 	c.Assert(err, qt.IsNil)
-	tree, err := NewTree(database, 140, hashFunc)
+	tree, err := NewTree(Config{Database: database, MaxLevels: 140,
+		HashFunction: hashFunc})
 	c.Assert(err, qt.IsNil)
 	defer tree.db.Close() //nolint:errcheck
 

vt.go

@@ -37,7 +37,7 @@ type kv struct {
 	v       []byte
 }
 
-func (p *params) keysValuesToKvs(ks, vs [][]byte) ([]kv, []Invalid, error) {
+func keysValuesToKvs(maxLevels int, ks, vs [][]byte) ([]kv, []Invalid, error) {
 	if len(ks) != len(vs) {
 		return nil, nil, fmt.Errorf("len(keys)!=len(values) (%d!=%d)",
 			len(ks), len(vs))
@@ -45,7 +45,7 @@ func (p *params) keysValuesToKvs(ks, vs [][]byte) ([]kv, []Invalid, error) {
 	var invalids []Invalid
 	var kvs []kv
 	for i := 0; i < len(ks); i++ {
-		keyPath, err := keyPathFromKey(p.maxLevels, ks[i])
+		keyPath, err := keyPathFromKey(maxLevels, ks[i])
 		if err != nil {
 			invalids = append(invalids, Invalid{i, err})
 			continue
@@ -101,7 +101,7 @@ func (t *vt) addBatch(ks, vs [][]byte) ([]Invalid, error) {
 
 	l := int(math.Log2(float64(nCPU)))
 
-	kvs, invalids, err := t.params.keysValuesToKvs(ks, vs)
+	kvs, invalids, err := keysValuesToKvs(t.params.maxLevels, ks, vs)
 	if err != nil {
 		return invalids, err
 	}
@@ -271,7 +271,7 @@ func upFromNodes(ns []*node) (*node, error) {
 			// when both sub nodes are empty, the parent is also empty
 			// or
 			// when 1st sub node is a leaf but the 2nd is empty, the
-			// leaf is used as parent
+			// leaf is used as 'parent'
 			res = append(res, ns[i])
 			continue
 		}

vt_test.go

@@ -7,6 +7,7 @@ import (
 	"testing"
 
 	qt "github.com/frankban/quicktest"
+	"go.vocdoni.io/dvote/db"
 	"go.vocdoni.io/dvote/db/badgerdb"
 )
 
@@ -16,9 +17,10 @@ func testVirtualTree(c *qt.C, maxLevels int, keys, values [][]byte) {
 	c.Assert(len(keys), qt.Equals, len(values))
 
 	// normal tree, to have an expected root value
-	database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
+	database, err := badgerdb.New(db.Options{Path: c.TempDir()})
 	c.Assert(err, qt.IsNil)
-	tree, err := NewTree(database, maxLevels, HashFunctionSha256)
+	tree, err := NewTree(Config{Database: database, MaxLevels: maxLevels,
+		HashFunction: HashFunctionSha256})
 	c.Assert(err, qt.IsNil)
 	for i := 0; i < len(keys); i++ {
 		err := tree.Add(keys[i], values[i])
@@ -121,9 +123,10 @@ func TestVirtualTreeAddBatch(t *testing.T) {
 	}
 
 	// normal tree, to have an expected root value
-	database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
+	database, err := badgerdb.New(db.Options{Path: c.TempDir()})
 	c.Assert(err, qt.IsNil)
-	tree, err := NewTree(database, maxLevels, HashFunctionBlake2b)
+	tree, err := NewTree(Config{Database: database, MaxLevels: maxLevels,
+		HashFunction: HashFunctionBlake2b})
 	c.Assert(err, qt.IsNil)
 	for i := 0; i < len(keys); i++ {
 		err := tree.Add(keys[i], values[i])