miden_assembly/ast/procedure/
procedure.rs

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

miden_assembly/ast/procedure/
procedure.rs
