miden_assembly/

procedure.rs

use alloc::sync::Arc;

use miden_assembly_syntax::{
    ast::{Path, Visibility, types::FunctionType},
    debuginfo::{SourceManager, SourceSpan, Spanned},
    diagnostics::Report,
};
use miden_core::Word;

use super::{
    GlobalItemIndex,
    assembler::{MAX_PROC_LOCALS, error::AssemblerError},
    mast_forest_builder::MastNodeRef,
};

// PROCEDURE CONTEXT
// ================================================================================================

/// Information about a procedure currently being compiled.
pub struct ProcedureContext {
    source_manager: Arc<dyn SourceManager>,
    gid: GlobalItemIndex,
    is_program_entrypoint: bool,
    span: SourceSpan,
    path: Arc<Path>,
    signature: Option<Arc<FunctionType>>,
    visibility: Visibility,
    is_kernel: bool,
    num_locals: u16,
}

// ------------------------------------------------------------------------------------------------
/// Constructors
impl ProcedureContext {
    pub fn new(
        gid: GlobalItemIndex,
        is_program_entrypoint: bool,
        path: Arc<Path>,
        visibility: Visibility,
        signature: Option<Arc<FunctionType>>,
        is_kernel: bool,
        source_manager: Arc<dyn SourceManager>,
    ) -> Self {
        Self {
            source_manager,
            gid,
            is_program_entrypoint,
            span: SourceSpan::UNKNOWN,
            path,
            visibility,
            signature,
            is_kernel,
            num_locals: 0,
        }
    }

    /// Sets the number of locals to allocate for the procedure.
    ///
    /// Returns an error if `num_locals` exceeds `MAX_PROC_LOCALS`, the largest count that
    /// stays representable in a `u16` once rounded up to a word boundary during frame-pointer
    /// codegen. The text parser enforces this on `@locals(..)`, but procedures built directly
    /// via the AST bypass the parser. So the limit is enforced here for all callers.
    ///
    /// Call [`Self::with_span`] first so the error can point at the procedure definition.
    pub fn with_num_locals(mut self, num_locals: u16) -> Result<Self, Report> {
        if num_locals > MAX_PROC_LOCALS {
            let source_file = self.source_manager.get(self.span.source_id()).ok();
            return Err(Report::new(AssemblerError::TooManyProcedureLocals {
                span: self.span,
                source_file,
                max_locals: MAX_PROC_LOCALS,
                num_locals,
            }));
        }
        self.num_locals = num_locals;
        Ok(self)
    }

    pub fn with_span(mut self, span: SourceSpan) -> Self {
        self.span = span;
        self
    }
}

// ------------------------------------------------------------------------------------------------
/// Public accessors
impl ProcedureContext {
    pub fn id(&self) -> GlobalItemIndex {
        self.gid
    }

    pub fn is_program_entrypoint(&self) -> bool {
        self.is_program_entrypoint
    }

    pub fn path(&self) -> &Arc<Path> {
        &self.path
    }

    pub fn signature(&self) -> Option<Arc<FunctionType>> {
        self.signature.clone()
    }

    pub fn set_signature(&mut self, signature: Option<Arc<FunctionType>>) {
        self.signature = signature;
    }

    pub fn num_locals(&self) -> u16 {
        self.num_locals
    }

    pub fn module(&self) -> &Path {
        self.path.parent().unwrap()
    }

    /// Returns true if the procedure is being assembled for a kernel.
    pub fn is_kernel(&self) -> bool {
        self.is_kernel
    }

    #[inline(always)]
    pub fn source_manager(&self) -> &dyn SourceManager {
        self.source_manager.as_ref()
    }
}

// ------------------------------------------------------------------------------------------------
/// State mutators
impl ProcedureContext {
    /// Transforms this procedure context into a [Procedure].
    ///
    /// The passed-in `mast_root` defines the MAST root of the procedure's body while `body_node`
    /// specifies the assembly-time reference to the procedure's body node.
    ///
    /// <div class="warning">
    /// `mast_root` and `body_node` must be consistent. That is, `body_node` must resolve to a MAST
    /// node whose digest equals `mast_root`.
    /// </div>
    pub(crate) fn into_procedure(self, mast_root: Word, body_node_ref: MastNodeRef) -> Procedure {
        let is_syscall = self.is_kernel && self.visibility.is_public();
        Procedure::new(
            self.path,
            self.visibility,
            self.signature,
            is_syscall,
            self.num_locals as u32,
            mast_root,
            body_node_ref,
        )
        .with_span(self.span)
    }
}

impl Spanned for ProcedureContext {
    fn span(&self) -> SourceSpan {
        self.span
    }
}

// PROCEDURE
// ================================================================================================

/// A compiled Miden Assembly procedure, consisting of MAST info and basic metadata.
///
/// Procedure metadata includes:
///
/// - Fully-qualified path of the procedure in Miden Assembly (if known).
/// - Number of procedure locals to allocate.
/// - The visibility of the procedure (e.g. public/private/syscall)
/// - The set of MAST roots invoked by this procedure.
/// - The original source span and file of the procedure (if available).
#[derive(Clone, Debug)]
pub struct Procedure {
    span: SourceSpan,
    path: Arc<Path>,
    signature: Option<Arc<FunctionType>>,
    visibility: Visibility,
    is_syscall: bool,
    num_locals: u32,
    /// The MAST root of the procedure.
    mast_root: Word,
    /// The assembly-time node reference which resolves to the above MAST root.
    body_node_ref: MastNodeRef,
}

// ------------------------------------------------------------------------------------------------
/// Constructors
impl Procedure {
    fn new(
        path: Arc<Path>,
        visibility: Visibility,
        signature: Option<Arc<FunctionType>>,
        is_syscall: bool,
        num_locals: u32,
        mast_root: Word,
        body_node_ref: MastNodeRef,
    ) -> Self {
        Self {
            span: SourceSpan::default(),
            path,
            visibility,
            signature,
            is_syscall,
            num_locals,
            mast_root,
            body_node_ref,
        }
    }

    pub(crate) fn with_span(mut self, span: SourceSpan) -> Self {
        self.span = span;
        self
    }
}

// ------------------------------------------------------------------------------------------------
/// Public accessors
impl Procedure {
    /// Returns source span of this procedure.
    pub fn span(&self) -> &SourceSpan {
        &self.span
    }

    /// Returns a reference to the fully-qualified name of this procedure
    pub fn path(&self) -> &Arc<Path> {
        &self.path
    }

    /// Returns true if this procedure is a syscallable procedure
    #[inline(always)]
    pub const fn is_syscall(&self) -> bool {
        self.is_syscall
    }

    /// Returns the visibility of this procedure as expressed in the original source code
    pub fn visibility(&self) -> Visibility {
        self.visibility
    }

    /// Returns a reference to the fully-qualified module path of this procedure
    pub fn module(&self) -> &Path {
        self.path.parent().unwrap()
    }

    /// Returns a reference to the type signature of this procedure
    pub fn signature(&self) -> Option<Arc<FunctionType>> {
        self.signature.clone()
    }

    /// Returns the number of memory locals reserved by the procedure.
    pub fn num_locals(&self) -> u32 {
        self.num_locals
    }

    /// Returns the root of this procedure's MAST.
    pub fn mast_root(&self) -> Word {
        self.mast_root
    }

    /// Returns the assembly-time node reference of this procedure.
    pub(crate) fn body_node_ref(&self) -> MastNodeRef {
        self.body_node_ref
    }
}

impl Spanned for Procedure {
    fn span(&self) -> SourceSpan {
        self.span
    }
}