artifact.rs

use alloc::sync::Arc;
use core::fmt;

use miden_assembly::Path;
use miden_core::Word;
use miden_mast_package::Package;
use midenc_hir::{constants::ConstantData, dialects::builtin, interner::Symbol};
use midenc_session::{Emit, OutputMode, OutputType, Session, Writer, diagnostics::Report};

use crate::{lower::NativePtr, masm};

mod project_support;

pub struct MasmComponent {
    pub id: Option<builtin::ComponentId>,
    /// The path of the root module for this component
    ///
    /// All components must have a canonical root module, even if empty
    pub root: Arc<Path>,
    /// Whether this component requires an initializer procedure.
    ///
    /// The initializer is responsible for initializing global variables and writing data segments
    /// into memory at program startup, and at cross-context call boundaries (in callee prologue).
    /// When set, a private root-local `init` procedure is generated and invoked by the lifted
    /// export wrappers (and the generated executable entrypoint) via a same-module symbol.
    pub requires_init: bool,
    /// The invocation target for the program entrypoint, if this component is executable.
    ///
    /// If unset, it indicates that the component is a library, even if it could be made executable.
    pub entrypoint: Option<masm::InvocationTarget>,
    /// The kernel library to link against
    pub kernel: Option<masm::KernelLibrary>,
    /// The rodata segments of this component keyed by the offset of the segment
    pub rodata: Vec<Rodata>,
    /// The address of the start of the global heap
    pub heap_base: u32,
    /// The address of the `__stack_pointer` global, if such a global has been defined
    pub stack_pointer: Option<u32>,
    /// Whether non-root support modules are implementation details of the component package.
    ///
    /// When true, support modules are statically linked into library packages instead of being
    /// surfaced as public package modules. This static linking — not procedure visibility — is
    /// what prunes the lowered core Wasm procedures from the package export surface in the common
    /// (non-synthetic-wrapper) component-library case, because MASM currently lowers `Internal`
    /// visibility to public.
    pub link_support_modules_privately: bool,
    /// The set of modules in this component
    pub modules: Vec<Arc<masm::Module>>,
}

impl Emit for MasmComponent {
    fn name(&self) -> Option<Symbol> {
        None
    }

    fn output_type(&self, _mode: OutputMode) -> OutputType {
        OutputType::Masm
    }

    fn write_to<W: Writer>(
        &self,
        mut writer: W,
        mode: OutputMode,
        _session: &Session,
    ) -> anyhow::Result<()> {
        if mode != OutputMode::Text {
            anyhow::bail!("masm emission does not support binary mode");
        }
        writer.write_fmt(core::format_args!("{self}"))?;
        Ok(())
    }
}

/// Represents a read-only data segment, combined with its content digest
#[derive(Clone, PartialEq, Eq)]
pub struct Rodata {
    /// The component to which this read-only data segment belongs
    pub component: builtin::ComponentId,
    /// The content digest computed for `data`
    pub digest: Word,
    /// The address at which the data for this segment begins
    pub start: NativePtr,
    /// The raw binary data for this segment
    pub data: Arc<ConstantData>,
}
impl fmt::Debug for Rodata {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("Rodata")
            .field("digest", &format_args!("{}", &self.digest))
            .field("start", &self.start)
            .field_with("data", |f| {
                f.debug_struct("ConstantData")
                    .field("len", &self.data.len())
                    .finish_non_exhaustive()
            })
            .finish()
    }
}
impl Rodata {
    pub fn size_in_bytes(&self) -> usize {
        self.data.len()
    }

    pub fn size_in_felts(&self) -> usize {
        self.data.len().div_ceil(4)
    }

    pub fn size_in_words(&self) -> usize {
        self.size_in_felts().div_ceil(4)
    }

    /// Attempt to convert this rodata object to its equivalent representation in felts
    ///
    /// See [Self::bytes_to_elements] for more details.
    pub fn to_elements(&self) -> Vec<miden_processor::Felt> {
        Self::bytes_to_elements(self.data.as_slice())
    }

    /// Attempt to convert the given bytes to their equivalent representation in felts
    ///
    /// The resulting felts will be in padded out to the nearest number of words, i.e. if the data
    /// only takes up 3 felts worth of bytes, then the resulting `Vec` will contain 4 felts, so that
    /// the total size is a valid number of words.
    pub fn bytes_to_elements(bytes: &[u8]) -> Vec<miden_processor::Felt> {
        use miden_processor::Felt;

        let mut felts = Vec::with_capacity(bytes.len() / 4);
        let mut iter = bytes.iter().copied().array_chunks::<4>();
        felts.extend(
            iter.by_ref().map(|chunk| Felt::new_unchecked(u32::from_le_bytes(chunk) as u64)),
        );
        let remainder = iter.into_remainder();
        if remainder.len() > 0 {
            let mut chunk = [0u8; 4];
            for (i, byte) in remainder.enumerate() {
                chunk[i] = byte;
            }
            felts.push(Felt::new_unchecked(u32::from_le_bytes(chunk) as u64));
        }

        let size_in_felts = bytes.len().div_ceil(4);
        let size_in_words = size_in_felts.div_ceil(4);
        let padding = (size_in_words * 4).abs_diff(felts.len());
        felts.resize(felts.len() + padding, Felt::ZERO);
        debug_assert_eq!(felts.len() % 4, 0, "expected to be a valid number of words");
        felts
    }
}

inventory::submit! {
    midenc_session::CompileFlag::new("test_harness")
        .long("test-harness")
        .action(midenc_session::FlagAction::SetTrue)
        .help("If present, causes the code generator to emit extra code for the VM test harness")
        .help_heading("Testing")
}

impl fmt::Display for MasmComponent {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        use crate::intrinsics::INTRINSICS_MODULE_NAMES;

        for module in self.modules.iter() {
            // Skip printing the standard library modules and intrinsics modules to focus on the
            // user-defined modules and avoid the
            // stack overflow error when printing large programs
            // https://github.com/0xMiden/miden-formatting/issues/4
            let module_name = module.path().as_str();
            let module_name_trimmed = module_name.trim_start_matches("::");
            if INTRINSICS_MODULE_NAMES.contains(&module_name) {
                continue;
            }
            if module.is_in_namespace(Path::new("std"))
                || module_name_trimmed.starts_with("miden::core")
                || module_name_trimmed.starts_with("miden::protocol")
            {
                continue;
            } else {
                writeln!(f, "# mod {}\n", &module_name)?;
                writeln!(f, "{module}")?;
            }
        }
        Ok(())
    }
}

impl MasmComponent {
    /// Get a mutable reference to the component root module (`modules[0]`).
    ///
    /// The root module is never cloned during lowering, so the unique reference is guaranteed.
    pub(crate) fn root_module_mut(&mut self) -> &mut masm::Module {
        debug_assert!(
            self.modules[0].path() == self.root.as_ref(),
            "modules[0] must be the component root module"
        );
        Arc::get_mut(&mut self.modules[0]).expect("expected unique reference")
    }

    /// Assemble this component into a Miden package.
    pub fn assemble(
        &self,
        account_component_metadata_bytes: Option<&[u8]>,
        session: &Session,
    ) -> Result<Arc<Package>, Report> {
        project_support::assemble(self, account_component_metadata_bytes, session)
    }

    /// Assemble this component into a Miden package using a pre-populated package registry.
    pub fn assemble_with_registry(
        &self,
        account_component_metadata_bytes: Option<&[u8]>,
        session: &Session,
        registry: &mut midenc_session::registry::HybridPackageRegistry,
    ) -> Result<Arc<Package>, Report> {
        project_support::assemble_with_registry(
            self,
            account_component_metadata_bytes,
            session,
            registry,
        )
    }
}

#[cfg(test)]
mod tests {
    use proptest::prelude::*;

    use super::*;

    fn validate_bytes_to_elements(bytes: &[u8]) {
        let result = Rodata::bytes_to_elements(bytes);

        // Each felt represents 4 bytes
        let expected_felts = bytes.len().div_ceil(4);
        // Felts should be padded to a multiple of 4 (1 word = 4 felts)
        let expected_total_felts = expected_felts.div_ceil(4) * 4;

        assert_eq!(
            result.len(),
            expected_total_felts,
            "For {} bytes, expected {} felts (padded from {} felts), but got {}",
            bytes.len(),
            expected_total_felts,
            expected_felts,
            result.len()
        );

        // Verify padding is zeros
        for (i, felt) in result.iter().enumerate().skip(expected_felts) {
            assert_eq!(*felt, miden_processor::Felt::ZERO, "Padding at index {i} should be zero");
        }
    }

    #[test]
    fn test_bytes_to_elements_edge_cases() {
        validate_bytes_to_elements(&[]);
        validate_bytes_to_elements(&[1]);
        validate_bytes_to_elements(&[0u8; 4]);
        validate_bytes_to_elements(&[0u8; 15]);
        validate_bytes_to_elements(&[0u8; 16]);
        validate_bytes_to_elements(&[0u8; 17]);
        validate_bytes_to_elements(&[0u8; 31]);
        validate_bytes_to_elements(&[0u8; 32]);
        validate_bytes_to_elements(&[0u8; 33]);
        validate_bytes_to_elements(&[0u8; 64]);
    }

    proptest! {
        #![proptest_config(ProptestConfig::with_cases(1000))]
        #[test]
        fn proptest_bytes_to_elements(bytes in prop::collection::vec(any::<u8>(), 0..=1000)) {
            validate_bytes_to_elements(&bytes);
        }

        #[test]
        fn proptest_bytes_to_elements_word_boundaries(size_factor in 0u32..=100) {
            // Test specifically around word boundaries
            // Test sizes around multiples of 16 (since 1 word = 4 felts = 16 bytes)
            let base_size = size_factor * 16;
            for offset in -2i32..=2 {
                let size = (base_size as i32 + offset).max(0) as usize;
                let bytes = vec![0u8; size];
                validate_bytes_to_elements(&bytes);
            }
        }
    }
}