miden-protocol 0.15.0

use miden_assembly::Library;
use miden_assembly::library::ProcedureExport;
use miden_processor::mast::{MastForest, MastNodeExt};

use crate::account::AccountProcedureRoot;
use crate::assembly::Path;
use crate::vm::AdviceMap;

// ACCOUNT COMPONENT CODE
// ================================================================================================

/// A [`Library`] that has been assembled for use as component code.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AccountComponentCode(Library);

impl AccountComponentCode {
    /// Returns a reference to the underlying [`Library`]
    pub fn as_library(&self) -> &Library {
        &self.0
    }

    /// Returns a reference to the code's [`MastForest`]
    pub fn mast_forest(&self) -> &MastForest {
        self.0.mast_forest().as_ref()
    }

    /// Consumes `self` and returns the underlying [`Library`]
    pub fn into_library(self) -> Library {
        self.0
    }

    /// Returns an iterator over the [`AccountProcedureRoot`]s of this component's exported
    /// procedures.
    pub fn procedure_roots(&self) -> impl Iterator<Item = AccountProcedureRoot> + '_ {
        self.0.exports().filter_map(|export| {
            export.as_procedure().map(|proc_export| {
                let digest = self.0.mast_forest()[proc_export.node].digest();
                AccountProcedureRoot::from_raw(digest)
            })
        })
    }

    /// Returns the procedure exports of this component.
    pub fn exports(&self) -> impl Iterator<Item = &ProcedureExport> + '_ {
        self.0.exports().filter_map(|export| export.as_procedure())
    }

    /// Returns the [`AccountProcedureRoot`] of the procedure with the specified path, or `None`
    /// if it was not found in this component's library.
    pub fn get_procedure_root_by_path(
        &self,
        proc_name: impl AsRef<Path>,
    ) -> Option<AccountProcedureRoot> {
        self.0.get_procedure_root_by_path(proc_name).map(AccountProcedureRoot::from_raw)
    }

    /// Returns a new [AccountComponentCode] with the provided advice map entries merged into the
    /// underlying [Library]'s [MastForest].
    ///
    /// This allows adding advice map entries to an already-compiled account component,
    /// which is useful when the entries are determined after compilation.
    pub fn with_advice_map(self, advice_map: AdviceMap) -> Self {
        if advice_map.is_empty() {
            return self;
        }

        Self(self.0.with_advice_map(advice_map))
    }
}

impl AsRef<Library> for AccountComponentCode {
    fn as_ref(&self) -> &Library {
        self.as_library()
    }
}

// CONVERSIONS
// ================================================================================================

impl From<Library> for AccountComponentCode {
    fn from(value: Library) -> Self {
        Self(value)
    }
}

impl From<AccountComponentCode> for Library {
    fn from(value: AccountComponentCode) -> Self {
        value.into_library()
    }
}

// TESTS
// ================================================================================================

#[cfg(test)]
mod tests {
    use alloc::string::ToString;
    use alloc::sync::Arc;

    use miden_core::{Felt, Word};

    use super::*;
    use crate::assembly::Assembler;

    #[test]
    fn test_account_component_code_with_advice_map() {
        let assembler = Assembler::default();
        let library = Arc::unwrap_or_clone(
            assembler
                .assemble_library(["pub proc test nop end"])
                .expect("failed to assemble library"),
        );
        let component_code = AccountComponentCode::from(library);

        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 component_code = component_code.with_advice_map(AdviceMap::default());
        assert_eq!(original_digest, component_code.as_library().digest());

        // Non-empty advice map should add entries
        let key = Word::from([10u32, 20, 30, 40]);
        let value = vec![Felt::from(200_u8)];
        let mut advice_map = AdviceMap::default();
        advice_map.insert(key, value.clone());

        let component_code = component_code.with_advice_map(advice_map);

        let mast = component_code.mast_forest();
        let stored = mast.advice_map().get(&key).expect("entry should be present");
        assert_eq!(stored.as_ref(), value.as_slice());
    }

    #[test]
    fn test_get_procedure_root_by_path() {
        let assembler = Assembler::default();
        let library = Arc::unwrap_or_clone(
            assembler
                .assemble_library(["pub proc test_proc nop end"])
                .expect("failed to assemble library"),
        );
        let component_code = AccountComponentCode::from(library);

        // The test library exports exactly one procedure.
        assert_eq!(component_code.procedure_roots().count(), 1);
        let expected = component_code.procedure_roots().next().expect("one procedure exported");

        let library_namespace = component_code
            .as_library()
            .module_infos()
            .next()
            .expect("library should have one module")
            .path()
            .to_string();
        let proc_path = alloc::format!("{library_namespace}::test_proc");

        let root = component_code
            .get_procedure_root_by_path(proc_path.as_str())
            .expect("test_proc should be present");
        assert_eq!(root, expected);

        assert!(component_code.get_procedure_root_by_path("bogus::missing").is_none());
    }
}