w3vm: added WithPrecompile (#200)

* adding option to allow precompile * update * update doc * set default precompiles if there are overwrites * fix dep --------- Co-authored-by: lmittmann <[email protected]>
lmittmann · Dec 12, 2024 · f0d8138 · f0d8138
1 parent f4b89b3
commit f0d8138
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 0 deletions.
diff --git a/docs/pages/vm-overview.mdx b/docs/pages/vm-overview.mdx
@@ -71,6 +71,7 @@ A new VM instance is created using the `w3vm.New` function, which accepts variou
 * `WithChainConfig(cfg *params.ChainConfig)`: Sets the chain configuration. If not provided, the VM defaults to the Mainnet configuration.
 * `WithNoBaseFee()`: Forces the EIP-1559 base fee to 0.
 * `WithBlockContext(ctx *vm.BlockContext)`: Sets the block context for the VM.
+* `WithPrecompile(addr common.Address, contract vm.PrecompiledContract)`: Registers a precompile contract at the given address in the VM.
 * `WithHeader(header *types.Header)`: Configures the block context for the VM using the provided header.
 * `WithState(state w3types.State)`: Sets the pre-state of the VM. When used with `WithFork`, the pre-state overrides the forked state.
 * `WithStateDB(db *state.StateDB)`: Specifies the state database for the VM, typically a snapshot from `VM.Snapshot`.

diff --git a/w3vm/vm.go b/w3vm/vm.go
@@ -112,6 +112,10 @@ func (v *VM) apply(msg *w3types.Message, isCall bool, hooks *tracing.Hooks) (*Re
 		NoBaseFee: v.opts.noBaseFee || isCall,
 	})
 
+	if len(v.opts.precompiles) > 0 {
+		evm.SetPrecompiles(v.opts.precompiles)
+	}
+
 	snap := v.db.Snapshot()
 
 	// apply the message to the evm
@@ -389,6 +393,8 @@ type options struct {
 	forkBlockNumber *big.Int
 	fetcher         Fetcher
 	tb              testing.TB
+
+	precompiles vm.PrecompiledContracts
 }
 
 func (opt *options) Signer() types.Signer {
@@ -446,6 +452,19 @@ func (opts *options) Init() error {
 			opts.blockCtx = defaultBlockContext()
 		}
 	}
+
+	// set precompiles
+	if len(opts.precompiles) > 0 {
+		rules := opts.chainConfig.Rules(opts.blockCtx.BlockNumber, opts.blockCtx.Random != nil, opts.blockCtx.Time)
+
+		// overwrite default precompiles
+		precompiles := vm.ActivePrecompiledContracts(rules)
+		for addr, contract := range opts.precompiles {
+			precompiles[addr] = contract
+		}
+		opts.precompiles = precompiles
+	}
+
 	return nil
 }
 
@@ -471,6 +490,13 @@ func WithBlockContext(ctx *vm.BlockContext) Option {
 	return func(vm *VM) { vm.opts.blockCtx = ctx }
 }
 
+// WithPrecompile registers a precompile contract at the given address in the VM.
+func WithPrecompile(addr common.Address, contract vm.PrecompiledContract) Option {
+	return func(vm *VM) {
+		vm.opts.precompiles[addr] = contract
+	}
+}
+
 // WithState sets the pre state of the VM.
 //
 // WithState can be used together with [WithFork] to only set the state of some