feat: support stateful precompile

core/vm/contract.go

@@ -44,6 +44,9 @@ type Contract struct {
 
 	Gas   uint64
 	value *uint256.Int
+
+	// Precompile and Code are mutually exclusive
+	Precompile PrecompiledContract
 }
 
 // NewContract returns a new contract environment for the execution of EVM.
@@ -163,3 +166,7 @@ func (c *Contract) SetCallCode(hash common.Hash, code []byte) {
 	c.Code = code
 	c.CodeHash = hash
 }
+
+func (c *Contract) SetPrecompile(precompile PrecompiledContract) {
+	c.Precompile = precompile
+}

core/vm/contracts.go

@@ -53,6 +53,13 @@ type PrecompiledContract interface {
 	Name() string
 }
 
+// StatefulPrecompiledContract has access to the EVM and Contract context when executed.
+type StatefulPrecompiledContract interface {
+	PrecompiledContract
+
+	RunWithEVM(evm *EVM, contract *Contract) ([]byte, error)
+}
+
 // PrecompiledContracts contains the precompiled contracts supported at the given fork.
 type PrecompiledContracts map[common.Address]PrecompiledContract
 

core/vm/evm.go

@@ -282,22 +282,23 @@ func (evm *EVM) Call(caller common.Address, addr common.Address, input []byte, g
 	}
 	evm.Context.Transfer(evm.StateDB, caller, addr, value)
 
+	// The contract is a scoped environment for this execution context only.
+	contract := NewContract(caller, addr, value, gas, evm.jumpDests)
+	contract.IsSystemCall = isSystemCall(caller)
 	if isPrecompile {
-		ret, gas, err = RunPrecompiledContract(p, input, gas, evm.Config.Tracer)
+		contract.SetPrecompile(p)
 	} else {
 		// Initialise a new contract and set the code that is to be used by the EVM.
 		code := evm.resolveCode(addr)
 		if len(code) == 0 {
-			ret, err = nil, nil // gas is unchanged
-		} else {
-			// The contract is a scoped environment for this execution context only.
-			contract := NewContract(caller, addr, value, gas, evm.jumpDests)
-			contract.IsSystemCall = isSystemCall(caller)
+			return nil, gas, nil // gas is unchanged
+		}
+
 		contract.SetCallCode(evm.resolveCodeHash(addr), code)
+	}
+
 	ret, err = evm.Run(contract, input, false)
 	gas = contract.Gas
-		}
-	}
 	// When an error was returned by the EVM or when setting the creation code
 	// above we revert to the snapshot and consume any gas remaining. Additionally,
 	// when we're in homestead this also counts for code storage gas errors.
@@ -345,17 +346,18 @@ func (evm *EVM) CallCode(caller common.Address, addr common.Address, input []byt
 	}
 	var snapshot = evm.StateDB.Snapshot()
 
-	// It is allowed to call precompiles, even via delegatecall
-	if p, isPrecompile := evm.precompile(addr); isPrecompile {
-		ret, gas, err = RunPrecompiledContract(p, input, gas, evm.Config.Tracer)
-	} else {
 	// Initialise a new contract and set the code that is to be used by the EVM.
 	// The contract is a scoped environment for this execution context only.
 	contract := NewContract(caller, caller, value, gas, evm.jumpDests)
+	if p, isPrecompile := evm.precompile(addr); isPrecompile {
+		// It is allowed to call precompiles, even via delegatecall
+		contract.SetPrecompile(p)
+	} else {
 		contract.SetCallCode(evm.resolveCodeHash(addr), evm.resolveCode(addr))
+	}
+
 	ret, err = evm.Run(contract, input, false)
 	gas = contract.Gas
-	}
 	if err != nil {
 		evm.StateDB.RevertToSnapshot(snapshot)
 		if err != ErrExecutionReverted {
@@ -388,18 +390,19 @@ func (evm *EVM) DelegateCall(originCaller common.Address, caller common.Address,
 	}
 	var snapshot = evm.StateDB.Snapshot()
 
-	// It is allowed to call precompiles, even via delegatecall
-	if p, isPrecompile := evm.precompile(addr); isPrecompile {
-		ret, gas, err = RunPrecompiledContract(p, input, gas, evm.Config.Tracer)
-	} else {
 	// Initialise a new contract and make initialise the delegate values
 	//
 	// Note: The value refers to the original value from the parent call.
 	contract := NewContract(originCaller, caller, value, gas, evm.jumpDests)
+	if p, isPrecompile := evm.precompile(addr); isPrecompile {
+		// It is allowed to call precompiles, even via delegatecall
+		contract.SetPrecompile(p)
+	} else {
 		contract.SetCallCode(evm.resolveCodeHash(addr), evm.resolveCode(addr))
+	}
+
 	ret, err = evm.Run(contract, input, false)
 	gas = contract.Gas
-	}
 	if err != nil {
 		evm.StateDB.RevertToSnapshot(snapshot)
 		if err != ErrExecutionReverted {
@@ -441,20 +444,20 @@ func (evm *EVM) StaticCall(caller common.Address, addr common.Address, input []b
 	// future scenarios
 	evm.StateDB.AddBalance(addr, new(uint256.Int), tracing.BalanceChangeTouchAccount)
 
-	if p, isPrecompile := evm.precompile(addr); isPrecompile {
-		ret, gas, err = RunPrecompiledContract(p, input, gas, evm.Config.Tracer)
-	} else {
 	// Initialise a new contract and set the code that is to be used by the EVM.
 	// The contract is a scoped environment for this execution context only.
 	contract := NewContract(caller, addr, new(uint256.Int), gas, evm.jumpDests)
+	if p, isPrecompile := evm.precompile(addr); isPrecompile {
+		contract.SetPrecompile(p)
+	} else {
 		contract.SetCallCode(evm.resolveCodeHash(addr), evm.resolveCode(addr))
+	}
 
 	// When an error was returned by the EVM or when setting the creation code
 	// above we revert to the snapshot and consume any gas remaining. Additionally
 	// when we're in Homestead this also counts for code storage gas errors.
 	ret, err = evm.Run(contract, input, true)
 	gas = contract.Gas
-	}
 	if err != nil {
 		evm.StateDB.RevertToSnapshot(snapshot)
 		if err != ErrExecutionReverted {

core/vm/interpreter.go

@@ -110,6 +110,17 @@ func (evm *EVM) Run(contract *Contract, input []byte, readOnly bool) (ret []byte
 	// as every returning call will return new data anyway.
 	evm.returnData = nil
 
+	contract.Input = input
+
+	if contract.Precompile != nil {
+		if stateful, ok := contract.Precompile.(StatefulPrecompiledContract); ok {
+			return stateful.RunWithEVM(evm, contract)
+		}
+
+		ret, contract.Gas, err = RunPrecompiledContract(contract.Precompile, input, contract.Gas, evm.Config.Tracer)
+		return ret, err
+	}
+
 	// Don't bother with the execution if there's no code.
 	if len(contract.Code) == 0 {
 		return nil, nil
@@ -145,7 +156,6 @@ func (evm *EVM) Run(contract *Contract, input []byte, readOnly bool) (ret []byte
 		returnStack(stack)
 		mem.Free()
 	}()
-	contract.Input = input
 
 	if debug {
 		defer func() { // this deferred method handles exit-with-error