use explicity structs and innervm calls

plugin/evm/atomic/sync/extender.go

@@ -16,26 +16,24 @@ import (
 	"github.com/ava-labs/libevm/log"
 )
 
-var _ sync.Extender = (*extender)(nil)
+var _ sync.Extender = (*Extender)(nil)
 
-type extender struct {
+// Extender is the sync extender for the atomic VM.
+type Extender struct {
 	backend     *state.AtomicBackend
 	trie        *state.AtomicTrie
 	requestSize uint16 // maximum number of leaves to sync in a single request
 }
 
-// Initialize initializes the sync extender with the backend and trie.
-func NewExtender() *extender {
-	return &extender{}
-}
-
-func (a *extender) Initialize(backend *state.AtomicBackend, trie *state.AtomicTrie, requestSize uint16) {
+// Initialize initializes the sync extender with the backend and trie and request size.
+func (a *Extender) Initialize(backend *state.AtomicBackend, trie *state.AtomicTrie, requestSize uint16) {
 	a.backend = backend
 	a.trie = trie
 	a.requestSize = requestSize
 }
 
-func (a *extender) Sync(ctx context.Context, client syncclient.LeafClient, verDB *versiondb.Database, summary message.Syncable) error {
+// Sync syncs the atomic summary with the given client and verDB.
+func (a *Extender) Sync(ctx context.Context, client syncclient.LeafClient, verDB *versiondb.Database, summary message.Syncable) error {
 	atomicSummary, ok := summary.(*Summary)
 	if !ok {
 		return fmt.Errorf("expected *Summary, got %T", summary)
@@ -60,7 +58,8 @@ func (a *extender) Sync(ctx context.Context, client syncclient.LeafClient, verDB
 	return err
 }
 
-func (a *extender) OnFinishBeforeCommit(lastAcceptedHeight uint64, Summary message.Syncable) error {
+// OnFinishBeforeCommit implements the sync.Extender interface by marking the previously last accepted block for the shared memory cursor.
+func (a *Extender) OnFinishBeforeCommit(lastAcceptedHeight uint64, Summary message.Syncable) error {
 	// Mark the previously last accepted block for the shared memory cursor, so that we will execute shared
 	// memory operations from the previously last accepted block when ApplyToSharedMemory
 	// is called.
@@ -71,7 +70,8 @@ func (a *extender) OnFinishBeforeCommit(lastAcceptedHeight uint64, Summary messa
 	return nil
 }
 
-func (a *extender) OnFinishAfterCommit(summaryHeight uint64) error {
+// OnFinishAfterCommit implements the sync.Extender interface by applying the atomic trie to the shared memory.
+func (a *Extender) OnFinishAfterCommit(summaryHeight uint64) error {
 	// the chain state is already restored, and, from this point on,
 	// the block synced to is the accepted block. The last operation
 	// is updating shared memory with the atomic trie.

plugin/evm/atomic/sync/leaf_handler.go

@@ -34,7 +34,7 @@ type leafHandler struct {
 	handlers.LeafRequestHandler
 }
 
-// NewAtomicLeafHandler returns a new uninitialzied atomicLeafHandler that can be later initialized
+// NewAtomicLeafHandler returns a new uninitialized leafHandler that can be later initialized
 func NewLeafHandler() *leafHandler {
 	return &leafHandler{
 		LeafRequestHandler: &uninitializedHandler{},

plugin/evm/atomic/sync/summary_provider.go

@@ -14,22 +14,20 @@ import (
 	"github.com/ava-labs/libevm/core/types"
 )
 
-var _ sync.SummaryProvider = (*summaryProvider)(nil)
+var _ sync.SummaryProvider = (*SummaryProvider)(nil)
 
-type summaryProvider struct {
+// SummaryProvider is the summary provider that provides the state summary for the atomic trie.
+type SummaryProvider struct {
 	trie *state.AtomicTrie
 }
 
-func NewSummaryProvider() *summaryProvider {
-	return &summaryProvider{}
-}
-
-func (a *summaryProvider) Initialize(trie *state.AtomicTrie) {
+// Initialize initializes the summary provider with the atomic trie.
+func (a *SummaryProvider) Initialize(trie *state.AtomicTrie) {
 	a.trie = trie
 }
 
 // StateSummaryAtBlock returns the block state summary at [blk] if valid.
-func (a *summaryProvider) StateSummaryAtBlock(blk *types.Block) (block.StateSummary, error) {
+func (a *SummaryProvider) StateSummaryAtBlock(blk *types.Block) (block.StateSummary, error) {
 	height := blk.NumberU64()
 	atomicRoot, err := a.trie.Root(height)
 	if err != nil {

plugin/evm/atomic/vm/vm.go

@@ -86,8 +86,8 @@ func (vm *VM) Initialize(
 	// Create the atomic extension structs
 	// some of them need to be initialized after the inner VM is initialized
 	blockExtender := newBlockExtender(extDataHashes, vm)
-	syncExtender := sync.NewExtender()
-	syncProvider := sync.NewSummaryProvider()
+	syncExtender := &sync.Extender{}
+	syncProvider := &sync.SummaryProvider{}
 	// Create and pass the leaf handler to the atomic extension
 	// it will be initialized after the inner VM is initialized
 	leafHandler := sync.NewLeafHandler()
@@ -105,7 +105,7 @@ func (vm *VM) Initialize(
 		ExtraSyncLeafHandlerConfig: atomicLeafTypeConfig,
 		Clock:                      &vm.clock,
 	}
-	if err := vm.SetExtensionConfig(extensionConfig); err != nil {
+	if err := vm.InnerVM.SetExtensionConfig(extensionConfig); err != nil {
 		return fmt.Errorf("failed to set extension config: %w", err)
 	}
 
@@ -125,13 +125,15 @@ func (vm *VM) Initialize(
 	}
 
 	// Atomic backend is available now, we can initialize structs that depend on it
-	syncProvider.Initialize(vm.AtomicBackend().AtomicTrie())
-	syncExtender.Initialize(vm.AtomicBackend(), vm.AtomicBackend().AtomicTrie(), vm.Config().StateSyncRequestSize)
-	leafHandler.Initialize(vm.AtomicBackend().AtomicTrie().TrieDB(), state.TrieKeyLength, message.Codec)
+	atomicBackend := vm.InnerVM.AtomicBackend()
+	atomicTrie := atomicBackend.AtomicTrie()
+	syncProvider.Initialize(atomicTrie)
+	syncExtender.Initialize(atomicBackend, atomicTrie, vm.Config().StateSyncRequestSize)
+	leafHandler.Initialize(atomicTrie.TrieDB(), state.TrieKeyLength, message.Codec)
 
 	return nil
 }
 
 func (vm *VM) chainConfigExtra() *extras.ChainConfig {
-	return params.GetExtra(vm.Ethereum().BlockChain().Config())
+	return params.GetExtra(vm.InnerVM.Ethereum().BlockChain().Config())
 }