Add thresholdNLeafs configurable

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.

addbatch_test.go

@@ -41,12 +41,14 @@ func debugTime(descr string, time1, time2 time.Duration) {
 func testInit(c *qt.C, n int) (*Tree, *Tree) {
 	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(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()
@@ -72,7 +74,8 @@ func TestAddBatchTreeEmpty(t *testing.T) {
 
 	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
 
@@ -95,7 +98,8 @@ func TestAddBatchTreeEmpty(t *testing.T) {
 
 	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()
@@ -122,7 +126,8 @@ func TestAddBatchTreeEmptyNotPowerOf2(t *testing.T) {
 
 	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
 
@@ -137,7 +142,8 @@ func TestAddBatchTreeEmptyNotPowerOf2(t *testing.T) {
 
 	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
 
@@ -169,13 +175,15 @@ func TestAddBatchTestVector1(t *testing.T) {
 	c := qt.New(t)
 	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(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
 
@@ -209,13 +217,15 @@ func TestAddBatchTestVector1(t *testing.T) {
 	// 2nd test vectors
 	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(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
 
@@ -257,13 +267,15 @@ func TestAddBatchTestVector2(t *testing.T) {
 
 	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(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
 
@@ -302,13 +314,15 @@ func TestAddBatchTestVector3(t *testing.T) {
 
 	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(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
 
@@ -351,13 +365,15 @@ func TestAddBatchTreeEmptyRandomKeys(t *testing.T) {
 
 	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(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
 
@@ -701,7 +717,8 @@ func TestAddBatchNotEmptyUnbalanced(t *testing.T) {
 
 	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()
@@ -778,7 +795,8 @@ func benchAdd(t *testing.T, ks, vs [][]byte) {
 
 	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
 
@@ -798,7 +816,8 @@ func benchAddBatch(t *testing.T, ks, vs [][]byte) {
 
 	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
 
@@ -831,7 +850,8 @@ func TestDbgStats(t *testing.T) {
 	// 1
 	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
 
@@ -845,7 +865,8 @@ func TestDbgStats(t *testing.T) {
 	// 2
 	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
 
@@ -858,7 +879,8 @@ func TestDbgStats(t *testing.T) {
 	// 3
 	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
 
@@ -893,7 +915,8 @@ func TestLoadVT(t *testing.T) {
 
 	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
 
@@ -929,7 +952,8 @@ func TestAddKeysWithEmptyValues(t *testing.T) {
 
 	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
 
@@ -950,7 +974,8 @@ func TestAddKeysWithEmptyValues(t *testing.T) {
 
 	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()
@@ -964,7 +989,8 @@ func TestAddKeysWithEmptyValues(t *testing.T) {
 	// use tree3 to add nil value array
 	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
 
@@ -1008,35 +1034,38 @@ func TestAddKeysWithEmptyValues(t *testing.T) {
 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(database1, 256, HashFunctionBlake2b)
+	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(database2, 256, HashFunctionBlake2b)
+	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(database3, 256, HashFunctionBlake2b)
+	tree3, err := NewTree(Config{database3, 256, testThresholdNLeafs,
+		HashFunctionBlake2b})
 	c.Assert(err, qt.IsNil)
 	defer tree3.db.Close() //nolint:errcheck
 
-	// customize thresholdNLeafs for the test
-	thresholdNLeafs = 1024
-
 	var keys, values [][]byte
-	for i := 0; i < 3*thresholdNLeafs; i++ {
+	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 < thresholdNLeafs+1 {
+		if i < testThresholdNLeafs+1 {
 			if err := tree2.Add(k, v); err != nil {
 				t.Fatal(err)
 			}
@@ -1046,7 +1075,7 @@ func TestAddBatchThresholdInDisk(t *testing.T) {
 		values = append(values, v)
 	}
 
-	invalids, err := tree2.AddBatch(keys[thresholdNLeafs+1:], values[thresholdNLeafs+1:])
+	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

circomproofs_test.go

@@ -14,7 +14,8 @@ func TestCircomVerifierProof(t *testing.T) {
 	c := qt.New(t)
 	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
 

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

@@ -90,12 +90,14 @@ func TestReadTreeDBG(t *testing.T) {
 
 	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(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

@@ -16,7 +16,8 @@ func TestGenerator(t *testing.T) {
 	c := qt.New(t)
 	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

@@ -46,11 +46,11 @@ const (
 )
 
 var (
-	// thresholdNLeafs defines the threshold number of leafs in the tree
-	// that determines if AddBatch will work in memory or in disk. Is
-	// defined as a var in order to have the ability to modify it for
-	// testing purposes.
-	thresholdNLeafs = 1024 // TODO define a reasonable value
+	// 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")
@@ -85,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
@@ -97,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
 	}
@@ -117,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)
@@ -220,7 +234,7 @@ func (t *Tree) AddBatchWithTx(wTx db.WriteTx, keys, values [][]byte) ([]Invalid,
 	if err != nil {
 		return nil, err
 	}
-	if nLeafs > thresholdNLeafs {
+	if nLeafs > t.thresholdNLeafs {
 		return t.addBatchInDisk(wTx, keys, values)
 	}
 	return t.addBatchInMemory(wTx, keys, values)
@@ -231,7 +245,7 @@ func (t *Tree) addBatchInDisk(wTx db.WriteTx, keys, values [][]byte) ([]Invalid,
 	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 {
+			if err := t.addWithTx(wTx, keys[i], values[i]); err != nil {
 				invalids = append(invalids, Invalid{i, err})
 			}
 		}
@@ -265,7 +279,8 @@ func (t *Tree) addBatchInDisk(wTx db.WriteTx, keys, values [][]byte) ([]Invalid,
 			// (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 {
+				if newRoot, err := t.add(wTx, root, 0,
+					buckets[i][j].k, buckets[i][j].v); err == nil {
 					inserted = j
 					root = newRoot
 					break
@@ -308,10 +323,15 @@ func (t *Tree) addBatchInDisk(wTx db.WriteTx, keys, values [][]byte) ([]Invalid,
 			// are done, iter over the cpuWTxs and copy their
 			// content into the main wTx
 			for j := 0; j < len(buckets[cpu]); j++ {
-				subRoots[cpu], err = t.add(txs[cpu], subRoots[cpu], l, buckets[cpu][j].k, buckets[cpu][j].v)
+				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})
+					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)
@@ -346,6 +366,7 @@ func (t *Tree) addBatchInDisk(wTx db.WriteTx, keys, values [][]byte) ([]Invalid,
 
 	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.
@@ -495,7 +516,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
@@ -1132,7 +1158,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
 		}
@@ -1157,8 +1184,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

@@ -72,7 +72,8 @@ func TestAddTestVectors(t *testing.T) {
 func testAdd(c *qt.C, hashFunc HashFunction, testVectors []string) {
 	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
 
@@ -104,7 +105,8 @@ func TestAddBatch(t *testing.T) {
 	c := qt.New(t)
 	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
 
@@ -122,7 +124,8 @@ func TestAddBatch(t *testing.T) {
 
 	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
 
@@ -145,7 +148,8 @@ func TestAddDifferentOrder(t *testing.T) {
 	c := qt.New(t)
 	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
 
@@ -160,7 +164,8 @@ func TestAddDifferentOrder(t *testing.T) {
 
 	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
 
@@ -185,7 +190,8 @@ func TestAddRepeatedIndex(t *testing.T) {
 	c := qt.New(t)
 	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
 
@@ -203,7 +209,8 @@ func TestUpdate(t *testing.T) {
 	c := qt.New(t)
 	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
 
@@ -256,7 +263,8 @@ func TestAux(t *testing.T) { // TODO split in proper tests
 	c := qt.New(t)
 	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
 
@@ -295,7 +303,8 @@ func TestGet(t *testing.T) {
 	c := qt.New(t)
 	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
 
@@ -319,7 +328,8 @@ func TestGenProofAndVerify(t *testing.T) {
 	c := qt.New(t)
 	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
 
@@ -351,7 +361,8 @@ func TestDumpAndImportDump(t *testing.T) {
 	c := qt.New(t)
 	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
 
@@ -369,7 +380,8 @@ func TestDumpAndImportDump(t *testing.T) {
 
 	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)
@@ -388,7 +400,8 @@ func TestRWMutex(t *testing.T) {
 	c := qt.New(t)
 	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
 
@@ -471,12 +484,14 @@ func TestAddBatchFullyUsed(t *testing.T) {
 
 	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(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
@@ -530,7 +545,8 @@ func TestSetRoot(t *testing.T) {
 	c := qt.New(t)
 	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"
@@ -586,7 +602,8 @@ func TestSnapshot(t *testing.T) {
 	c := qt.New(t)
 	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
@@ -636,7 +653,8 @@ func TestGetFromSnapshotExpectArboErrKeyNotFound(t *testing.T) {
 
 	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
 
@@ -658,7 +676,8 @@ func TestKeyLen(t *testing.T) {
 	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
@@ -690,7 +709,8 @@ func TestKeyLen(t *testing.T) {
 	maxLevels = 32
 	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
@@ -760,7 +780,8 @@ func TestKeyLenBiggerThan32(t *testing.T) {
 	maxLevels := 264
 	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
@@ -802,7 +823,8 @@ func benchmarkAdd(b *testing.B, hashFunc HashFunction, ks, vs [][]byte) {
 	c := qt.New(b)
 	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_test.go

@@ -19,7 +19,8 @@ func testVirtualTree(c *qt.C, maxLevels int, keys, values [][]byte) {
 	// normal tree, to have an expected root value
 	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])
@@ -124,7 +125,8 @@ func TestVirtualTreeAddBatch(t *testing.T) {
 	// normal tree, to have an expected root value
 	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])