refactor: minor improvements to AccountComponentCode

CHANGELOG.md

@@ -40,6 +40,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/component/code.rs

@@ -1,33 +1,64 @@
-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::{MastForest, MastNodeExt};
 
 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 MAST forest and procedure exports.
+    pub fn new(mast: Arc<MastForest>, exports: Vec<ProcedureExport>) -> Self {
+        Self { mast, exports }
     }
 
-    /// Returns a reference to the code's [`MastForest`]
+    // PUBLIC ACCESSORS
+    // --------------------------------------------------------------------------------------------
+
+    /// Returns a reference to the code's [`MastForest`].
     pub fn mast_forest(&self) -> &MastForest {
-        self.0.mast_forest().as_ref()
+        self.mast.as_ref()
     }
 
-    /// Consumes `self` and returns the underlying [`Library`]
-    pub fn into_library(self) -> Library {
-        self.0
+    /// Returns the [`MastForest`] wrapped in an [`Arc`].
+    pub fn mast(&self) -> Arc<MastForest> {
+        self.mast.clone()
     }
 
-    /// Returns a new [AccountComponentCode] with the provided advice map entries merged into the
-    /// underlying [Library]'s [MastForest].
+    /// Returns the procedure exports of this component.
+    pub fn exports(&self) -> &[ProcedureExport] {
+        &self.exports
+    }
+
+    /// 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())
+    }
+
+    /// 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 +67,46 @@ 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);
 
-impl AsRef<Library> for AccountComponentCode {
-    fn as_ref(&self) -> &Library {
-        self.as_library()
+        Self {
+            mast: Arc::new(mast),
+            exports: self.exports,
+        }
     }
 }
 
 // CONVERSIONS
 // ================================================================================================
 
 impl From<Library> for AccountComponentCode {
-    fn from(value: Library) -> Self {
-        Self(value)
+    fn from(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 }
     }
 }
 
 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")
+    }
+}
+
+impl From<&AccountComponentCode> for Library {
+    fn from(value: &AccountComponentCode) -> Self {
+        let exports: BTreeMap<_, _> =
+            value.exports.iter().cloned().map(|e| (e.path.clone(), e.into())).collect();
+
+        Library::new(value.mast.clone(), exports)
+            .expect("AccountComponentCode should have at least one export")
     }
 }
 
@@ -82,10 +131,9 @@ mod tests {
         assert!(component_code.mast_forest().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]);

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

@@ -196,31 +196,26 @@ 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 (digest, 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 get_procedures(&self) -> impl Iterator<Item = (Word, bool)> + '_ {
+        let mast = self.code.mast_forest();
+        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);
+            (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/code_builder/mod.rs

@@ -70,7 +70,7 @@ use crate::standards_lib::StandardsLib;
 /// let script = CodeBuilder::default()
 ///     .with_linked_module("my::module", module_code).context("failed to link module")?
 ///     .with_statically_linked_library(&my_lib).context("failed to link static library")?
-///     .with_dynamically_linked_library(&fpi_lib).context("failed to link dynamic library")?  // For FPI calls
+///     .with_dynamically_linked_library(fpi_lib).context("failed to link dynamic library")?  // For FPI calls
 ///     .compile_tx_script(script_code).context("failed to parse tx script")?;
 /// # Ok(())
 /// # }
@@ -223,9 +223,10 @@ impl CodeBuilder {
     /// Returns an error if the library cannot be added to the assembler
     pub fn with_dynamically_linked_library(
         mut self,
-        library: impl AsRef<Library>,
+        library: impl Into<Library>,
     ) -> Result<Self, CodeBuilderError> {
-        self.link_dynamic_library(library.as_ref())?;
+        let library = library.into();
+        self.link_dynamic_library(&library)?;
         Ok(self)
     }
 
@@ -697,7 +698,7 @@ mod tests {
         let builder = CodeBuilder::default()
             .with_statically_linked_library(&static_lib)
             .context("failed to link static library")?
-            .with_dynamically_linked_library(&dynamic_lib)
+            .with_dynamically_linked_library(dynamic_lib)
             .context("failed to link dynamic library")?;
 
         builder

crates/miden-standards/src/testing/account_component/incr_nonce.rs

@@ -18,7 +18,7 @@ static INCR_NONCE_AUTH_LIBRARY: LazyLock<Library> = LazyLock::new(|| {
     CodeBuilder::default()
         .compile_component_code("incr_nonce", INCR_NONCE_AUTH_CODE)
         .expect("incr nonce code should be valid")
-        .into_library()
+        .into()
 });
 
 /// Creates a mock authentication [`AccountComponent`] for testing purposes under the "incr_nonce"

crates/miden-standards/src/testing/mock_account_code.rs

@@ -142,14 +142,14 @@ static MOCK_FAUCET_LIBRARY: LazyLock<Library> = LazyLock::new(|| {
     CodeBuilder::default()
         .compile_component_code("mock::faucet", MOCK_FAUCET_CODE)
         .expect("mock faucet code should be valid")
-        .into_library()
+        .into()
 });
 
 static MOCK_ACCOUNT_LIBRARY: LazyLock<Library> = LazyLock::new(|| {
     CodeBuilder::default()
         .compile_component_code("mock::account", MOCK_ACCOUNT_CODE)
         .expect("mock account code should be valid")
-        .into_library()
+        .into()
 });
 
 // MOCK ACCOUNT CODE EXT

crates/miden-testing/src/kernel_tests/tx/test_account.rs

@@ -1571,7 +1571,7 @@ async fn transaction_executor_account_code_using_custom_library() -> anyhow::Res
         .build_existing()?;
 
     let tx_script = CodeBuilder::default()
-        .with_dynamically_linked_library(&account_component_lib)?
+        .with_dynamically_linked_library(account_component_lib.clone())?
         .compile_tx_script(tx_script_src)?;
 
     let tx_context = TransactionContextBuilder::new(native_account.clone())

crates/miden-testing/src/kernel_tests/tx/test_array.rs

@@ -126,7 +126,7 @@ async fn test_array_get_and_set() -> anyhow::Result<()> {
 
     // Compile the transaction script with the wrapper library linked
     let tx_script = CodeBuilder::default()
-        .with_dynamically_linked_library(&wrapper_library)?
+        .with_dynamically_linked_library(wrapper_library.clone())?
         .compile_tx_script(tx_script_code)?;
 
     // Create transaction context and execute
@@ -243,7 +243,7 @@ async fn test_double_word_array_get_and_set() -> anyhow::Result<()> {
     );
 
     let tx_script = CodeBuilder::default()
-        .with_dynamically_linked_library(&wrapper_library)?
+        .with_dynamically_linked_library(wrapper_library.clone())?
         .compile_tx_script(tx_script_code)?;
 
     let tx_context = TransactionContextBuilder::new(account).tx_script(tx_script).build()?;

crates/miden-testing/src/kernel_tests/tx/test_tx.rs

@@ -796,7 +796,7 @@ async fn inputs_created_correctly() -> anyhow::Result<()> {
         "#;
 
     let tx_script = CodeBuilder::default()
-        .with_dynamically_linked_library(component_code.as_library())?
+        .with_dynamically_linked_library(component_code)?
         .compile_tx_script(script)?;
 
     assert!(tx_script.mast().advice_map().get(&Word::try_from([1u64, 2, 3, 4])?).is_some());