move block extender to vm

plugin/evm/atomic_block_extension.go → plugin/evm/atomic/vm/block_extension.go

@@ -1,7 +1,7 @@
 // Copyright (C) 2019-2025, Ava Labs, Inc. All rights reserved.
 // See the file LICENSE for licensing terms.
 
-package evm
+package vm
 
 import (
 	"errors"
@@ -29,7 +29,8 @@ var (
 
 var (
 	errNilEthBlock  = errors.New("nil ethBlock")
-	errMissingUTXOs = errors.New("missing UTXOs")
+	ErrMissingUTXOs = errors.New("missing UTXOs")
+	ErrEmptyBlock   = errors.New("empty block")
 )
 
 type blockExtender struct {
@@ -131,7 +132,7 @@ func (be *blockExtension) SyntacticVerify(rules extras.Rules) error {
 	txs := ethBlock.Transactions()
 	atomicTxs := be.atomicTxs
 	if len(txs) == 0 && len(atomicTxs) == 0 {
-		return errEmptyBlock
+		return ErrEmptyBlock
 	}
 
 	// If we are in ApricotPhase4, ensure that ExtDataGasUsed is populated correctly.
@@ -193,12 +194,12 @@ func (be *blockExtension) Accept(acceptedBatch database.Batch) error {
 	vm := be.blockExtender.vm
 	for _, tx := range be.atomicTxs {
 		// Remove the accepted transaction from the mempool
-		vm.mempool.RemoveTx(tx)
+		vm.AtomicMempool().RemoveTx(tx)
 	}
 
 	// Update VM state for atomic txs in this block. This includes updating the
 	// atomic tx repo, atomic trie, and shared memory.
-	atomicState, err := vm.atomicBackend.GetVerifiedAtomicState(common.Hash(be.block.ID()))
+	atomicState, err := vm.AtomicBackend().GetVerifiedAtomicState(common.Hash(be.block.ID()))
 	if err != nil {
 		// should never occur since [b] must be verified before calling Accept
 		return err
@@ -213,12 +214,12 @@ func (be *blockExtension) Reject() error {
 	vm := be.blockExtender.vm
 	for _, tx := range be.atomicTxs {
 		// Re-issue the transaction in the mempool, continue even if it fails
-		vm.mempool.RemoveTx(tx)
-		if err := vm.mempool.AddRemoteTx(tx); err != nil {
+		vm.AtomicMempool().RemoveTx(tx)
+		if err := vm.AtomicMempool().AddRemoteTx(tx); err != nil {
 			log.Debug("Failed to re-issue transaction in rejected block", "txID", tx.ID(), "err", err)
 		}
 	}
-	atomicState, err := vm.atomicBackend.GetVerifiedAtomicState(common.Hash(be.block.ID()))
+	atomicState, err := vm.AtomicBackend().GetVerifiedAtomicState(common.Hash(be.block.ID()))
 	if err != nil {
 		// should never occur since [b] must be verified before calling Reject
 		return err
@@ -229,9 +230,7 @@ func (be *blockExtension) Reject() error {
 // CleanupVerified is called when the block is cleaned up after a failed insertion.
 func (be *blockExtension) CleanupVerified() {
 	vm := be.blockExtender.vm
-	// If an error occurred inserting the block into the chain or if we are not pinning to memory,
-	// unpin the atomic trie changes from memory (if they were pinned).
-	if atomicState, err := vm.atomicBackend.GetVerifiedAtomicState(be.block.GetEthBlock().Hash()); err == nil {
+	if atomicState, err := vm.AtomicBackend().GetVerifiedAtomicState(be.block.GetEthBlock().Hash()); err == nil {
 		atomicState.Reject()
 	}
 }
@@ -247,7 +246,7 @@ func (be *blockExtension) verifyUTXOsPresent(atomicTxs []*atomic.Tx) error {
 	b := be.block
 	blockHash := common.Hash(b.ID())
 	vm := be.blockExtender.vm
-	if vm.atomicBackend.IsBonus(b.Height(), blockHash) {
+	if vm.AtomicBackend().IsBonus(b.Height(), blockHash) {
 		log.Info("skipping atomic tx verification on bonus block", "block", blockHash)
 		return nil
 	}
@@ -260,7 +259,7 @@ func (be *blockExtension) verifyUTXOsPresent(atomicTxs []*atomic.Tx) error {
 			return err
 		}
 		if _, err := vm.ctx.SharedMemory.Get(chainID, requests.RemoveRequests); err != nil {
-			return fmt.Errorf("%w: %s", errMissingUTXOs, err)
+			return fmt.Errorf("%w: %s", ErrMissingUTXOs, err)
 		}
 	}
 	return nil

plugin/evm/ext_data_hashes.go → plugin/evm/atomic/vm/ext_data_hashes.go

@@ -1,7 +1,7 @@
 // Copyright (C) 2019-2025, Ava Labs, Inc. All rights reserved.
 // See the file LICENSE for licensing terms.
 
-package evm
+package vm
 
 import (
 	_ "embed"

plugin/evm/atomic/tx_semantic_verifier.go → plugin/evm/atomic/vm/tx_semantic_verifier.go

@@ -1,7 +1,7 @@
 // Copyright (C) 2019-2025, Ava Labs, Inc. All rights reserved.
 // See the file LICENSE for licensing terms.
 
-package atomic
+package vm
 
 import (
 	"context"
@@ -17,11 +17,12 @@ import (
 	"github.com/ava-labs/avalanchego/vms/platformvm/fx"
 	"github.com/ava-labs/avalanchego/vms/secp256k1fx"
 	"github.com/ava-labs/coreth/params/extras"
+	"github.com/ava-labs/coreth/plugin/evm/atomic"
 	"github.com/ava-labs/coreth/plugin/evm/extension"
 	"github.com/ava-labs/coreth/plugin/evm/upgrade/ap0"
 )
 
-var _ Visitor = (*SemanticVerifier)(nil)
+var _ atomic.Visitor = (*SemanticVerifier)(nil)
 
 var (
 	ErrAssetIDMismatch            = errors.New("asset IDs in the input don't match the utxo")
@@ -49,13 +50,13 @@ type VerifierBackend struct {
 // SemanticVerifier is a visitor that checks the semantic validity of atomic transactions.
 type SemanticVerifier struct {
 	Backend *VerifierBackend
-	Tx      *Tx
+	Tx      *atomic.Tx
 	Parent  extension.ExtendedBlock
 	BaseFee *big.Int
 }
 
 // ImportTx verifies this transaction is valid.
-func (s *SemanticVerifier) ImportTx(utx *UnsignedImportTx) error {
+func (s *SemanticVerifier) ImportTx(utx *atomic.UnsignedImportTx) error {
 	backend := s.Backend
 	ctx := backend.Ctx
 	rules := backend.Rules
@@ -73,7 +74,7 @@ func (s *SemanticVerifier) ImportTx(utx *UnsignedImportTx) error {
 		if err != nil {
 			return err
 		}
-		txFee, err := CalculateDynamicFee(gasUsed, s.BaseFee)
+		txFee, err := atomic.CalculateDynamicFee(gasUsed, s.BaseFee)
 		if err != nil {
 			return err
 		}
@@ -118,7 +119,7 @@ func (s *SemanticVerifier) ImportTx(utx *UnsignedImportTx) error {
 		utxoBytes := allUTXOBytes[i]
 
 		utxo := &avax.UTXO{}
-		if _, err := Codec.Unmarshal(utxoBytes, utxo); err != nil {
+		if _, err := atomic.Codec.Unmarshal(utxoBytes, utxo); err != nil {
 			return fmt.Errorf("failed to unmarshal UTXO: %w", err)
 		}
 
@@ -148,7 +149,7 @@ func conflicts(backend *VerifierBackend, inputs set.Set[ids.ID], ancestor extens
 	lastAcceptedHeight := lastAcceptedBlock.Height()
 	for ancestor.Height() > lastAcceptedHeight {
 		ancestorExtIntf := ancestor.GetBlockExtension()
-		ancestorExt, ok := ancestorExtIntf.(AtomicBlockContext)
+		ancestorExt, ok := ancestorExtIntf.(atomic.AtomicBlockContext)
 		if !ok {
 			return fmt.Errorf("expected block extension to be AtomicBlockContext but got %T", ancestorExtIntf)
 		}
@@ -180,7 +181,7 @@ func conflicts(backend *VerifierBackend, inputs set.Set[ids.ID], ancestor extens
 }
 
 // ExportTx verifies this transaction is valid.
-func (s *SemanticVerifier) ExportTx(utx *UnsignedExportTx) error {
+func (s *SemanticVerifier) ExportTx(utx *atomic.UnsignedExportTx) error {
 	backend := s.Backend
 	ctx := backend.Ctx
 	rules := backend.Rules
@@ -198,7 +199,7 @@ func (s *SemanticVerifier) ExportTx(utx *UnsignedExportTx) error {
 		if err != nil {
 			return err
 		}
-		txFee, err := CalculateDynamicFee(gasUsed, s.BaseFee)
+		txFee, err := atomic.CalculateDynamicFee(gasUsed, s.BaseFee)
 		if err != nil {
 			return err
 		}

plugin/evm/atomic/vm/vm.go

@@ -0,0 +1,111 @@
+// Copyright (C) 2019-2025, Ava Labs, Inc. All rights reserved.
+// See the file LICENSE for licensing terms.
+
+package vm
+
+import (
+	"context"
+	"fmt"
+
+	avalanchedatabase "github.com/ava-labs/avalanchego/database"
+	"github.com/ava-labs/avalanchego/snow"
+	avalanchecommon "github.com/ava-labs/avalanchego/snow/engine/common"
+	"github.com/ava-labs/avalanchego/snow/engine/snowman/block"
+	"github.com/ava-labs/avalanchego/utils/constants"
+
+	"github.com/ava-labs/coreth/params"
+	"github.com/ava-labs/coreth/params/extras"
+	"github.com/ava-labs/coreth/plugin/evm/atomic/state"
+	"github.com/ava-labs/coreth/plugin/evm/atomic/txpool"
+	"github.com/ava-labs/coreth/plugin/evm/extension"
+
+	"github.com/ava-labs/libevm/common"
+)
+
+var (
+	_ block.ChainVM                      = (*VM)(nil)
+	_ block.BuildBlockWithContextChainVM = (*VM)(nil)
+	_ block.StateSyncableVM              = (*VM)(nil)
+)
+
+// TODO: remove this
+// InnerVM is the interface that must be implemented by the VM
+// that's being wrapped by the extension
+type InnerVM interface {
+	extension.ExtensibleVM
+	avalanchecommon.VM
+	block.ChainVM
+	block.BuildBlockWithContextChainVM
+	block.StateSyncableVM
+
+	// TODO: remove these
+	AtomicBackend() *state.AtomicBackend
+	AtomicMempool() *txpool.Mempool
+}
+
+type VM struct {
+	InnerVM
+	ctx *snow.Context
+}
+
+func WrapVM(vm InnerVM) *VM {
+	return &VM{InnerVM: vm}
+}
+
+// Initialize implements the snowman.ChainVM interface
+func (vm *VM) Initialize(
+	ctx context.Context,
+	chainCtx *snow.Context,
+	db avalanchedatabase.Database,
+	genesisBytes []byte,
+	upgradeBytes []byte,
+	configBytes []byte,
+	toEngine chan<- avalanchecommon.Message,
+	fxs []*avalanchecommon.Fx,
+	appSender avalanchecommon.AppSender,
+) error {
+	vm.ctx = chainCtx
+
+	var extDataHashes map[common.Hash]common.Hash
+	// Set the chain config for mainnet/fuji chain IDs
+	switch chainCtx.NetworkID {
+	case constants.MainnetID:
+		extDataHashes = mainnetExtDataHashes
+	case constants.FujiID:
+		extDataHashes = fujiExtDataHashes
+	}
+	// Free the memory of the extDataHash map
+	fujiExtDataHashes = nil
+	mainnetExtDataHashes = nil
+
+	// Create the atomic extension structs
+	// some of them need to be initialized after the inner VM is initialized
+	blockExtender := newBlockExtender(extDataHashes, vm)
+
+	extensionConfig := &extension.Config{
+		BlockExtender: blockExtender,
+	}
+	if err := vm.SetExtensionConfig(extensionConfig); err != nil {
+		return fmt.Errorf("failed to set extension config: %w", err)
+	}
+
+	// Initialize inner vm with the provided parameters
+	if err := vm.InnerVM.Initialize(
+		ctx,
+		chainCtx,
+		db,
+		genesisBytes,
+		upgradeBytes,
+		configBytes,
+		toEngine,
+		fxs,
+		appSender,
+	); err != nil {
+		return fmt.Errorf("failed to initialize inner VM: %w", err)
+	}
+	return nil
+}
+
+func (vm *VM) chainConfigExtra() *extras.ChainConfig {
+	return params.GetExtra(vm.Ethereum().BlockChain().Config())
+}

plugin/evm/export_tx_test.go

@@ -22,6 +22,7 @@ import (
 	"github.com/ava-labs/avalanchego/vms/secp256k1fx"
 	"github.com/ava-labs/coreth/params/extras"
 	"github.com/ava-labs/coreth/plugin/evm/atomic"
+	atomicvm "github.com/ava-labs/coreth/plugin/evm/atomic/vm"
 	"github.com/ava-labs/coreth/utils"
 	"github.com/ava-labs/libevm/common"
 	"github.com/holiman/uint256"
@@ -922,7 +923,7 @@ func TestExportTxSemanticVerify(t *testing.T) {
 			t.Fatal(err)
 		}
 
-		backend := &atomic.VerifierBackend{
+		backend := &atomicvm.VerifierBackend{
 			Ctx:          vm.ctx,
 			Fx:           &vm.fx,
 			Rules:        test.rules,
@@ -935,7 +936,7 @@ func TestExportTxSemanticVerify(t *testing.T) {
 			tx := test.tx
 			exportTx := tx.UnsignedAtomicTx
 
-			err := exportTx.Visit(&atomic.SemanticVerifier{
+			err := exportTx.Visit(&atomicvm.SemanticVerifier{
 				Backend: backend,
 				Tx:      tx,
 				Parent:  parent,
@@ -1783,7 +1784,7 @@ func TestNewExportTx(t *testing.T) {
 
 			exportTx := tx.UnsignedAtomicTx
 
-			backend := &atomic.VerifierBackend{
+			backend := &atomicvm.VerifierBackend{
 				Ctx:          tvm.vm.ctx,
 				Fx:           &tvm.vm.fx,
 				Rules:        tvm.vm.currentRules(),
@@ -1792,7 +1793,7 @@ func TestNewExportTx(t *testing.T) {
 				SecpCache:    tvm.vm.secpCache,
 			}
 
-			if err := exportTx.Visit(&atomic.SemanticVerifier{
+			if err := exportTx.Visit(&atomicvm.SemanticVerifier{
 				Backend: backend,
 				Tx:      tx,
 				Parent:  parent,
@@ -1987,7 +1988,7 @@ func TestNewExportTxMulticoin(t *testing.T) {
 			}
 
 			exportTx := tx.UnsignedAtomicTx
-			backend := &atomic.VerifierBackend{
+			backend := &atomicvm.VerifierBackend{
 				Ctx:          tvm.vm.ctx,
 				Fx:           &tvm.vm.fx,
 				Rules:        tvm.vm.currentRules(),
@@ -1996,7 +1997,7 @@ func TestNewExportTxMulticoin(t *testing.T) {
 				SecpCache:    tvm.vm.secpCache,
 			}
 
-			if err := exportTx.Visit(&atomic.SemanticVerifier{
+			if err := exportTx.Visit(&atomicvm.SemanticVerifier{
 				Backend: backend,
 				Tx:      tx,
 				Parent:  parent,