From 1fd35aa5f37fbb5802118e28291b5337f436dbb6 Mon Sep 17 00:00:00 2001
From: codchen <codchen03@gmail.com>
Date: Mon, 15 Apr 2024 14:46:10 +0800
Subject: [PATCH] Refactor BuyGas

---
 core/state_transition.go                      | 30 ++++++++++++-------
 .../internal/tracetest/calltrace_test.go      |  2 +-
 eth/tracers/tracers_test.go                   |  2 +-
 3 files changed, 22 insertions(+), 12 deletions(-)

diff --git a/core/state_transition.go b/core/state_transition.go
index e901077cfc..400b81462d 100644
--- a/core/state_transition.go
+++ b/core/state_transition.go
@@ -180,7 +180,7 @@ func TransactionToMessage(tx *types.Transaction, s types.Signer, baseFee *big.In
 // indicates a core error meaning that the message would always fail for that particular
 // state and would never be accepted within a block.
 func ApplyMessage(evm *vm.EVM, msg *Message, gp *GasPool) (*ExecutionResult, error) {
-	return NewStateTransition(evm, msg, gp).TransitionDb()
+	return NewStateTransition(evm, msg, gp, false).TransitionDb()
 }
 
 // StateTransition represents a state transition.
@@ -212,15 +212,17 @@ type StateTransition struct {
 	initialGas   uint64
 	state        vm.StateDB
 	evm          *vm.EVM
+	feeCharged   bool
 }
 
 // NewStateTransition initialises and returns a new state transition object.
-func NewStateTransition(evm *vm.EVM, msg *Message, gp *GasPool) *StateTransition {
+func NewStateTransition(evm *vm.EVM, msg *Message, gp *GasPool, feeCharged bool) *StateTransition {
 	return &StateTransition{
-		gp:    gp,
-		evm:   evm,
-		msg:   msg,
-		state: evm.StateDB,
+		gp:         gp,
+		evm:        evm,
+		msg:        msg,
+		state:      evm.StateDB,
+		feeCharged: feeCharged,
 	}
 }
 
@@ -232,7 +234,7 @@ func (st *StateTransition) to() common.Address {
 	return *st.msg.To
 }
 
-func (st *StateTransition) buyGas() error {
+func (st *StateTransition) BuyGas() error {
 	mgval := new(big.Int).SetUint64(st.msg.GasLimit)
 	mgval = mgval.Mul(mgval, st.msg.GasPrice)
 	balanceCheck := new(big.Int).Set(mgval)
@@ -256,10 +258,14 @@ func (st *StateTransition) buyGas() error {
 	if have, want := st.state.GetBalance(st.msg.From), balanceCheck; have.Cmp(want) < 0 {
 		return fmt.Errorf("%w: address %v have %v want %v", ErrInsufficientFunds, st.msg.From.Hex(), have, want)
 	}
+	st.state.SubBalance(st.msg.From, mgval, tracing.BalanceDecreaseGasBuy)
+	return nil
+}
+
+func (st *StateTransition) initGas() error {
 	if err := st.gp.SubGas(st.msg.GasLimit); err != nil {
 		return err
 	}
-
 	if st.evm.Config.Tracer != nil && st.evm.Config.Tracer.OnGasChange != nil {
 		st.evm.Config.Tracer.OnGasChange(0, st.msg.GasLimit, tracing.GasChangeTxInitialBalance)
 	}
@@ -267,7 +273,6 @@ func (st *StateTransition) buyGas() error {
 	st.gasRemaining += st.msg.GasLimit
 
 	st.initialGas = st.msg.GasLimit
-	st.state.SubBalance(st.msg.From, mgval, tracing.BalanceDecreaseGasBuy)
 	return nil
 }
 
@@ -346,7 +351,12 @@ func (st *StateTransition) preCheck() error {
 			}
 		}
 	}
-	return st.buyGas()
+	if !st.feeCharged {
+		if err := st.BuyGas(); err != nil {
+			return err
+		}
+	}
+	return st.initGas()
 }
 
 // TransitionDb will transition the state by applying the current message and
diff --git a/eth/tracers/internal/tracetest/calltrace_test.go b/eth/tracers/internal/tracetest/calltrace_test.go
index 9a115c7235..00ac53f315 100644
--- a/eth/tracers/internal/tracetest/calltrace_test.go
+++ b/eth/tracers/internal/tracetest/calltrace_test.go
@@ -253,7 +253,7 @@ func benchTracer(tracerName string, test *callTracerTest, b *testing.B) {
 		}
 		evm := vm.NewEVM(context, txContext, statedb, test.Genesis.Config, vm.Config{Tracer: tracer.Hooks})
 		snap := statedb.Snapshot()
-		st := core.NewStateTransition(evm, msg, new(core.GasPool).AddGas(tx.Gas()))
+		st := core.NewStateTransition(evm, msg, new(core.GasPool).AddGas(tx.Gas()), false)
 		if _, err = st.TransitionDb(); err != nil {
 			b.Fatalf("failed to execute transaction: %v", err)
 		}
diff --git a/eth/tracers/tracers_test.go b/eth/tracers/tracers_test.go
index 505db4fa4e..4c85f28af7 100644
--- a/eth/tracers/tracers_test.go
+++ b/eth/tracers/tracers_test.go
@@ -99,7 +99,7 @@ func BenchmarkTransactionTrace(b *testing.B) {
 
 	for i := 0; i < b.N; i++ {
 		snap := statedb.Snapshot()
-		st := core.NewStateTransition(evm, msg, new(core.GasPool).AddGas(tx.Gas()))
+		st := core.NewStateTransition(evm, msg, new(core.GasPool).AddGas(tx.Gas()), false)
 		_, err = st.TransitionDb()
 		if err != nil {
 			b.Fatal(err)