Migrate kv db to go.vocdoni.io/dvote/db interface

Case tree empty, AddBatch was 10.95x times faster than without AddBatch
	nCPU: 4, nLeafs: 1024, hash: Poseidon, db: memory
	dbgStats(hash: 2.047k, dbGet: 1, dbPut: 2.049k)
Case tree not empty w/ few leafs, AddBatch was 7.28x times faster than without AddBatch
	nCPU: 4, nLeafs: 1024, hash: Poseidon, db: memory
	dbgStats(hash: 2.047k, dbGet: 198, dbPut: 2.049k)
Case tree not empty w/ enough leafs, AddBatch was 5.94x times faster than without AddBatch
	nCPU: 4, nLeafs: 1024, hash: Poseidon, db: memory
	dbgStats(hash: 2.047k, dbGet: 1.000k, dbPut: 2.049k)
Case tree not empty, AddBatch was 9.27x times faster than without AddBatch
	nCPU: 4, nLeafs: 4096, hash: Poseidon, db: memory
	dbgStats(hash: 8.191k, dbGet: 1.800k, dbPut: 8.193k)
Case tree not empty & unbalanced, AddBatch was 10.67x times faster than without AddBatch
	nCPU: 4, nLeafs: 4096, hash: Poseidon, db: memory
	dbgStats(hash: 10.409k, dbGet: 2.668k, dbPut: 10.861k)
TestAddBatchBench: nCPU: 4, nLeafs: 50000, hash: Blake2b, db: badgerdb
	Add loop:	10.10829114s
	AddBatch:	732.030263ms
		dbgStats(hash: 122.518k, dbGet: 1, dbPut: 122.520k)
TestDbgStats
	add in loop in emptyTree dbgStats(hash: 141.721k, dbGet: 134.596k, dbPut: 161.721k)
	addbatch caseEmptyTree dbgStats(hash: 24.402k, dbGet: 1, dbPut: 24.404k)
	addbatch caseNotEmptyTree dbgStats(hash: 26.868k, dbGet: 2.468k, dbPut: 26.872k)

addbatch_test.go

@@ -5,15 +5,13 @@ import (
 	"encoding/hex"
 	"fmt"
 	"math/big"
-	"os"
 	"runtime"
 	"sort"
 	"testing"
 	"time"
 
 	qt "github.com/frankban/quicktest"
-	"github.com/iden3/go-merkletree/db/leveldb"
-	"github.com/iden3/go-merkletree/db/memory"
+	"go.vocdoni.io/dvote/db"
 )
 
 var debug = true
@@ -39,13 +37,15 @@ func debugTime(descr string, time1, time2 time.Duration) {
 }
 
 func testInit(c *qt.C, n int) (*Tree, *Tree) {
-	tree1, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
+	database1, err := db.NewBadgerDB(c.TempDir())
+	c.Assert(err, qt.IsNil)
+	tree1, err := NewTree(database1, 100, HashFunctionPoseidon)
 	c.Assert(err, qt.IsNil)
-	defer tree1.db.Close()
 
-	tree2, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
+	database2, err := db.NewBadgerDB(c.TempDir())
+	c.Assert(err, qt.IsNil)
+	tree2, err := NewTree(database2, 100, HashFunctionPoseidon)
 	c.Assert(err, qt.IsNil)
-	defer tree2.db.Close()
 
 	bLen := HashFunctionPoseidon.Len()
 	// add the initial leafs to fill a bit the trees before calling the
@@ -68,9 +68,11 @@ func TestAddBatchTreeEmpty(t *testing.T) {
 
 	nLeafs := 1024
 
-	tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
+	database, err := db.NewBadgerDB(c.TempDir())
+	c.Assert(err, qt.IsNil)
+	tree, err := NewTree(database, 100, HashFunctionPoseidon)
 	c.Assert(err, qt.IsNil)
-	defer tree.db.Close()
+	defer tree.db.Close() //nolint:errcheck //nolint:errcheck
 
 	bLen := tree.HashFunction().Len()
 	start := time.Now()
@@ -83,9 +85,11 @@ func TestAddBatchTreeEmpty(t *testing.T) {
 	}
 	time1 := time.Since(start)
 
-	tree2, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
+	database2, err := db.NewBadgerDB(c.TempDir())
 	c.Assert(err, qt.IsNil)
-	defer tree2.db.Close()
+	tree2, err := NewTree(database2, 100, HashFunctionPoseidon)
+	c.Assert(err, qt.IsNil)
+	defer tree2.db.Close() //nolint:errcheck //nolint:errcheck
 	tree2.dbgInit()
 
 	var keys, values [][]byte
@@ -115,9 +119,11 @@ func TestAddBatchTreeEmptyNotPowerOf2(t *testing.T) {
 
 	nLeafs := 1027
 
-	tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
+	database, err := db.NewBadgerDB(c.TempDir())
+	c.Assert(err, qt.IsNil)
+	tree, err := NewTree(database, 100, HashFunctionPoseidon)
 	c.Assert(err, qt.IsNil)
-	defer tree.db.Close()
+	defer tree.db.Close() //nolint:errcheck
 
 	bLen := tree.HashFunction().Len()
 	for i := 0; i < nLeafs; i++ {
@@ -128,9 +134,11 @@ func TestAddBatchTreeEmptyNotPowerOf2(t *testing.T) {
 		}
 	}
 
-	tree2, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
+	database2, err := db.NewBadgerDB(c.TempDir())
 	c.Assert(err, qt.IsNil)
-	defer tree2.db.Close()
+	tree2, err := NewTree(database2, 100, HashFunctionPoseidon)
+	c.Assert(err, qt.IsNil)
+	defer tree2.db.Close() //nolint:errcheck
 
 	var keys, values [][]byte
 	for i := 0; i < nLeafs; i++ {
@@ -158,13 +166,17 @@ func randomBytes(n int) []byte {
 
 func TestAddBatchTreeEmptyTestVector(t *testing.T) {
 	c := qt.New(t)
-	tree1, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionBlake2b)
+	database1, err := db.NewBadgerDB(c.TempDir())
+	c.Assert(err, qt.IsNil)
+	tree1, err := NewTree(database1, 100, HashFunctionBlake2b)
 	c.Assert(err, qt.IsNil)
-	defer tree1.db.Close()
+	defer tree1.db.Close() //nolint:errcheck
 
-	tree2, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionBlake2b)
+	database2, err := db.NewBadgerDB(c.TempDir())
 	c.Assert(err, qt.IsNil)
-	defer tree2.db.Close()
+	tree2, err := NewTree(database2, 100, HashFunctionBlake2b)
+	c.Assert(err, qt.IsNil)
+	defer tree2.db.Close() //nolint:errcheck
 
 	// leafs in 2nd level subtrees: [ 6, 0, 1, 1]
 	testvectorKeys := []string{
@@ -194,13 +206,17 @@ func TestAddBatchTreeEmptyTestVector(t *testing.T) {
 	c.Check(tree2.Root(), qt.DeepEquals, tree1.Root())
 
 	// 2nd test vectors
-	tree1, err = NewTree(memory.NewMemoryStorage(), 100, HashFunctionBlake2b)
+	database1, err = db.NewBadgerDB(c.TempDir())
+	c.Assert(err, qt.IsNil)
+	tree1, err = NewTree(database1, 100, HashFunctionBlake2b)
 	c.Assert(err, qt.IsNil)
-	defer tree1.db.Close()
+	defer tree1.db.Close() //nolint:errcheck
 
-	tree2, err = NewTree(memory.NewMemoryStorage(), 100, HashFunctionBlake2b)
+	database2, err = db.NewBadgerDB(c.TempDir())
 	c.Assert(err, qt.IsNil)
-	defer tree2.db.Close()
+	tree2, err = NewTree(database2, 100, HashFunctionBlake2b)
+	c.Assert(err, qt.IsNil)
+	defer tree2.db.Close() //nolint:errcheck
 
 	testvectorKeys = []string{
 		"1c7c2265e368314ca58ed2e1f33a326f1220e234a566d55c3605439dbe411642",
@@ -239,13 +255,17 @@ func TestAddBatchTreeEmptyRandomKeys(t *testing.T) {
 
 	nLeafs := 8
 
-	tree1, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionBlake2b)
+	database1, err := db.NewBadgerDB(c.TempDir())
+	c.Assert(err, qt.IsNil)
+	tree1, err := NewTree(database1, 100, HashFunctionBlake2b)
 	c.Assert(err, qt.IsNil)
-	defer tree1.db.Close()
+	defer tree1.db.Close() //nolint:errcheck
 
-	tree2, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionBlake2b)
+	database2, err := db.NewBadgerDB(c.TempDir())
 	c.Assert(err, qt.IsNil)
-	defer tree2.db.Close()
+	tree2, err := NewTree(database2, 100, HashFunctionBlake2b)
+	c.Assert(err, qt.IsNil)
+	defer tree2.db.Close() //nolint:errcheck
 
 	var keys, values [][]byte
 	for i := 0; i < nLeafs; i++ {
@@ -585,9 +605,11 @@ func TestAddBatchNotEmptyUnbalanced(t *testing.T) {
 	}
 	time1 := time.Since(start)
 
-	tree2, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
+	database2, err := db.NewBadgerDB(c.TempDir())
+	c.Assert(err, qt.IsNil)
+	tree2, err := NewTree(database2, 100, HashFunctionPoseidon)
 	c.Assert(err, qt.IsNil)
-	defer tree2.db.Close()
+	defer tree2.db.Close() //nolint:errcheck
 	tree2.dbgInit()
 
 	var keys, values [][]byte
@@ -641,7 +663,7 @@ func TestFlp2(t *testing.T) {
 
 func TestAddBatchBench(t *testing.T) {
 	nLeafs := 50_000
-	printTestContext("TestAddBatchBench: ", nLeafs, "Blake2b", "leveldb")
+	printTestContext("TestAddBatchBench: ", nLeafs, "Blake2b", "badgerdb")
 
 	// prepare inputs
 	var ks, vs [][]byte
@@ -660,13 +682,11 @@ func TestAddBatchBench(t *testing.T) {
 func benchAdd(t *testing.T, ks, vs [][]byte) {
 	c := qt.New(t)
 
-	dbDir := t.TempDir()
-	defer os.RemoveAll(dbDir) //nolint:errcheck
-	// storage, err := pebble.NewPebbleStorage(dbDir, false)
-	storage, err := leveldb.NewLevelDbStorage(dbDir, false)
+	database, err := db.NewBadgerDB(c.TempDir())
 	c.Assert(err, qt.IsNil)
-	tree, err := NewTree(storage, 140, HashFunctionBlake2b)
+	tree, err := NewTree(database, 140, HashFunctionBlake2b)
 	c.Assert(err, qt.IsNil)
+	defer tree.db.Close() //nolint:errcheck
 
 	start := time.Now()
 	for i := 0; i < len(ks); i++ {
@@ -682,13 +702,11 @@ func benchAdd(t *testing.T, ks, vs [][]byte) {
 func benchAddBatch(t *testing.T, ks, vs [][]byte) {
 	c := qt.New(t)
 
-	dbDir := t.TempDir()
-	defer os.RemoveAll(dbDir) //nolint:errcheck
-	// storage, err := pebble.NewPebbleStorage(dbDir, false)
-	storage, err := leveldb.NewLevelDbStorage(dbDir, false)
+	database, err := db.NewBadgerDB(c.TempDir())
 	c.Assert(err, qt.IsNil)
-	tree, err := NewTree(storage, 140, HashFunctionBlake2b)
+	tree, err := NewTree(database, 140, HashFunctionBlake2b)
 	c.Assert(err, qt.IsNil)
+	defer tree.db.Close() //nolint:errcheck
 
 	tree.dbgInit()
 
@@ -717,9 +735,11 @@ func TestDbgStats(t *testing.T) {
 	}
 
 	// 1
-	tree1, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionBlake2b)
+	database1, err := db.NewBadgerDB(c.TempDir())
 	c.Assert(err, qt.IsNil)
-	defer tree1.db.Close()
+	tree1, err := NewTree(database1, 100, HashFunctionBlake2b)
+	c.Assert(err, qt.IsNil)
+	defer tree1.db.Close() //nolint:errcheck
 
 	tree1.dbgInit()
 
@@ -729,9 +749,11 @@ func TestDbgStats(t *testing.T) {
 	}
 
 	// 2
-	tree2, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionBlake2b)
+	database2, err := db.NewBadgerDB(c.TempDir())
+	c.Assert(err, qt.IsNil)
+	tree2, err := NewTree(database2, 100, HashFunctionBlake2b)
 	c.Assert(err, qt.IsNil)
-	defer tree2.db.Close()
+	defer tree2.db.Close() //nolint:errcheck
 
 	tree2.dbgInit()
 
@@ -740,9 +762,11 @@ func TestDbgStats(t *testing.T) {
 	c.Assert(len(invalids), qt.Equals, 0)
 
 	// 3
-	tree3, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionBlake2b)
+	database3, err := db.NewBadgerDB(c.TempDir())
 	c.Assert(err, qt.IsNil)
-	defer tree3.db.Close()
+	tree3, err := NewTree(database3, 100, HashFunctionBlake2b)
+	c.Assert(err, qt.IsNil)
+	defer tree3.db.Close() //nolint:errcheck
 
 	tree3.dbgInit()
 
@@ -773,9 +797,11 @@ func TestLoadVT(t *testing.T) {
 
 	nLeafs := 1024
 
-	tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
+	database, err := db.NewBadgerDB(c.TempDir())
+	c.Assert(err, qt.IsNil)
+	tree, err := NewTree(database, 100, HashFunctionPoseidon)
 	c.Assert(err, qt.IsNil)
-	defer tree.db.Close()
+	defer tree.db.Close() //nolint:errcheck
 
 	var keys, values [][]byte
 	for i := 0; i < nLeafs; i++ {
@@ -798,3 +824,6 @@ func TestLoadVT(t *testing.T) {
 }
 
 // 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
+// recreate the case

go.mod

@@ -5,6 +5,6 @@ go 1.16
 require (
 	github.com/frankban/quicktest v1.11.3
 	github.com/iden3/go-iden3-crypto v0.0.6-0.20210308142348-8f85683b2cef
-	github.com/iden3/go-merkletree v0.0.0-20210308143313-8b63ca866189
+	go.vocdoni.io/dvote v1.0.0
 	golang.org/x/crypto v0.0.0-20210317152858-513c2a44f670
 )

tree.go

@@ -13,6 +13,7 @@ package arbo
 
 import (
 	"bytes"
+	"crypto/sha256"
 	"encoding/binary"
 	"encoding/hex"
 	"fmt"
@@ -20,7 +21,7 @@ import (
 	"math"
 	"sync"
 
-	"github.com/iden3/go-merkletree/db"
+	"go.vocdoni.io/dvote/db"
 )
 
 const (
@@ -47,14 +48,17 @@ var (
 	dbKeyNLeafs = []byte("nleafs")
 	emptyValue  = []byte{0}
 
+	// ErrKeyNotFound is used when a key is not found in the db neither in
+	// the current db Batch.
+	ErrKeyNotFound = fmt.Errorf("key not found")
 	// ErrKeyAlreadyExists is used when trying to add a key as leaf to the
 	// tree that already exists.
 	ErrKeyAlreadyExists = fmt.Errorf("key already exists")
 	// ErrInvalidValuePrefix is used when going down into the tree, a value
 	// is read from the db and has an unrecognized prefix.
 	ErrInvalidValuePrefix = fmt.Errorf("invalid value prefix")
-	// ErrDBNoTx is used when trying to use Tree.dbPut but Tree.tx==nil
-	ErrDBNoTx = fmt.Errorf("dbPut error: no db Tx")
+	// ErrDBNoTx is used when trying to use Tree.dbPut but Tree.dbBatch==nil
+	ErrDBNoTx = fmt.Errorf("dbPut error: no db Batch")
 	// ErrMaxLevel indicates when going down into the tree, the max level is
 	// reached
 	ErrMaxLevel = fmt.Errorf("max level reached")
@@ -66,10 +70,11 @@ var (
 // Tree defines the struct that implements the MerkleTree functionalities
 type Tree struct {
 	sync.RWMutex
-	tx        db.Tx
-	db        db.Storage
-	maxLevels int
-	root      []byte
+	dbBatch     db.Batch
+	batchMemory kvMap // TODO TMP
+	db          db.Database
+	maxLevels   int
+	root        []byte
 
 	hashFunction HashFunction
 	// TODO in the methods that use it, check if emptyHash param is len>0
@@ -79,19 +84,33 @@ type Tree struct {
 	dbg *dbgStats
 }
 
-// NewTree returns a new Tree, if there is a Tree still in the given storage, it
+const bmSize = sha256.Size
+
+// TMP
+type kvMap map[[bmSize]byte]kv
+
+// Get retreives the value respective to a key from the KvMap
+func (m kvMap) Get(k []byte) ([]byte, bool) {
+	v, ok := m[sha256.Sum256(k)]
+	return v.v, ok
+}
+
+// Put stores a key and a value in the KvMap
+func (m kvMap) Put(k, v []byte) {
+	m[sha256.Sum256(k)] = kv{k: k, v: v}
+}
+
+// NewTree returns a new Tree, if there is a Tree still in the given database, it
 // will load it.
-func NewTree(storage db.Storage, maxLevels int, hash HashFunction) (*Tree, error) {
-	t := Tree{db: storage, maxLevels: maxLevels, hashFunction: hash}
+func NewTree(database db.Database, maxLevels int, hash HashFunction) (*Tree, error) {
+	t := Tree{db: database, maxLevels: maxLevels, hashFunction: hash}
 	t.emptyHash = make([]byte, t.hashFunction.Len()) // empty
 
 	root, err := t.dbGet(dbKeyRoot)
-	if err == db.ErrNotFound {
+	if err == ErrKeyNotFound {
 		// store new root 0
-		t.tx, err = t.db.NewTx()
-		if err != nil {
-			return nil, err
-		}
+		t.dbBatch = t.db.NewBatch()
+		t.batchMemory = make(map[[bmSize]byte]kv) // TODO TMP
 		t.root = t.emptyHash
 		if err = t.dbPut(dbKeyRoot, t.root); err != nil {
 			return nil, err
@@ -99,7 +118,7 @@ func NewTree(storage db.Storage, maxLevels int, hash HashFunction) (*Tree, error
 		if err = t.setNLeafs(0); err != nil {
 			return nil, err
 		}
-		if err = t.tx.Commit(); err != nil {
+		if err = t.dbBatch.Write(); err != nil {
 			return nil, err
 		}
 		return &t, err
@@ -147,10 +166,8 @@ func (t *Tree) AddBatch(keys, values [][]byte) ([]int, error) {
 	t.root = vt.root.h
 
 	// store pairs in db
-	t.tx, err = t.db.NewTx()
-	if err != nil {
-		return nil, err
-	}
+	t.dbBatch = t.db.NewBatch()
+	t.batchMemory = make(map[[bmSize]byte]kv) // TODO TMP
 	for i := 0; i < len(pairs); i++ {
 		if err := t.dbPut(pairs[i][0], pairs[i][1]); err != nil {
 			return nil, err
@@ -167,8 +184,8 @@ func (t *Tree) AddBatch(keys, values [][]byte) ([]int, error) {
 		return nil, err
 	}
 
-	// commit db tx
-	if err := t.tx.Commit(); err != nil {
+	// commit db dbBatch
+	if err := t.dbBatch.Write(); err != nil {
 		return nil, err
 	}
 	return invalids, nil
@@ -200,10 +217,8 @@ func (t *Tree) Add(k, v []byte) error {
 	defer t.Unlock()
 
 	var err error
-	t.tx, err = t.db.NewTx()
-	if err != nil {
-		return err
-	}
+	t.dbBatch = t.db.NewBatch()
+	t.batchMemory = make(map[[bmSize]byte]kv) // TODO TMP
 
 	// TODO check validity of key & value for Tree.hashFunction
 
@@ -219,7 +234,7 @@ func (t *Tree) Add(k, v []byte) error {
 	if err = t.incNLeafs(1); err != nil {
 		return err
 	}
-	return t.tx.Commit()
+	return t.dbBatch.Write()
 }
 
 func (t *Tree) add(fromLvl int, k, v []byte) error {
@@ -476,10 +491,8 @@ func (t *Tree) Update(k, v []byte) error {
 	defer t.Unlock()
 
 	var err error
-	t.tx, err = t.db.NewTx()
-	if err != nil {
-		return err
-	}
+	t.dbBatch = t.db.NewBatch()
+	t.batchMemory = make(map[[bmSize]byte]kv) // TODO TMP
 
 	keyPath := make([]byte, t.hashFunction.Len())
 	copy(keyPath[:], k)
@@ -507,7 +520,7 @@ func (t *Tree) Update(k, v []byte) error {
 	// go up to the root
 	if len(siblings) == 0 {
 		t.root = leafKey
-		return t.tx.Commit()
+		return t.dbBatch.Write()
 	}
 	root, err := t.up(leafKey, siblings, path, len(siblings)-1, 0)
 	if err != nil {
@@ -519,7 +532,7 @@ func (t *Tree) Update(k, v []byte) error {
 	if err := t.dbPut(dbKeyRoot, t.root); err != nil {
 		return err
 	}
-	return t.tx.Commit()
+	return t.dbBatch.Write()
 }
 
 // GenProof generates a MerkleTree proof for the given key. If the key exists in
@@ -681,11 +694,12 @@ func CheckProof(hashFunc HashFunction, k, v, root, packedSiblings []byte) (bool,
 }
 
 func (t *Tree) dbPut(k, v []byte) error {
-	if t.tx == nil {
+	if t.dbBatch == nil {
 		return ErrDBNoTx
 	}
 	t.dbg.incDbPut()
-	return t.tx.Put(k, v)
+	t.batchMemory.Put(k, v) // TODO TMP
+	return t.dbBatch.Put(k, v)
 }
 
 func (t *Tree) dbGet(k []byte) ([]byte, error) {
@@ -699,14 +713,20 @@ func (t *Tree) dbGet(k []byte) ([]byte, error) {
 	if err == nil {
 		return v, nil
 	}
-	if t.tx != nil {
-		return t.tx.Get(k)
+	if t.dbBatch != nil {
+		// TODO TMP
+		v, ok := t.batchMemory.Get(k)
+		if !ok {
+			return nil, ErrKeyNotFound
+		}
+		// /TMP
+		return v, nil
 	}
-	return nil, db.ErrNotFound
+	return nil, ErrKeyNotFound
 }
 
-// Warning: should be called with a Tree.tx created, and with a Tree.tx.Commit
-// after the setNLeafs call.
+// Warning: should be called with a Tree.dbBatch created, and with a
+// Tree.dbBatch.Write after the incNLeafs call.
 func (t *Tree) incNLeafs(nLeafs int) error {
 	oldNLeafs, err := t.GetNLeafs()
 	if err != nil {
@@ -716,8 +736,8 @@ func (t *Tree) incNLeafs(nLeafs int) error {
 	return t.setNLeafs(newNLeafs)
 }
 
-// Warning: should be called with a Tree.tx created, and with a Tree.tx.Commit
-// after the setNLeafs call.
+// Warning: should be called with a Tree.dbBatch created, and with a
+// Tree.dbBatch.Write after the setNLeafs call.
 func (t *Tree) setNLeafs(nLeafs int) error {
 	b := make([]byte, 8)
 	binary.LittleEndian.PutUint64(b, uint64(nLeafs))

tree_test.go

@@ -7,9 +7,19 @@ import (
 	"time"
 
 	qt "github.com/frankban/quicktest"
-	"github.com/iden3/go-merkletree/db/memory"
+	"go.vocdoni.io/dvote/db"
 )
 
+func TestDBBatch(t *testing.T) {
+	c := qt.New(t)
+
+	database, err := db.NewBadgerDB(c.TempDir())
+	c.Assert(err, qt.IsNil)
+	b := database.NewBatch()
+	err = b.Put([]byte("a"), []byte("b"))
+	c.Assert(err, qt.IsNil)
+}
+
 func TestAddTestVectors(t *testing.T) {
 	c := qt.New(t)
 
@@ -32,9 +42,11 @@ func TestAddTestVectors(t *testing.T) {
 }
 
 func testAdd(c *qt.C, hashFunc HashFunction, testVectors []string) {
-	tree, err := NewTree(memory.NewMemoryStorage(), 10, hashFunc)
+	database, err := db.NewBadgerDB(c.TempDir())
 	c.Assert(err, qt.IsNil)
-	defer tree.db.Close()
+	tree, err := NewTree(database, 10, hashFunc)
+	c.Assert(err, qt.IsNil)
+	defer tree.db.Close() //nolint:errcheck
 
 	c.Check(hex.EncodeToString(tree.Root()), qt.Equals, testVectors[0])
 
@@ -63,9 +75,11 @@ func testAdd(c *qt.C, hashFunc HashFunction, testVectors []string) {
 
 func TestAddBatch(t *testing.T) {
 	c := qt.New(t)
-	tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
+	database, err := db.NewBadgerDB(c.TempDir())
+	c.Assert(err, qt.IsNil)
+	tree, err := NewTree(database, 100, HashFunctionPoseidon)
 	c.Assert(err, qt.IsNil)
-	defer tree.db.Close()
+	defer tree.db.Close() //nolint:errcheck
 
 	bLen := tree.HashFunction().Len()
 	for i := 0; i < 1000; i++ {
@@ -80,9 +94,11 @@ func TestAddBatch(t *testing.T) {
 	c.Check(rootBI.String(), qt.Equals,
 		"296519252211642170490407814696803112091039265640052570497930797516015811235")
 
-	tree2, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
+	database, err = db.NewBadgerDB(c.TempDir())
+	c.Assert(err, qt.IsNil)
+	tree2, err := NewTree(database, 100, HashFunctionPoseidon)
 	c.Assert(err, qt.IsNil)
-	defer tree2.db.Close()
+	defer tree2.db.Close() //nolint:errcheck
 
 	var keys, values [][]byte
 	for i := 0; i < 1000; i++ {
@@ -102,9 +118,11 @@ func TestAddBatch(t *testing.T) {
 
 func TestAddDifferentOrder(t *testing.T) {
 	c := qt.New(t)
-	tree1, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
+	database1, err := db.NewBadgerDB(c.TempDir())
 	c.Assert(err, qt.IsNil)
-	defer tree1.db.Close()
+	tree1, err := NewTree(database1, 100, HashFunctionPoseidon)
+	c.Assert(err, qt.IsNil)
+	defer tree1.db.Close() //nolint:errcheck
 
 	bLen := tree1.HashFunction().Len()
 	for i := 0; i < 16; i++ {
@@ -115,9 +133,11 @@ func TestAddDifferentOrder(t *testing.T) {
 		}
 	}
 
-	tree2, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
+	database2, err := db.NewBadgerDB(c.TempDir())
+	c.Assert(err, qt.IsNil)
+	tree2, err := NewTree(database2, 100, HashFunctionPoseidon)
 	c.Assert(err, qt.IsNil)
-	defer tree2.db.Close()
+	defer tree2.db.Close() //nolint:errcheck
 
 	for i := 16 - 1; i >= 0; i-- {
 		k := BigIntToBytes(bLen, big.NewInt(int64(i)))
@@ -134,9 +154,11 @@ func TestAddDifferentOrder(t *testing.T) {
 
 func TestAddRepeatedIndex(t *testing.T) {
 	c := qt.New(t)
-	tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
+	database, err := db.NewBadgerDB(c.TempDir())
+	c.Assert(err, qt.IsNil)
+	tree, err := NewTree(database, 100, HashFunctionPoseidon)
 	c.Assert(err, qt.IsNil)
-	defer tree.db.Close()
+	defer tree.db.Close() //nolint:errcheck
 
 	bLen := tree.HashFunction().Len()
 	k := BigIntToBytes(bLen, big.NewInt(int64(3)))
@@ -150,9 +172,11 @@ func TestAddRepeatedIndex(t *testing.T) {
 
 func TestUpdate(t *testing.T) {
 	c := qt.New(t)
-	tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
+	database, err := db.NewBadgerDB(c.TempDir())
 	c.Assert(err, qt.IsNil)
-	defer tree.db.Close()
+	tree, err := NewTree(database, 100, HashFunctionPoseidon)
+	c.Assert(err, qt.IsNil)
+	defer tree.db.Close() //nolint:errcheck
 
 	bLen := tree.HashFunction().Len()
 	k := BigIntToBytes(bLen, big.NewInt(int64(20)))
@@ -201,9 +225,11 @@ func TestUpdate(t *testing.T) {
 
 func TestAux(t *testing.T) { // TODO split in proper tests
 	c := qt.New(t)
-	tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
+	database, err := db.NewBadgerDB(c.TempDir())
+	c.Assert(err, qt.IsNil)
+	tree, err := NewTree(database, 100, HashFunctionPoseidon)
 	c.Assert(err, qt.IsNil)
-	defer tree.db.Close()
+	defer tree.db.Close() //nolint:errcheck
 
 	bLen := tree.HashFunction().Len()
 	k := BigIntToBytes(bLen, big.NewInt(int64(1)))
@@ -238,9 +264,11 @@ func TestAux(t *testing.T) { // TODO split in proper tests
 
 func TestGet(t *testing.T) {
 	c := qt.New(t)
-	tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
+	database, err := db.NewBadgerDB(c.TempDir())
+	c.Assert(err, qt.IsNil)
+	tree, err := NewTree(database, 100, HashFunctionPoseidon)
 	c.Assert(err, qt.IsNil)
-	defer tree.db.Close()
+	defer tree.db.Close() //nolint:errcheck
 
 	bLen := tree.HashFunction().Len()
 	for i := 0; i < 10; i++ {
@@ -260,9 +288,11 @@ func TestGet(t *testing.T) {
 
 func TestGenProofAndVerify(t *testing.T) {
 	c := qt.New(t)
-	tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
+	database, err := db.NewBadgerDB(c.TempDir())
 	c.Assert(err, qt.IsNil)
-	defer tree.db.Close()
+	tree, err := NewTree(database, 100, HashFunctionPoseidon)
+	c.Assert(err, qt.IsNil)
+	defer tree.db.Close() //nolint:errcheck
 
 	bLen := tree.HashFunction().Len()
 	for i := 0; i < 10; i++ {
@@ -286,9 +316,11 @@ func TestGenProofAndVerify(t *testing.T) {
 
 func TestDumpAndImportDump(t *testing.T) {
 	c := qt.New(t)
-	tree1, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
+	database1, err := db.NewBadgerDB(c.TempDir())
+	c.Assert(err, qt.IsNil)
+	tree1, err := NewTree(database1, 100, HashFunctionPoseidon)
 	c.Assert(err, qt.IsNil)
-	defer tree1.db.Close()
+	defer tree1.db.Close() //nolint:errcheck
 
 	bLen := tree1.HashFunction().Len()
 	for i := 0; i < 16; i++ {
@@ -302,9 +334,11 @@ func TestDumpAndImportDump(t *testing.T) {
 	e, err := tree1.Dump(nil)
 	c.Assert(err, qt.IsNil)
 
-	tree2, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
+	database2, err := db.NewBadgerDB(c.TempDir())
+	c.Assert(err, qt.IsNil)
+	tree2, err := NewTree(database2, 100, HashFunctionPoseidon)
 	c.Assert(err, qt.IsNil)
-	defer tree2.db.Close()
+	defer tree2.db.Close() //nolint:errcheck
 	err = tree2.ImportDump(e)
 	c.Assert(err, qt.IsNil)
 	c.Check(tree2.Root(), qt.DeepEquals, tree1.Root())
@@ -314,9 +348,11 @@ func TestDumpAndImportDump(t *testing.T) {
 
 func TestRWMutex(t *testing.T) {
 	c := qt.New(t)
-	tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
+	database, err := db.NewBadgerDB(c.TempDir())
 	c.Assert(err, qt.IsNil)
-	defer tree.db.Close()
+	tree, err := NewTree(database, 100, HashFunctionPoseidon)
+	c.Assert(err, qt.IsNil)
+	defer tree.db.Close() //nolint:errcheck
 
 	bLen := tree.HashFunction().Len()
 	var keys, values [][]byte
@@ -343,17 +379,18 @@ func TestRWMutex(t *testing.T) {
 
 func TestSetGetNLeafs(t *testing.T) {
 	c := qt.New(t)
-	tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
+	database, err := db.NewBadgerDB(c.TempDir())
+	c.Assert(err, qt.IsNil)
+	tree, err := NewTree(database, 100, HashFunctionPoseidon)
 	c.Assert(err, qt.IsNil)
 
 	// 0
-	tree.tx, err = tree.db.NewTx()
-	c.Assert(err, qt.IsNil)
+	tree.dbBatch = tree.db.NewBatch()
 
 	err = tree.setNLeafs(0)
 	c.Assert(err, qt.IsNil)
 
-	err = tree.tx.Commit()
+	err = tree.dbBatch.Write()
 	c.Assert(err, qt.IsNil)
 
 	n, err := tree.GetNLeafs()
@@ -361,13 +398,12 @@ func TestSetGetNLeafs(t *testing.T) {
 	c.Assert(n, qt.Equals, 0)
 
 	// 1024
-	tree.tx, err = tree.db.NewTx()
-	c.Assert(err, qt.IsNil)
+	tree.dbBatch = tree.db.NewBatch()
 
 	err = tree.setNLeafs(1024)
 	c.Assert(err, qt.IsNil)
 
-	err = tree.tx.Commit()
+	err = tree.dbBatch.Write()
 	c.Assert(err, qt.IsNil)
 
 	n, err = tree.GetNLeafs()
@@ -375,8 +411,7 @@ func TestSetGetNLeafs(t *testing.T) {
 	c.Assert(n, qt.Equals, 1024)
 
 	// 2**64 -1
-	tree.tx, err = tree.db.NewTx()
-	c.Assert(err, qt.IsNil)
+	tree.dbBatch = tree.db.NewBatch()
 
 	maxUint := ^uint(0)
 	maxInt := int(maxUint >> 1)
@@ -384,7 +419,7 @@ func TestSetGetNLeafs(t *testing.T) {
 	err = tree.setNLeafs(maxInt)
 	c.Assert(err, qt.IsNil)
 
-	err = tree.tx.Commit()
+	err = tree.dbBatch.Write()
 	c.Assert(err, qt.IsNil)
 
 	n, err = tree.GetNLeafs()
@@ -413,9 +448,11 @@ func BenchmarkAdd(b *testing.B) {
 
 func benchmarkAdd(b *testing.B, hashFunc HashFunction, ks, vs [][]byte) {
 	c := qt.New(b)
-	tree, err := NewTree(memory.NewMemoryStorage(), 140, hashFunc)
+	database, err := db.NewBadgerDB(c.TempDir())
+	c.Assert(err, qt.IsNil)
+	tree, err := NewTree(database, 140, hashFunc)
 	c.Assert(err, qt.IsNil)
-	defer tree.db.Close()
+	defer tree.db.Close() //nolint:errcheck
 
 	for i := 0; i < len(ks); i++ {
 		if err := tree.Add(ks[i], vs[i]); err != nil {

vt_test.go

@@ -6,7 +6,7 @@ import (
 	"testing"
 
 	qt "github.com/frankban/quicktest"
-	"github.com/iden3/go-merkletree/db/memory"
+	"go.vocdoni.io/dvote/db"
 )
 
 func TestVirtualTreeTestVectors(t *testing.T) {
@@ -76,7 +76,9 @@ 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
-	tree, err := NewTree(memory.NewMemoryStorage(), maxLevels, HashFunctionSha256)
+	database, err := db.NewBadgerDB(c.TempDir())
+	c.Assert(err, qt.IsNil)
+	tree, err := NewTree(database, maxLevels, HashFunctionSha256)
 	c.Assert(err, qt.IsNil)
 	for i := 0; i < len(keys); i++ {
 		err := tree.Add(keys[i], values[i])
@@ -113,7 +115,9 @@ func TestVirtualTreeAddBatch(t *testing.T) {
 	}
 
 	// normal tree, to have an expected root value
-	tree, err := NewTree(memory.NewMemoryStorage(), maxLevels, HashFunctionBlake2b)
+	database, err := db.NewBadgerDB(c.TempDir())
+	c.Assert(err, qt.IsNil)
+	tree, err := NewTree(database, maxLevels, HashFunctionBlake2b)
 	c.Assert(err, qt.IsNil)
 	for i := 0; i < len(keys); i++ {
 		err := tree.Add(keys[i], values[i])