miden_assembly_syntax/ast/procedure/
procedure.rs1use 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#[derive(Default, Debug, Copy, Clone, PartialEq, Eq)]
14#[repr(u8)]
15pub enum Visibility {
16 Public = 0,
18 Syscall = 1,
20 #[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 pub fn is_exported(&self) -> bool {
38 matches!(self, Self::Public | Self::Syscall)
39 }
40
41 pub fn is_syscall(&self) -> bool {
43 matches!(self, Self::Syscall)
44 }
45}
46
47#[derive(Clone)]
52pub struct Procedure {
53 span: SourceSpan,
55 docs: Option<DocString>,
57 attrs: AttributeSet,
59 name: ProcedureName,
61 visibility: Visibility,
63 num_locals: u16,
65 body: Block,
67 pub(super) invoked: BTreeSet<Invoke>,
69}
70
71impl Procedure {
73 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 pub fn with_docs(mut self, docs: Option<Span<String>>) -> Self {
96 self.docs = docs.map(DocString::new);
97 self
98 }
99
100 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 pub(crate) fn set_visibility(&mut self, visibility: Visibility) {
116 self.visibility = visibility;
117 }
118}
119
120impl Procedure {
122 pub fn name(&self) -> &ProcedureName {
124 &self.name
125 }
126
127 pub fn visibility(&self) -> Visibility {
129 self.visibility
130 }
131
132 pub fn num_locals(&self) -> u16 {
134 self.num_locals
135 }
136
137 pub fn is_entrypoint(&self) -> bool {
140 self.name.is_main()
141 }
142
143 pub fn docs(&self) -> Option<Span<&str>> {
145 self.docs.as_ref().map(|docstring| docstring.as_spanned_str())
146 }
147
148 #[inline]
150 pub fn attributes(&self) -> &AttributeSet {
151 &self.attrs
152 }
153
154 #[inline]
156 pub fn attributes_mut(&mut self) -> &mut AttributeSet {
157 &mut self.attrs
158 }
159
160 #[inline]
162 pub fn has_attribute(&self, name: impl AsRef<str>) -> bool {
163 self.attrs.has(name)
164 }
165
166 #[inline]
168 pub fn get_attribute(&self, name: impl AsRef<str>) -> Option<&Attribute> {
169 self.attrs.get(name)
170 }
171
172 pub fn body(&self) -> &Block {
174 &self.body
175 }
176
177 pub fn body_mut(&mut self) -> &mut Block {
179 &mut self.body
180 }
181
182 pub fn iter(&self) -> core::slice::Iter<'_, crate::ast::Op> {
184 self.body.iter()
185 }
186
187 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 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}