Skip to main content

miden_assembly/
procedure.rs

1use alloc::sync::Arc;
2
3use miden_assembly_syntax::{
4    ast::{Attribute, AttributeSet, MetaExpr, Path, PathBuf, Visibility, types::FunctionType},
5    debuginfo::{SourceManager, SourceSpan, Spanned},
6    diagnostics::Report,
7};
8use miden_core::Word;
9
10use super::{
11    GlobalItemIndex,
12    assembler::{MAX_PROC_LOCALS, error::AssemblerError},
13    mast_forest_builder::{MastNodeRef, MastNodeUse, SourceNodeRef},
14};
15
16// PROCEDURE CONTEXT
17// ================================================================================================
18
19/// Information about a procedure currently being compiled.
20pub struct ProcedureContext {
21    source_manager: Arc<dyn SourceManager>,
22    gid: GlobalItemIndex,
23    is_program_entrypoint: bool,
24    span: SourceSpan,
25    path: Arc<Path>,
26    signature: Option<Arc<FunctionType>>,
27    attributes: AttributeSet,
28    visibility: Visibility,
29    is_kernel: bool,
30    num_locals: u16,
31}
32
33// ------------------------------------------------------------------------------------------------
34/// Constructors
35impl ProcedureContext {
36    pub fn new(
37        gid: GlobalItemIndex,
38        is_program_entrypoint: bool,
39        path: Arc<Path>,
40        visibility: Visibility,
41        signature: Option<Arc<FunctionType>>,
42        is_kernel: bool,
43        source_manager: Arc<dyn SourceManager>,
44    ) -> Self {
45        Self {
46            source_manager,
47            gid,
48            is_program_entrypoint,
49            span: SourceSpan::UNKNOWN,
50            path,
51            visibility,
52            signature,
53            attributes: Default::default(),
54            is_kernel,
55            num_locals: 0,
56        }
57    }
58
59    /// Sets the number of locals to allocate for the procedure.
60    ///
61    /// Returns an error if `num_locals` exceeds `MAX_PROC_LOCALS`, the largest count that
62    /// stays representable in a `u16` once rounded up to a word boundary during frame-pointer
63    /// codegen. The text parser enforces this on `@locals(..)`, but procedures built directly
64    /// via the AST bypass the parser. So the limit is enforced here for all callers.
65    ///
66    /// Call [`Self::with_span`] first so the error can point at the procedure definition.
67    pub fn with_num_locals(mut self, num_locals: u16) -> Result<Self, Report> {
68        if num_locals > MAX_PROC_LOCALS {
69            let source_file = self.source_manager.get(self.span.source_id()).ok();
70            return Err(Report::new(AssemblerError::TooManyProcedureLocals {
71                span: self.span,
72                source_file,
73                max_locals: MAX_PROC_LOCALS,
74                num_locals,
75            }));
76        }
77        self.num_locals = num_locals;
78        Ok(self)
79    }
80
81    pub fn with_span(mut self, span: SourceSpan) -> Self {
82        self.span = span;
83        self
84    }
85
86    /// Sets the attributes attached to this procedure.
87    pub fn with_attributes(mut self, attributes: AttributeSet) -> Self {
88        self.attributes = attributes;
89        self
90    }
91}
92
93// ------------------------------------------------------------------------------------------------
94/// Public accessors
95impl ProcedureContext {
96    pub fn id(&self) -> GlobalItemIndex {
97        self.gid
98    }
99
100    pub fn is_program_entrypoint(&self) -> bool {
101        self.is_program_entrypoint
102    }
103
104    pub fn path(&self) -> &Arc<Path> {
105        &self.path
106    }
107
108    pub fn signature(&self) -> Option<Arc<FunctionType>> {
109        self.signature.clone()
110    }
111
112    pub fn set_signature(&mut self, signature: Option<Arc<FunctionType>>) {
113        self.signature = signature;
114    }
115
116    pub fn num_locals(&self) -> u16 {
117        self.num_locals
118    }
119
120    pub fn module(&self) -> &Path {
121        self.path.parent().unwrap()
122    }
123
124    /// Returns true if the procedure is being assembled for a kernel.
125    pub fn is_kernel(&self) -> bool {
126        self.is_kernel
127    }
128
129    #[inline(always)]
130    pub fn source_manager(&self) -> &dyn SourceManager {
131        self.source_manager.as_ref()
132    }
133}
134
135// ------------------------------------------------------------------------------------------------
136/// State mutators
137impl ProcedureContext {
138    /// Transforms this procedure context into a [Procedure].
139    ///
140    /// The passed-in `mast_root` defines the MAST root of the procedure's body while `body_node`
141    /// specifies the assembly-time reference to the procedure's body node.
142    ///
143    /// <div class="warning">
144    /// `mast_root` and `body_node` must be consistent. That is, `body_node` must resolve to a MAST
145    /// node whose digest equals `mast_root`.
146    /// </div>
147    pub(crate) fn into_procedure(self, mast_root: Word, body_node: MastNodeUse) -> Procedure {
148        let is_syscall = self.is_kernel && self.visibility.is_public();
149        Procedure::new(
150            self.path,
151            self.visibility,
152            self.signature,
153            self.attributes,
154            is_syscall,
155            self.num_locals as u32,
156            mast_root,
157            body_node,
158        )
159        .with_span(self.span)
160    }
161}
162
163impl Spanned for ProcedureContext {
164    fn span(&self) -> SourceSpan {
165        self.span
166    }
167}
168
169// PROCEDURE
170// ================================================================================================
171
172/// A compiled Miden Assembly procedure, consisting of MAST info and basic metadata.
173///
174/// Procedure metadata includes:
175///
176/// - Fully-qualified path of the procedure in Miden Assembly (if known).
177/// - Number of procedure locals to allocate.
178/// - The visibility of the procedure (e.g. public/private/syscall)
179/// - The attributes attached to the procedure.
180/// - The set of MAST roots invoked by this procedure.
181/// - The original source span and file of the procedure (if available).
182#[derive(Clone, Debug)]
183pub struct Procedure {
184    span: SourceSpan,
185    path: Arc<Path>,
186    signature: Option<Arc<FunctionType>>,
187    attributes: AttributeSet,
188    visibility: Visibility,
189    is_syscall: bool,
190    num_locals: u32,
191    /// The MAST root of the procedure.
192    mast_root: Word,
193    /// The assembly-time node reference which resolves to the above MAST root.
194    body_node_ref: MastNodeRef,
195    /// The exact source/debug occurrence for this procedure body.
196    body_source_ref: SourceNodeRef,
197}
198
199// ------------------------------------------------------------------------------------------------
200/// Constructors
201impl Procedure {
202    fn new(
203        path: Arc<Path>,
204        visibility: Visibility,
205        signature: Option<Arc<FunctionType>>,
206        attributes: AttributeSet,
207        is_syscall: bool,
208        num_locals: u32,
209        mast_root: Word,
210        body_node: MastNodeUse,
211    ) -> Self {
212        Self {
213            span: SourceSpan::default(),
214            path,
215            visibility,
216            signature,
217            attributes,
218            is_syscall,
219            num_locals,
220            mast_root,
221            body_node_ref: body_node.node_ref(),
222            body_source_ref: body_node.source_ref(),
223        }
224    }
225
226    pub(crate) fn with_span(mut self, span: SourceSpan) -> Self {
227        self.span = span;
228        self
229    }
230}
231
232// ------------------------------------------------------------------------------------------------
233/// Public accessors
234impl Procedure {
235    /// Returns source span of this procedure.
236    pub fn span(&self) -> &SourceSpan {
237        &self.span
238    }
239
240    /// Returns a reference to the fully-qualified name of this procedure
241    pub fn path(&self) -> &Arc<Path> {
242        &self.path
243    }
244
245    /// Returns true if this procedure is a syscallable procedure
246    #[inline(always)]
247    pub const fn is_syscall(&self) -> bool {
248        self.is_syscall
249    }
250
251    /// Returns the visibility of this procedure as expressed in the original source code
252    pub fn visibility(&self) -> Visibility {
253        self.visibility
254    }
255
256    /// Returns a reference to the fully-qualified module path of this procedure
257    pub fn module(&self) -> &Path {
258        self.path.parent().unwrap()
259    }
260
261    /// Returns a reference to the type signature of this procedure
262    pub fn signature(&self) -> Option<Arc<FunctionType>> {
263        self.signature.clone()
264    }
265
266    /// Returns the attributes attached to this procedure.
267    pub fn attributes(&self) -> &AttributeSet {
268        &self.attributes
269    }
270
271    /// Returns the fully-qualified `@source_name`, if present.
272    ///
273    /// The returned path is formed by joining the procedure's module path and `@source_name`.
274    ///
275    /// # `@source_name` specification
276    ///
277    /// The attribute must contain exactly one quoted string.
278    ///
279    /// `@source_name` allows a producer to preserve a source-level function name while giving the
280    /// emitted Miden Assembly procedure a distinct symbol. This is useful when multiple source
281    /// functions share a name but require unique assembler symbols.
282    ///
283    /// Producers emitting this attribute must respect these requirements:
284    ///
285    /// - Give every emitted procedure whose source-level name is duplicated a distinct,
286    ///   deterministic assembler symbol.
287    /// - Attach `@source_name("original name")` to every such procedure, using the original
288    ///   source-level name.
289    /// - Do not attach `@source_name` to procedures with a unique source-level name or without a
290    ///   source-level name.
291    ///
292    /// Duplicate `@source_name` values are valid. They identify the source-facing name; the
293    /// separately recorded unique linkage name identifies the corresponding assembler procedure.
294    ///
295    /// # Errors
296    ///
297    /// Returns an error if a `@source_name` attribute is present, but is not of the form
298    /// `@source_name("...")`.
299    pub fn source_name_fully_qualified(
300        &self,
301        source_manager: &dyn SourceManager,
302    ) -> Result<Option<PathBuf>, Report> {
303        let Some(attribute) = self.attributes.get("source_name") else {
304            return Ok(None);
305        };
306
307        if let Attribute::List(list) = attribute
308            && let [MetaExpr::String(name)] = list.as_slice()
309        {
310            return Ok(Some(self.path.parent().unwrap().join(name)));
311        }
312
313        let span = attribute.span();
314        Err(Report::new(AssemblerError::InvalidSourceNameAttribute {
315            span,
316            source_file: source_manager.get(span.source_id()).ok(),
317        }))
318    }
319
320    /// Returns the number of memory locals reserved by the procedure.
321    pub fn num_locals(&self) -> u32 {
322        self.num_locals
323    }
324
325    /// Returns the root of this procedure's MAST.
326    pub fn mast_root(&self) -> Word {
327        self.mast_root
328    }
329
330    /// Returns the assembly-time node reference of this procedure.
331    pub(crate) fn body_node_ref(&self) -> MastNodeRef {
332        self.body_node_ref
333    }
334
335    pub(crate) fn body_node_use(&self) -> MastNodeUse {
336        MastNodeUse::new(self.body_node_ref, self.body_source_ref)
337    }
338
339    pub(crate) fn body_source_ref(&self) -> SourceNodeRef {
340        self.body_source_ref
341    }
342}
343
344impl Spanned for Procedure {
345    fn span(&self) -> SourceSpan {
346        self.span
347    }
348}
349
350#[cfg(test)]
351mod tests {
352    use alloc::{sync::Arc, vec};
353
354    use miden_assembly_syntax::{
355        PathBuf,
356        ast::{Attribute, Ident, MetaExpr},
357        debuginfo::{DefaultSourceManager, SourceLanguage, Uri},
358    };
359
360    use super::*;
361
362    /// Constructs a [Procedure] with `attrs` attached, for testing attribute accessors.
363    fn procedure_with_attributes(attrs: vec::IntoIter<Attribute>) -> Procedure {
364        Procedure::new(
365            Arc::from(PathBuf::new("::test::module::foo").unwrap()),
366            Visibility::Private,
367            None,
368            AttributeSet::new(attrs),
369            false,
370            0,
371            Word::default(),
372            MastNodeUse::new(MastNodeRef::from(0), SourceNodeRef::from(0)),
373        )
374    }
375
376    #[test]
377    fn source_name_fully_qualified_is_none_without_attribute() {
378        let source_manager = DefaultSourceManager::default();
379        let procedure = procedure_with_attributes(vec![].into_iter());
380
381        assert_eq!(procedure.source_name_fully_qualified(&source_manager).unwrap(), None);
382    }
383
384    #[test]
385    fn source_name_fully_qualified_returns_quoted_string_joined_to_module_path() {
386        let source_manager = DefaultSourceManager::default();
387        let attribute = Attribute::from_iter(
388            Ident::new("source_name").unwrap(),
389            [MetaExpr::String(Ident::new("bar").unwrap())],
390        );
391        let procedure = procedure_with_attributes(vec![attribute].into_iter());
392
393        assert_eq!(
394            procedure.source_name_fully_qualified(&source_manager).unwrap(),
395            Some(PathBuf::new("::test::module::bar").unwrap()),
396        );
397    }
398
399    #[test]
400    fn malformed_source_name_attributes_are_rejected() {
401        let source_manager = DefaultSourceManager::default();
402        let file = source_manager.load(
403            SourceLanguage::Masm,
404            Uri::new("test.masm"),
405            "@source_name(unquoted)".into(),
406        );
407        let span = SourceSpan::new(file.id(), 0..19);
408
409        // Generate some malformed `@source_name` attributes
410        let malformed = vec![
411            Attribute::Marker(Ident::new("source_name").unwrap()),
412            // `@source_name(unquoted)`
413            Attribute::from_iter(
414                Ident::new("source_name").unwrap(),
415                [MetaExpr::Ident(Ident::new("unquoted").unwrap())],
416            ),
417            // `@source_name("one", "two")`
418            Attribute::from_iter(
419                Ident::new("source_name").unwrap(),
420                [
421                    MetaExpr::String(Ident::new("one").unwrap()),
422                    MetaExpr::String(Ident::new("two").unwrap()),
423                ],
424            ),
425            // `@source_name(value = "named")`
426            Attribute::from_iter(
427                Ident::new("source_name").unwrap(),
428                [(Ident::new("value").unwrap(), MetaExpr::String(Ident::new("named").unwrap()))],
429            ),
430        ];
431
432        for attribute in malformed {
433            let procedure = procedure_with_attributes(vec![attribute.with_span(span)].into_iter());
434            let error = procedure.source_name_fully_qualified(&source_manager).unwrap_err();
435
436            match error.downcast_ref::<AssemblerError>() {
437                Some(AssemblerError::InvalidSourceNameAttribute { source_file, .. }) => {
438                    // The error must be attributed to the file in which the attribute occurred
439                    assert_eq!(source_file.as_ref(), Some(&file));
440                },
441                unexpected => panic!("expected InvalidSourceNameAttribute, got {unexpected:?}"),
442            }
443        }
444    }
445}