refactor!: gas consumption for stateful precompiles

Author: ARR4N

core/vm/contracts.go

@@ -174,8 +174,7 @@ func (args *evmCallArgs) RunPrecompiledContract(p PrecompiledContract, input []b
 		return nil, 0, ErrOutOfGas
 	}
 	suppliedGas -= gasCost
-	output, err := args.run(p, input)
-	return output, suppliedGas, err
+	return args.run(p, input, suppliedGas)
 }
 
 // ECRECOVER implemented as a native contract.

core/vm/contracts.libevm.go

@@ -50,42 +50,39 @@ const (
 
 // run runs the [PrecompiledContract], differentiating between stateful and
 // regular types.
-func (args *evmCallArgs) run(p PrecompiledContract, input []byte) (ret []byte, err error) {
+func (args *evmCallArgs) run(p PrecompiledContract, input []byte, suppliedGas uint64) (ret []byte, remainingGas uint64, err error) {
 	if p, ok := p.(statefulPrecompile); ok {
-		return p.run(args, input)
+		return p(args, input, suppliedGas)
 	}
-	return p.Run(input)
+	// Gas consumption for regular precompiles was already handled by the native
+	// RunPrecompiledContract(), which called this method.
+	ret, err = p.Run(input)
+	return ret, suppliedGas, err
 }
 
-// PrecompiledStatefulRun is the stateful equivalent of the Run() method of a
+// PrecompiledStatefulContract is the stateful equivalent of a
 // [PrecompiledContract].
-type PrecompiledStatefulRun func(env PrecompileEnvironment, input []byte) ([]byte, error)
+type PrecompiledStatefulContract func(env PrecompileEnvironment, input []byte, suppliedGas uint64) (ret []byte, remainingGas uint64, err error)
 
 // NewStatefulPrecompile constructs a new PrecompiledContract that can be used
 // via an [EVM] instance but MUST NOT be called directly; a direct call to Run()
 // reserves the right to panic. See other requirements defined in the comments
 // on [PrecompiledContract].
-func NewStatefulPrecompile(run PrecompiledStatefulRun, requiredGas func([]byte) uint64) PrecompiledContract {
-	return statefulPrecompile{
-		gas: requiredGas,
-		run: run,
-	}
+func NewStatefulPrecompile(run PrecompiledStatefulContract) PrecompiledContract {
+	return statefulPrecompile(run)
 }
 
-type statefulPrecompile struct {
-	gas func([]byte) uint64
-	run PrecompiledStatefulRun
-}
+type statefulPrecompile PrecompiledStatefulContract
 
-func (p statefulPrecompile) RequiredGas(input []byte) uint64 {
-	return p.gas(input)
-}
+// RequiredGas always returns zero as this gas is consumed before the contract
+// is run.
+func (p statefulPrecompile) RequiredGas(input []byte) uint64 { return 0 }
 
 func (p statefulPrecompile) Run([]byte) ([]byte, error) {
 	// https://google.github.io/styleguide/go/best-practices.html#when-to-panic
 	// This would indicate an API misuse and would occur in tests, not in
 	// production.
-	panic(fmt.Sprintf("BUG: call to %T.Run(); MUST call %T", p, p.run))
+	panic(fmt.Sprintf("BUG: call to %T.Run(); MUST call %T itself", p, p))
 }
 
 // A PrecompileEnvironment provides information about the context in which a

core/vm/contracts.libevm_test.go

@@ -93,23 +93,18 @@ func TestNewStatefulPrecompile(t *testing.T) {
 			caller, self, stateVal, readOnly, input,
 		))
 	}
-	run := func(env vm.PrecompileEnvironment, input []byte) ([]byte, error) {
+	run := func(env vm.PrecompileEnvironment, input []byte, suppliedGas uint64) ([]byte, uint64, error) {
 		if got, want := env.StateDB() != nil, !env.ReadOnly(); got != want {
-			return nil, fmt.Errorf("PrecompileEnvironment().StateDB() must be non-nil i.f.f. not read-only; got non-nil? %t; want %t", got, want)
+			return nil, 0, fmt.Errorf("PrecompileEnvironment().StateDB() must be non-nil i.f.f. not read-only; got non-nil? %t; want %t", got, want)
 		}
 
 		addrs := env.Addresses()
 		val := env.ReadOnlyState().GetState(precompile, slot)
-		return makeOutput(addrs.Caller, addrs.Self, input, val, env.ReadOnly()), nil
+		return makeOutput(addrs.Caller, addrs.Self, input, val, env.ReadOnly()), suppliedGas - gasCost, nil
 	}
 	hooks := &hookstest.Stub{
 		PrecompileOverrides: map[common.Address]libevm.PrecompiledContract{
-			precompile: vm.NewStatefulPrecompile(
-				run,
-				func(b []byte) uint64 {
-					return gasCost
-				},
-			),
+			precompile: vm.NewStatefulPrecompile(run),
 		},
 	}
 	hooks.Register(t)
@@ -204,13 +199,12 @@ func TestInheritReadOnly(t *testing.T) {
 	hooks := &hookstest.Stub{
 		PrecompileOverrides: map[common.Address]libevm.PrecompiledContract{
 			precompile: vm.NewStatefulPrecompile(
-				func(env vm.PrecompileEnvironment, input []byte) ([]byte, error) {
+				func(env vm.PrecompileEnvironment, input []byte, suppliedGas uint64) ([]byte, uint64, error) {
 					if env.ReadOnly() {
-						return []byte{ifReadOnly}, nil
+						return []byte{ifReadOnly}, suppliedGas, nil
 					}
-					return []byte{ifNotReadOnly}, nil
+					return []byte{ifNotReadOnly}, suppliedGas, nil
 				},
-				func([]byte) uint64 { return 0 },
 			),
 		},
 	}