Add statedb ExitTree implementation

batchbuilder/batchbuilder.go

@@ -52,6 +52,6 @@ func (bb *BatchBuilder) Reset(batchNum uint64, fromSynchronizer bool) error {
 
 // BuildBatch takes the transactions and returns the common.ZKInputs of the next batch
 func (bb *BatchBuilder) BuildBatch(configBatch *ConfigBatch, l1usertxs, l1coordinatortxs []*common.L1Tx, l2txs []*common.L2Tx, tokenIDs []common.TokenID) (*common.ZKInputs, error) {
-	zkInputs, _, err := bb.localStateDB.ProcessTxs(l1usertxs, l1coordinatortxs, l2txs)
+	zkInputs, _, err := bb.localStateDB.ProcessTxs(false, l1usertxs, l1coordinatortxs, l2txs)
 	return zkInputs, err
 }

db/statedb/statedb.go

@@ -179,16 +179,33 @@ func (s *StateDB) Reset(batchNum uint64) error {
 	}
 	// idx is obtained from the statedb reset
 	s.idx, err = s.getIdx()
-	return err
+	if err != nil {
+		return err
+	}
+
+	// open the MT for the current s.db
+	mt, err := merkletree.NewMerkleTree(s.db, s.mt.MaxLevels())
+	if err != nil {
+		return err
+	}
+	s.mt = mt
+
+	return nil
 }
 
 // GetAccount returns the account for the given Idx
 func (s *StateDB) GetAccount(idx common.Idx) (*common.Account, error) {
-	vBytes, err := s.db.Get(idx.Bytes())
+	return getAccountInTreeDB(s.db, idx)
+}
+
+// getAccountInTreeDB is abstracted from StateDB to be used from StateDB and
+// from ExitTree.  GetAccount returns the account for the given Idx
+func getAccountInTreeDB(sto db.Storage, idx common.Idx) (*common.Account, error) {
+	vBytes, err := sto.Get(idx.Bytes())
 	if err != nil {
 		return nil, err
 	}
-	accBytes, err := s.db.Get(vBytes)
+	accBytes, err := sto.Get(vBytes)
 	if err != nil {
 		return nil, err
 	}
@@ -201,6 +218,14 @@ func (s *StateDB) GetAccount(idx common.Idx) (*common.Account, error) {
 // StateDB.mt==nil, MerkleTree is not affected, otherwise updates the
 // MerkleTree, returning a CircomProcessorProof.
 func (s *StateDB) CreateAccount(idx common.Idx, account *common.Account) (*merkletree.CircomProcessorProof, error) {
+	return createAccountInTreeDB(s.db, s.mt, idx, account)
+}
+
+// createAccountInTreeDB is abstracted from StateDB to be used from StateDB and
+// from ExitTree.  Creates a new Account in the StateDB for the given Idx.  If
+// StateDB.mt==nil, MerkleTree is not affected, otherwise updates the
+// MerkleTree, returning a CircomProcessorProof.
+func createAccountInTreeDB(sto db.Storage, mt *merkletree.MerkleTree, idx common.Idx, account *common.Account) (*merkletree.CircomProcessorProof, error) {
 	// store at the DB the key: v, and value: leaf.Bytes()
 	v, err := account.HashValue()
 	if err != nil {
@@ -212,7 +237,7 @@ func (s *StateDB) CreateAccount(idx common.Idx, account *common.Account) (*merkl
 	}
 
 	// store the Leaf value
-	tx, err := s.db.NewTx()
+	tx, err := sto.NewTx()
 	if err != nil {
 		return nil, err
 	}
@@ -229,9 +254,10 @@ func (s *StateDB) CreateAccount(idx common.Idx, account *common.Account) (*merkl
 		return nil, err
 	}
 
-	if s.mt != nil {
-		return s.mt.AddAndGetCircomProof(idx.BigInt(), v)
+	if mt != nil {
+		return mt.AddAndGetCircomProof(idx.BigInt(), v)
 	}
+
 	return nil, nil
 }
 
@@ -239,7 +265,15 @@ func (s *StateDB) CreateAccount(idx common.Idx, account *common.Account) (*merkl
 // StateDB.mt==nil, MerkleTree is not affected, otherwise updates the
 // MerkleTree, returning a CircomProcessorProof.
 func (s *StateDB) UpdateAccount(idx common.Idx, account *common.Account) (*merkletree.CircomProcessorProof, error) {
-	// store at the DB the key: v, and value: leaf.Bytes()
+	return updateAccountInTreeDB(s.db, s.mt, idx, account)
+}
+
+// updateAccountInTreeDB is abstracted from StateDB to be used from StateDB and
+// from ExitTree.  Updates the Account in the StateDB for the given Idx.  If
+// StateDB.mt==nil, MerkleTree is not affected, otherwise updates the
+// MerkleTree, returning a CircomProcessorProof.
+func updateAccountInTreeDB(sto db.Storage, mt *merkletree.MerkleTree, idx common.Idx, account *common.Account) (*merkletree.CircomProcessorProof, error) {
+	// store at the DB the key: v, and value: account.Bytes()
 	v, err := account.HashValue()
 	if err != nil {
 		return nil, err
@@ -249,7 +283,7 @@ func (s *StateDB) UpdateAccount(idx common.Idx, account *common.Account) (*merkl
 		return nil, err
 	}
 
-	tx, err := s.db.NewTx()
+	tx, err := sto.NewTx()
 	if err != nil {
 		return nil, err
 	}
@@ -260,8 +294,8 @@ func (s *StateDB) UpdateAccount(idx common.Idx, account *common.Account) (*merkl
 		return nil, err
 	}
 
-	if s.mt != nil {
-		return s.mt.Update(idx.BigInt(), v)
+	if mt != nil {
+		return mt.Update(idx.BigInt(), v)
 	}
 	return nil, nil
 }
@@ -344,6 +378,13 @@ func (l *LocalStateDB) Reset(batchNum uint64, fromSynchronizer bool) error {
 		if err != nil {
 			return err
 		}
+		// open the MT for the current s.db
+		mt, err := merkletree.NewMerkleTree(l.db, l.mt.MaxLevels())
+		if err != nil {
+			return err
+		}
+		l.mt = mt
+
 		return nil
 	}
 	// use checkpoint from LocalStateDB

db/statedb/txprocessors.go

@@ -4,105 +4,185 @@ import (
 	"math/big"
 
 	"github.com/hermeznetwork/hermez-node/common"
+	"github.com/iden3/go-iden3-crypto/poseidon"
+	"github.com/iden3/go-merkletree"
 	"github.com/iden3/go-merkletree/db"
+	"github.com/iden3/go-merkletree/db/memory"
 )
 
 // keyidx is used as key in the db to store the current Idx
 var keyidx = []byte("idx")
 
-// ProcessTxs process the given L1Txs & L2Txs applying the needed updates
-// to the StateDB depending on the transaction Type. Returns the
-// common.ZKInputs to generate the SnarkProof later used by the BatchBuilder,
-// and returns common.ExitTreeLeaf that is later used by the Synchronizer to
-// update the HistoryDB.
-func (s *StateDB) ProcessTxs(l1usertxs, l1coordinatortxs []*common.L1Tx, l2txs []*common.L2Tx) (*common.ZKInputs, []*common.ExitTreeLeaf, error) {
+// FUTURE This will be used from common once pending PR is merged
+type ExitInfo struct {
+	Idx       *common.Idx
+	Proof     *merkletree.CircomVerifierProof
+	Nullifier *big.Int
+	Balance   *big.Int
+}
+
+// ProcessTxs process the given L1Txs & L2Txs applying the needed updates to
+// the StateDB depending on the transaction Type. Returns the common.ZKInputs
+// to generate the SnarkProof later used by the BatchBuilder, and if
+// cmpExitTree is set to true, returns common.ExitTreeLeaf that is later used
+// by the Synchronizer to update the HistoryDB.
+func (s *StateDB) ProcessTxs(cmpExitTree bool, l1usertxs, l1coordinatortxs []*common.L1Tx, l2txs []*common.L2Tx) (*common.ZKInputs, []*ExitInfo, error) {
+	var err error
+	var exitTree *merkletree.MerkleTree
+	exits := make(map[common.Idx]common.Account)
+	if cmpExitTree {
+		// TBD if ExitTree is only in memory or stored in disk, for the moment
+		// only needed in memory
+		exitTree, err = merkletree.NewMerkleTree(memory.NewMemoryStorage(), s.mt.MaxLevels())
+		if err != nil {
+			return nil, nil, err
+		}
+	}
+
 	for _, tx := range l1usertxs {
-		err := s.processL1Tx(tx)
+		exitIdx, exitAccount, err := s.processL1Tx(exitTree, tx)
 		if err != nil {
 			return nil, nil, err
 		}
+		if exitIdx != nil && cmpExitTree {
+			exits[*exitIdx] = *exitAccount
+		}
 	}
 	for _, tx := range l1coordinatortxs {
-		err := s.processL1Tx(tx)
+		exitIdx, exitAccount, err := s.processL1Tx(exitTree, tx)
 		if err != nil {
 			return nil, nil, err
 		}
+		if exitIdx != nil && cmpExitTree {
+			exits[*exitIdx] = *exitAccount
+		}
 	}
 	for _, tx := range l2txs {
-		err := s.processL2Tx(tx)
+		exitIdx, exitAccount, err := s.processL2Tx(exitTree, tx)
 		if err != nil {
 			return nil, nil, err
 		}
+		if exitIdx != nil && cmpExitTree {
+			exits[*exitIdx] = *exitAccount
+		}
 	}
 
-	return nil, nil, nil
+	if !cmpExitTree {
+		return nil, nil, nil
+	}
+
+	// once all txs processed (exitTree root frozen), for each leaf
+	// generate ExitInfo data
+	var exitInfos []*ExitInfo
+	for exitIdx, exitAccount := range exits {
+		// 0. generate MerkleProof
+		p, err := exitTree.GenerateCircomVerifierProof(exitIdx.BigInt(), nil)
+		if err != nil {
+			return nil, nil, err
+		}
+		// 1. compute nullifier
+		exitAccStateValue, err := exitAccount.HashValue()
+		if err != nil {
+			return nil, nil, err
+		}
+		nullifier, err := poseidon.Hash([]*big.Int{
+			exitAccStateValue,
+			big.NewInt(int64(s.currentBatch)),
+			exitTree.Root().BigInt(),
+		})
+		if err != nil {
+			return nil, nil, err
+		}
+		// 2. generate ExitInfo
+		ei := &ExitInfo{
+			Idx:       &exitIdx,
+			Proof:     p,
+			Nullifier: nullifier,
+			Balance:   exitAccount.Balance,
+		}
+		exitInfos = append(exitInfos, ei)
+	}
+
+	// return exitInfos, so Synchronizer will be able to store it into
+	// HistoryDB for the concrete BatchNum
+	return nil, exitInfos, nil
 }
 
 // processL2Tx process the given L2Tx applying the needed updates to
 // the StateDB depending on the transaction Type.
-func (s *StateDB) processL2Tx(tx *common.L2Tx) error {
+func (s *StateDB) processL2Tx(exitTree *merkletree.MerkleTree, tx *common.L2Tx) (*common.Idx, *common.Account, error) {
 	switch tx.Type {
 	case common.TxTypeTransfer:
 		// go to the MT account of sender and receiver, and update
 		// balance & nonce
 		err := s.applyTransfer(tx.Tx())
 		if err != nil {
-			return err
+			return nil, nil, err
 		}
 	case common.TxTypeExit:
 		// execute exit flow
+		exitAccount, err := s.applyExit(exitTree, tx.Tx())
+		if err != nil {
+			return nil, nil, err
+		}
+		return &tx.FromIdx, exitAccount, nil
 	default:
 	}
-	return nil
+	return nil, nil, nil
 }
 
 // processL1Tx process the given L1Tx applying the needed updates to the
 // StateDB depending on the transaction Type.
-func (s *StateDB) processL1Tx(tx *common.L1Tx) error {
+func (s *StateDB) processL1Tx(exitTree *merkletree.MerkleTree, tx *common.L1Tx) (*common.Idx, *common.Account, error) {
 	switch tx.Type {
 	case common.TxTypeForceTransfer, common.TxTypeTransfer:
 		// go to the MT account of sender and receiver, and update balance
 		// & nonce
 		err := s.applyTransfer(tx.Tx())
 		if err != nil {
-			return err
+			return nil, nil, err
 		}
 	case common.TxTypeCreateAccountDeposit:
 		// add new account to the MT, update balance of the MT account
 		err := s.applyCreateAccount(tx)
 		if err != nil {
-			return err
+			return nil, nil, err
 		}
 	case common.TxTypeDeposit:
 		// update balance of the MT account
 		err := s.applyDeposit(tx, false)
 		if err != nil {
-			return err
+			return nil, nil, err
 		}
 	case common.TxTypeDepositTransfer:
 		// update balance in MT account, update balance & nonce of sender
 		// & receiver
 		err := s.applyDeposit(tx, true)
 		if err != nil {
-			return err
+			return nil, nil, err
 		}
 	case common.TxTypeCreateAccountDepositTransfer:
 		// add new account to the merkletree, update balance in MT account,
 		// update balance & nonce of sender & receiver
 		err := s.applyCreateAccount(tx)
 		if err != nil {
-			return err
+			return nil, nil, err
 		}
 		err = s.applyTransfer(tx.Tx())
 		if err != nil {
-			return err
+			return nil, nil, err
 		}
 	case common.TxTypeExit:
 		// execute exit flow
+		exitAccount, err := s.applyExit(exitTree, tx.Tx())
+		if err != nil {
+			return nil, nil, err
+		}
+		return &tx.FromIdx, exitAccount, nil
 	default:
 	}
 
-	return nil
+	return nil, nil, nil
 }
 
 // applyCreateAccount creates a new account in the account of the depositer, it
@@ -195,6 +275,43 @@ func (s *StateDB) applyTransfer(tx *common.Tx) error {
 	return nil
 }
 
+func (s *StateDB) applyExit(exitTree *merkletree.MerkleTree, tx *common.Tx) (*common.Account, error) {
+	// 0. substract tx.Amount from current Account in StateMT
+	// add the tx.Amount into the Account (tx.FromIdx) in the ExitMT
+	acc, err := s.GetAccount(tx.FromIdx)
+	if err != nil {
+		return nil, err
+	}
+	acc.Balance = new(big.Int).Sub(acc.Balance, tx.Amount)
+	_, err = s.UpdateAccount(tx.FromIdx, acc)
+	if err != nil {
+		return nil, err
+	}
+
+	exitAccount, err := getAccountInTreeDB(exitTree.DB(), tx.FromIdx)
+	if err == db.ErrNotFound {
+		// 1a. if idx does not exist in exitTree:
+		// add new leaf 'ExitTreeLeaf', where ExitTreeLeaf.Balance = exitAmount (exitAmount=tx.Amount)
+		exitAccount := &common.Account{
+			TokenID:   acc.TokenID,
+			Nonce:     common.Nonce(1),
+			Balance:   tx.Amount,
+			PublicKey: acc.PublicKey,
+			EthAddr:   acc.EthAddr,
+		}
+		_, err = createAccountInTreeDB(exitTree.DB(), exitTree, tx.FromIdx, exitAccount)
+		return exitAccount, err
+	} else if err != nil {
+		return exitAccount, err
+	}
+
+	// 1b. if idx already exist in exitTree:
+	// update account, where account.Balance += exitAmount
+	exitAccount.Balance = new(big.Int).Add(exitAccount.Balance, tx.Amount)
+	_, err = updateAccountInTreeDB(exitTree.DB(), exitTree, tx.FromIdx, exitAccount)
+	return exitAccount, err
+}
+
 // getIdx returns the stored Idx from the localStateDB, which is the last Idx
 // used for an Account in the localStateDB.
 func (s *StateDB) getIdx() (common.Idx, error) {

go.mod

@@ -9,7 +9,7 @@ require (
 	github.com/gobuffalo/packr/v2 v2.8.0
 	github.com/iden3/go-iden3-core v0.0.8
 	github.com/iden3/go-iden3-crypto v0.0.6-0.20200823174058-e04ca5764a15
-	github.com/iden3/go-merkletree v0.0.0-20200819092443-dc656fdd32fc
+	github.com/iden3/go-merkletree v0.0.0-20200825093552-a4b68208bb41
 	github.com/jinzhu/gorm v1.9.15
 	github.com/jmoiron/sqlx v1.2.0
 	github.com/lib/pq v1.8.0

go.sum

@@ -306,6 +306,8 @@ github.com/iden3/go-merkletree v0.0.0-20200819075941-77df3096f7ea h1:0MtMnN42ULN
 github.com/iden3/go-merkletree v0.0.0-20200819075941-77df3096f7ea/go.mod h1:MRe6i0mi2oDVUzgBIHsNRE6XAg8EBuqIQZMsd+do+dU=
 github.com/iden3/go-merkletree v0.0.0-20200819092443-dc656fdd32fc h1:VnRP7JCp5TJHnTC+r8NrXGkRBvp62tRGqeA3FcdEq+Q=
 github.com/iden3/go-merkletree v0.0.0-20200819092443-dc656fdd32fc/go.mod h1:MRe6i0mi2oDVUzgBIHsNRE6XAg8EBuqIQZMsd+do+dU=
+github.com/iden3/go-merkletree v0.0.0-20200825093552-a4b68208bb41 h1:mCOMMQ/YmL9ST9kk7ifT961chESkB2GFFEp8osF0Jw8=
+github.com/iden3/go-merkletree v0.0.0-20200825093552-a4b68208bb41/go.mod h1:MRe6i0mi2oDVUzgBIHsNRE6XAg8EBuqIQZMsd+do+dU=
 github.com/iden3/go-wasm3 v0.0.1/go.mod h1:j+TcAB94Dfrjlu5kJt83h2OqAU+oyNUTwNZnQyII1sI=
 github.com/inconshreveable/mousetrap v1.0.0/go.mod h1:PxqpIevigyE2G7u3NXJIT2ANytuPF1OarO4DADm73n8=
 github.com/influxdata/influxdb v1.2.3-0.20180221223340-01288bdb0883/go.mod h1:qZna6X/4elxqT3yI9iZYdZrWWdeFOOprn86kgg4+IzY=