Add more methods to MultiGas struct

arbitrum/multigas/resources.go

@@ -1,5 +1,7 @@
 package multigas
 
+import "github.com/ethereum/go-ethereum/common/math"
+
 // ResourceKind represents a dimension for the multi-dimensional gas.
 type ResourceKind uint8
 
@@ -13,4 +15,88 @@ const (
 )
 
 // MultiGas tracks gas for each resource separately.
-type MultiGas [NumResourceKind]uint64
+type MultiGas struct {
+	gas    [NumResourceKind]uint64
+	refund uint64
+}
+
+func ZeroGas() *MultiGas {
+	return &MultiGas{}
+}
+
+func NewMultiGas(kind ResourceKind, amount uint64) *MultiGas {
+	mg := ZeroGas()
+	mg.gas[kind] = amount
+	return mg
+}
+
+func ComputationGas(amount uint64) *MultiGas {
+	return NewMultiGas(ResourceKindComputation, amount)
+}
+
+func HistoryGrowthGas(amount uint64) *MultiGas {
+	return NewMultiGas(ResourceKindHistoryGrowth, amount)
+}
+
+func StorageAccessGas(amount uint64) *MultiGas {
+	return NewMultiGas(ResourceKindStorageAccess, amount)
+}
+
+func StorageGrowthGas(amount uint64) *MultiGas {
+	return NewMultiGas(ResourceKindStorageGrowth, amount)
+}
+
+func (z *MultiGas) Get(kind ResourceKind) uint64 {
+	return z.gas[kind]
+}
+
+func (z *MultiGas) Set(kind ResourceKind, gas uint64) *MultiGas {
+	z.gas[kind] = gas
+	return z
+}
+
+// GetRefund gets the SSTORE refund computed at the end of the transaction.
+func (z *MultiGas) GetRefund() uint64 {
+	return z.refund
+}
+
+// SetRefund sets the SSTORE refund computed at the end of the transaction.
+func (z *MultiGas) SetRefund(amount uint64) *MultiGas {
+	z.refund = amount
+	return z
+}
+
+// SafeAdd sets z to the sum x+y and returns z and checks for overflow.
+func (z *MultiGas) SafeAdd(x *MultiGas, y *MultiGas) (*MultiGas, bool) {
+	for i := ResourceKindUnknown; i < NumResourceKind; i++ {
+		var overflow bool
+		z.gas[i], overflow = math.SafeAdd(x.gas[i], y.gas[i])
+		if overflow {
+			return z, overflow
+		}
+	}
+	return z, false
+}
+
+// SafeIncrement increments the given resource kind by the amount of gas and checks for overflow.
+func (z *MultiGas) SafeIncrement(kind ResourceKind, gas uint64) bool {
+	result, overflow := math.SafeAdd(z.gas[kind], gas)
+	if overflow {
+		return overflow
+	}
+	z.gas[kind] = result
+	return false
+}
+
+// SingleGas converts the multi-gas to single-dimensional gas checking for overflow.
+func (z *MultiGas) SingleGas() (uint64, bool) {
+	var sum uint64
+	for _, value := range z.gas {
+		var overflow bool
+		sum, overflow = math.SafeAdd(sum, value)
+		if overflow {
+			return 0, overflow
+		}
+	}
+	return sum, false
+}

arbitrum/multigas/resources_test.go

@@ -0,0 +1,63 @@
+package multigas
+
+import (
+	"math"
+	"testing"
+)
+
+func TestMultiGas(t *testing.T) {
+	// Test SafeAdd
+	gas, overflow := new(MultiGas).SafeAdd(ComputationGas(10), HistoryGrowthGas(20))
+	if overflow {
+		t.Errorf("unexpected overflow: got %v, want %v", overflow, false)
+	}
+	if got, want := gas.Get(ResourceKindComputation), uint64(10); got != want {
+		t.Errorf("unexpected computation gas: got %v, want %v", got, want)
+	}
+	if got, want := gas.Get(ResourceKindHistoryGrowth), uint64(20); got != want {
+		t.Errorf("unexpected history growth gas: got %v, want %v", got, want)
+	}
+	if got, want := gas.Get(ResourceKindStorageAccess), uint64(0); got != want {
+		t.Errorf("unexpected storage access gas: got %v, want %v", got, want)
+	}
+	if got, want := gas.Get(ResourceKindStorageGrowth), uint64(0); got != want {
+		t.Errorf("unexpected storage growth gas: got %v, want %v", got, want)
+	}
+
+	// Test SafeAdd checks for overflow
+	_, overflow = new(MultiGas).SafeAdd(ComputationGas(math.MaxUint64), ComputationGas(1))
+	if !overflow {
+		t.Errorf("unexpected overflow: got %v, want %v", overflow, true)
+	}
+
+	// Test SafeIncrement
+	overflow = gas.SafeIncrement(ResourceKindComputation, 11)
+	if overflow {
+		t.Errorf("unexpected overflow: got %v, want %v", overflow, false)
+	}
+	if got, want := gas.Get(ResourceKindComputation), uint64(21); got != want {
+		t.Errorf("unexpected computation gas: got %v, want %v", got, want)
+	}
+
+	// Test SafeIncrement checks for overflow
+	overflow = gas.SafeIncrement(ResourceKindComputation, math.MaxUint64)
+	if !overflow {
+		t.Errorf("unexpected overflow: got %v, want %v", overflow, true)
+	}
+
+	// Test SingleGas
+	singleGas, overflow := gas.SingleGas()
+	if overflow {
+		t.Errorf("unexpected overflow: got %v, want %v", overflow, false)
+	}
+	if want := uint64(41); singleGas != want {
+		t.Errorf("unexpected storage growth gas: got %v, want %v", singleGas, want)
+	}
+
+	// Test SingleGas checks for overflow
+	gas.Set(ResourceKindComputation, math.MaxUint64)
+	_, overflow = gas.SingleGas()
+	if !overflow {
+		t.Errorf("unexpected overflow: got %v, want %v", overflow, true)
+	}
+}