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, 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("export")
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 number of locals to allocate for this procedure
64    num_locals: u16,
65    /// The body of the procedure
66    body: Block,
67    /// The set of callees for any call-like instruction in the procedure body.
68    pub(super) invoked: BTreeSet<Invoke>,
69}
70
71/// Construction
72impl Procedure {
73    /// Creates a new [Procedure] from the given source span, visibility, name, number of locals,
74    /// and code block.
75    pub fn new(
76        span: SourceSpan,
77        visibility: Visibility,
78        name: ProcedureName,
79        num_locals: u16,
80        body: Block,
81    ) -> Self {
82        Self {
83            span,
84            docs: None,
85            attrs: Default::default(),
86            name,
87            visibility,
88            num_locals,
89            invoked: Default::default(),
90            body,
91        }
92    }
93
94    /// Adds documentation to this procedure definition
95    pub fn with_docs(mut self, docs: Option<Span<String>>) -> Self {
96        self.docs = docs.map(DocString::new);
97        self
98    }
99
100    /// Adds attributes to this procedure definition
101    pub fn with_attributes<I>(mut self, attrs: I) -> Self
102    where
103        I: IntoIterator<Item = Attribute>,
104    {
105        self.attrs.extend(attrs);
106        self
107    }
108
109    /// Modifies the visibility of this procedure.
110    ///
111    /// This is made crate-local as the visibility of a procedure is virtually always determined
112    /// by the source code from which it was derived; the only exception being kernel modules,
113    /// where exported procedures take on syscall visibility once the module is identified as
114    /// a kernel.
115    pub(crate) fn set_visibility(&mut self, visibility: Visibility) {
116        self.visibility = visibility;
117    }
118}
119
120/// Metadata
121impl Procedure {
122    /// Returns the name of this procedure within its containing module.
123    pub fn name(&self) -> &ProcedureName {
124        &self.name
125    }
126
127    /// Returns the visibility of this procedure
128    pub fn visibility(&self) -> Visibility {
129        self.visibility
130    }
131
132    /// Returns the number of locals allocated by this procedure.
133    pub fn num_locals(&self) -> u16 {
134        self.num_locals
135    }
136
137    /// Returns true if this procedure corresponds to the `begin`..`end` block of an executable
138    /// module.
139    pub fn is_entrypoint(&self) -> bool {
140        self.name.is_main()
141    }
142
143    /// Returns the documentation for this procedure, if present.
144    pub fn docs(&self) -> Option<Span<&str>> {
145        self.docs.as_ref().map(|docstring| docstring.as_spanned_str())
146    }
147
148    /// Get the attributes attached to this procedure
149    #[inline]
150    pub fn attributes(&self) -> &AttributeSet {
151        &self.attrs
152    }
153
154    /// Get the attributes attached to this procedure, mutably
155    #[inline]
156    pub fn attributes_mut(&mut self) -> &mut AttributeSet {
157        &mut self.attrs
158    }
159
160    /// Returns true if this procedure has an attribute named `name`
161    #[inline]
162    pub fn has_attribute(&self, name: impl AsRef<str>) -> bool {
163        self.attrs.has(name)
164    }
165
166    /// Returns the attribute named `name`, if present
167    #[inline]
168    pub fn get_attribute(&self, name: impl AsRef<str>) -> Option<&Attribute> {
169        self.attrs.get(name)
170    }
171
172    /// Returns a reference to the [Block] containing the body of this procedure.
173    pub fn body(&self) -> &Block {
174        &self.body
175    }
176
177    /// Returns a mutable reference to the [Block] containing the body of this procedure.
178    pub fn body_mut(&mut self) -> &mut Block {
179        &mut self.body
180    }
181
182    /// Returns an iterator over the operations of the top-level [Block] of this procedure.
183    pub fn iter(&self) -> core::slice::Iter<'_, crate::ast::Op> {
184        self.body.iter()
185    }
186
187    /// Returns an iterator over the set of invocation targets of this procedure, i.e. the callees
188    /// of any call instructions in the body of this procedure.
189    pub fn invoked<'a, 'b: 'a>(&'b self) -> impl Iterator<Item = &'a Invoke> + 'a {
190        if self.invoked.is_empty() {
191            InvokedIter::Empty
192        } else {
193            InvokedIter::NonEmpty(self.invoked.iter())
194        }
195    }
196
197    /// Extends the set of procedures known to be invoked by this procedure.
198    ///
199    /// This is for internal use only, and is called during semantic analysis once we've identified
200    /// the set of invoked procedures for a given definition.
201    pub fn extend_invoked<I>(&mut self, iter: I)
202    where
203        I: IntoIterator<Item = Invoke>,
204    {
205        self.invoked.extend(iter);
206    }
207}
208
209#[doc(hidden)]
210pub(crate) enum InvokedIter<'a, I: Iterator<Item = &'a Invoke> + 'a> {
211    Empty,
212    NonEmpty(I),
213}
214
215impl<'a, I> Iterator for InvokedIter<'a, I>
216where
217    I: Iterator<Item = &'a Invoke> + 'a,
218{
219    type Item = <I as Iterator>::Item;
220
221    fn next(&mut self) -> Option<Self::Item> {
222        match self {
223            Self::Empty => None,
224            Self::NonEmpty(iter) => {
225                let result = iter.next();
226                if result.is_none() {
227                    *self = Self::Empty;
228                }
229                result
230            },
231        }
232    }
233}
234
235impl Spanned for Procedure {
236    fn span(&self) -> SourceSpan {
237        self.span
238    }
239}
240
241impl crate::prettier::PrettyPrint for Procedure {
242    fn render(&self) -> crate::prettier::Document {
243        use crate::prettier::*;
244
245        let mut doc = self
246            .docs
247            .as_ref()
248            .map(|docstring| docstring.render())
249            .unwrap_or(Document::Empty);
250
251        if !self.attrs.is_empty() {
252            doc += self
253                .attrs
254                .iter()
255                .map(|attr| attr.render())
256                .reduce(|acc, attr| acc + nl() + attr)
257                .unwrap_or(Document::Empty);
258        }
259
260        if self.is_entrypoint() {
261            doc += const_text("begin");
262        } else {
263            doc += display(self.visibility) + const_text(".") + display(&self.name);
264            if self.num_locals > 0 {
265                doc += const_text(".") + display(self.num_locals);
266            }
267        }
268
269        doc + self.body.render() + const_text("end") + nl()
270    }
271}
272
273impl fmt::Debug for Procedure {
274    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
275        f.debug_struct("Procedure")
276            .field("docs", &self.docs)
277            .field("attrs", &self.attrs)
278            .field("name", &self.name)
279            .field("visibility", &self.visibility)
280            .field("num_locals", &self.num_locals)
281            .field("body", &self.body)
282            .field("invoked", &self.invoked)
283            .finish()
284    }
285}
286
287impl Eq for Procedure {}
288
289impl PartialEq for Procedure {
290    fn eq(&self, other: &Self) -> bool {
291        self.name == other.name
292            && self.visibility == other.visibility
293            && self.num_locals == other.num_locals
294            && self.body == other.body
295            && self.attrs == other.attrs
296            && self.docs == other.docs
297    }
298}