miden_assembly_syntax/ast/procedure/
procedure.rs

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

miden_assembly_syntax/ast/procedure/
procedure.rs
