impl SubprotocolManager trait to pass to each stage

Author: prajwolrg

crates/asm/common/src/lib.rs

@@ -8,9 +8,9 @@ mod state;
 mod subprotocol;
 mod tx;
 
-pub use error::AsmError;
-pub use msg::{InterprotoMsg, Log, NullMsg};
-pub use spec::{AsmSpec, Stage};
-pub use state::{AnchorState, ChainViewState, SectionState};
-pub use subprotocol::{MsgRelayer, Subprotocol, SubprotocolId};
-pub use tx::{TagPayload, TxInput};
+pub use error::*;
+pub use msg::*;
+pub use spec::*;
+pub use state::*;
+pub use subprotocol::*;
+pub use tx::*;

crates/asm/common/src/spec.rs

@@ -1,4 +1,4 @@
-use crate::Subprotocol;
+use crate::{Subprotocol, SubprotocolManager};
 
 /// Specification for a concrete ASM instantiation describing the subprotocols we
 /// want to invoke and in what order.
@@ -8,11 +8,11 @@ use crate::Subprotocol;
 pub trait AsmSpec {
     /// Function that calls the loader with each subprotocol we intend to
     /// process, in the order we intend to process them.
-    fn call_subprotocols(stage: &mut impl Stage);
+    fn call_subprotocols(stage: &mut impl Stage, manager: &mut impl SubprotocolManager);
 }
 
 /// Implementation of a subprotocol handling stage.
 pub trait Stage {
     /// Invoked by the ASM spec to perform logic relating to a specific subprotocol.
-    fn process_subprotocol<S: Subprotocol>(&mut self);
+    fn process_subprotocol<S: Subprotocol>(&mut self, manager: &mut impl SubprotocolManager);
 }

crates/asm/common/src/subprotocol.rs

@@ -8,7 +8,7 @@ use std::any::Any;
 
 use borsh::{BorshDeserialize, BorshSerialize};
 
-use crate::{Log, TxInput, msg::InterprotoMsg};
+use crate::{AsmError, Log, SectionState, TxInput, msg::InterprotoMsg};
 
 /// Identifier for a subprotocol.
 pub type SubprotocolId = u8;
@@ -57,3 +57,93 @@ pub trait MsgRelayer: Any {
     /// Gets this msg relayer as a `&dyn Any`.
     fn as_mut_any(&mut self) -> &mut dyn Any;
 }
+
+/// Subprotocol handler trait for a loaded subprotocol.
+pub trait SubprotoHandler {
+    /// Processes transactions that were previously collected.
+    fn process_txs(&mut self, txs: &[TxInput<'_>], relayer: &mut dyn MsgRelayer);
+
+    /// Accepts a message.  This is called while processing other subprotocols.
+    /// These should not be processed until we do the finalization.
+    ///
+    /// This MUST NOT act on any messages that were accepted before this was
+    /// called.
+    ///
+    /// # Panics
+    ///
+    /// If an mismatched message type (behind the `dyn`) is provided.
+    fn accept_msg(&mut self, msg: &dyn InterprotoMsg);
+
+    /// Processes the messages received.
+    fn process_msgs(&mut self);
+
+    /// Repacks the state into a [`SectionState`] instance.
+    fn to_section(&self) -> SectionState;
+}
+
+/// Manages the lifecycle and execution of subprotocol handlers in the Anchor State Machine (ASM).
+///
+/// Implementors of this trait maintain a collection of subprotocol handlers and
+/// provide methods to insert, remove, lookup, and drive execution (transactions and messages),
+/// as well as extract the final `SectionState`.
+pub trait SubprotocolManager: MsgRelayer + Sized {
+    /// Inserts a new subprotocol by consuming its initial state and creating its handler.
+    fn insert_subproto<S: Subprotocol>(&mut self, state: S::State);
+
+    /// Inserts a boxed handler directly.
+    fn insert_handler<S: Subprotocol>(&mut self, handler: Box<dyn SubprotoHandler>);
+
+    /// Removes and returns the handler for the given `id`.
+    ///
+    /// # Errors
+    ///
+    /// Returns `AsmError::InvalidSubprotocol(id)` if no handler with that ID is present.
+    fn remove_handler(&mut self, id: SubprotocolId) -> Result<Box<dyn SubprotoHandler>, AsmError>;
+
+    /// Retrieves an immutable reference to the handler for the given `id`.
+    ///
+    /// # Errors
+    ///
+    /// Returns `AsmError::InvalidSubprotocol(id)` if no handler with that ID is present.
+    fn get_handler(&self, id: SubprotocolId) -> Result<&dyn SubprotoHandler, AsmError>;
+
+    /// Retrieves a mutable reference to the handler for the given `id`.
+    ///
+    /// # Errors
+    ///
+    /// Returns `AsmError::InvalidSubprotocol(id)` if no handler with that ID is present.
+    fn get_handler_mut(
+        &mut self,
+        id: SubprotocolId,
+    ) -> Result<&mut Box<dyn SubprotoHandler>, AsmError>;
+
+    /// Dispatches transaction processing to the appropriate handler.
+    ///
+    /// This default implementation temporarily removes the handler to satisfy
+    /// borrow-checker constraints, invokes `process_txs` with `self` as the relayer,
+    /// and then reinserts the handler.
+    fn invoke_process_txs<S: Subprotocol>(&mut self, txs: &[TxInput<'_>]) {
+        // We temporarily take the handler out of the map so we can call
+        // `process_txs` with `self` as the relayer without violating the
+        // borrow checker.
+        let mut h = self
+            .remove_handler(S::ID)
+            .expect("asm: unloaded subprotocol");
+        h.process_txs(txs, self);
+        self.insert_handler::<S>(h);
+    }
+
+    /// Dispatches buffered inter-protocol message processing to the handler.
+    fn invoke_process_msgs<S: Subprotocol>(&mut self) {
+        let h = self
+            .get_handler_mut(S::ID)
+            .expect("asm: unloaded subprotocol");
+        h.process_msgs()
+    }
+
+    /// Extracts the finalized `SectionState` from the handler.
+    fn to_section_state<S: Subprotocol>(&self) -> SectionState {
+        let h = self.get_handler(S::ID).expect("asm: unloaded subprotocol");
+        h.to_section()
+    }
+}

crates/asm/stf/src/handler.rs

@@ -3,32 +3,10 @@
 use std::{any::Any, collections::BTreeMap};
 
 use strata_asm_common::{
-    AsmError, InterprotoMsg, Log, MsgRelayer, SectionState, Subprotocol, SubprotocolId, TxInput,
+    AsmError, InterprotoMsg, Log, MsgRelayer, SectionState, SubprotoHandler, Subprotocol,
+    SubprotocolId, SubprotocolManager, TxInput,
 };
 
-/// Subprotocol handler trait for a loaded subprotocol.
-pub(crate) trait SubprotoHandler {
-    /// Processes transactions that were previously collected.
-    fn process_txs(&mut self, txs: &[TxInput<'_>], relayer: &mut dyn MsgRelayer);
-
-    /// Accepts a message.  This is called while processing other subprotocols.
-    /// These should not be processed until we do the finalization.
-    ///
-    /// This MUST NOT act on any messages that were accepted before this was
-    /// called.
-    ///
-    /// # Panics
-    ///
-    /// If an mismatched message type (behind the `dyn`) is provided.
-    fn accept_msg(&mut self, msg: &dyn InterprotoMsg);
-
-    /// Processes the messages received.
-    fn process_msgs(&mut self);
-
-    /// Repacks the state into a [`SectionState`] instance.
-    fn to_section(&self) -> SectionState;
-}
-
 /// Wrapper around the common subprotocol interface that handles the common
 /// buffering logic for interproto messages.
 pub(crate) struct HandlerImpl<S: Subprotocol, R> {
@@ -84,60 +62,48 @@ pub(crate) struct HandlerRelayer {
     logs: Vec<Log>,
 }
 
-impl HandlerRelayer {
-    pub(crate) fn new() -> Self {
-        Self {
-            handlers: BTreeMap::new(),
-            logs: Vec::new(),
-        }
-    }
-
+impl SubprotocolManager for HandlerRelayer {
     /// Inserts a subproto by creating a handler for it.
-    pub(crate) fn insert_subproto<S: Subprotocol>(&mut self, state: S::State) {
+    fn insert_subproto<S: Subprotocol>(&mut self, state: S::State) {
         let handler = HandlerImpl::<S, Self>::from_state(state);
         if self.handlers.insert(S::ID, Box::new(handler)).is_some() {
             panic!("asm: loaded state twice");
         }
     }
 
-    pub(crate) fn get_handler(&self, id: SubprotocolId) -> Result<&dyn SubprotoHandler, AsmError> {
+    fn insert_handler<S: Subprotocol>(&mut self, handler: Box<dyn SubprotoHandler>) {
+        self.handlers.insert(S::ID, handler);
+    }
+
+    fn remove_handler(&mut self, id: SubprotocolId) -> Result<Box<dyn SubprotoHandler>, AsmError> {
+        self.handlers
+            .remove(&id)
+            .ok_or(AsmError::InvalidSubprotocol(id))
+    }
+
+    fn get_handler(&self, id: SubprotocolId) -> Result<&dyn SubprotoHandler, AsmError> {
         self.handlers
             .get(&id)
             .map(Box::as_ref)
             .ok_or(AsmError::InvalidSubprotocol(id))
     }
 
-    pub(crate) fn get_handler_mut(
+    fn get_handler_mut(
         &mut self,
         id: SubprotocolId,
     ) -> Result<&mut Box<dyn SubprotoHandler>, AsmError> {
         self.handlers
             .get_mut(&id)
             .ok_or(AsmError::InvalidSubprotocol(id))
     }
+}
 
-    pub(crate) fn invoke_process_txs<S: Subprotocol>(&mut self, txs: &[TxInput<'_>]) {
-        // We temporarily take the handler out of the map so we can call
-        // `process_txs` with `self` as the relayer without violating the
-        // borrow checker.
-        let mut h = self
-            .handlers
-            .remove(&S::ID)
-            .expect("asm: unloaded subprotocol");
-        h.process_txs(txs, self);
-        self.handlers.insert(S::ID, h);
-    }
-
-    pub(crate) fn invoke_process_msgs<S: Subprotocol>(&mut self) {
-        let h = self
-            .get_handler_mut(S::ID)
-            .expect("asm: unloaded subprotocol");
-        h.process_msgs()
-    }
-
-    pub(crate) fn to_section_state<S: Subprotocol>(&self) -> SectionState {
-        let h = self.get_handler(S::ID).expect("asm: unloaded subprotocol");
-        h.to_section()
+impl HandlerRelayer {
+    pub(crate) fn new() -> Self {
+        Self {
+            handlers: BTreeMap::new(),
+            logs: Vec::new(),
+        }
     }
 }
 

crates/asm/stf/src/stage.rs

@@ -3,31 +3,23 @@
 
 use std::collections::BTreeMap;
 
-use strata_asm_common::{AnchorState, SectionState, Stage, Subprotocol, SubprotocolId, TxInput};
-
-use crate::handler::HandlerRelayer;
+use strata_asm_common::{
+    AnchorState, SectionState, Stage, Subprotocol, SubprotocolId, SubprotocolManager, TxInput,
+};
 
 /// Stage that loads each subprotocol from the anchor state we're basing off of.
 pub(crate) struct SubprotoLoaderStage<'a> {
     anchor_state: &'a AnchorState,
-    handler: HandlerRelayer,
 }
 
 impl<'a> SubprotoLoaderStage<'a> {
     pub(crate) fn new(anchor_state: &'a AnchorState) -> Self {
-        Self {
-            anchor_state,
-            handler: HandlerRelayer::new(),
-        }
-    }
-
-    pub(crate) fn into_handler(self) -> HandlerRelayer {
-        self.handler
+        Self { anchor_state }
     }
 }
 
 impl Stage for SubprotoLoaderStage<'_> {
-    fn process_subprotocol<S: Subprotocol>(&mut self) {
+    fn process_subprotocol<S: Subprotocol>(&mut self, manager: &mut impl SubprotocolManager) {
         // Load or create the subprotocol state.
         // OPTIMIZE: Linear scan is done every time to find the section
         let state = match self.anchor_state.find_section(S::ID) {
@@ -37,50 +29,41 @@ impl Stage for SubprotoLoaderStage<'_> {
             None => S::init(),
         };
 
-        self.handler.insert_subproto::<S>(state);
+        manager.insert_subproto::<S>(state);
     }
 }
 
 /// Stage to process txs pre-extracted from the block for each subprotocol.
 pub(crate) struct ProcessStage<'b> {
     tx_bufs: BTreeMap<SubprotocolId, Vec<TxInput<'b>>>,
-    handler: HandlerRelayer,
 }
 
 impl<'b> ProcessStage<'b> {
-    pub(crate) fn new(
-        tx_bufs: BTreeMap<SubprotocolId, Vec<TxInput<'b>>>,
-        handler: HandlerRelayer,
-    ) -> Self {
-        Self { tx_bufs, handler }
-    }
-
-    pub(crate) fn into_handler(self) -> HandlerRelayer {
-        self.handler
+    pub(crate) fn new(tx_bufs: BTreeMap<SubprotocolId, Vec<TxInput<'b>>>) -> Self {
+        Self { tx_bufs }
     }
 }
 
 impl Stage for ProcessStage<'_> {
-    fn process_subprotocol<S: Subprotocol>(&mut self) {
+    fn process_subprotocol<S: Subprotocol>(&mut self, manager: &mut impl SubprotocolManager) {
         let txs = self
             .tx_bufs
             .get(&S::ID)
             .map(|v| v.as_slice())
             .unwrap_or(&[]);
-        self.handler.invoke_process_txs::<S>(txs);
+        manager.invoke_process_txs::<S>(txs);
     }
 }
 
 /// Stage to handle messages exchanged between subprotocols in execution.
 pub(crate) struct FinishStage {
-    handler: HandlerRelayer,
     sections: Vec<SectionState>,
 }
 
 impl FinishStage {
-    pub(crate) fn new(handler: HandlerRelayer) -> Self {
+    pub(crate) fn new() -> Self {
         let sections = Vec::new();
-        Self { handler, sections }
+        Self { sections }
     }
 
     pub(crate) fn into_sections(self) -> Vec<SectionState> {
@@ -89,9 +72,9 @@ impl FinishStage {
 }
 
 impl Stage for FinishStage {
-    fn process_subprotocol<S: Subprotocol>(&mut self) {
-        self.handler.invoke_process_msgs::<S>();
-        let section = self.handler.to_section_state::<S>();
+    fn process_subprotocol<S: Subprotocol>(&mut self, manager: &mut impl SubprotocolManager) {
+        manager.invoke_process_msgs::<S>();
+        let section = manager.to_section_state::<S>();
         self.sections.push(section);
     }
 }