Skip to main content

miden_assembly_syntax/ast/
type.rs

1use alloc::{boxed::Box, string::String, sync::Arc, vec::Vec};
2
3use miden_debug_types::{SourceManager, SourceSpan, Span, Spanned};
4use midenc_hir_type::{AddressSpace, Type, TypeRepr};
5
6use super::{
7    ConstantExpr, DocString, GlobalItemIndex, Ident, ItemIndex, Path, SymbolResolution,
8    SymbolResolutionError, Visibility, types,
9};
10
11/// Maximum allowed nesting depth of type expressions during resolution.
12///
13/// This limit is intended to prevent stack overflows from maliciously deep type expressions while
14/// remaining far above typical type nesting in real programs.
15const MAX_TYPE_EXPR_NESTING: usize = 256;
16
17/// Abstracts over resolving an item to a concrete [Type], using one of:
18///
19/// * A [GlobalItemIndex]
20/// * An [ItemIndex]
21/// * A [Path]
22/// * A [TypeExpr]
23///
24/// Since type resolution happens in two different contexts during assembly, this abstraction allows
25/// us to share more of the resolution logic in both places.
26///
27/// NOTE: Most methods of this trait take a mutable reference to the resolver, so that the resolver
28/// can mutate its own state as necessary during resolution (e.g. to manage a cache, or other side
29/// table-like data structures).
30pub trait TypeResolver<E> {
31    fn source_manager(&self) -> Arc<dyn SourceManager>;
32    /// Should be called by consumers of this resolver to convert a [SymbolResolutionError] to the
33    /// error type used by the [TypeResolver] implementation.
34    fn resolve_local_failed(&self, err: SymbolResolutionError) -> E;
35    /// Get the [Type] corresponding to the item given by `gid`
36    fn get_type(&mut self, context: SourceSpan, gid: GlobalItemIndex) -> Result<Type, E>;
37    /// Get the [Type] corresponding to the item in the current module given by `id`
38    fn get_local_type(&mut self, context: SourceSpan, id: ItemIndex) -> Result<Option<Type>, E>;
39    /// Attempt to resolve a symbol path, given by a `TypeExpr::Ref`, to an item
40    fn resolve_type_ref(&mut self, ty: Span<&Path>) -> Result<SymbolResolution, E>;
41    /// Resolve a [TypeExpr] to a concrete [Type]
42    fn resolve(&mut self, ty: &TypeExpr) -> Result<Option<Type>, E> {
43        ty.resolve_type(self)
44    }
45}
46
47// TYPE DECLARATION
48// ================================================================================================
49
50/// An abstraction over the different types of type declarations allowed in Miden Assembly
51#[derive(Debug, Clone, PartialEq, Eq)]
52pub enum TypeDecl {
53    /// A named type, i.e. a type alias
54    Alias(TypeAlias),
55    /// A C-like enumeration type with associated constants
56    Enum(EnumType),
57}
58
59impl TypeDecl {
60    /// Adds documentation to this type alias
61    pub fn with_docs(self, docs: Option<Span<String>>) -> Self {
62        match self {
63            Self::Alias(ty) => Self::Alias(ty.with_docs(docs)),
64            Self::Enum(ty) => Self::Enum(ty.with_docs(docs)),
65        }
66    }
67
68    /// Get the name assigned to this type declaration
69    pub fn name(&self) -> &Ident {
70        match self {
71            Self::Alias(ty) => &ty.name,
72            Self::Enum(ty) => &ty.name,
73        }
74    }
75
76    /// Get the visibility of this type declaration
77    pub const fn visibility(&self) -> Visibility {
78        match self {
79            Self::Alias(ty) => ty.visibility,
80            Self::Enum(ty) => ty.visibility,
81        }
82    }
83
84    /// Get the documentation of this enum type
85    pub fn docs(&self) -> Option<Span<&str>> {
86        match self {
87            Self::Alias(ty) => ty.docs(),
88            Self::Enum(ty) => ty.docs(),
89        }
90    }
91
92    /// Get the type expression associated with this declaration
93    pub fn ty(&self) -> TypeExpr {
94        match self {
95            Self::Alias(ty) => ty.ty.clone(),
96            Self::Enum(ty) => TypeExpr::Primitive(Span::new(ty.span, ty.ty.clone())),
97        }
98    }
99}
100
101impl Spanned for TypeDecl {
102    fn span(&self) -> SourceSpan {
103        match self {
104            Self::Alias(spanned) => spanned.span,
105            Self::Enum(spanned) => spanned.span,
106        }
107    }
108}
109
110impl From<TypeAlias> for TypeDecl {
111    fn from(value: TypeAlias) -> Self {
112        Self::Alias(value)
113    }
114}
115
116impl From<EnumType> for TypeDecl {
117    fn from(value: EnumType) -> Self {
118        Self::Enum(value)
119    }
120}
121
122impl crate::prettier::PrettyPrint for TypeDecl {
123    fn render(&self) -> crate::prettier::Document {
124        match self {
125            Self::Alias(ty) => ty.render(),
126            Self::Enum(ty) => ty.render(),
127        }
128    }
129}
130
131// FUNCTION TYPE
132// ================================================================================================
133
134/// A procedure type signature
135#[derive(Debug, Clone)]
136pub struct FunctionType {
137    pub span: SourceSpan,
138    pub cc: types::CallConv,
139    pub args: Vec<TypeExpr>,
140    pub results: Vec<TypeExpr>,
141}
142
143impl Eq for FunctionType {}
144
145impl PartialEq for FunctionType {
146    fn eq(&self, other: &Self) -> bool {
147        self.cc == other.cc && self.args == other.args && self.results == other.results
148    }
149}
150
151impl core::hash::Hash for FunctionType {
152    fn hash<H: core::hash::Hasher>(&self, state: &mut H) {
153        self.cc.hash(state);
154        self.args.hash(state);
155        self.results.hash(state);
156    }
157}
158
159impl Spanned for FunctionType {
160    fn span(&self) -> SourceSpan {
161        self.span
162    }
163}
164
165impl FunctionType {
166    pub fn new(cc: types::CallConv, args: Vec<TypeExpr>, results: Vec<TypeExpr>) -> Self {
167        Self {
168            span: SourceSpan::UNKNOWN,
169            cc,
170            args,
171            results,
172        }
173    }
174
175    /// Override the default source span
176    #[inline]
177    pub fn with_span(mut self, span: SourceSpan) -> Self {
178        self.span = span;
179        self
180    }
181}
182
183impl crate::prettier::PrettyPrint for FunctionType {
184    fn render(&self) -> crate::prettier::Document {
185        use crate::prettier::*;
186
187        let singleline_args = self
188            .args
189            .iter()
190            .map(PrettyPrint::render)
191            .reduce(|acc, arg| acc + const_text(", ") + arg)
192            .unwrap_or(Document::Empty);
193        let multiline_args = indent(
194            4,
195            nl() + self
196                .args
197                .iter()
198                .map(PrettyPrint::render)
199                .reduce(|acc, arg| acc + const_text(",") + nl() + arg)
200                .unwrap_or(Document::Empty),
201        ) + nl();
202        let args = singleline_args | multiline_args;
203        let args = const_text("(") + args + const_text(")");
204
205        match self.results.len() {
206            0 => args,
207            1 => args + const_text(" -> ") + self.results[0].render(),
208            _ => {
209                let results = self
210                    .results
211                    .iter()
212                    .map(PrettyPrint::render)
213                    .reduce(|acc, r| acc + const_text(", ") + r)
214                    .unwrap_or(Document::Empty);
215                args + const_text(" -> ") + const_text("(") + results + const_text(")")
216            },
217        }
218    }
219}
220
221// TYPE EXPRESSION
222// ================================================================================================
223
224/// A syntax-level type expression (i.e. primitive type, reference to nominal type, etc.)
225#[derive(Debug, Clone, Eq, PartialEq, Hash)]
226pub enum TypeExpr {
227    /// A primitive integral type, e.g. `i1`, `u16`
228    Primitive(Span<Type>),
229    /// A pointer type expression, e.g. `*u8`
230    Ptr(PointerType),
231    /// An array type expression, e.g. `[u8; 32]`
232    Array(ArrayType),
233    /// A struct type expression, e.g. `struct { a: u32 }`
234    Struct(StructType),
235    /// A reference to a type aliased by name, e.g. `Foo`
236    Ref(Span<Arc<Path>>),
237}
238
239impl TypeExpr {
240    /// Set the name associated with this type expression, if applicable.
241    ///
242    /// Currently this just sets the name of struct types, but if we add other types with names in
243    /// the future, we can support them here.
244    pub fn set_name(&mut self, name: Ident) {
245        match self {
246            Self::Struct(struct_ty) => {
247                struct_ty.name = Some(name);
248            },
249            Self::Primitive(_) | Self::Ptr(_) | Self::Array(_) | Self::Ref(_) => (),
250        }
251    }
252
253    /// Get any references to other types present in this expression
254    pub fn references(&self) -> Vec<Span<Arc<Path>>> {
255        use alloc::collections::BTreeSet;
256
257        let mut worklist = smallvec::SmallVec::<[_; 4]>::from_slice(&[self]);
258        let mut references = BTreeSet::new();
259
260        while let Some(ty) = worklist.pop() {
261            match ty {
262                Self::Primitive(_) => {},
263                Self::Ptr(ty) => {
264                    worklist.push(&ty.pointee);
265                },
266                Self::Array(ty) => {
267                    worklist.push(&ty.elem);
268                },
269                Self::Struct(ty) => {
270                    for field in ty.fields.iter() {
271                        worklist.push(&field.ty);
272                    }
273                },
274                Self::Ref(ty) => {
275                    references.insert(ty.clone());
276                },
277            }
278        }
279
280        references.into_iter().collect()
281    }
282
283    /// Resolve this type expression to a concrete type, using `resolver`
284    pub fn resolve_type<E, R>(&self, resolver: &mut R) -> Result<Option<Type>, E>
285    where
286        R: ?Sized + TypeResolver<E>,
287    {
288        self.resolve_type_with_depth(resolver, 0)
289    }
290
291    fn resolve_type_with_depth<E, R>(
292        &self,
293        resolver: &mut R,
294        depth: usize,
295    ) -> Result<Option<Type>, E>
296    where
297        R: ?Sized + TypeResolver<E>,
298    {
299        if depth > MAX_TYPE_EXPR_NESTING {
300            let source_manager = resolver.source_manager();
301            return Err(resolver.resolve_local_failed(
302                SymbolResolutionError::type_expression_depth_exceeded(
303                    self.span(),
304                    MAX_TYPE_EXPR_NESTING,
305                    source_manager.as_ref(),
306                ),
307            ));
308        }
309
310        match self {
311            TypeExpr::Ref(path) => {
312                let mut current_path = path.clone();
313                loop {
314                    match resolver.resolve_type_ref(current_path.as_deref())? {
315                        SymbolResolution::Local(item) => {
316                            return resolver.get_local_type(current_path.span(), item.into_inner());
317                        },
318                        SymbolResolution::External(path) => {
319                            // We don't have a definition for this type yet
320                            if path == current_path {
321                                break Ok(None);
322                            }
323                            current_path = path;
324                        },
325                        SymbolResolution::Exact { gid, .. } => {
326                            return resolver.get_type(current_path.span(), gid).map(Some);
327                        },
328                        SymbolResolution::Module { path: module_path, .. } => {
329                            break Err(resolver.resolve_local_failed(
330                                SymbolResolutionError::invalid_symbol_type(
331                                    path.span(),
332                                    "type",
333                                    module_path.span(),
334                                    &resolver.source_manager(),
335                                ),
336                            ));
337                        },
338                        SymbolResolution::MastRoot(item) => {
339                            break Err(resolver.resolve_local_failed(
340                                SymbolResolutionError::invalid_symbol_type(
341                                    path.span(),
342                                    "type",
343                                    item.span(),
344                                    &resolver.source_manager(),
345                                ),
346                            ));
347                        },
348                    }
349                }
350            },
351            TypeExpr::Primitive(t) => Ok(Some(t.inner().clone())),
352            TypeExpr::Array(t) => Ok(t
353                .elem
354                .resolve_type_with_depth(resolver, depth + 1)?
355                .map(|elem| Type::Array(Arc::new(types::ArrayType::new(elem, t.arity))))),
356            TypeExpr::Ptr(ty) => Ok(ty
357                .pointee
358                .resolve_type_with_depth(resolver, depth + 1)?
359                .map(|pointee| Type::Ptr(Arc::new(types::PointerType::new(pointee))))),
360            TypeExpr::Struct(t) => {
361                let mut fields = Vec::with_capacity(t.fields.len());
362                for field in t.fields.iter() {
363                    let field_ty = field.ty.resolve_type_with_depth(resolver, depth + 1)?;
364                    if let Some(field_ty) = field_ty {
365                        fields.push((field.name.clone().into_inner(), field_ty));
366                    } else {
367                        return Ok(None);
368                    }
369                }
370                Ok(Some(Type::Struct(Arc::new(types::StructType::from_parts(
371                    t.name.clone().map(Ident::into_inner),
372                    t.repr.into_inner(),
373                    fields,
374                )))))
375            },
376        }
377    }
378}
379
380impl From<Type> for TypeExpr {
381    fn from(ty: Type) -> Self {
382        match ty {
383            Type::Array(t) => Self::Array(ArrayType::new(t.element_type().clone().into(), t.len())),
384            Type::Struct(t) => {
385                let name = t.name().and_then(|name| Ident::new(name.as_ref()).ok());
386                let fields = t.fields().iter().enumerate().map(|(i, ft)| {
387                    let name = ft
388                        .name
389                        .as_deref()
390                        .map(Ident::new)
391                        .and_then(Result::ok)
392                        .unwrap_or_else(|| Ident::new(format!("field{i}")).unwrap());
393                    StructField {
394                        span: SourceSpan::UNKNOWN,
395                        name,
396                        ty: ft.ty.clone().into(),
397                    }
398                });
399                Self::Struct(
400                    StructType::new(name, fields)
401                        .with_repr(Span::unknown(t.repr()))
402                        .with_span(SourceSpan::UNKNOWN),
403                )
404            },
405            Type::Ptr(t) => Self::Ptr((*t).clone().into()),
406            Type::Function(_) => {
407                Self::Ptr(PointerType::new(TypeExpr::Primitive(Span::unknown(Type::Felt))))
408            },
409            Type::List(t) => Self::Ptr(
410                PointerType::new((*t).clone().into()).with_address_space(AddressSpace::Byte),
411            ),
412            Type::Unknown | Type::Never | Type::F64 => panic!("unrepresentable type value: {ty}"),
413            ty => Self::Primitive(Span::unknown(ty)),
414        }
415    }
416}
417
418impl Spanned for TypeExpr {
419    fn span(&self) -> SourceSpan {
420        match self {
421            Self::Primitive(spanned) => spanned.span(),
422            Self::Ptr(spanned) => spanned.span(),
423            Self::Array(spanned) => spanned.span(),
424            Self::Struct(spanned) => spanned.span(),
425            Self::Ref(spanned) => spanned.span(),
426        }
427    }
428}
429
430impl crate::prettier::PrettyPrint for TypeExpr {
431    fn render(&self) -> crate::prettier::Document {
432        use crate::prettier::*;
433
434        match self {
435            Self::Primitive(ty) => display(ty),
436            Self::Ptr(ty) => ty.render(),
437            Self::Array(ty) => ty.render(),
438            Self::Struct(ty) => ty.render(),
439            Self::Ref(ty) => display(ty),
440        }
441    }
442}
443
444// POINTER TYPE
445// ================================================================================================
446
447#[derive(Debug, Clone)]
448pub struct PointerType {
449    pub span: SourceSpan,
450    pub pointee: Box<TypeExpr>,
451    addrspace: Option<AddressSpace>,
452}
453
454impl From<types::PointerType> for PointerType {
455    fn from(ty: types::PointerType) -> Self {
456        let types::PointerType { addrspace, pointee } = ty;
457        let pointee = Box::new(TypeExpr::from(pointee));
458        Self {
459            span: SourceSpan::UNKNOWN,
460            pointee,
461            addrspace: Some(addrspace),
462        }
463    }
464}
465
466impl Eq for PointerType {}
467
468impl PartialEq for PointerType {
469    fn eq(&self, other: &Self) -> bool {
470        self.address_space() == other.address_space() && self.pointee == other.pointee
471    }
472}
473
474impl core::hash::Hash for PointerType {
475    fn hash<H: core::hash::Hasher>(&self, state: &mut H) {
476        self.pointee.hash(state);
477        self.address_space().hash(state);
478    }
479}
480
481impl Spanned for PointerType {
482    fn span(&self) -> SourceSpan {
483        self.span
484    }
485}
486
487impl PointerType {
488    pub fn new(pointee: TypeExpr) -> Self {
489        Self {
490            span: SourceSpan::UNKNOWN,
491            pointee: Box::new(pointee),
492            addrspace: None,
493        }
494    }
495
496    /// Override the default source span
497    #[inline]
498    pub fn with_span(mut self, span: SourceSpan) -> Self {
499        self.span = span;
500        self
501    }
502
503    /// Override the default address space
504    #[inline]
505    pub fn with_address_space(mut self, addrspace: AddressSpace) -> Self {
506        self.addrspace = Some(addrspace);
507        self
508    }
509
510    /// Get the address space of this pointer type
511    #[inline]
512    pub fn address_space(&self) -> AddressSpace {
513        self.addrspace.unwrap_or(AddressSpace::Element)
514    }
515}
516
517impl crate::prettier::PrettyPrint for PointerType {
518    fn render(&self) -> crate::prettier::Document {
519        use crate::prettier::*;
520
521        let doc = const_text("ptr<") + self.pointee.render();
522        if let Some(addrspace) = self.addrspace.as_ref() {
523            doc + const_text(", ") + text(format!("addrspace({addrspace})")) + const_text(">")
524        } else {
525            doc + const_text(">")
526        }
527    }
528}
529
530// ARRAY TYPE
531// ================================================================================================
532
533#[derive(Debug, Clone)]
534pub struct ArrayType {
535    pub span: SourceSpan,
536    pub elem: Box<TypeExpr>,
537    pub arity: usize,
538}
539
540impl Eq for ArrayType {}
541
542impl PartialEq for ArrayType {
543    fn eq(&self, other: &Self) -> bool {
544        self.arity == other.arity && self.elem == other.elem
545    }
546}
547
548impl core::hash::Hash for ArrayType {
549    fn hash<H: core::hash::Hasher>(&self, state: &mut H) {
550        self.elem.hash(state);
551        self.arity.hash(state);
552    }
553}
554
555impl Spanned for ArrayType {
556    fn span(&self) -> SourceSpan {
557        self.span
558    }
559}
560
561impl ArrayType {
562    pub fn new(elem: TypeExpr, arity: usize) -> Self {
563        Self {
564            span: SourceSpan::UNKNOWN,
565            elem: Box::new(elem),
566            arity,
567        }
568    }
569
570    /// Override the default source span
571    #[inline]
572    pub fn with_span(mut self, span: SourceSpan) -> Self {
573        self.span = span;
574        self
575    }
576}
577
578impl crate::prettier::PrettyPrint for ArrayType {
579    fn render(&self) -> crate::prettier::Document {
580        use crate::prettier::*;
581
582        const_text("[")
583            + self.elem.render()
584            + const_text("; ")
585            + display(self.arity)
586            + const_text("]")
587    }
588}
589
590// STRUCT TYPE
591// ================================================================================================
592
593#[derive(Debug, Clone)]
594pub struct StructType {
595    pub span: SourceSpan,
596    pub name: Option<Ident>,
597    pub repr: Span<TypeRepr>,
598    pub fields: Vec<StructField>,
599}
600
601impl Eq for StructType {}
602
603impl PartialEq for StructType {
604    fn eq(&self, other: &Self) -> bool {
605        self.name == other.name && self.repr == other.repr && self.fields == other.fields
606    }
607}
608
609impl core::hash::Hash for StructType {
610    fn hash<H: core::hash::Hasher>(&self, state: &mut H) {
611        self.name.hash(state);
612        self.repr.hash(state);
613        self.fields.hash(state);
614    }
615}
616
617impl Spanned for StructType {
618    fn span(&self) -> SourceSpan {
619        self.span
620    }
621}
622
623impl StructType {
624    pub fn new(name: Option<Ident>, fields: impl IntoIterator<Item = StructField>) -> Self {
625        Self {
626            span: SourceSpan::UNKNOWN,
627            name,
628            repr: Span::unknown(TypeRepr::Default),
629            fields: fields.into_iter().collect(),
630        }
631    }
632
633    /// Override the default struct representation
634    #[inline]
635    pub fn with_repr(mut self, repr: Span<TypeRepr>) -> Self {
636        self.repr = repr;
637        self
638    }
639
640    /// Override the default source span
641    #[inline]
642    pub fn with_span(mut self, span: SourceSpan) -> Self {
643        self.span = span;
644        self
645    }
646}
647
648impl crate::prettier::PrettyPrint for StructType {
649    fn render(&self) -> crate::prettier::Document {
650        use crate::prettier::*;
651
652        let repr = match &*self.repr {
653            TypeRepr::Default => Document::Empty,
654            TypeRepr::BigEndian => const_text(" @bigendian"),
655            repr @ (TypeRepr::Align(_) | TypeRepr::Packed(_) | TypeRepr::Transparent) => {
656                text(format!(" @{repr}"))
657            },
658        };
659
660        let singleline_body = self
661            .fields
662            .iter()
663            .map(PrettyPrint::render)
664            .reduce(|acc, field| acc + const_text(", ") + field)
665            .unwrap_or(Document::Empty);
666        let multiline_body = indent(
667            4,
668            nl() + self
669                .fields
670                .iter()
671                .map(PrettyPrint::render)
672                .reduce(|acc, field| acc + const_text(",") + nl() + field)
673                .unwrap_or(Document::Empty),
674        ) + nl();
675        let body = singleline_body | multiline_body;
676
677        const_text("struct") + repr + const_text(" { ") + body + const_text(" }")
678    }
679}
680
681// STRUCT FIELD
682// ================================================================================================
683
684#[derive(Debug, Clone)]
685pub struct StructField {
686    pub span: SourceSpan,
687    pub name: Ident,
688    pub ty: TypeExpr,
689}
690
691impl Eq for StructField {}
692
693impl PartialEq for StructField {
694    fn eq(&self, other: &Self) -> bool {
695        self.name == other.name && self.ty == other.ty
696    }
697}
698
699impl core::hash::Hash for StructField {
700    fn hash<H: core::hash::Hasher>(&self, state: &mut H) {
701        self.name.hash(state);
702        self.ty.hash(state);
703    }
704}
705
706impl Spanned for StructField {
707    fn span(&self) -> SourceSpan {
708        self.span
709    }
710}
711
712impl crate::prettier::PrettyPrint for StructField {
713    fn render(&self) -> crate::prettier::Document {
714        use crate::prettier::*;
715
716        display(&self.name) + const_text(": ") + self.ty.render()
717    }
718}
719
720// TYPE ALIAS
721// ================================================================================================
722
723/// A [TypeAlias] represents a named [Type].
724///
725/// Type aliases correspond to type declarations in Miden Assembly source files. They are called
726/// aliases, rather than declarations, as the type system for Miden Assembly is structural, rather
727/// than nominal, and so two aliases with the same underlying type are considered equivalent.
728#[derive(Debug, Clone)]
729pub struct TypeAlias {
730    span: SourceSpan,
731    /// The documentation string attached to this definition.
732    docs: Option<DocString>,
733    /// The visibility of this type alias
734    pub visibility: Visibility,
735    /// The name of this type alias
736    pub name: Ident,
737    /// The concrete underlying type
738    pub ty: TypeExpr,
739}
740
741impl TypeAlias {
742    /// Create a new type alias from a name and type
743    pub fn new(visibility: Visibility, name: Ident, ty: TypeExpr) -> Self {
744        Self {
745            span: name.span(),
746            docs: None,
747            visibility,
748            name,
749            ty,
750        }
751    }
752
753    /// Adds documentation to this type alias
754    pub fn with_docs(mut self, docs: Option<Span<String>>) -> Self {
755        self.docs = docs.map(DocString::new);
756        self
757    }
758
759    /// Override the default source span
760    #[inline]
761    pub fn with_span(mut self, span: SourceSpan) -> Self {
762        self.span = span;
763        self
764    }
765
766    /// Set the source span
767    #[inline]
768    pub fn set_span(&mut self, span: SourceSpan) {
769        self.span = span;
770    }
771
772    /// Returns the documentation associated with this item.
773    pub fn docs(&self) -> Option<Span<&str>> {
774        self.docs.as_ref().map(|docstring| docstring.as_spanned_str())
775    }
776
777    /// Get the name of this type alias
778    pub fn name(&self) -> &Ident {
779        &self.name
780    }
781
782    /// Get the visibility of this type alias
783    #[inline]
784    pub const fn visibility(&self) -> Visibility {
785        self.visibility
786    }
787}
788
789impl Eq for TypeAlias {}
790
791impl PartialEq for TypeAlias {
792    fn eq(&self, other: &Self) -> bool {
793        self.visibility == other.visibility
794            && self.name == other.name
795            && self.docs == other.docs
796            && self.ty == other.ty
797    }
798}
799
800impl core::hash::Hash for TypeAlias {
801    fn hash<H: core::hash::Hasher>(&self, state: &mut H) {
802        let Self { span: _, docs, visibility, name, ty } = self;
803        docs.hash(state);
804        visibility.hash(state);
805        name.hash(state);
806        ty.hash(state);
807    }
808}
809
810impl Spanned for TypeAlias {
811    fn span(&self) -> SourceSpan {
812        self.span
813    }
814}
815
816impl crate::prettier::PrettyPrint for TypeAlias {
817    fn render(&self) -> crate::prettier::Document {
818        use crate::prettier::*;
819
820        let mut doc = self.docs.as_ref().map(PrettyPrint::render).unwrap_or(Document::Empty);
821
822        if self.visibility.is_public() {
823            doc += display(self.visibility) + const_text(" ");
824        }
825
826        doc + const_text("type")
827            + const_text(" ")
828            + display(&self.name)
829            + const_text(" = ")
830            + self.ty.render()
831    }
832}
833
834// ENUM TYPE
835// ================================================================================================
836
837/// A combined type alias and constant declaration corresponding to a C-like enumeration.
838///
839/// C-style enumerations are effectively a type alias for an integer type with a limited set of
840/// valid values with associated names (referred to as _variants_ of the enum type).
841///
842/// In Miden Assembly, these provide a means for a procedure to declare that it expects an argument
843/// of the underlying integral type, but that values other than those of the declared variants are
844/// illegal/invalid. Currently, these are unchecked, and are only used to convey semantic
845/// information. In the future, we may perform static analysis to try and identify invalid instances
846/// of the enumeration when derived from a constant.
847#[derive(Debug, Clone)]
848pub struct EnumType {
849    span: SourceSpan,
850    /// The documentation string attached to this definition.
851    docs: Option<DocString>,
852    /// The visibility of this enum type
853    visibility: Visibility,
854    /// The enum name
855    name: Ident,
856    /// The type of the discriminant value used for this enum's variants
857    ///
858    /// NOTE: The type must be an integral value, and this is enforced by [`Self::new`].
859    ty: Type,
860    /// The enum variants
861    variants: Vec<Variant>,
862}
863
864impl EnumType {
865    /// Construct a new enum type with the given name and variants
866    ///
867    /// The caller is assumed to have already validated that `ty` is an integral type, and this
868    /// function will assert that this is the case.
869    pub fn new(
870        visibility: Visibility,
871        name: Ident,
872        ty: Type,
873        variants: impl IntoIterator<Item = Variant>,
874    ) -> Self {
875        assert!(ty.is_integer(), "only integer types are allowed in enum type definitions");
876        Self {
877            span: name.span(),
878            docs: None,
879            visibility,
880            name,
881            ty,
882            variants: Vec::from_iter(variants),
883        }
884    }
885
886    /// Adds documentation to this enum declaration.
887    pub fn with_docs(mut self, docs: Option<Span<String>>) -> Self {
888        self.docs = docs.map(DocString::new);
889        self
890    }
891
892    /// Override the default source span
893    pub fn with_span(mut self, span: SourceSpan) -> Self {
894        self.span = span;
895        self
896    }
897
898    /// Returns true if this is a C-style enum where the discriminant is the value
899    pub fn is_c_like(&self) -> bool {
900        !self.variants.is_empty() && self.variants.iter().all(|v| v.value_ty.is_none())
901    }
902
903    /// Set the source span
904    pub fn set_span(&mut self, span: SourceSpan) {
905        self.span = span;
906    }
907
908    /// Get the name of this enum type
909    pub fn name(&self) -> &Ident {
910        &self.name
911    }
912
913    /// Get the visibility of this enum type
914    pub const fn visibility(&self) -> Visibility {
915        self.visibility
916    }
917
918    /// Returns the documentation associated with this item.
919    pub fn docs(&self) -> Option<Span<&str>> {
920        self.docs.as_ref().map(|docstring| docstring.as_spanned_str())
921    }
922
923    /// Get the concrete type of this enum's variants
924    pub fn ty(&self) -> &Type {
925        &self.ty
926    }
927
928    /// Get the variants of this enum type
929    pub fn variants(&self) -> &[Variant] {
930        &self.variants
931    }
932
933    /// Get the variants of this enum type, mutably
934    pub fn variants_mut(&mut self) -> &mut Vec<Variant> {
935        &mut self.variants
936    }
937
938    /// Split this definition into its type alias and variant parts
939    pub fn into_parts(self) -> (TypeAlias, Vec<Variant>) {
940        let Self {
941            span,
942            docs,
943            visibility,
944            name,
945            ty,
946            variants,
947        } = self;
948        let alias = TypeAlias {
949            span,
950            docs,
951            visibility,
952            name,
953            ty: TypeExpr::Primitive(Span::new(span, ty)),
954        };
955        (alias, variants)
956    }
957}
958
959impl Spanned for EnumType {
960    fn span(&self) -> SourceSpan {
961        self.span
962    }
963}
964
965impl Eq for EnumType {}
966
967impl PartialEq for EnumType {
968    fn eq(&self, other: &Self) -> bool {
969        self.visibility == other.visibility
970            && self.name == other.name
971            && self.docs == other.docs
972            && self.ty == other.ty
973            && self.variants == other.variants
974    }
975}
976
977impl core::hash::Hash for EnumType {
978    fn hash<H: core::hash::Hasher>(&self, state: &mut H) {
979        let Self {
980            span: _,
981            docs,
982            visibility,
983            name,
984            ty,
985            variants,
986        } = self;
987        docs.hash(state);
988        visibility.hash(state);
989        name.hash(state);
990        ty.hash(state);
991        variants.hash(state);
992    }
993}
994
995impl crate::prettier::PrettyPrint for EnumType {
996    fn render(&self) -> crate::prettier::Document {
997        use crate::prettier::*;
998
999        let mut doc = self.docs.as_ref().map(PrettyPrint::render).unwrap_or(Document::Empty);
1000
1001        let variants = self
1002            .variants
1003            .iter()
1004            .map(PrettyPrint::render)
1005            .reduce(|acc, v| acc + const_text(",") + nl() + v)
1006            .unwrap_or(Document::Empty);
1007
1008        if self.visibility.is_public() {
1009            doc += display(self.visibility) + const_text(" ");
1010        }
1011
1012        doc + const_text("enum")
1013            + const_text(" ")
1014            + display(&self.name)
1015            + const_text(" : ")
1016            + self.ty.render()
1017            + const_text(" {")
1018            + nl()
1019            + variants
1020            + const_text("}")
1021    }
1022}
1023
1024// ENUM VARIANT
1025// ================================================================================================
1026
1027/// A variant of an [EnumType].
1028///
1029/// See the [EnumType] docs for more information.
1030#[derive(Debug, Clone)]
1031pub struct Variant {
1032    pub span: SourceSpan,
1033    /// The documentation string attached to the constant derived from this variant.
1034    pub docs: Option<DocString>,
1035    /// The name of this enum variant
1036    pub name: Ident,
1037    /// The payload value type of this variant
1038    ///
1039    /// NOTE: This is not supported in Miden Assembly text format yet, but can be set when lowering
1040    /// directly to the AST.
1041    pub value_ty: Option<TypeExpr>,
1042    /// The discriminant value associated with this variant
1043    pub discriminant: ConstantExpr,
1044}
1045
1046impl Variant {
1047    /// Construct a new variant of an [EnumType], with the given name and discriminant value.
1048    pub fn new(name: Ident, discriminant: ConstantExpr, payload: Option<TypeExpr>) -> Self {
1049        Self {
1050            span: name.span(),
1051            docs: None,
1052            name,
1053            value_ty: payload,
1054            discriminant,
1055        }
1056    }
1057
1058    /// Override the span for this variant
1059    pub fn with_span(mut self, span: SourceSpan) -> Self {
1060        self.span = span;
1061        self
1062    }
1063
1064    /// Adds documentation to this variant
1065    pub fn with_docs(mut self, docs: Option<Span<String>>) -> Self {
1066        self.docs = docs.map(DocString::new);
1067        self
1068    }
1069
1070    /// Used to validate that this variant's discriminant value is an instance of `ty`,
1071    /// which must be a type valid for use as the underlying representation for an enum, i.e. an
1072    /// integer type up to 64 bits in size.
1073    ///
1074    /// It is expected that the discriminant expression has been folded to an integer value by the
1075    /// time this is called. If the discriminant has not been fully folded, then an error will be
1076    /// returned.
1077    pub fn assert_instance_of(&self, ty: &Type) -> Result<(), crate::SemanticAnalysisError> {
1078        use crate::{FIELD_MODULUS, SemanticAnalysisError};
1079
1080        let value = match &self.discriminant {
1081            ConstantExpr::Int(value) => value.as_int(),
1082            _ => {
1083                return Err(SemanticAnalysisError::InvalidEnumDiscriminant {
1084                    span: self.discriminant.span(),
1085                    repr: ty.clone(),
1086                });
1087            },
1088        };
1089
1090        match ty {
1091            Type::Felt if value >= FIELD_MODULUS => {
1092                Err(SemanticAnalysisError::InvalidEnumDiscriminant {
1093                    span: self.discriminant.span(),
1094                    repr: ty.clone(),
1095                })
1096            },
1097            // IntValue is represented as an unsigned integer, so negative discriminants
1098            // are rejected during constant evaluation.
1099            Type::Felt => Ok(()),
1100            Type::I1 if value > 1 => Err(SemanticAnalysisError::InvalidEnumDiscriminant {
1101                span: self.discriminant.span(),
1102                repr: ty.clone(),
1103            }),
1104            Type::I1 => Ok(()),
1105            Type::I8 | Type::U8 if value > u8::MAX as u64 => {
1106                Err(SemanticAnalysisError::InvalidEnumDiscriminant {
1107                    span: self.discriminant.span(),
1108                    repr: ty.clone(),
1109                })
1110            },
1111            Type::I8 | Type::U8 => Ok(()),
1112            Type::I16 | Type::U16 if value > u16::MAX as u64 => {
1113                Err(SemanticAnalysisError::InvalidEnumDiscriminant {
1114                    span: self.discriminant.span(),
1115                    repr: ty.clone(),
1116                })
1117            },
1118            Type::I16 | Type::U16 => Ok(()),
1119            Type::I32 | Type::U32 if value > u32::MAX as u64 => {
1120                Err(SemanticAnalysisError::InvalidEnumDiscriminant {
1121                    span: self.discriminant.span(),
1122                    repr: ty.clone(),
1123                })
1124            },
1125            Type::I32 | Type::U32 => Ok(()),
1126            Type::I64 | Type::U64 if value >= FIELD_MODULUS => {
1127                Err(SemanticAnalysisError::InvalidEnumDiscriminant {
1128                    span: self.discriminant.span(),
1129                    repr: ty.clone(),
1130                })
1131            },
1132            _ => Err(SemanticAnalysisError::InvalidEnumRepr { span: self.span }),
1133        }
1134    }
1135}
1136
1137impl Spanned for Variant {
1138    fn span(&self) -> SourceSpan {
1139        self.span
1140    }
1141}
1142
1143impl Eq for Variant {}
1144
1145impl PartialEq for Variant {
1146    fn eq(&self, other: &Self) -> bool {
1147        self.name == other.name
1148            && self.value_ty == other.value_ty
1149            && self.discriminant == other.discriminant
1150            && self.docs == other.docs
1151    }
1152}
1153
1154impl core::hash::Hash for Variant {
1155    fn hash<H: core::hash::Hasher>(&self, state: &mut H) {
1156        let Self {
1157            span: _,
1158            docs,
1159            name,
1160            value_ty,
1161            discriminant,
1162        } = self;
1163        docs.hash(state);
1164        name.hash(state);
1165        value_ty.hash(state);
1166        discriminant.hash(state);
1167    }
1168}
1169
1170impl crate::prettier::PrettyPrint for Variant {
1171    fn render(&self) -> crate::prettier::Document {
1172        use crate::prettier::*;
1173
1174        let doc = self.docs.as_ref().map(PrettyPrint::render).unwrap_or(Document::Empty);
1175
1176        let name = display(&self.name);
1177        let name_and_payload = if let Some(value_ty) = self.value_ty.as_ref() {
1178            name + const_text("(") + value_ty.render() + const_text(")")
1179        } else {
1180            name
1181        };
1182        doc + name_and_payload + const_text(" = ") + self.discriminant.render()
1183    }
1184}
1185
1186#[cfg(test)]
1187mod tests {
1188    use alloc::{string::ToString, sync::Arc};
1189    use core::str::FromStr;
1190
1191    use miden_debug_types::{DefaultSourceManager, SourceFile, SourceId, SourceLanguage, Uri};
1192
1193    use super::*;
1194    use crate::{ast::Form, prettier::PrettyPrint};
1195
1196    struct DummyResolver {
1197        source_manager: Arc<dyn SourceManager>,
1198    }
1199
1200    impl DummyResolver {
1201        fn new() -> Self {
1202            Self {
1203                source_manager: Arc::new(DefaultSourceManager::default()),
1204            }
1205        }
1206    }
1207
1208    impl TypeResolver<SymbolResolutionError> for DummyResolver {
1209        fn source_manager(&self) -> Arc<dyn SourceManager> {
1210            self.source_manager.clone()
1211        }
1212
1213        fn resolve_local_failed(&self, err: SymbolResolutionError) -> SymbolResolutionError {
1214            err
1215        }
1216
1217        fn get_type(
1218            &mut self,
1219            context: SourceSpan,
1220            _gid: GlobalItemIndex,
1221        ) -> Result<Type, SymbolResolutionError> {
1222            Err(SymbolResolutionError::undefined(context, self.source_manager.as_ref()))
1223        }
1224
1225        fn get_local_type(
1226            &mut self,
1227            _context: SourceSpan,
1228            _id: ItemIndex,
1229        ) -> Result<Option<Type>, SymbolResolutionError> {
1230            Ok(None)
1231        }
1232
1233        fn resolve_type_ref(
1234            &mut self,
1235            ty: Span<&Path>,
1236        ) -> Result<SymbolResolution, SymbolResolutionError> {
1237            Err(SymbolResolutionError::undefined(ty.span(), self.source_manager.as_ref()))
1238        }
1239    }
1240
1241    fn nested_type_expr(depth: usize) -> TypeExpr {
1242        let mut expr = TypeExpr::Primitive(Span::unknown(Type::Felt));
1243        for i in 0..depth {
1244            expr = match i % 3 {
1245                0 => TypeExpr::Ptr(PointerType::new(expr)),
1246                1 => TypeExpr::Array(ArrayType::new(expr, 1)),
1247                _ => {
1248                    let field = StructField {
1249                        span: SourceSpan::UNKNOWN,
1250                        name: Ident::from_str("field").expect("valid ident"),
1251                        ty: expr,
1252                    };
1253                    TypeExpr::Struct(StructType::new(None, [field]))
1254                },
1255            };
1256        }
1257        expr
1258    }
1259
1260    fn test_source_file(source: &str) -> Arc<SourceFile> {
1261        Arc::new(SourceFile::new(
1262            SourceId::default(),
1263            SourceLanguage::Masm,
1264            Uri::new("memory:///type-expr-test.masm"),
1265            source.to_string().into_boxed_str(),
1266        ))
1267    }
1268
1269    fn parse_type_alias_expr(source: &str) -> TypeExpr {
1270        let mut forms =
1271            crate::parser::parse_forms(test_source_file(source)).expect("type alias should parse");
1272        assert_eq!(forms.len(), 1, "expected exactly one parsed form");
1273        match forms.pop().expect("expected parsed form") {
1274            Form::Type(alias) => alias.ty,
1275            form => panic!("expected type alias form, got {form:?}"),
1276        }
1277    }
1278
1279    fn repr_round_trip_struct(repr: TypeRepr) -> TypeExpr {
1280        TypeExpr::Struct(
1281            StructType::new(
1282                None,
1283                [
1284                    StructField {
1285                        span: SourceSpan::UNKNOWN,
1286                        name: Ident::from_str("prefix").expect("valid ident"),
1287                        ty: TypeExpr::Primitive(Span::unknown(Type::Felt)),
1288                    },
1289                    StructField {
1290                        span: SourceSpan::UNKNOWN,
1291                        name: Ident::from_str("suffix").expect("valid ident"),
1292                        ty: TypeExpr::Primitive(Span::unknown(Type::U32)),
1293                    },
1294                ],
1295            )
1296            .with_repr(Span::unknown(repr)),
1297        )
1298    }
1299
1300    #[test]
1301    fn type_expr_depth_boundary() {
1302        let mut resolver = DummyResolver::new();
1303
1304        let ok_expr = nested_type_expr(MAX_TYPE_EXPR_NESTING);
1305        assert!(ok_expr.resolve_type(&mut resolver).is_ok());
1306
1307        let err_expr = nested_type_expr(MAX_TYPE_EXPR_NESTING + 1);
1308        let err = err_expr.resolve_type(&mut resolver).expect_err("expected depth-exceeded error");
1309        assert!(
1310            matches!(err, SymbolResolutionError::TypeExpressionDepthExceeded { max_depth, .. }
1311                if max_depth == MAX_TYPE_EXPR_NESTING)
1312        );
1313    }
1314
1315    #[test]
1316    fn struct_type_expr_render_round_trips_non_default_reprs() {
1317        for repr in [
1318            TypeRepr::BigEndian,
1319            TypeRepr::align(16),
1320            TypeRepr::packed(1),
1321            TypeRepr::packed(2),
1322            TypeRepr::Transparent,
1323        ] {
1324            let rendered = repr_round_trip_struct(repr).to_pretty_string();
1325            assert!(
1326                rendered.starts_with("struct @"),
1327                "non-default struct repr should render after `struct`: {rendered}"
1328            );
1329
1330            let parsed = parse_type_alias_expr(&format!("type RoundTrip = {rendered}\n"));
1331            let TypeExpr::Struct(parsed) = parsed else {
1332                panic!("expected rendered type to parse back as a struct");
1333            };
1334            assert_eq!(*parsed.repr, repr);
1335            assert_eq!(parsed.fields[0].name.as_str(), "prefix");
1336            assert_eq!(parsed.fields[1].name.as_str(), "suffix");
1337        }
1338    }
1339
1340    #[test]
1341    fn type_expr_from_type_preserves_wide_integer_primitives() {
1342        for ty in [Type::I64, Type::U64, Type::I128, Type::U128] {
1343            let expr = TypeExpr::from(ty.clone());
1344            let TypeExpr::Primitive(actual) = expr else {
1345                panic!("expected primitive type expression for {ty}, got {expr:?}");
1346            };
1347            assert_eq!(actual.into_inner(), ty);
1348        }
1349    }
1350
1351    #[test]
1352    fn type_expr_from_type_preserves_struct_metadata() {
1353        let ty = Type::Struct(Arc::new(types::StructType::from_parts(
1354            Some(Arc::from("miden:base/core-types@1.0.0/account-id")),
1355            TypeRepr::BigEndian,
1356            [
1357                (Arc::<str>::from("prefix"), Type::Felt),
1358                (Arc::<str>::from("suffix"), Type::Felt),
1359            ],
1360        )));
1361
1362        let TypeExpr::Struct(actual) = TypeExpr::from(ty) else {
1363            panic!("expected struct type expression");
1364        };
1365        assert_eq!(
1366            actual.name.as_ref().map(Ident::as_str),
1367            Some("miden:base/core-types@1.0.0/account-id"),
1368        );
1369        assert_eq!(*actual.repr, TypeRepr::BigEndian);
1370        assert_eq!(actual.fields[0].name.as_str(), "prefix");
1371        assert_eq!(actual.fields[1].name.as_str(), "suffix");
1372    }
1373
1374    #[test]
1375    fn parsed_struct_type_preserves_field_names_through_resolution() {
1376        let expr = parse_type_alias_expr(
1377            "type AccountId = struct @bigendian { prefix: felt, suffix: felt }\n",
1378        );
1379
1380        let mut resolver = DummyResolver::new();
1381        let resolved = expr
1382            .resolve_type(&mut resolver)
1383            .expect("struct type should resolve")
1384            .expect("struct type should be concrete");
1385        let Type::Struct(resolved_struct) = &resolved else {
1386            panic!("expected resolved struct type, got {resolved:?}");
1387        };
1388        assert_eq!(resolved_struct.repr(), TypeRepr::BigEndian);
1389        assert_eq!(resolved_struct.fields()[0].name.as_deref(), Some("prefix"));
1390        assert_eq!(resolved_struct.fields()[1].name.as_deref(), Some("suffix"));
1391
1392        let TypeExpr::Struct(converted) = TypeExpr::from(resolved) else {
1393            panic!("expected concrete struct to convert back to struct type expression");
1394        };
1395        assert_eq!(*converted.repr, TypeRepr::BigEndian);
1396        assert_eq!(converted.fields[0].name.as_str(), "prefix");
1397        assert_eq!(converted.fields[1].name.as_str(), "suffix");
1398    }
1399}