refactor(asm): clean up types and finalization dataflow some more

Author: delbonis

crates/asm/common/src/subprotocol.rs

@@ -8,7 +8,7 @@ use std::any::Any;
 
 use borsh::{BorshDeserialize, BorshSerialize};
 
-use crate::{AsmError, Log, SectionState, TxInput, msg::InterprotoMsg};
+use crate::{Log, SectionState, TxInput, msg::InterprotoMsg};
 
 /// Identifier for a subprotocol.
 pub type SubprotocolId = u8;
@@ -40,13 +40,17 @@ pub trait Subprotocol: 'static {
     /// Update it's own state and output a list of InterProtoMsg addressed to other subprotocols
     fn process_txs(state: &mut Self::State, txs: &[TxInput<'_>], relayer: &mut impl MsgRelayer);
 
-    /// Use the msg other subprotocols to update its state. Also generate the event
-    /// logs that is later needed for introspection. Return the commitment of the events. The actual
+    /// Use the msgs other subprotocols to update its state.
+    ///
+    /// TODO:
+    /// Also generate the event logs that is later needed for other components
+    /// to read ASM activity. Return the commitment of the events. The actual
     /// event is defined by the subprotocol and is not visible to the ASM.
-    fn finalize_state(state: &mut Self::State, msgs: &[Self::Msg]);
+    fn process_msgs(state: &mut Self::State, msgs: &[Self::Msg]);
 }
 
-/// Generic message relayer interface.
+/// Generic message relayer interface which subprotocols can use to interact
+/// with each other and the outside world.
 pub trait MsgRelayer: Any {
     /// Relays a message to the destination subprotocol.
     fn relay_msg(&mut self, m: &dyn InterprotoMsg);
@@ -78,8 +82,8 @@ pub trait SubprotoHandler {
     /// 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);
+    /// Processes the buffered messages stored in the handler.
+    fn process_buffered_msgs(&mut self);
 
     /// Repacks the state into a [`SectionState`] instance.
     fn to_section(&self) -> SectionState;

crates/asm/stf/src/manager.rs

@@ -67,7 +67,6 @@ pub(crate) struct SubprotoManager {
     logs: Vec<Log>,
 }
 
-/* SubprotocolManager for */
 impl SubprotoManager {
     /// Inserts a subproto by creating a handler for it, wrapping a tstate.
     pub(crate) fn insert_subproto<S: Subprotocol>(&mut self, state: S::State) {
@@ -122,6 +121,7 @@ impl SubprotoManager {
             .ok_or(AsmError::InvalidSubprotocol(id))
     }
 
+    #[allow(unused)]
     fn get_handler(&self, id: SubprotocolId) -> Result<&dyn SubprotoHandler, AsmError> {
         self.handlers
             .get(&id)
@@ -139,10 +139,29 @@ impl SubprotoManager {
     }
 
     /// Extracts the section state for a subprotocol.
+    #[allow(unused)]
     pub(crate) fn to_section_state<S: Subprotocol>(&self) -> SectionState {
         let h = self.get_handler(S::ID).expect("asm: unloaded subprotocol");
         h.to_section()
     }
+
+    /// Exports each handler as a section we can use when constructing the final
+    /// `AnchorState`.  Consumes the manager.
+    pub(crate) fn export_sections(self) -> Vec<SectionState> {
+        let sections = self
+            .handlers
+            .into_values()
+            .map(|h| h.to_section())
+            .collect::<Vec<_>>();
+
+        // sanity check
+        assert!(
+            sections.is_sorted_by_key(|s| s.id),
+            "asm: sections not sorted on export"
+        );
+
+        sections
+    }
 }
 
 impl SubprotoManager {

crates/asm/stf/src/stage.rs

@@ -3,7 +3,7 @@
 
 use std::collections::BTreeMap;
 
-use strata_asm_common::{AnchorState, SectionState, Stage, Subprotocol, SubprotocolId, TxInput};
+use strata_asm_common::{AnchorState, Stage, Subprotocol, SubprotocolId, TxInput};
 
 use crate::manager::SubprotoManager;
 
@@ -65,26 +65,17 @@ impl Stage for ProcessStage<'_, '_> {
 
 /// Stage to handle messages exchanged between subprotocols in execution.
 pub(crate) struct FinishStage<'m> {
-    sections: Vec<SectionState>,
     manager: &'m mut SubprotoManager,
 }
 
 impl<'m> FinishStage<'m> {
     pub(crate) fn new(manager: &'m mut SubprotoManager) -> Self {
-        let sections = Vec::new();
-        Self { sections, manager }
-    }
-
-    pub(crate) fn into_sorted_sections(mut self) -> Vec<SectionState> {
-        self.sections.sort_by_key(|state| state.id);
-        self.sections
+        Self { manager }
     }
 }
 
 impl Stage for FinishStage<'_> {
     fn process_subprotocol<S: Subprotocol>(&mut self) {
         self.manager.invoke_process_msgs::<S>();
-        let section = self.manager.to_section_state::<S>();
-        self.sections.push(section);
     }
 }

crates/asm/stf/src/transition.rs

@@ -26,10 +26,9 @@ impl AsmSpec for StrataAsmSpec {
 
 /// Computes the next AnchorState by applying the Anchor State Machine (ASM) state transition
 /// function (STF) to the given previous state and new L1 block.
-pub fn asm_stf<S: AsmSpec>(pre_state: AnchorState, block: Block) -> Result<AnchorState, AsmError> {
+pub fn asm_stf<S: AsmSpec>(pre_state: AnchorState, block: &Block) -> Result<AnchorState, AsmError> {
+    // 1. Validate and update PoW header continuity for the new block.
     let mut pow_state = pre_state.chain_view.pow_state.clone();
-
-    // 1. Validate and update PoW header continuity for the new block
     pow_state
         .check_and_update_continuity(&block.header, &Params::MAINNET)
         .map_err(AsmError::InvalidL1Header)?;
@@ -39,19 +38,20 @@ pub fn asm_stf<S: AsmSpec>(pre_state: AnchorState, block: Block) -> Result<Ancho
 
     let mut manager = SubprotoManager::new();
 
-    // 3. LOAD: bring each subprotocol into a HandlerRelayer
+    // 3. LOAD: Bring each subprotocol into the subproto manager.
     let mut loader_stage = SubprotoLoaderStage::new(&pre_state, &mut manager);
     S::call_subprotocols(&mut loader_stage);
 
-    // 4. PROCESS: feed each subprotocol its slice of txs
+    // 4. PROCESS: Feed each subprotocol its slice of txs.
     let mut process_stage = ProcessStage::new(all_relevant_transactions, &mut manager);
     S::call_subprotocols(&mut process_stage);
 
-    // 5. FINISH: let each subprotocol process its buffered interproto messages
+    // 5. FINISH: Let each subprotocol process its buffered interproto messages.
     let mut finish_stage = FinishStage::new(&mut manager);
     S::call_subprotocols(&mut finish_stage);
 
-    let sections = finish_stage.into_sorted_sections();
+    // 6. Construct the final `AnchorState` we return.
+    let sections = manager.export_sections();
     let chain_view = ChainViewState { pow_state };
     Ok(AnchorState {
         chain_view,