Merge pull request #261 from hermeznetwork/feature/til-extra

Move logic to extend til blocks to til

test/til/txs.go

@@ -21,7 +21,10 @@ func newBatchData(batchNum int) common.BatchData {
 		L2Txs:            []common.L2Tx{},
 		Batch: common.Batch{
 			BatchNum:  common.BatchNum(batchNum),
-			StateRoot: big.NewInt(0), ExitRoot: big.NewInt(0)},
+			StateRoot: big.NewInt(0), ExitRoot: big.NewInt(0),
+			FeeIdxsCoordinator: make([]common.Idx, 0),
+			CollectedFees:      make(map[common.TokenID]*big.Int),
+		},
 	}
 }
 
@@ -36,13 +39,22 @@ func newBlock(blockNum int64) common.BlockData {
 	}
 }
 
+type contextExtra struct {
+	openToForge     int64
+	toForgeL1TxsNum int64
+	nonces          map[common.Idx]common.Nonce
+	idx             int
+}
+
 // Context contains the data of the test
 type Context struct {
 	Instructions          []instruction
 	userNames             []string
-	Users                 map[string]*User
+	Users                 map[string]*User // Name -> *User
+	usersByIdx            map[int]*User
+	accountsByIdx         map[int]*Account
 	LastRegisteredTokenID common.TokenID
-	l1CreatedAccounts     map[string]*Account
+	l1CreatedAccounts     map[string]*Account // (Name, TokenID) -> *Account
 
 	// rollupConstMaxL1UserTx Maximum L1-user transactions allowed to be queued in a batch
 	rollupConstMaxL1UserTx int
@@ -59,6 +71,8 @@ type Context struct {
 		l2Txs            []L2Tx
 	}
 	blockNum int64
+
+	extra contextExtra
 }
 
 // NewContext returns a new Context
@@ -67,6 +81,8 @@ func NewContext(rollupConstMaxL1UserTx int) *Context {
 	return &Context{
 		Users:                 make(map[string]*User),
 		l1CreatedAccounts:     make(map[string]*Account),
+		usersByIdx:            make(map[int]*User),
+		accountsByIdx:         make(map[int]*Account),
 		LastRegisteredTokenID: 0,
 
 		rollupConstMaxL1UserTx: rollupConstMaxL1UserTx,
@@ -82,17 +98,26 @@ func NewContext(rollupConstMaxL1UserTx int) *Context {
 		openToForge: 1,
 		//nolint:gomnd
 		blockNum: 2, // rollup genesis blockNum
+		extra: contextExtra{
+			openToForge:     0,
+			toForgeL1TxsNum: 0,
+			nonces:          make(map[common.Idx]common.Nonce),
+			idx:             common.UserThreshold,
+		},
 	}
 }
 
 // Account contains the data related to the account for a specific TokenID of a User
 type Account struct {
-	Idx   common.Idx
-	Nonce common.Nonce
+	Idx      common.Idx
+	TokenID  common.TokenID
+	Nonce    common.Nonce
+	BatchNum int
 }
 
 // User contains the data related to a testing user
 type User struct {
+	Name     string
 	BJJ      *babyjub.PrivateKey
 	Addr     ethCommon.Address
 	Accounts map[common.TokenID]*Account
@@ -361,10 +386,14 @@ func (tc *Context) calculateIdxForL1Txs(isCoordinatorTxs bool, txs []L1Tx) error
 				return fmt.Errorf("Can not create same account twice (same User (%s) & same TokenID (%d)) (this is a design property of Til)", tx.fromIdxName, tx.L1Tx.TokenID)
 			}
 			tc.Users[tx.fromIdxName].Accounts[tx.L1Tx.TokenID] = &Account{
-				Idx:   common.Idx(tc.idx),
-				Nonce: common.Nonce(0),
+				Idx:      common.Idx(tc.idx),
+				TokenID:  tx.L1Tx.TokenID,
+				Nonce:    common.Nonce(0),
+				BatchNum: tc.currBatchNum,
 			}
 			tc.l1CreatedAccounts[idxTokenIDToString(tx.fromIdxName, tx.L1Tx.TokenID)] = tc.Users[tx.fromIdxName].Accounts[tx.L1Tx.TokenID]
+			tc.accountsByIdx[tc.idx] = tc.Users[tx.fromIdxName].Accounts[tx.L1Tx.TokenID]
+			tc.usersByIdx[tc.idx] = tc.Users[tx.fromIdxName]
 			tc.idx++
 		}
 		if isCoordinatorTxs {
@@ -606,6 +635,7 @@ func (tc *Context) generateKeys(userNames []string) {
 		addr := ethCrypto.PubkeyToAddress(key.PublicKey)
 
 		u := User{
+			Name:     userNames[i-1],
 			BJJ:      &sk,
 			Addr:     addr,
 			Accounts: make(map[common.TokenID]*Account),
@@ -622,3 +652,205 @@ func L1TxsToCommonL1Txs(l1 []L1Tx) []common.L1Tx {
 	}
 	return r
 }
+
+// ConfigExtra is the configuration used in FillBlocksExtra to extend the
+// blocks returned by til.
+type ConfigExtra struct {
+	// Address to set as forger for each batch
+	BootCoordAddr ethCommon.Address
+	// Coordinator user name used to select the corresponding accounts to
+	// collect coordinator fees
+	CoordUser string
+}
+
+// FillBlocksExtra fills extra fields not generated by til in each block, so
+// that the blockData is closer to what the HistoryDB stores.  The filled fields are:
+// - blocks[].Rollup.L1UserTxs[].BatchNum
+// - blocks[].Rollup.Batch.EthBlockNum
+// - blocks[].Rollup.Batch.ForgerAddr
+// - blocks[].Rollup.Batch.ForgeL1TxsNum
+// - blocks[].Rollup.Batch.L1CoordinatorTxs[].TxID
+// - blocks[].Rollup.Batch.L1CoordinatorTxs[].BatchNum
+// - blocks[].Rollup.Batch.L1CoordinatorTxs[].EthBlockNum
+// - blocks[].Rollup.Batch.L1CoordinatorTxs[].Position
+// - blocks[].Rollup.Batch.L2Txs[].TxID
+// - blocks[].Rollup.Batch.L2Txs[].Position
+// - blocks[].Rollup.Batch.L2Txs[].Nonce
+// - blocks[].Rollup.Batch.ExitTree
+// - blocks[].Rollup.Batch.CreatedAccounts
+// - blocks[].Rollup.Batch.FeeIdxCoordinator
+// - blocks[].Rollup.Batch.CollectedFees
+func (tc *Context) FillBlocksExtra(blocks []common.BlockData, cfg *ConfigExtra) error {
+	// Fill extra fields not generated by til in til block
+	for i := range blocks {
+		block := &blocks[i]
+		for j := range block.Rollup.Batches {
+			batch := &block.Rollup.Batches[j]
+			if batch.L1Batch {
+				// Set BatchNum for forged L1UserTxs to til blocks
+				bn := batch.Batch.BatchNum
+				for k := range blocks {
+					block := &blocks[k]
+					for l := range block.Rollup.L1UserTxs {
+						tx := &block.Rollup.L1UserTxs[l]
+						if *tx.ToForgeL1TxsNum == tc.extra.openToForge {
+							tx.BatchNum = &bn
+						}
+					}
+				}
+				tc.extra.openToForge++
+			}
+
+			batch.Batch.EthBlockNum = block.Block.EthBlockNum
+			// til doesn't fill the batch forger addr
+			batch.Batch.ForgerAddr = cfg.BootCoordAddr
+			if batch.L1Batch {
+				toForgeL1TxsNumCpy := tc.extra.toForgeL1TxsNum
+				// til doesn't fill the ForgeL1TxsNum
+				batch.Batch.ForgeL1TxsNum = &toForgeL1TxsNumCpy
+				tc.extra.toForgeL1TxsNum++
+			}
+
+			batchNum := batch.Batch.BatchNum
+			for k := range batch.L1CoordinatorTxs {
+				tx := &batch.L1CoordinatorTxs[k]
+				tx.BatchNum = &batchNum
+				tx.EthBlockNum = batch.Batch.EthBlockNum
+			}
+		}
+	}
+
+	// Fill CreatedAccounts
+	for i := range blocks {
+		block := &blocks[i]
+		for j := range block.Rollup.Batches {
+			batch := &block.Rollup.Batches[j]
+			l1Txs := []common.L1Tx{}
+			if batch.L1Batch {
+				for _, tx := range tc.Queues[*batch.Batch.ForgeL1TxsNum] {
+					l1Txs = append(l1Txs, tx.L1Tx)
+				}
+			}
+			l1Txs = append(l1Txs, batch.L1CoordinatorTxs...)
+			for k := range l1Txs {
+				tx := &l1Txs[k]
+				if tx.Type == common.TxTypeCreateAccountDeposit ||
+					tx.Type == common.TxTypeCreateAccountDepositTransfer {
+					user, ok := tc.usersByIdx[tc.extra.idx]
+					if !ok {
+						return fmt.Errorf("Created account with idx: %v not found", tc.extra.idx)
+					}
+					batch.CreatedAccounts = append(batch.CreatedAccounts,
+						common.Account{
+							Idx:       common.Idx(tc.extra.idx),
+							TokenID:   tx.TokenID,
+							BatchNum:  batch.Batch.BatchNum,
+							PublicKey: user.BJJ.Public(),
+							EthAddr:   user.Addr,
+							Nonce:     0,
+							Balance:   big.NewInt(0),
+						})
+					tc.extra.idx++
+				}
+			}
+		}
+	}
+
+	// Fill expected positions in L1CoordinatorTxs and L2Txs
+	for i := range blocks {
+		block := &blocks[i]
+		for j := range block.Rollup.Batches {
+			batch := &block.Rollup.Batches[j]
+			position := 0
+			if batch.L1Batch {
+				position = len(tc.Queues[*batch.Batch.ForgeL1TxsNum])
+			}
+			for k := range batch.L1CoordinatorTxs {
+				tx := &batch.L1CoordinatorTxs[k]
+				tx.Position = position
+				position++
+				nTx, err := common.NewL1Tx(tx)
+				if err != nil {
+					return err
+				}
+				*tx = *nTx
+			}
+			for k := range batch.L2Txs {
+				tx := &batch.L2Txs[k]
+				tx.Position = position
+				position++
+				tc.extra.nonces[tx.FromIdx]++
+				tx.Nonce = tc.extra.nonces[tx.FromIdx]
+				nTx, err := common.NewL2Tx(tx)
+				if err != nil {
+					return err
+				}
+				*tx = *nTx
+			}
+		}
+	}
+
+	// Fill ExitTree (only AccountIdx and Balance)
+	for i := range blocks {
+		block := &blocks[i]
+		for j := range block.Rollup.Batches {
+			batch := &block.Rollup.Batches[j]
+			if batch.L1Batch {
+				for _, _tx := range tc.Queues[*batch.Batch.ForgeL1TxsNum] {
+					tx := _tx.L1Tx
+					if tx.Type == common.TxTypeForceExit {
+						batch.ExitTree =
+							append(batch.ExitTree,
+								common.ExitInfo{
+									AccountIdx: tx.FromIdx,
+									Balance:    tx.Amount,
+								})
+					}
+				}
+			}
+			for k := range batch.L2Txs {
+				tx := &batch.L2Txs[k]
+				if tx.Type == common.TxTypeExit {
+					batch.ExitTree = append(batch.ExitTree, common.ExitInfo{
+						AccountIdx: tx.FromIdx,
+						Balance:    tx.Amount,
+					})
+				}
+				fee, err := common.CalcFeeAmount(tx.Amount, tx.Fee)
+				if err != nil {
+					return err
+				}
+
+				// Find the TokenID of the tx
+				fromAcc, ok := tc.accountsByIdx[int(tx.FromIdx)]
+				if !ok {
+					return fmt.Errorf("L2tx.FromIdx idx: %v not found", tx.FromIdx)
+				}
+
+				// Find the idx of the CoordUser for the
+				// TokenID, and if it exists, add the fee to
+				// the collectedFees.  Only consider the
+				// coordinator account to receive fee if it was
+				// created in this or a previous batch
+				if acc, ok := tc.l1CreatedAccounts[idxTokenIDToString(cfg.CoordUser, fromAcc.TokenID)]; ok &&
+					common.BatchNum(acc.BatchNum) <= batch.Batch.BatchNum {
+					found := false
+					for _, idx := range batch.Batch.FeeIdxsCoordinator {
+						if idx == common.Idx(acc.Idx) {
+							found = true
+							break
+						}
+					}
+					if !found {
+						batch.Batch.FeeIdxsCoordinator = append(batch.Batch.FeeIdxsCoordinator,
+							common.Idx(acc.Idx))
+						batch.Batch.CollectedFees[fromAcc.TokenID] = big.NewInt(0)
+					}
+					collected := batch.Batch.CollectedFees[fromAcc.TokenID]
+					collected.Add(collected, fee)
+				}
+			}
+		}
+	}
+	return nil
+}