From f0d8138457516f5477d4cabaa173cc214581c274 Mon Sep 17 00:00:00 2001
From: Dean Eigenmann <7621705+decanus@users.noreply.github.com>
Date: Thu, 12 Dec 2024 08:29:55 +0100
Subject: [PATCH] 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 <lmittmann@users.noreply.github.com>
---
 docs/pages/vm-overview.mdx |  1 +
 w3vm/vm.go                 | 26 ++++++++++++++++++++++++++
 2 files changed, 27 insertions(+)

diff --git a/docs/pages/vm-overview.mdx b/docs/pages/vm-overview.mdx
index 724372a..181defc 100644
--- 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
index ecdee6f..414685b 100644
--- 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