AddBatch tests abstract code reusage

addbatch.go

@@ -179,10 +179,7 @@ func (t *Tree) AddBatchOpt(keys, values [][]byte) ([]int, error) {
 			return nil, err
 		}
 
-		if err = t.finalizeAddBatch(); err != nil {
-			return nil, err
-		}
-		return invalids, nil
+		return t.finalizeAddBatch(len(keys), invalids)
 	}
 
 	// CASE B: if nLeafs<nBuckets
@@ -204,10 +201,7 @@ func (t *Tree) AddBatchOpt(keys, values [][]byte) ([]int, error) {
 			}
 		}
 
-		if err = t.finalizeAddBatch(); err != nil {
-			return nil, err
-		}
-		return invalids, nil
+		return t.finalizeAddBatch(len(keys), invalids)
 	}
 
 	keysAtL, err := t.getKeysAtLevel(l + 1)
@@ -225,10 +219,7 @@ func (t *Tree) AddBatchOpt(keys, values [][]byte) ([]int, error) {
 			return nil, err
 		}
 
-		if err = t.finalizeAddBatch(); err != nil {
-			return nil, err
-		}
-		return invalids, nil
+		return t.finalizeAddBatch(len(keys), invalids)
 	}
 
 	// CASE E
@@ -266,27 +257,28 @@ func (t *Tree) AddBatchOpt(keys, values [][]byte) ([]int, error) {
 		}
 		invalids = append(invalids, invalidsCaseD...)
 
-		if err = t.finalizeAddBatch(); err != nil {
-			return nil, err
-		}
-		return invalids, nil
+		return t.finalizeAddBatch(len(keys), invalids)
 	}
 
-	// TODO update NLeafs from DB
-
 	return nil, fmt.Errorf("UNIMPLEMENTED")
 }
 
-func (t *Tree) finalizeAddBatch() error {
+func (t *Tree) finalizeAddBatch(nKeys int, invalids []int) ([]int, error) {
 	// store root to db
 	if err := t.tx.Put(dbKeyRoot, t.root); err != nil {
-		return err
+		return nil, err
 	}
+
+	// update nLeafs
+	if err := t.incNLeafs(nKeys - len(invalids)); err != nil {
+		return nil, err
+	}
+
 	// commit db tx
 	if err := t.tx.Commit(); err != nil {
-		return err
+		return nil, err
 	}
-	return nil
+	return invalids, nil
 }
 
 func (t *Tree) caseA(nCPU int, kvs []kv) ([]int, error) {
@@ -395,9 +387,8 @@ func (t *Tree) caseC(nCPU, l int, keysAtL [][]byte, kvs []kv) ([]int, error) {
 	// add the key-values that have not been used yet
 	var invalids []int
 	for i := 0; i < len(excedents); i++ {
-		// Add until the level L
 		if err = t.add(0, excedents[i].k, excedents[i].v); err != nil {
-			invalids = append(invalids, excedents[i].pos) // TODO WIP
+			invalids = append(invalids, excedents[i].pos)
 		}
 	}
 	return invalids, nil

addbatch_test.go

@@ -11,54 +11,34 @@ import (
 	"github.com/iden3/go-merkletree/db/memory"
 )
 
-func TestBatchAux(t *testing.T) { // TODO TMP this test will be delted
-	c := qt.New(t)
-
-	nLeafs := 16
-
-	tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
+func testInit(c *qt.C, n int) (*Tree, *Tree) {
+	tree1, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
 	c.Assert(err, qt.IsNil)
-	defer tree.db.Close()
-
-	start := time.Now()
-	for i := 0; i < nLeafs; i++ {
-		k := BigIntToBytes(big.NewInt(int64(i)))
-		v := BigIntToBytes(big.NewInt(int64(i * 2)))
-		if err := tree.Add(k, v); err != nil {
-			t.Fatal(err)
-		}
-	}
-	fmt.Println(time.Since(start))
+	defer tree1.db.Close()
 
 	tree2, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
 	c.Assert(err, qt.IsNil)
 	defer tree2.db.Close()
 
-	for i := 0; i < 8; i++ {
+	// add the initial leafs to fill a bit the trees before calling the
+	// AddBatch method
+	for i := 0; i < n; i++ {
 		k := BigIntToBytes(big.NewInt(int64(i)))
 		v := BigIntToBytes(big.NewInt(int64(i * 2)))
+		if err := tree1.Add(k, v); err != nil {
+			c.Fatal(err)
+		}
 		if err := tree2.Add(k, v); err != nil {
-			t.Fatal(err)
+			c.Fatal(err)
 		}
 	}
-	// tree.PrintGraphviz(nil)
-	// tree2.PrintGraphviz(nil)
+	return tree1, tree2
+}
 
-	var keys, values [][]byte
-	for i := 8; i < nLeafs; i++ {
-		k := BigIntToBytes(big.NewInt(int64(i)))
-		v := BigIntToBytes(big.NewInt(int64(i * 2)))
-		keys = append(keys, k)
-		values = append(values, v)
-	}
-	start = time.Now()
-	indexes, err := tree2.AddBatchOpt(keys, values)
-	c.Assert(err, qt.IsNil)
-	fmt.Println(time.Since(start))
-	c.Check(len(indexes), qt.Equals, 0)
-
-	// check that both trees roots are equal
-	c.Check(tree2.Root(), qt.DeepEquals, tree.Root())
+func ratio(t1, t2 time.Duration) float64 {
+	a := float64(t1)
+	b := float64(t2)
+	return (a / b)
 }
 
 func TestAddBatchCaseA(t *testing.T) {
@@ -78,7 +58,7 @@ func TestAddBatchCaseA(t *testing.T) {
 			t.Fatal(err)
 		}
 	}
-	fmt.Println("time elapsed without CASE A:	", time.Since(start))
+	time1 := time.Since(start)
 
 	tree2, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
 	c.Assert(err, qt.IsNil)
@@ -94,7 +74,9 @@ func TestAddBatchCaseA(t *testing.T) {
 	start = time.Now()
 	indexes, err := tree2.AddBatchOpt(keys, values)
 	c.Assert(err, qt.IsNil)
-	fmt.Println("time elapsed with CASE A:	", time.Since(start))
+	time2 := time.Since(start)
+	fmt.Printf("CASE A, AddBatch was %f times faster than without AddBatch\n",
+		ratio(time1, time2))
 	c.Check(len(indexes), qt.Equals, 0)
 
 	// check that both trees roots are equal
@@ -141,37 +123,23 @@ func TestAddBatchCaseB(t *testing.T) {
 	c := qt.New(t)
 
 	nLeafs := 1024
+	initialNLeafs := 99 // TMP TODO use const minLeafsThreshold-1 once ready
 
-	tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
-	c.Assert(err, qt.IsNil)
-	defer tree.db.Close()
+	tree1, tree2 := testInit(c, initialNLeafs)
 
 	start := time.Now()
-	for i := 0; i < nLeafs; i++ {
+	for i := initialNLeafs; i < nLeafs; i++ {
 		k := BigIntToBytes(big.NewInt(int64(i)))
 		v := BigIntToBytes(big.NewInt(int64(i * 2)))
-		if err := tree.Add(k, v); err != nil {
-			t.Fatal(err)
-		}
-	}
-	fmt.Println("time elapsed without CASE B:	", time.Since(start))
-
-	tree2, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
-	c.Assert(err, qt.IsNil)
-	defer tree2.db.Close()
-
-	// add the initial leafs to fill a bit the tree before calling the
-	// AddBatch method
-	for i := 0; i < 99; i++ { // TMP TODO use const minLeafsThreshold-1 once ready
-		k := BigIntToBytes(big.NewInt(int64(i)))
-		v := BigIntToBytes(big.NewInt(int64(i * 2)))
-		if err := tree2.Add(k, v); err != nil {
+		if err := tree1.Add(k, v); err != nil {
 			t.Fatal(err)
 		}
 	}
+	time1 := time.Since(start)
 
+	// prepare the key-values to be added
 	var keys, values [][]byte
-	for i := 99; i < nLeafs; i++ {
+	for i := initialNLeafs; i < nLeafs; i++ {
 		k := BigIntToBytes(big.NewInt(int64(i)))
 		v := BigIntToBytes(big.NewInt(int64(i * 2)))
 		keys = append(keys, k)
@@ -180,11 +148,13 @@ func TestAddBatchCaseB(t *testing.T) {
 	start = time.Now()
 	indexes, err := tree2.AddBatchOpt(keys, values)
 	c.Assert(err, qt.IsNil)
-	fmt.Println("time elapsed with CASE B:	", time.Since(start))
+	time2 := time.Since(start)
+	fmt.Printf("CASE B, AddBatch was %f times faster than without AddBatch\n",
+		ratio(time1, time2))
 	c.Check(len(indexes), qt.Equals, 0)
 
 	// check that both trees roots are equal
-	c.Check(tree2.Root(), qt.DeepEquals, tree.Root())
+	c.Check(tree2.Root(), qt.DeepEquals, tree1.Root())
 }
 
 func TestGetKeysAtLevel(t *testing.T) {
@@ -319,37 +289,23 @@ func TestAddBatchCaseC(t *testing.T) {
 	c := qt.New(t)
 
 	nLeafs := 1024
+	initialNLeafs := 101 // TMP TODO use const minLeafsThreshold+1 once ready
 
-	tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
-	c.Assert(err, qt.IsNil)
-	defer tree.db.Close()
+	tree1, tree2 := testInit(c, initialNLeafs)
 
 	start := time.Now()
-	for i := 0; i < nLeafs; i++ {
+	for i := initialNLeafs; i < nLeafs; i++ {
 		k := BigIntToBytes(big.NewInt(int64(i)))
 		v := BigIntToBytes(big.NewInt(int64(i * 2)))
-		if err := tree.Add(k, v); err != nil {
-			t.Fatal(err)
-		}
-	}
-	fmt.Println("time elapsed without CASE C:	", time.Since(start))
-
-	tree2, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
-	c.Assert(err, qt.IsNil)
-	defer tree2.db.Close()
-
-	// add the initial leafs to fill a bit the tree before calling the
-	// AddBatch method
-	for i := 0; i < 101; i++ { // TMP TODO use const minLeafsThreshold-1 once ready
-		k := BigIntToBytes(big.NewInt(int64(i)))
-		v := BigIntToBytes(big.NewInt(int64(i * 2)))
-		if err := tree2.Add(k, v); err != nil {
+		if err := tree1.Add(k, v); err != nil {
 			t.Fatal(err)
 		}
 	}
+	time1 := time.Since(start)
 
+	// prepare the key-values to be added
 	var keys, values [][]byte
-	for i := 101; i < nLeafs; i++ {
+	for i := initialNLeafs; i < nLeafs; i++ {
 		k := BigIntToBytes(big.NewInt(int64(i)))
 		v := BigIntToBytes(big.NewInt(int64(i * 2)))
 		keys = append(keys, k)
@@ -358,48 +314,36 @@ func TestAddBatchCaseC(t *testing.T) {
 	start = time.Now()
 	indexes, err := tree2.AddBatchOpt(keys, values)
 	c.Assert(err, qt.IsNil)
-	fmt.Println("time elapsed with CASE C:	", time.Since(start))
+	time2 := time.Since(start)
+	fmt.Printf("CASE C, AddBatch was %f times faster than without AddBatch\n",
+		ratio(time1, time2))
 	c.Check(len(indexes), qt.Equals, 0)
 
 	// check that both trees roots are equal
-	c.Check(tree2.Root(), qt.DeepEquals, tree.Root())
+	c.Check(tree2.Root(), qt.DeepEquals, tree1.Root())
 }
 
 func TestAddBatchCaseD(t *testing.T) {
 	c := qt.New(t)
 
 	nLeafs := 4096
+	initialNLeafs := 900
 
-	tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
-	c.Assert(err, qt.IsNil)
-	defer tree.db.Close()
+	tree1, tree2 := testInit(c, initialNLeafs)
 
 	start := time.Now()
-	for i := 0; i < nLeafs; i++ {
+	for i := initialNLeafs; i < nLeafs; i++ {
 		k := BigIntToBytes(big.NewInt(int64(i)))
 		v := BigIntToBytes(big.NewInt(int64(i * 2)))
-		if err := tree.Add(k, v); err != nil {
-			t.Fatal(err)
-		}
-	}
-	fmt.Println("time elapsed without CASE D:	", time.Since(start))
-
-	tree2, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
-	c.Assert(err, qt.IsNil)
-	defer tree2.db.Close()
-
-	// add the initial leafs to fill a bit the tree before calling the
-	// AddBatch method
-	for i := 0; i < 900; i++ { // TMP TODO use const minLeafsThreshold+1 once ready
-		k := BigIntToBytes(big.NewInt(int64(i)))
-		v := BigIntToBytes(big.NewInt(int64(i * 2)))
-		if err := tree2.Add(k, v); err != nil {
+		if err := tree1.Add(k, v); err != nil {
 			t.Fatal(err)
 		}
 	}
+	time1 := time.Since(start)
 
+	// prepare the key-values to be added
 	var keys, values [][]byte
-	for i := 900; i < nLeafs; i++ {
+	for i := initialNLeafs; i < nLeafs; i++ {
 		k := BigIntToBytes(big.NewInt(int64(i)))
 		v := BigIntToBytes(big.NewInt(int64(i * 2)))
 		keys = append(keys, k)
@@ -408,31 +352,32 @@ func TestAddBatchCaseD(t *testing.T) {
 	start = time.Now()
 	indexes, err := tree2.AddBatchOpt(keys, values)
 	c.Assert(err, qt.IsNil)
-	fmt.Println("time elapsed with CASE D:	", time.Since(start))
+	time2 := time.Since(start)
+	fmt.Printf("CASE D, AddBatch was %f times faster than without AddBatch\n",
+		ratio(time1, time2))
 	c.Check(len(indexes), qt.Equals, 0)
 
 	// check that both trees roots are equal
-	c.Check(tree2.Root(), qt.DeepEquals, tree.Root())
+	c.Check(tree2.Root(), qt.DeepEquals, tree1.Root())
 }
 
 func TestAddBatchCaseE(t *testing.T) {
 	c := qt.New(t)
 
 	nLeafs := 4096
+	initialNLeafs := 900
 
-	tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
-	c.Assert(err, qt.IsNil)
-	defer tree.db.Close()
+	tree1, _ := testInit(c, initialNLeafs)
 
 	start := time.Now()
-	for i := 0; i < nLeafs; i++ {
+	for i := initialNLeafs; i < nLeafs; i++ {
 		k := BigIntToBytes(big.NewInt(int64(i)))
 		v := BigIntToBytes(big.NewInt(int64(i * 2)))
-		if err := tree.Add(k, v); err != nil {
+		if err := tree1.Add(k, v); err != nil {
 			t.Fatal(err)
 		}
 	}
-	fmt.Println("time elapsed without CASE E:	", time.Since(start))
+	time1 := time.Since(start)
 
 	tree2, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
 	c.Assert(err, qt.IsNil)
@@ -441,7 +386,7 @@ func TestAddBatchCaseE(t *testing.T) {
 	var keys, values [][]byte
 	// add the initial leafs to fill a bit the tree before calling the
 	// AddBatch method
-	for i := 0; i < 900; i++ { // TMP TODO use const minLeafsThreshold+1 once ready
+	for i := 0; i < initialNLeafs; i++ {
 		k := BigIntToBytes(big.NewInt(int64(i)))
 		v := BigIntToBytes(big.NewInt(int64(i * 2)))
 		// use only the keys of one bucket, store the not used ones for
@@ -456,7 +401,7 @@ func TestAddBatchCaseE(t *testing.T) {
 		}
 	}
 
-	for i := 900; i < nLeafs; i++ {
+	for i := initialNLeafs; i < nLeafs; i++ {
 		k := BigIntToBytes(big.NewInt(int64(i)))
 		v := BigIntToBytes(big.NewInt(int64(i * 2)))
 		keys = append(keys, k)
@@ -465,11 +410,13 @@ func TestAddBatchCaseE(t *testing.T) {
 	start = time.Now()
 	indexes, err := tree2.AddBatchOpt(keys, values)
 	c.Assert(err, qt.IsNil)
-	fmt.Println("time elapsed with CASE E:	", time.Since(start))
+	time2 := time.Since(start)
+	fmt.Printf("CASE E, AddBatch was %f times faster than without AddBatch\n",
+		ratio(time1, time2))
 	c.Check(len(indexes), qt.Equals, 0)
 
 	// check that both trees roots are equal
-	c.Check(tree2.Root(), qt.DeepEquals, tree.Root())
+	c.Check(tree2.Root(), qt.DeepEquals, tree1.Root())
 }
 
 func TestHighestPowerOfTwo(t *testing.T) {
@@ -513,8 +460,5 @@ func TestHighestPowerOfTwo(t *testing.T) {
 // less keys than nBuckets (so CASE C could be applied if first few leafs are
 // added to balance the tree)
 
-// TODO for Cases tests, add initial keys, do snapshot, and then measure time of
-// adding the rest of keys with loop over normal Add, and with AddBatch
-
 // TODO test adding batch with repeated keys in the batch
 // TODO test adding batch with multiple invalid keys