Add multigas to EndTxHook and safeDecrement

Author: MishkaRogachev

arbitrum/multigas/resources.go

@@ -322,6 +322,26 @@ func (z MultiGas) SaturatingIncrement(kind ResourceKind, gas uint64) MultiGas {
 	return res
 }
 
+// SaturatingDecrement returns a copy of z with the given resource kind
+// and the total decremented by gas. On underflow, the field(s) are clamped to 0.
+func (z MultiGas) SaturatingDecrement(kind ResourceKind, gas uint64) MultiGas {
+	res := z
+
+	if v, c := bits.Sub64(res.gas[kind], gas, 0); c != 0 {
+		res.gas[kind] = 0 // clamp
+	} else {
+		res.gas[kind] = v
+	}
+
+	if t, c := bits.Sub64(res.total, gas, 0); c != 0 {
+		res.total = 0 // clamp
+	} else {
+		res.total = t
+	}
+
+	return res
+}
+
 // SaturatingIncrementInto increments the given resource kind and the total
 // in place by gas. On overflow, the affected field(s) are clamped to MaxUint64.
 // Unlike SaturatingIncrement, this method mutates the receiver directly and

arbitrum/multigas/resources_test.go

@@ -388,6 +388,43 @@ func TestSaturatingIncrementIntoClampsOnOverflow(t *testing.T) {
 	}
 }
 
+func TestSaturatingDecrement(t *testing.T) {
+	// normal decrement
+	gas := ComputationGas(10)
+	newGas := gas.SaturatingDecrement(ResourceKindComputation, 5)
+	if got, want := newGas.Get(ResourceKindComputation), uint64(5); got != want {
+		t.Errorf("unexpected computation gas: got %v, want %v", got, want)
+	}
+	if got, want := newGas.SingleGas(), uint64(5); got != want {
+		t.Errorf("unexpected single gas: got %v, want %v", got, want)
+	}
+
+	// saturating decrement on kind
+	gas = ComputationGas(0)
+	newGas = gas.SaturatingDecrement(ResourceKindComputation, 1)
+	if got, want := newGas.Get(ResourceKindComputation), uint64(0); got != want {
+		t.Errorf("expected computation gas to clamp to zero: got %v, want %v", got, want)
+	}
+	if got, want := newGas.SingleGas(), uint64(0); got != want {
+		t.Errorf("expected total to clamp to zero: got %v, want %v", got, want)
+	}
+
+	// total-only decrement case
+	gas = MultiGasFromPairs(
+		Pair{ResourceKindComputation, math.MaxUint64 - 1},
+		Pair{ResourceKindHistoryGrowth, 1},
+	)
+
+	newGas = gas.SaturatingDecrement(ResourceKindHistoryGrowth, 1)
+	if got, want := newGas.Get(ResourceKindHistoryGrowth), uint64(0); got != want {
+		t.Errorf("unexpected history growth gas: got %v, want %v", got, want)
+	}
+
+	if got, want := newGas.SingleGas(), uint64(math.MaxUint64-1); got != want {
+		t.Errorf("unexpected total gas: got %v, want %v", got, want)
+	}
+}
+
 func TestMultiGasSingleGasTracking(t *testing.T) {
 	g := ZeroGas()
 	if got := g.SingleGas(); got != 0 {

core/state_transition.go

@@ -774,7 +774,7 @@ func (st *stateTransition) execute() (*ExecutionResult, error) {
 		tracer.CaptureArbitrumTransfer(nil, &tipReceipient, tipAmount, false, tracing.BalanceIncreaseRewardTransactionFee)
 	}
 
-	st.evm.ProcessingHook.EndTxHook(st.gasRemaining, vmerr == nil)
+	st.evm.ProcessingHook.EndTxHook(st.gasRemaining, usedMultiGas, vmerr == nil)
 
 	// Arbitrum: record self destructs
 	if tracer := st.evm.Config.Tracer; tracer != nil && tracer.CaptureArbitrumTransfer != nil {

core/vm/evm_arbitrum.go

@@ -50,7 +50,7 @@ type TxProcessingHook interface {
 	HeldGas() uint64
 	NonrefundableGas() uint64
 	DropTip() bool
-	EndTxHook(totalGasUsed uint64, evmSuccess bool)
+	EndTxHook(totalGasUsed uint64, usedMultiGas multigas.MultiGas, evmSuccess bool)
 	ScheduledTxes() types.Transactions
 	L1BlockNumber(blockCtx BlockContext) (uint64, error)
 	L1BlockHash(blockCtx BlockContext, l1BlocKNumber uint64) (common.Hash, error)
@@ -83,7 +83,7 @@ func (p DefaultTxProcessor) NonrefundableGas() uint64 { return 0 }
 
 func (p DefaultTxProcessor) DropTip() bool { return false }
 
-func (p DefaultTxProcessor) EndTxHook(_ uint64, _ bool) {}
+func (p DefaultTxProcessor) EndTxHook(_ uint64, _ multigas.MultiGas, _ bool) {}
 
 func (p DefaultTxProcessor) ScheduledTxes() types.Transactions {
 	return types.Transactions{}