Integrate BatchBuilder with StateDB

batchbuilder/batchbuilder.go

@@ -1,11 +1,11 @@
 package batchbuilder
 
 import (
+	"math/big"
+
 	ethCommon "github.com/ethereum/go-ethereum/common"
 	"github.com/hermeznetwork/hermez-node/common"
-	"github.com/iden3/go-merkletree"
-	"github.com/iden3/go-merkletree/db"
-	"github.com/iden3/go-merkletree/db/memory"
+	"github.com/hermeznetwork/hermez-node/db/statedb"
 )
 
 // ConfigCircuit contains the circuit configuration
@@ -15,11 +15,12 @@ type ConfigCircuit struct {
 	SMTLevelsMax uint64
 }
 
-// BatchBuilder implements the batch builder type, which contains the functionallities
+// BatchBuilder implements the batch builder type, which contains the
+// functionalities
 type BatchBuilder struct {
-	StateDB        db.Storage // where the MTs will be stored by the Synchronizer
+	// idx holds the current Idx that the BatchBuilder is using
 	idx            uint64
-	mt             *merkletree.MerkleTree
+	localStateDB   *statedb.LocalStateDB
 	configCircuits []ConfigCircuit
 }
 
@@ -30,19 +31,19 @@ type ConfigBatch struct {
 
 // NewBatchBuilder constructs a new BatchBuilder, and executes the bb.Reset
 // method
-func NewBatchBuilder(stateDB db.Storage, configCircuits []ConfigCircuit, batchNum int, idx, nLevels uint64) (*BatchBuilder, error) {
-	localMt, err := merkletree.NewMerkleTree(memory.NewMemoryStorage(), int(nLevels))
+func NewBatchBuilder(synchronizerStateDB *statedb.StateDB, configCircuits []ConfigCircuit, batchNum int, idx, nLevels uint64) (*BatchBuilder, error) {
+	localStateDB, err := statedb.NewLocalStateDB(synchronizerStateDB, true, int(nLevels))
 	if err != nil {
 		return nil, err
 	}
+
 	bb := BatchBuilder{
-		StateDB:        stateDB,
 		idx:            idx,
-		mt:             localMt,
+		localStateDB:   localStateDB,
 		configCircuits: configCircuits,
 	}
 
-	err = bb.Reset(batchNum, idx, true)
+	err = bb.Reset(batchNum, true)
 	return &bb, err
 }
 
@@ -50,8 +51,12 @@ func NewBatchBuilder(stateDB db.Storage, configCircuits []ConfigCircuit, batchNu
 // `batchNum`.  If `fromSynchronizer` is true, the BatchBuilder must take a
 // copy of the rollup state from the Synchronizer at that `batchNum`, otherwise
 // it can just roll back the internal copy.
-func (bb *BatchBuilder) Reset(batchNum int, idx uint64, fromSynchronizer bool) error {
-	// TODO
+func (bb *BatchBuilder) Reset(batchNum int, fromSynchronizer bool) error {
+	err := bb.localStateDB.Reset(batchNum, fromSynchronizer)
+	if err != nil {
+		return err
+	}
+	// bb.idx = idx // TODO idx will be obtained from the statedb reset
 	return nil
 }
 
@@ -98,7 +103,7 @@ func (bb *BatchBuilder) processL1Tx(tx common.L1Tx) error {
 		}
 	case common.TxTypeCreateAccountDeposit:
 		// add new account to the MT, update balance of the MT account
-		err := bb.applyCreateLeaf(tx)
+		err := bb.applyCreateAccount(tx)
 		if err != nil {
 			return err
 		}
@@ -118,7 +123,7 @@ func (bb *BatchBuilder) processL1Tx(tx common.L1Tx) error {
 	case common.TxTypeCreateAccountDepositAndTransfer:
 		// add new account to the merkletree, update balance in MT account,
 		// update balance & nonce of sender & receiver
-		err := bb.applyCreateLeaf(tx)
+		err := bb.applyCreateAccount(tx)
 		if err != nil {
 			return err
 		}
@@ -134,10 +139,10 @@ func (bb *BatchBuilder) processL1Tx(tx common.L1Tx) error {
 	return nil
 }
 
-// applyCreateLeaf creates a new account in the account of the depositer, it stores
+// applyCreateAccount creates a new account in the account of the depositer, it stores
 // the deposit value
-func (bb *BatchBuilder) applyCreateLeaf(tx common.L1Tx) error {
-	account := common.Account{
+func (bb *BatchBuilder) applyCreateAccount(tx common.L1Tx) error {
+	account := &common.Account{
 		TokenID:   tx.TokenID,
 		Nonce:     0,
 		Balance:   tx.LoadAmount,
@@ -145,85 +150,78 @@ func (bb *BatchBuilder) applyCreateLeaf(tx common.L1Tx) error {
 		EthAddr:   tx.FromEthAddr,
 	}
 
-	v, err := account.HashValue()
-	if err != nil {
-		return err
-	}
-	dbTx, err := bb.mt.DB().NewTx()
-	if err != nil {
-		return err
-	}
-
-	err = bb.CreateBalance(dbTx, common.Idx(bb.idx+1), account)
-	if err != nil {
-		return err
-	}
-	accountBytes, err := account.Bytes()
+	err := bb.localStateDB.CreateAccount(common.Idx(bb.idx+1), account)
 	if err != nil {
 		return err
 	}
-	dbTx.Put(v.Bytes(), accountBytes[:])
 
-	// if everything is fine, do dbTx & increment idx
-	if err := dbTx.Commit(); err != nil {
-		return err
-	}
 	bb.idx = bb.idx + 1
 	return nil
 }
 
-// applyDeposit updates the balance in the account of the depositer, if andTransfer parameter is set to true, the method will also apply the Transfer of the L1Tx/DepositAndTransfer
-func (bb *BatchBuilder) applyDeposit(tx common.L1Tx, andTransfer bool) error {
-	dbTx, err := bb.mt.DB().NewTx()
-	if err != nil {
-		return err
-	}
-
-	// deposit
-	err = bb.UpdateBalance(dbTx, tx.FromIdx, tx.LoadAmount, false)
+// applyDeposit updates the balance in the account of the depositer, if
+// andTransfer parameter is set to true, the method will also apply the
+// Transfer of the L1Tx/DepositAndTransfer
+func (bb *BatchBuilder) applyDeposit(tx common.L1Tx, transfer bool) error {
+	// deposit the tx.LoadAmount into the sender account
+	accSender, err := bb.localStateDB.GetAccount(tx.FromIdx)
 	if err != nil {
 		return err
 	}
+	accSender.Balance = new(big.Int).Add(accSender.Balance, tx.LoadAmount)
 
 	// in case that the tx is a L1Tx>DepositAndTransfer
-	if andTransfer {
-		// transact
-		err = bb.UpdateBalance(dbTx, tx.FromIdx, tx.Tx.Amount, true)
+	if transfer {
+		accReceiver, err := bb.localStateDB.GetAccount(tx.ToIdx)
 		if err != nil {
 			return err
 		}
-		err = bb.UpdateBalance(dbTx, tx.ToIdx, tx.Tx.Amount, false)
+		// substract amount to the sender
+		accSender.Balance = new(big.Int).Sub(accSender.Balance, tx.Tx.Amount)
+		// add amount to the receiver
+		accReceiver.Balance = new(big.Int).Add(accReceiver.Balance, tx.Tx.Amount)
+		// update receiver account in localStateDB
+		err = bb.localStateDB.UpdateAccount(tx.ToIdx, accReceiver)
 		if err != nil {
 			return err
 		}
 	}
-
-	if err := dbTx.Commit(); err != nil {
+	// update sender account in localStateDB
+	err = bb.localStateDB.UpdateAccount(tx.FromIdx, accSender)
+	if err != nil {
 		return err
 	}
 	return nil
 }
 
-// applyTransfer updates the balance & nonce in the account of the sender, and the
-// balance in the account of the receiver
+// applyTransfer updates the balance & nonce in the account of the sender, and
+// the balance in the account of the receiver
 func (bb *BatchBuilder) applyTransfer(tx common.Tx) error {
-	dbTx, err := bb.mt.DB().NewTx()
+	// get sender and receiver accounts from localStateDB
+	accSender, err := bb.localStateDB.GetAccount(tx.FromIdx)
 	if err != nil {
 		return err
 	}
-
-	// transact
-	err = bb.UpdateBalance(dbTx, tx.FromIdx, tx.Amount, true)
+	accReceiver, err := bb.localStateDB.GetAccount(tx.ToIdx)
 	if err != nil {
 		return err
 	}
-	err = bb.UpdateBalance(dbTx, tx.ToIdx, tx.Amount, false)
+
+	// substract amount to the sender
+	accSender.Balance = new(big.Int).Sub(accSender.Balance, tx.Amount)
+	// add amount to the receiver
+	accReceiver.Balance = new(big.Int).Add(accReceiver.Balance, tx.Amount)
+
+	// update receiver account in localStateDB
+	err = bb.localStateDB.UpdateAccount(tx.ToIdx, accReceiver)
 	if err != nil {
 		return err
 	}
-
-	if err := dbTx.Commit(); err != nil {
+	// update sender account in localStateDB
+	err = bb.localStateDB.UpdateAccount(tx.FromIdx, accSender)
+	if err != nil {
 		return err
 	}
+
 	return nil
 }

batchbuilder/batchbuilder_test.go

@@ -2,28 +2,22 @@ package batchbuilder
 
 import (
 	"fmt"
+	"io/ioutil"
 	"testing"
 
-	"github.com/iden3/go-merkletree/db/memory"
+	"github.com/hermeznetwork/hermez-node/db/statedb"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
-// func genTxs(n int) common.Tx {
-//         return common.Tx{}
-// }
-// func genL1Txs(n int) []common.L1Tx {
-//
-//         return nil
-// }
-// func genL2Txs(n int) []common.L2Tx {
-//
-//         return nil
-// }
-
 func TestBatchBuilder(t *testing.T) {
-	stateDB := memory.NewMemoryStorage()
+	dir, err := ioutil.TempDir("", "tmpdb")
+	require.Nil(t, err)
+
+	sdb, err := statedb.NewStateDB(dir, false, false, 0)
+	assert.Nil(t, err)
 
-	bb, err := NewBatchBuilder(stateDB, nil, 0, 0, 32)
+	bb, err := NewBatchBuilder(sdb, nil, 0, 0, 32)
 	assert.Nil(t, err)
 	fmt.Println(bb)
 }

batchbuilder/state.go

@@ -1,85 +0,0 @@
-package batchbuilder
-
-import (
-	"math/big"
-
-	"github.com/hermeznetwork/hermez-node/common"
-	"github.com/iden3/go-merkletree/db"
-)
-
-// TODO next iteration move the methods of this file into StateDB, which Synchronizer will use in the disk DB, and BatchBuilder will use with the MemoryDB
-
-// GetBalance returns the balance for a given Idx from the DB
-func (bb *BatchBuilder) GetBalance(tx db.Tx, idx common.Idx) (*common.Account, error) {
-	idxBytes := idx.Bytes()
-	vBytes, err := tx.Get(idxBytes[:])
-	if err != nil {
-		return nil, err
-	}
-	var b [32 * common.NLEAFELEMS]byte
-	copy(b[:], vBytes)
-	leaf, err := common.AccountFromBytes(b)
-	if err != nil {
-		return nil, err
-	}
-	return leaf, nil
-}
-
-// CreateBalance stores the Account into the Idx position in the MerkleTree, also adds db entry for the Account value
-func (bb *BatchBuilder) CreateBalance(tx db.Tx, idx common.Idx, leaf common.Account) error {
-	// store at the DB the key: v, and value: leaf.Bytes()
-	v, err := leaf.HashValue()
-	if err != nil {
-		return err
-	}
-	leafBytes, err := leaf.Bytes()
-	if err != nil {
-		return err
-	}
-
-	// store the Account value
-	tx.Put(v.Bytes(), leafBytes[:])
-	// Add k & v into the MT
-	err = bb.mt.Add(idx.BigInt(), v)
-	if err != nil {
-		return err
-	}
-
-	return nil
-}
-
-// UpdateBalance updates the balance of the leaf of a given Idx.
-// If sending==true: will substract the amount, if sending==false will add the ammount
-func (bb *BatchBuilder) UpdateBalance(tx db.Tx, idx common.Idx, amount *big.Int, sending bool) error {
-	leaf, err := bb.GetBalance(tx, idx)
-	if err != nil {
-		return err
-	}
-
-	// TODO add checks that the numbers are correct and there is no missing value neither impossible values
-	if sending {
-		leaf.Balance = new(big.Int).Sub(leaf.Balance, amount)
-	} else {
-		leaf.Balance = new(big.Int).Add(leaf.Balance, amount)
-	}
-
-	// store at the DB the key: v, and value: leaf.Bytes()
-	v, err := leaf.HashValue()
-	if err != nil {
-		return err
-	}
-	leafBytes, err := leaf.Bytes()
-	if err != nil {
-		return err
-	}
-
-	// store the Account value
-	tx.Put(v.Bytes(), leafBytes[:])
-	// Add k & v into the MT
-	_, err = bb.mt.Update(idx.BigInt(), v)
-	if err != nil {
-		return err
-	}
-
-	return nil
-}