Fix tx indexes

Refactor duplicated logic
iotaledger · Jan 16, 2025 · eb89081 · eb89081
1 parent 13a11e5
commit eb89081
Show file tree

Hide file tree

Showing 4 changed files with 81 additions and 108 deletions.
diff --git a/packages/evm/jsonrpc/evmchain.go b/packages/evm/jsonrpc/evmchain.go
@@ -725,17 +725,11 @@ func (e *EVMChain) debugTraceBlock(config *tracers.TraceConfig, block *types.Blo
 	blockNumber := uint64(iscBlock.BlockIndex())
 
 	blockTxs := block.Transactions()
-	fakeTxs := make([]*types.Transaction, 0, len(blockTxs))
-	for _, tx := range blockTxs {
-		if evmutil.IsFakeTransaction(tx) {
-			fakeTxs = append(fakeTxs, tx)
-		}
-	}
 
 	tracer, err := newTracer(tracerType, &tracers.Context{
 		BlockHash:   block.Hash(),
 		BlockNumber: new(big.Int).SetUint64(blockNumber),
-	}, config.TracerConfig, true, types.Transactions(fakeTxs))
+	}, config.TracerConfig, true, blockTxs)
 	if err != nil {
 		return nil, err
 	}

diff --git a/packages/evm/jsonrpc/tracer.go b/packages/evm/jsonrpc/tracer.go
@@ -6,6 +6,7 @@ import (
 
 	"github.com/ethereum/go-ethereum/core/types"
 	"github.com/ethereum/go-ethereum/eth/tracers"
+	"github.com/iotaledger/wasp/packages/evm/evmutil"
 )
 
 type Tracer struct {
@@ -28,3 +29,46 @@ func newTracer(tracerType string, ctx *tracers.Context, cfg json.RawMessage, tra
 	}
 	return fn(ctx, cfg, traceBlock, initValue)
 }
+
+func GetTraceResults(
+	blockTxs []*types.Transaction,
+	traceBlock bool,
+	getFakeTxTrace func(tx *types.Transaction) (json.RawMessage, error),
+	getTxTrace func(tx *types.Transaction) (json.RawMessage, error),
+	getSingleTxTrace func() (json.RawMessage, error),
+	reason error,
+) (json.RawMessage, error) {
+	var traceResult []byte
+	var err error
+	if traceBlock {
+		results := make([]TxTraceResult, 0, len(blockTxs))
+		var jsResult json.RawMessage
+		for _, tx := range blockTxs {
+			if evmutil.IsFakeTransaction(tx) {
+				jsResult, err = getFakeTxTrace(tx)
+				if err != nil {
+					return nil, err
+				}
+			} else {
+				jsResult, err = getTxTrace(tx)
+				if err != nil {
+					return nil, err
+				}
+			}
+
+			results = append(results, TxTraceResult{TxHash: tx.Hash(), Result: jsResult})
+		}
+
+		traceResult, err = json.Marshal(results)
+		if err != nil {
+			return nil, err
+		}
+	} else {
+		traceResult, err = getSingleTxTrace()
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	return traceResult, reason
+}
diff --git a/packages/evm/jsonrpc/tracer_call.go b/packages/evm/jsonrpc/tracer_call.go
@@ -121,14 +121,14 @@ type TxTraceResult struct {
 
 type callTracer struct {
 	txToStack  map[common.Hash][]CallFrame
-	currentTx  common.Hash
 	config     callTracerConfig
 	gasLimit   uint64
 	depth      int
 	interrupt  atomic.Bool // Atomic flag to signal execution interruption
 	reason     error       // Textual reason for the interruption
+	currentTx  common.Hash
 	traceBlock bool
-	fakeTxs    []*types.Transaction
+	blockTxs   []*types.Transaction
 }
 
 type callTracerConfig struct {
@@ -172,7 +172,7 @@ func newCallTracer(ctx *tracers.Context, cfg json.RawMessage, traceBlock bool, i
 	}, nil
 }
 
-func newCallTracerObject(_ *tracers.Context, cfg json.RawMessage, traceBlock bool, fakeTxs []*types.Transaction) (*callTracer, error) {
+func newCallTracerObject(_ *tracers.Context, cfg json.RawMessage, traceBlock bool, blockTxs []*types.Transaction) (*callTracer, error) {
 	var config callTracerConfig
 	if cfg != nil {
 		if err := json.Unmarshal(cfg, &config); err != nil {
@@ -181,7 +181,7 @@ func newCallTracerObject(_ *tracers.Context, cfg json.RawMessage, traceBlock boo
 	}
 	// First callframe contains tx context info
 	// and is populated on start and end.
-	return &callTracer{txToStack: make(map[common.Hash][]CallFrame), currentTx: common.Hash{}, config: config, traceBlock: traceBlock, fakeTxs: fakeTxs}, nil
+	return &callTracer{txToStack: make(map[common.Hash][]CallFrame), currentTx: common.Hash{}, config: config, traceBlock: traceBlock, blockTxs: blockTxs}, nil
 }
 
 // OnEnter is called when EVM enters a new scope (via call, create or selfdestruct).
@@ -290,45 +290,21 @@ var ErrIncorrectTopLevelCalls = errors.New("incorrect number of top-level calls"
 // GetResult returns the json-encoded nested list of call traces, and any
 // error arising from the encoding or forceful termination (via `Stop`).
 func (t *callTracer) GetResult() (json.RawMessage, error) {
-	if t.traceBlock {
-		// otherwise return all call frames
-		results := make([]TxTraceResult, 0, len(t.txToStack))
-		for txHash, stack := range t.txToStack {
-			csJSON, err := json.Marshal(stack[0])
-			if err != nil {
-				return nil, err
-			}
-			results = append(results, TxTraceResult{
-				TxHash: txHash,
-				Result: csJSON,
-			})
-		}
-
-		for _, tx := range t.fakeTxs {
-			csJSON, err := t.TraceFakeTx(tx)
-			if err != nil {
-				return nil, err
+	return GetTraceResults(
+		t.blockTxs,
+		t.traceBlock,
+		t.TraceFakeTx,
+		func(tx *types.Transaction) (json.RawMessage, error) {
+			stack, ok := t.txToStack[tx.Hash()]
+			if !ok {
+				return nil, fmt.Errorf("no call stack for tx %s", tx.Hash().Hex())
 			}
-			results = append(results, TxTraceResult{TxHash: tx.Hash(), Result: csJSON})
-		}
-
-		res, err := json.Marshal(results)
-		if err != nil {
-			return nil, err
-		}
-		return res, t.reason
-	}
-
-	if len(t.txToStack) != 1 {
-		return nil, ErrIncorrectTopLevelCalls
-	}
-
-	res, err := json.Marshal(t.txToStack[t.currentTx][0])
-	if err != nil {
-		return nil, err
-	}
-
-	return res, t.reason
+			return json.Marshal(stack[0])
+		},
+		func() (json.RawMessage, error) {
+			return json.Marshal(t.txToStack[t.currentTx][0])
+		},
+		t.reason)
 }
 
 // Stop terminates execution of the tracer at the first opportune moment.

diff --git a/packages/evm/jsonrpc/tracer_prestate.go b/packages/evm/jsonrpc/tracer_prestate.go
@@ -58,23 +58,23 @@ type prestateTracer struct {
 	interrupt     atomic.Bool // Atomic flag to signal execution interruption
 	reason        error       // Textual reason for the interruption
 	traceBlock    bool
-	fakeTxs       types.Transactions
+	blockTxs      types.Transactions
 }
 
 type prestateTracerConfig struct {
 	DiffMode bool `json:"diffMode"` // If true, this tracer will return state modifications
 }
 
 func newPrestateTracer(ctx *tracers.Context, cfg json.RawMessage, traceBlock bool, initValue any) (*Tracer, error) {
-	var fakeTxs types.Transactions
+	var blockTxs types.Transactions
 
 	if initValue == nil && traceBlock {
 		return nil, fmt.Errorf("initValue with block transactions is required for block tracing")
 	}
 
 	if initValue != nil {
 		var ok bool
-		fakeTxs, ok = initValue.(types.Transactions)
+		blockTxs, ok = initValue.(types.Transactions)
 		if !ok {
 			return nil, fmt.Errorf("invalid init value type for prestateTracer: %T", initValue)
 		}
@@ -87,7 +87,7 @@ func newPrestateTracer(ctx *tracers.Context, cfg json.RawMessage, traceBlock boo
 		config:     config,
 		traceBlock: traceBlock,
 		states:     make(map[common.Hash]*PrestateTxValue),
-		fakeTxs:    fakeTxs,
+		blockTxs:   blockTxs,
 	}
 	return &Tracer{
 		Tracer: &tracers.Tracer{
@@ -197,63 +197,22 @@ func (t *prestateTracer) OnTxEnd(receipt *types.Receipt, err error) {
 // GetResult returns the json-encoded nested list of call traces, and any
 // error arising from the encoding or forceful termination (via `Stop`).
 func (t *prestateTracer) GetResult() (json.RawMessage, error) {
-	var res []byte
-	var err error
-	if t.traceBlock {
-		if t.config.DiffMode {
-			result := []TxTraceResult{}
-			for txHash, txState := range t.states {
-				var diffResult json.RawMessage
-				diffResult, err = json.Marshal(PrestateDiffResult{txState.Post, txState.Pre})
-				if err != nil {
-					return nil, err
-				}
-				result = append(result, TxTraceResult{TxHash: txHash, Result: diffResult})
-
-				for _, tx := range t.fakeTxs {
-					var csJSON json.RawMessage
-					csJSON, err = t.TraceFakeTx(tx)
-					if err != nil {
-						return nil, err
-					}
-					result = append(result, TxTraceResult{TxHash: tx.Hash(), Result: csJSON})
-				}
+	return GetTraceResults(
+		t.blockTxs,
+		t.traceBlock,
+		t.TraceFakeTx,
+		func(tx *types.Transaction) (json.RawMessage, error) {
+			txState := t.states[tx.Hash()]
+			if t.config.DiffMode {
+				return json.Marshal(PrestateDiffResult{txState.Post, txState.Pre})
 			}
-			res, err = json.Marshal(result)
-		} else {
-			result := []TxTraceResult{}
-			for txHash, txState := range t.states {
-				var preState json.RawMessage
-				preState, err = json.Marshal(txState.Pre)
-				if err != nil {
-					return nil, err
-				}
-				result = append(result, TxTraceResult{TxHash: txHash, Result: preState})
+			return json.Marshal(txState.Pre)
+		}, func() (json.RawMessage, error) {
+			if t.config.DiffMode {
+				return json.Marshal(PrestateDiffResult{t.states[t.currentTxHash].Post, t.states[t.currentTxHash].Pre})
 			}
-
-			for _, tx := range t.fakeTxs {
-				var csJSON json.RawMessage
-				csJSON, err = t.TraceFakeTx(tx)
-				if err != nil {
-					return nil, err
-				}
-				result = append(result, TxTraceResult{TxHash: tx.Hash(), Result: csJSON})
-			}
-
-			res, err = json.Marshal(result)
-		}
-	} else {
-		if t.config.DiffMode {
-			res, err = json.Marshal(PrestateDiffResult{t.states[t.currentTxHash].Post, t.states[t.currentTxHash].Pre})
-		} else {
-			res, err = json.Marshal(t.states[t.currentTxHash].Pre)
-		}
-	}
-
-	if err != nil {
-		return nil, err
-	}
-	return json.RawMessage(res), t.reason
+			return json.Marshal(t.states[t.currentTxHash].Pre)
+		}, t.reason)
 }
 
 // Stop terminates execution of the tracer at the first opportune moment.