miden_assembly_syntax/ast/procedure/
procedure.rs

1use alloc::{collections::BTreeSet, string::String};
2use core::fmt;
3
4use miden_debug_types::{SourceSpan, Span, Spanned};
5
6use super::ProcedureName;
7use crate::ast::{Attribute, AttributeSet, Block, DocString, FunctionType, Invoke};
8
9// PROCEDURE VISIBILITY
10// ================================================================================================
11
12/// Represents the visibility of a procedure globally.
13#[derive(Default, Debug, Copy, Clone, PartialEq, Eq)]
14#[repr(u8)]
15pub enum Visibility {
16    /// The procedure is visible for call/exec.
17    Public = 0,
18    /// The procedure is visible to syscalls only.
19    Syscall = 1,
20    /// The procedure is visible only locally to the exec instruction.
21    #[default]
22    Private = 2,
23}
24
25impl fmt::Display for Visibility {
26    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
27        if self.is_exported() {
28            f.write_str("pub proc")
29        } else {
30            f.write_str("proc")
31        }
32    }
33}
34
35impl Visibility {
36    /// Returns true if the procedure has been explicitly exported
37    pub fn is_exported(&self) -> bool {
38        matches!(self, Self::Public | Self::Syscall)
39    }
40
41    /// Returns true if the procedure is a syscall export
42    pub fn is_syscall(&self) -> bool {
43        matches!(self, Self::Syscall)
44    }
45}
46
47// PROCEDURE
48// ================================================================================================
49
50/// Represents a concrete procedure definition in Miden Assembly syntax
51#[derive(Clone)]
52pub struct Procedure {
53    /// The source span of the full procedure body
54    span: SourceSpan,
55    /// The documentation attached to this procedure
56    docs: Option<DocString>,
57    /// The attributes attached to this procedure
58    attrs: AttributeSet,
59    /// The local name of this procedure
60    name: ProcedureName,
61    /// The visibility of this procedure (i.e. whether it is exported or not)
62    visibility: Visibility,
63    /// The type signature of this procedure, if known
64    ty: Option<FunctionType>,
65    /// The number of locals to allocate for this procedure
66    num_locals: u16,
67    /// The body of the procedure
68    body: Block,
69    /// The set of callees for any call-like instruction in the procedure body.
70    pub(super) invoked: BTreeSet<Invoke>,
71}
72
73/// Construction
74impl Procedure {
75    /// Creates a new [Procedure] from the given source span, visibility, name, number of locals,
76    /// and code block.
77    pub fn new(
78        span: SourceSpan,
79        visibility: Visibility,
80        name: ProcedureName,
81        num_locals: u16,
82        body: Block,
83    ) -> Self {
84        Self {
85            span,
86            docs: None,
87            attrs: Default::default(),
88            name,
89            visibility,
90            ty: None,
91            num_locals,
92            invoked: Default::default(),
93            body,
94        }
95    }
96
97    /// Specify the type signature of this procedure
98    pub fn with_signature(mut self, ty: FunctionType) -> Self {
99        self.ty = Some(ty);
100        self
101    }
102
103    /// Adds documentation to this procedure definition
104    pub fn with_docs(mut self, docs: Option<Span<String>>) -> Self {
105        self.docs = docs.map(DocString::new);
106        self
107    }
108
109    /// Adds attributes to this procedure definition
110    pub fn with_attributes<I>(mut self, attrs: I) -> Self
111    where
112        I: IntoIterator<Item = Attribute>,
113    {
114        self.attrs.extend(attrs);
115        self
116    }
117
118    /// Override the visibility of this procedure.
119    pub fn set_visibility(&mut self, visibility: Visibility) {
120        self.visibility = visibility;
121    }
122
123    /// Override the type signature of this procedure.
124    pub fn set_signature(&mut self, signature: FunctionType) {
125        self.ty = Some(signature);
126    }
127}
128
129/// Metadata
130impl Procedure {
131    /// Returns the name of this procedure within its containing module.
132    pub fn name(&self) -> &ProcedureName {
133        &self.name
134    }
135
136    /// Returns the visibility of this procedure
137    pub fn visibility(&self) -> Visibility {
138        self.visibility
139    }
140
141    /// Get the type signature of this procedure, if known
142    pub fn signature(&self) -> Option<&FunctionType> {
143        self.ty.as_ref()
144    }
145
146    /// Get the type signature of this procedure mutably, if known
147    pub fn signature_mut(&mut self) -> Option<&mut FunctionType> {
148        self.ty.as_mut()
149    }
150
151    /// Returns the number of locals allocated by this procedure.
152    pub fn num_locals(&self) -> u16 {
153        self.num_locals
154    }
155
156    /// Returns true if this procedure corresponds to the `begin`..`end` block of an executable
157    /// module.
158    pub fn is_entrypoint(&self) -> bool {
159        self.name.is_main()
160    }
161
162    /// Returns the documentation for this procedure, if present.
163    pub fn docs(&self) -> Option<Span<&str>> {
164        self.docs.as_ref().map(|docstring| docstring.as_spanned_str())
165    }
166
167    /// Get the attributes attached to this procedure
168    #[inline]
169    pub fn attributes(&self) -> &AttributeSet {
170        &self.attrs
171    }
172
173    /// Get the attributes attached to this procedure, mutably
174    #[inline]
175    pub fn attributes_mut(&mut self) -> &mut AttributeSet {
176        &mut self.attrs
177    }
178
179    /// Returns true if this procedure has an attribute named `name`
180    #[inline]
181    pub fn has_attribute(&self, name: impl AsRef<str>) -> bool {
182        self.attrs.has(name)
183    }
184
185    /// Returns the attribute named `name`, if present
186    #[inline]
187    pub fn get_attribute(&self, name: impl AsRef<str>) -> Option<&Attribute> {
188        self.attrs.get(name)
189    }
190
191    /// Returns a reference to the [Block] containing the body of this procedure.
192    pub fn body(&self) -> &Block {
193        &self.body
194    }
195
196    /// Returns a mutable reference to the [Block] containing the body of this procedure.
197    pub fn body_mut(&mut self) -> &mut Block {
198        &mut self.body
199    }
200
201    /// Returns an iterator over the operations of the top-level [Block] of this procedure.
202    pub fn iter(&self) -> core::slice::Iter<'_, crate::ast::Op> {
203        self.body.iter()
204    }
205
206    /// Returns an iterator over the set of invocation targets of this procedure, i.e. the callees
207    /// of any call instructions in the body of this procedure.
208    pub fn invoked<'a, 'b: 'a>(&'b self) -> impl Iterator<Item = &'a Invoke> + 'a {
209        if self.invoked.is_empty() {
210            InvokedIter::Empty
211        } else {
212            InvokedIter::NonEmpty(self.invoked.iter())
213        }
214    }
215
216    /// Extends the set of procedures known to be invoked by this procedure.
217    ///
218    /// This is for internal use only, and is called during semantic analysis once we've identified
219    /// the set of invoked procedures for a given definition.
220    pub fn extend_invoked<I>(&mut self, iter: I)
221    where
222        I: IntoIterator<Item = Invoke>,
223    {
224        self.invoked.extend(iter);
225    }
226}
227
228#[doc(hidden)]
229pub(crate) enum InvokedIter<'a, I: Iterator<Item = &'a Invoke> + 'a> {
230    Empty,
231    NonEmpty(I),
232}
233
234impl<'a, I> Iterator for InvokedIter<'a, I>
235where
236    I: Iterator<Item = &'a Invoke> + 'a,
237{
238    type Item = <I as Iterator>::Item;
239
240    fn next(&mut self) -> Option<Self::Item> {
241        match self {
242            Self::Empty => None,
243            Self::NonEmpty(iter) => {
244                let result = iter.next();
245                if result.is_none() {
246                    *self = Self::Empty;
247                }
248                result
249            },
250        }
251    }
252}
253
254impl Spanned for Procedure {
255    fn span(&self) -> SourceSpan {
256        self.span
257    }
258}
259
260impl crate::prettier::PrettyPrint for Procedure {
261    fn render(&self) -> crate::prettier::Document {
262        use crate::prettier::*;
263
264        let mut doc = self
265            .docs
266            .as_ref()
267            .map(|docstring| docstring.render())
268            .unwrap_or(Document::Empty);
269
270        if !self.attrs.is_empty() {
271            doc += self
272                .attrs
273                .iter()
274                .map(|attr| attr.render())
275                .reduce(|acc, attr| acc + nl() + attr)
276                .unwrap_or(Document::Empty);
277        }
278
279        if self.is_entrypoint() {
280            doc += const_text("begin");
281        } else {
282            match self.signature() {
283                Some(sig) if sig.cc != crate::ast::types::CallConv::Fast => {
284                    doc += text(format!("@callconv(\"{}\")", &sig.cc)) + nl();
285                },
286                _ => (),
287            }
288            doc += display(self.visibility) + const_text(" ") + display(&self.name);
289            if self.num_locals > 0 {
290                doc += const_text(".") + display(self.num_locals);
291            }
292            if let Some(sig) = self.signature() {
293                doc += sig.render();
294            }
295        }
296
297        doc + self.body.render() + const_text("end") + nl()
298    }
299}
300
301impl fmt::Debug for Procedure {
302    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
303        f.debug_struct("Procedure")
304            .field("docs", &self.docs)
305            .field("attrs", &self.attrs)
306            .field("name", &self.name)
307            .field("visibility", &self.visibility)
308            .field("num_locals", &self.num_locals)
309            .field("ty", &self.ty)
310            .field("body", &self.body)
311            .field("invoked", &self.invoked)
312            .finish()
313    }
314}
315
316impl Eq for Procedure {}
317
318impl PartialEq for Procedure {
319    fn eq(&self, other: &Self) -> bool {
320        self.name == other.name
321            && self.visibility == other.visibility
322            && self.num_locals == other.num_locals
323            && self.ty == other.ty
324            && self.body == other.body
325            && self.attrs == other.attrs
326            && self.docs == other.docs
327    }
328}
miden_assembly_syntax/ast/procedure/procedure.rs

miden_assembly_syntax/ast/procedure/
procedure.rs
