From af23b9c63ffb2c0cb46c921986cce6044d0b090d Mon Sep 17 00:00:00 2001
From: arnaucube <root@arnaucube.com>
Date: Tue, 31 Aug 2021 10:38:07 +0200
Subject: [PATCH] Fix level when going down & in computeHashes

Fix level when going down & in computeHashes, this affected the methods
Get, GenProof & AddBatch.
---
 tree.go      |  4 ++--
 tree_test.go | 49 +++++++++++++++++++++++++++++++++++++++++++++++++
 vt.go        |  4 ++--
 3 files changed, 53 insertions(+), 4 deletions(-)

diff --git a/tree.go b/tree.go
index 5695fdb..9b8c6ec 100644
--- a/tree.go
+++ b/tree.go
@@ -340,7 +340,7 @@ func (t *Tree) add(wTx db.WriteTx, root []byte, fromLvl int, k, v []byte) ([]byt
 func (t *Tree) down(rTx db.ReadTx, newKey, currKey []byte, siblings [][]byte,
 	path []bool, currLvl int, getLeaf bool) (
 	[]byte, []byte, [][]byte, error) {
-	if currLvl > t.maxLevels-1 {
+	if currLvl > t.maxLevels {
 		return nil, nil, nil, ErrMaxLevel
 	}
 
@@ -413,7 +413,7 @@ func (t *Tree) down(rTx db.ReadTx, newKey, currKey []byte, siblings [][]byte,
 func (t *Tree) downVirtually(siblings [][]byte, oldKey, newKey []byte, oldPath,
 	newPath []bool, currLvl int) ([][]byte, error) {
 	var err error
-	if currLvl > t.maxLevels-1 {
+	if currLvl > t.maxLevels {
 		return nil, ErrMaxVirtualLevel
 	}
 
diff --git a/tree_test.go b/tree_test.go
index 7cdb7e3..976ff34 100644
--- a/tree_test.go
+++ b/tree_test.go
@@ -454,6 +454,55 @@ func TestRWMutex(t *testing.T) {
 //         c.Assert(n, qt.Equals, maxInt)
 // }
 
+func TestAddBatchFullyUsed(t *testing.T) {
+	c := qt.New(t)
+
+	database1, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
+	c.Assert(err, qt.IsNil)
+	tree1, err := NewTree(database1, 4, HashFunctionPoseidon)
+	c.Assert(err, qt.IsNil)
+
+	database2, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
+	c.Assert(err, qt.IsNil)
+	tree2, err := NewTree(database2, 4, HashFunctionPoseidon)
+	c.Assert(err, qt.IsNil)
+
+	var keys, values [][]byte
+	for i := 0; i < 16; i++ {
+		k := BigIntToBytes(32, big.NewInt(int64(i)))
+		v := k
+
+		keys = append(keys, k)
+		values = append(values, v)
+
+		// add one by one expecting no error
+		err := tree1.Add(k, v)
+		c.Assert(err, qt.IsNil)
+	}
+
+	invalids, err := tree2.AddBatch(keys, values)
+	c.Assert(err, qt.IsNil)
+	c.Assert(0, qt.Equals, len(invalids))
+
+	root1, err := tree1.Root()
+	c.Assert(err, qt.IsNil)
+	root2, err := tree2.Root()
+	c.Assert(err, qt.IsNil)
+	c.Assert(root1, qt.DeepEquals, root2)
+
+	// get all key-values and check that are equal between both trees
+	for i := 0; i < 16; i++ {
+		auxK1, auxV1, err := tree1.Get(BigIntToBytes(32, big.NewInt(int64(i))))
+		c.Assert(err, qt.IsNil)
+
+		auxK2, auxV2, err := tree2.Get(BigIntToBytes(32, big.NewInt(int64(i))))
+		c.Assert(err, qt.IsNil)
+
+		c.Assert(auxK1, qt.DeepEquals, auxK2)
+		c.Assert(auxV1, qt.DeepEquals, auxV2)
+	}
+}
+
 func TestSnapshot(t *testing.T) {
 	c := qt.New(t)
 	database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
diff --git a/vt.go b/vt.go
index b79ded9..a18bbc8 100644
--- a/vt.go
+++ b/vt.go
@@ -318,11 +318,11 @@ func (t *vt) computeHashes() ([][2][]byte, error) {
 	wg.Add(nCPU)
 	for i := 0; i < nCPU; i++ {
 		go func(cpu int) {
-			bucketVT := newVT(t.params.maxLevels-l, t.params.hashFunction)
+			bucketVT := newVT(t.params.maxLevels, t.params.hashFunction)
 			bucketVT.params.dbg = newDbgStats()
 			bucketVT.root = nodesAtL[cpu]
 
-			bucketPairs[cpu], err = bucketVT.root.computeHashes(l,
+			bucketPairs[cpu], err = bucketVT.root.computeHashes(l-1,
 				t.params.maxLevels, bucketVT.params, bucketPairs[cpu])
 			if err != nil {
 				errs[cpu] = err