feat: vm.JumpTable override

core/vm/contracts.libevm.go

@@ -18,14 +18,10 @@ package vm
 
 import (
 	"fmt"
-	"math/big"
 
 	"github.com/holiman/uint256"
 
 	"github.com/ava-labs/libevm/common"
-	"github.com/ava-labs/libevm/core/types"
-	"github.com/ava-labs/libevm/libevm"
-	"github.com/ava-labs/libevm/params"
 )
 
 // evmCallArgs mirrors the parameters of the [EVM] methods Call(), CallCode(),
@@ -124,20 +120,9 @@ func (p statefulPrecompile) Run([]byte) ([]byte, error) {
 // precompiled contract is being run; and (b) a means of calling other
 // contracts.
 type PrecompileEnvironment interface {
-	ChainConfig() *params.ChainConfig
-	Rules() params.Rules
-	ReadOnly() bool
-	// StateDB will be non-nil i.f.f !ReadOnly().
-	StateDB() StateDB
-	// ReadOnlyState will always be non-nil.
-	ReadOnlyState() libevm.StateReader
-	Addresses() *libevm.AddressContext
-	IncomingCallType() CallType
-
-	BlockHeader() (types.Header, error)
-	BlockNumber() *big.Int
-	BlockTime() uint64
+	Environment
 
+	IncomingCallType() CallType
 	// Call is equivalent to [EVM.Call] except that the `caller` argument is
 	// removed and automatically determined according to the type of call that
 	// invoked the precompile.

core/vm/contracts.libevm_test.go

@@ -145,7 +145,7 @@ func TestNewStatefulPrecompile(t *testing.T) {
 
 	run := func(env vm.PrecompileEnvironment, input []byte, suppliedGas uint64) ([]byte, uint64, error) {
 		if got, want := env.StateDB() != nil, !env.ReadOnly(); 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)
+			return nil, 0, fmt.Errorf("Environment().StateDB() must be non-nil i.f.f. not read-only; got non-nil? %t; want %t", got, want)
 		}
 		hdr, err := env.BlockHeader()
 		if err != nil {

core/vm/environment.libevm.go

@@ -28,6 +28,23 @@ import (
 	"github.com/ava-labs/libevm/params"
 )
 
+// An Environment provides information about the context in which an instruction
+// is being executed.
+type Environment interface {
+	ChainConfig() *params.ChainConfig
+	Rules() params.Rules
+	ReadOnly() bool
+	// StateDB will be non-nil i.f.f !ReadOnly().
+	StateDB() StateDB
+	// ReadOnlyState will always be non-nil.
+	ReadOnlyState() libevm.StateReader
+	Addresses() *libevm.AddressContext
+
+	BlockHeader() (types.Header, error)
+	BlockNumber() *big.Int
+	BlockTime() uint64
+}
+
 var _ PrecompileEnvironment = (*environment)(nil)
 
 type environment struct {

core/vm/evm.libevm_test.go

@@ -38,6 +38,8 @@ func (o *evmArgOverrider) OverrideNewEVMArgs(args *NewEVMArgs) *NewEVMArgs {
 	return args
 }
 
+func (evmArgOverrider) OverrideJumpTable(_ params.Rules, jt *JumpTable) *JumpTable { return jt }
+
 func (o *evmArgOverrider) OverrideEVMResetArgs(r params.Rules, _ *EVMResetArgs) *EVMResetArgs {
 	o.gotResetChainID = r.ChainID
 	return &EVMResetArgs{

core/vm/hooks.libevm.go

@@ -34,6 +34,7 @@ var libevmHooks Hooks
 type Hooks interface {
 	OverrideNewEVMArgs(*NewEVMArgs) *NewEVMArgs
 	OverrideEVMResetArgs(params.Rules, *EVMResetArgs) *EVMResetArgs
+	OverrideJumpTable(params.Rules, *JumpTable) *JumpTable
 }
 
 // NewEVMArgs are the arguments received by [NewEVM], available for override
@@ -74,3 +75,10 @@ func (evm *EVM) overrideEVMResetArgs(txCtx TxContext, statedb StateDB) (TxContex
 	args := libevmHooks.OverrideEVMResetArgs(evm.chainRules, &EVMResetArgs{txCtx, statedb})
 	return args.TxContext, args.StateDB
 }
+
+func overrideJumpTable(r params.Rules, jt *JumpTable) *JumpTable {
+	if libevmHooks == nil {
+		return jt
+	}
+	return libevmHooks.OverrideJumpTable(r, jt)
+}

core/vm/interpreter.go

@@ -95,6 +95,7 @@ func NewEVMInterpreter(evm *EVM) *EVMInterpreter {
 		}
 	}
 	evm.Config.ExtraEips = extraEips
+	table = overrideJumpTable(evm.chainRules, table)
 	return &EVMInterpreter{evm: evm, table: table}
 }
 

core/vm/jump_table.libevm.go

@@ -0,0 +1,60 @@
+// Copyright 2024 the libevm authors.
+//
+// The libevm additions to go-ethereum are free software: you can redistribute
+// them and/or modify them under the terms of the GNU Lesser General Public License
+// as published by the Free Software Foundation, either version 3 of the License,
+// or (at your option) any later version.
+//
+// The libevm additions are distributed in the hope that they will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU Lesser
+// General Public License for more details.
+//
+// You should have received a copy of the GNU Lesser General Public License
+// along with the go-ethereum library. If not, see
+// <http://www.gnu.org/licenses/>.
+
+package vm
+
+// An OperationBuilder is a factory for a new operations to include in a
+// [JumpTable].
+type OperationBuilder struct {
+	Execute            OperationFunc
+	ConstantGas        uint64
+	DynamicGas         func(_ *EVM, _ *Contract, _ *Stack, _ *Memory, requestedMemorySize uint64) (uint64, error)
+	MinStack, MaxStack int
+	MemorySize         func(s *Stack) (size uint64, overflow bool)
+}
+
+// Build constructs the operation.
+func (b OperationBuilder) Build() *operation {
+	o := &operation{
+		execute:     b.Execute.internal(),
+		constantGas: b.ConstantGas,
+		dynamicGas:  b.DynamicGas,
+		minStack:    b.MinStack,
+		maxStack:    b.MaxStack,
+		memorySize:  b.MemorySize,
+	}
+	return o
+}
+
+// An OperationFunc is the execution function of a custom instruction.
+type OperationFunc func(_ Environment, pc *uint64, _ *EVMInterpreter, _ *ScopeContext) ([]byte, error)
+
+// internal converts an exported [OperationFunc] into an un-exported
+// [executionFunc] as required to build an [operation].
+func (fn OperationFunc) internal() executionFunc {
+	return func(pc *uint64, interpreter *EVMInterpreter, scope *ScopeContext) ([]byte, error) {
+		env := &environment{
+			evm:  interpreter.evm,
+			self: scope.Contract,
+			// The CallType isn't exposed by an instruction's [Environment] and,
+			// although [UnknownCallType] is the default value, it's explicitly
+			// set to avoid future accidental setting without proper
+			// justification.
+			callType: UnknownCallType,
+		}
+		return fn(env, pc, interpreter, scope)
+	}
+}

core/vm/jump_table.libevm_test.go

@@ -0,0 +1,110 @@
+// Copyright 2024 the libevm authors.
+//
+// The libevm additions to go-ethereum are free software: you can redistribute
+// them and/or modify them under the terms of the GNU Lesser General Public License
+// as published by the Free Software Foundation, either version 3 of the License,
+// or (at your option) any later version.
+//
+// The libevm additions are distributed in the hope that they will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU Lesser
+// General Public License for more details.
+//
+// You should have received a copy of the GNU Lesser General Public License
+// along with the go-ethereum library. If not, see
+// <http://www.gnu.org/licenses/>.
+
+package vm_test
+
+import (
+	"fmt"
+	"reflect"
+	"testing"
+
+	"github.com/holiman/uint256"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/ava-labs/libevm/common"
+	"github.com/ava-labs/libevm/core/vm"
+	"github.com/ava-labs/libevm/libevm/ethtest"
+	"github.com/ava-labs/libevm/params"
+)
+
+type vmHooksStub struct {
+	replacement *vm.JumpTable
+	overridden  bool
+}
+
+var _ vm.Hooks = (*vmHooksStub)(nil)
+
+// OverrideJumpTable overrides all non-nil operations from s.replacement .
+func (s *vmHooksStub) OverrideJumpTable(_ params.Rules, jt *vm.JumpTable) *vm.JumpTable {
+	s.overridden = true
+	for op, instr := range s.replacement {
+		if instr != nil {
+			fmt.Println(op, instr)
+			jt[op] = instr
+		}
+	}
+	return jt
+}
+
+func (*vmHooksStub) OverrideNewEVMArgs(a *vm.NewEVMArgs) *vm.NewEVMArgs { return a }
+
+func (*vmHooksStub) OverrideEVMResetArgs(r params.Rules, a *vm.EVMResetArgs) *vm.EVMResetArgs {
+	return a
+}
+
+// An opRecorder is an instruction that records its inputs.
+type opRecorder struct {
+	stateVal common.Hash
+}
+
+func (op *opRecorder) execute(env vm.Environment, pc *uint64, interpreter *vm.EVMInterpreter, scope *vm.ScopeContext) ([]byte, error) {
+	op.stateVal = env.StateDB().GetState(scope.Contract.Address(), common.Hash{})
+	return nil, nil
+}
+
+func TestOverrideJumpTable(t *testing.T) {
+	const (
+		opcode          = 1
+		gasLimit uint64 = 1e6
+	)
+	rng := ethtest.NewPseudoRand(142857)
+	gasCost := 1 + rng.Uint64n(gasLimit)
+	spy := &opRecorder{}
+
+	vmHooks := &vmHooksStub{
+		replacement: &vm.JumpTable{
+			opcode: vm.OperationBuilder{
+				Execute:     spy.execute,
+				ConstantGas: gasCost,
+				MemorySize: func(s *vm.Stack) (size uint64, overflow bool) {
+					return 0, false
+				},
+			}.Build(),
+		},
+	}
+	vm.RegisterHooks(vmHooks)
+
+	state, evm := ethtest.NewZeroEVM(t)
+
+	contract := rng.Address()
+	state.CreateAccount(contract)
+	state.SetCode(contract, []byte{opcode})
+	value := rng.Hash()
+	state.SetState(contract, common.Hash{}, value)
+
+	_, gasRemaining, err := evm.Call(vm.AccountRef(rng.Address()), contract, []byte{}, gasLimit, uint256.NewInt(0))
+	require.NoError(t, err, "evm.Call([contract with overridden opcode])")
+	assert.Equal(t, gasLimit-gasCost, gasRemaining, "gas remaining")
+	assert.Equal(t, spy.stateVal, value, "StateDB propagated")
+}
+
+func TestOperationFieldCount(t *testing.T) {
+	// The libevm OperationBuilder assumes that the 6 struct fields are the only
+	// ones.
+	op := vm.OperationBuilder{}.Build()
+	require.Equalf(t, 6, reflect.TypeOf(*op).NumField(), "number of fields in %T struct", *op)
+}

core/vm/jump_table_export.go

@@ -24,7 +24,12 @@ import (
 
 // LookupInstructionSet returns the instruction set for the fork configured by
 // the rules.
-func LookupInstructionSet(rules params.Rules) (JumpTable, error) {
+func LookupInstructionSet(rules params.Rules) (jt JumpTable, err error) {
+	defer func() {
+		if err == nil { // NOTE `err ==` NOT !=
+			jt = *overrideJumpTable(rules, &jt)
+		}
+	}()
 	switch {
 	case rules.IsVerkle:
 		return newCancunInstructionSet(), errors.New("verkle-fork not defined yet")