refactor: minor improvements to AccountComponentCode

CHANGELOG.md

@@ -41,6 +41,7 @@
 
 ### Changes
 
+- [BREAKING] Deconstructed `AccountComponentCode` into `Arc<MastForest>` + `Vec<ProcedureExport>` fields, replacing the `Library` newtype wrapper ([#2597](https://github.com/0xMiden/protocol/pull/2597)).
 - [BREAKING] Removed `NoteAssets::add_asset`; `OutputNoteBuilder` now accumulates assets in a `Vec` and computes the commitment only when `build()` is called, avoiding rehashing on every asset addition. ([#2577](https://github.com/0xMiden/protocol/pull/2577))
 - [BREAKING] Made `supported_types` a required parameter of `AccountComponentMetadata::new()`; removed `with_supported_type`, `with_supported_types`, `with_supports_all_types`, and `with_supports_regular_types` builder methods; added `AccountType::all()` and `AccountType::regular()` helpers ([#2554](https://github.com/0xMiden/protocol/pull/2554)).
 - [BREAKING] Migrated to miden-vm 0.21 and miden-crypto 0.22 ([#2508](https://github.com/0xMiden/miden-base/pull/2508)).

crates/miden-protocol/src/account/code/mod.rs

@@ -93,8 +93,9 @@ impl AccountCode {
     pub(super) fn from_components_unchecked(
         components: &[AccountComponent],
     ) -> Result<Self, AccountError> {
+        let mast_forests: Vec<_> = components.iter().map(|component| component.mast()).collect();
         let (merged_mast_forest, _) =
-            MastForest::merge(components.iter().map(|component| component.mast_forest()))
+            MastForest::merge(mast_forests.iter().map(|mast| mast.as_ref()))
                 .map_err(AccountError::AccountComponentMastForestMergeError)?;
 
         let mut builder = AccountProcedureBuilder::new();
@@ -338,7 +339,7 @@ impl AccountProcedureBuilder {
     fn add_auth_component(&mut self, component: &AccountComponent) -> Result<(), AccountError> {
         let mut auth_proc_count = 0;
 
-        for (proc_root, is_auth) in component.get_procedures() {
+        for (proc_root, is_auth) in component.procedures() {
             self.add_procedure(proc_root);
 
             if is_auth {
@@ -358,20 +359,19 @@ impl AccountProcedureBuilder {
     }
 
     fn add_component(&mut self, component: &AccountComponent) -> Result<(), AccountError> {
-        for (proc_mast_root, is_auth) in component.get_procedures() {
+        for (proc_root, is_auth) in component.procedures() {
             if is_auth {
                 return Err(AccountError::AccountCodeMultipleAuthComponents);
             }
-            self.add_procedure(proc_mast_root);
+            self.add_procedure(proc_root);
         }
 
         Ok(())
     }
 
-    fn add_procedure(&mut self, proc_mast_root: Word) {
+    fn add_procedure(&mut self, proc_root: AccountProcedureRoot) {
         // Allow procedures with the same MAST root from different components, but only add them
         // once.
-        let proc_root = AccountProcedureRoot::from_raw(proc_mast_root);
         if !self.procedures.contains(&proc_root) {
             self.procedures.push(proc_root);
         }

crates/miden-protocol/src/account/component/code.rs

@@ -1,33 +1,125 @@
-use miden_assembly::Library;
-use miden_processor::mast::MastForest;
+use alloc::collections::BTreeMap;
+use alloc::sync::Arc;
+use alloc::vec::Vec;
 
+use miden_assembly::library::ProcedureExport;
+use miden_assembly::{Library, Path};
+use miden_core::Word;
+use miden_core::mast::{ExternalNodeBuilder, MastForest, MastForestContributor, MastNodeExt};
+
+use crate::account::AccountProcedureRoot;
 use crate::vm::AdviceMap;
 
 // ACCOUNT COMPONENT CODE
 // ================================================================================================
 
-/// A [`Library`] that has been assembled for use as component code.
+/// The code associated with an account component, consisting of a [`MastForest`] and the set of
+/// procedures exported by the component.
 #[derive(Debug, Clone, PartialEq, Eq)]
-pub struct AccountComponentCode(Library);
+pub struct AccountComponentCode {
+    mast: Arc<MastForest>,
+    exports: Vec<ProcedureExport>,
+}
 
 impl AccountComponentCode {
-    /// Returns a reference to the underlying [`Library`]
-    pub fn as_library(&self) -> &Library {
-        &self.0
+    // CONSTRUCTORS
+    // --------------------------------------------------------------------------------------------
+
+    /// Creates a new [`AccountComponentCode`] from the provided [`Library`].
+    pub fn from_library(library: Library) -> Self {
+        let mast = library.mast_forest().clone();
+        let exports: Vec<ProcedureExport> =
+            library.exports().filter_map(|export| export.as_procedure().cloned()).collect();
+
+        Self { mast, exports }
     }
 
-    /// Returns a reference to the code's [`MastForest`]
-    pub fn mast_forest(&self) -> &MastForest {
-        self.0.mast_forest().as_ref()
+    /// Returns a new [`AccountComponentCode`] containing only external node references to the
+    /// procedures in the provided [`Library`].
+    ///
+    /// Note: This method creates a minimal [`MastForest`] where each exported procedure is
+    /// represented by an external node referencing its digest, rather than copying the entire
+    /// library's MAST forest. The actual procedure code will be resolved at runtime via the
+    /// `MastForestStore`.
+    pub fn from_library_reference(library: &Library) -> Self {
+        let mut mast = MastForest::new();
+        let mut exports = Vec::new();
+
+        for export in library.exports() {
+            if let Some(proc_export) = export.as_procedure() {
+                // Get the digest of the procedure from the library
+                let digest = library.mast_forest()[proc_export.node].digest();
+
+                // Create an external node referencing the digest
+                let node_id = ExternalNodeBuilder::new(digest)
+                    .add_to_forest(&mut mast)
+                    .expect("adding external node to forest should not fail");
+                mast.make_root(node_id);
+
+                exports.push(ProcedureExport {
+                    node: node_id,
+                    path: proc_export.path.clone(),
+                    signature: proc_export.signature.clone(),
+                    attributes: proc_export.attributes.clone(),
+                });
+            }
+        }
+
+        Self { mast: Arc::new(mast), exports }
     }
 
-    /// Consumes `self` and returns the underlying [`Library`]
-    pub fn into_library(self) -> Library {
-        self.0
+    /// Creates a new [`AccountComponentCode`] from the provided MAST forest and procedure exports.
+    ///
+    /// # Panics
+    ///
+    /// Panics if any of the exported procedure node IDs are not found in the MAST forest.
+    pub fn from_parts(mast: Arc<MastForest>, exports: Vec<ProcedureExport>) -> Self {
+        for export in &exports {
+            assert!(
+                mast.get_node_by_id(export.node).is_some(),
+                "exported procedure node not found in the MAST forest"
+            );
+        }
+
+        Self { mast, exports }
     }
 
-    /// Returns a new [AccountComponentCode] with the provided advice map entries merged into the
-    /// underlying [Library]'s [MastForest].
+    // PUBLIC ACCESSORS
+    // --------------------------------------------------------------------------------------------
+
+    /// Returns the [`MastForest`] wrapped in an [`Arc`].
+    pub fn mast(&self) -> Arc<MastForest> {
+        self.mast.clone()
+    }
+
+    /// Returns the procedure exports of this component.
+    pub fn exports(&self) -> &[ProcedureExport] {
+        &self.exports
+    }
+
+    /// Returns an iterator over the [`AccountProcedureRoot`]s of this component's exported
+    /// procedures.
+    pub fn procedure_roots(&self) -> impl Iterator<Item = AccountProcedureRoot> + '_ {
+        self.exports
+            .iter()
+            .map(|export| AccountProcedureRoot::from_raw(self.mast[export.node].digest()))
+    }
+
+    /// Returns the digest of the procedure with the specified path, or `None` if it was not found
+    /// in this component.
+    pub fn get_procedure_root_by_path(&self, path: impl AsRef<Path>) -> Option<Word> {
+        let path = path.as_ref().to_absolute();
+        self.exports
+            .iter()
+            .find(|export| export.path.as_ref() == path.as_ref())
+            .map(|export| self.mast[export.node].digest())
+    }
+
+    // MUTATORS
+    // --------------------------------------------------------------------------------------------
+
+    /// Returns a new [`AccountComponentCode`] with the provided advice map entries merged into the
+    /// underlying [`MastForest`].
     ///
     /// This allows adding advice map entries to an already-compiled account component,
     /// which is useful when the entries are determined after compilation.
@@ -36,28 +128,39 @@ impl AccountComponentCode {
             return self;
         }
 
-        Self(self.0.with_advice_map(advice_map))
+        let mut mast = Arc::unwrap_or_clone(self.mast);
+        mast.advice_map_mut().extend(advice_map);
+
+        Self {
+            mast: Arc::new(mast),
+            exports: self.exports,
+        }
     }
-}
 
-impl AsRef<Library> for AccountComponentCode {
-    fn as_ref(&self) -> &Library {
-        self.as_library()
+    /// Clears the debug info from the underlying [`MastForest`].
+    pub fn clear_debug_info(&mut self) {
+        let mut mast = self.mast.clone();
+        Arc::make_mut(&mut mast).clear_debug_info();
+        self.mast = mast;
     }
 }
 
 // CONVERSIONS
 // ================================================================================================
 
 impl From<Library> for AccountComponentCode {
-    fn from(value: Library) -> Self {
-        Self(value)
+    fn from(library: Library) -> Self {
+        Self::from_library(library)
     }
 }
 
 impl From<AccountComponentCode> for Library {
     fn from(value: AccountComponentCode) -> Self {
-        value.into_library()
+        let exports: BTreeMap<_, _> =
+            value.exports.into_iter().map(|e| (e.path.clone(), e.into())).collect();
+
+        Library::new(value.mast, exports)
+            .expect("AccountComponentCode should have at least one export")
     }
 }
 
@@ -79,13 +182,12 @@ mod tests {
             .expect("failed to assemble library");
         let component_code = AccountComponentCode::from(library);
 
-        assert!(component_code.mast_forest().advice_map().is_empty());
+        assert!(component_code.mast().advice_map().is_empty());
 
         // Empty advice map should be a no-op (digest stays the same)
-        let cloned = component_code.clone();
-        let original_digest = cloned.as_library().digest();
+        let original_digest = *Library::from(component_code.clone()).digest();
         let component_code = component_code.with_advice_map(AdviceMap::default());
-        assert_eq!(original_digest, component_code.as_library().digest());
+        assert_eq!(&original_digest, Library::from(component_code.clone()).digest());
 
         // Non-empty advice map should add entries
         let key = Word::from([10u32, 20, 30, 40]);
@@ -95,7 +197,7 @@ mod tests {
 
         let component_code = component_code.with_advice_map(advice_map);
 
-        let mast = component_code.mast_forest();
+        let mast = component_code.mast();
         let stored = mast.advice_map().get(&key).expect("entry should be present");
         assert_eq!(stored.as_ref(), value.as_slice());
     }

crates/miden-protocol/src/account/component/mod.rs

@@ -13,7 +13,7 @@ pub use storage::*;
 mod code;
 pub use code::AccountComponentCode;
 
-use crate::account::{AccountType, StorageSlot};
+use crate::account::{AccountProcedureRoot, AccountType, StorageSlot};
 use crate::assembly::Path;
 use crate::errors::AccountError;
 use crate::{MastForest, Word};
@@ -166,9 +166,9 @@ impl AccountComponent {
         &self.code
     }
 
-    /// Returns a reference to the underlying [`MastForest`] of this component.
-    pub fn mast_forest(&self) -> &MastForest {
-        self.code.mast_forest()
+    /// Returns the [`MastForest`] of this component wrapped in an [`Arc`](alloc::sync::Arc).
+    pub fn mast(&self) -> alloc::sync::Arc<MastForest> {
+        self.code.mast()
     }
 
     /// Returns a slice of the underlying [`StorageSlot`]s of this component.
@@ -196,31 +196,27 @@ impl AccountComponent {
         self.metadata.supported_types().contains(&account_type)
     }
 
-    /// Returns a vector of tuples (digest, is_auth) for all procedures in this component.
+    /// Returns an iterator over ([`AccountProcedureRoot`], is_auth) for all procedures in this
+    /// component.
     ///
     /// A procedure is considered an authentication procedure if it has the `@auth_script`
     /// attribute.
-    pub fn get_procedures(&self) -> Vec<(Word, bool)> {
-        let library = self.code.as_library();
-        let mut procedures = Vec::new();
-        for export in library.exports() {
-            if let Some(proc_export) = export.as_procedure() {
-                let digest = library
-                    .mast_forest()
-                    .get_node_by_id(proc_export.node)
-                    .expect("export node not in the forest")
-                    .digest();
-                let is_auth = proc_export.attributes.has(AUTH_SCRIPT_ATTRIBUTE);
-                procedures.push((digest, is_auth));
-            }
-        }
-        procedures
+    pub fn procedures(&self) -> impl Iterator<Item = (AccountProcedureRoot, bool)> + '_ {
+        let mast = self.code.mast();
+        self.code.exports().iter().map(move |proc_export| {
+            let digest = mast
+                .get_node_by_id(proc_export.node)
+                .expect("export node not in the forest")
+                .digest();
+            let is_auth = proc_export.attributes.has(AUTH_SCRIPT_ATTRIBUTE);
+            (AccountProcedureRoot::from_raw(digest), is_auth)
+        })
     }
 
     /// Returns the digest of the procedure with the specified path, or `None` if it was not found
-    /// in this component's library or its library path is malformed.
+    /// in this component.
     pub fn get_procedure_root_by_path(&self, proc_name: impl AsRef<Path>) -> Option<Word> {
-        self.code.as_library().get_procedure_root_by_path(proc_name)
+        self.code.get_procedure_root_by_path(proc_name)
     }
 }
 

crates/miden-protocol/src/lib.rs

@@ -37,7 +37,7 @@ pub use protocol::ProtocolLib;
 pub mod assembly {
     pub use miden_assembly::ast::{Module, ModuleKind, ProcedureName, QualifiedProcedureName};
     pub use miden_assembly::debuginfo::SourceManagerSync;
-    pub use miden_assembly::library::LibraryExport;
+    pub use miden_assembly::library::{LibraryExport, ProcedureExport as LibraryProcedureExport};
     pub use miden_assembly::{
         Assembler,
         DefaultSourceManager,

crates/miden-standards/src/account/interface/test.rs

@@ -461,7 +461,7 @@ fn test_custom_account_custom_notes() {
         end
     ";
     let note_script = CodeBuilder::default()
-        .with_dynamically_linked_library(account_component.component_code())
+        .with_dynamically_linked_library(account_component.component_code().clone())
         .unwrap()
         .compile_note_script(compatible_source_code)
         .unwrap();
@@ -488,7 +488,7 @@ fn test_custom_account_custom_notes() {
         end
     ";
     let note_script = CodeBuilder::default()
-        .with_dynamically_linked_library(account_component.component_code())
+        .with_dynamically_linked_library(account_component.component_code().clone())
         .unwrap()
         .compile_note_script(incompatible_source_code)
         .unwrap();
@@ -571,7 +571,7 @@ fn test_custom_account_multiple_components_custom_notes() {
         end
     ";
     let note_script = CodeBuilder::default()
-        .with_dynamically_linked_library(custom_component.component_code())
+        .with_dynamically_linked_library(custom_component.component_code().clone())
         .unwrap()
         .compile_note_script(compatible_source_code)
         .unwrap();
@@ -609,7 +609,7 @@ fn test_custom_account_multiple_components_custom_notes() {
         end
     ";
     let note_script = CodeBuilder::default()
-        .with_dynamically_linked_library(custom_component.component_code())
+        .with_dynamically_linked_library(custom_component.component_code().clone())
         .unwrap()
         .compile_note_script(incompatible_source_code)
         .unwrap();