Skip to main content

shape_ast/ast/
types.rs

1//! Type system definitions for Shape AST
2
3use super::DocComment;
4use super::functions::Annotation;
5use super::span::Span;
6use super::type_path::TypePath;
7use serde::{Deserialize, Serialize};
8
9#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
10pub enum TypeAnnotation {
11    /// Basic types: number, string, bool, row, pattern, etc.
12    Basic(String),
13    /// Vector type: T[] or Vec<T>
14    Array(Box<TypeAnnotation>),
15    /// Tuple type: [T1, T2, T3]
16    Tuple(Vec<TypeAnnotation>),
17    /// Object type: { field1: T1, field2?: T2 }
18    Object(Vec<ObjectTypeField>),
19    /// Function type: (T1, T2) => T3
20    Function {
21        params: Vec<FunctionParam>,
22        returns: Box<TypeAnnotation>,
23    },
24    /// Union type: T1 | T2 | T3 (discriminated union - value is ONE of the types)
25    Union(Vec<TypeAnnotation>),
26    /// Intersection type: T1 + T2 (structural merge - value has ALL fields from both types)
27    /// Only valid for object/trait types. Field collisions are compile-time errors.
28    Intersection(Vec<TypeAnnotation>),
29    /// Generic type: Map<K, V>
30    Generic {
31        name: TypePath,
32        args: Vec<TypeAnnotation>,
33    },
34    /// Type reference (custom type or type alias)
35    Reference(TypePath),
36    /// Void type
37    Void,
38    /// Never type
39    Never,
40    /// Null type
41    Null,
42    /// Undefined type
43    Undefined,
44    /// Trait object type: dyn Trait1 + Trait2
45    /// Represents a type-erased value that implements the given traits
46    Dyn(Vec<TypePath>),
47}
48
49impl TypeAnnotation {
50    pub fn option(inner: TypeAnnotation) -> Self {
51        TypeAnnotation::Generic {
52            name: TypePath::simple("Option"),
53            args: vec![inner],
54        }
55    }
56
57    pub fn option_inner(&self) -> Option<&TypeAnnotation> {
58        match self {
59            TypeAnnotation::Generic { name, args }
60                if name.as_str() == "Option" && args.len() == 1 =>
61            {
62                args.first()
63            }
64            _ => None,
65        }
66    }
67
68    pub fn into_option_inner(self) -> Option<TypeAnnotation> {
69        match self {
70            TypeAnnotation::Generic { name, mut args }
71                if name.as_str() == "Option" && args.len() == 1 =>
72            {
73                Some(args.remove(0))
74            }
75            _ => None,
76        }
77    }
78
79    pub fn is_option(&self) -> bool {
80        self.option_inner().is_some()
81    }
82
83    /// Extract a simple type name if this is a Reference or Basic type
84    ///
85    /// Returns `Some(type_name)` for:
86    /// - `TypeAnnotation::Reference(path)` - e.g., `Currency`, `foo::MyType`
87    /// - `TypeAnnotation::Basic(name)` - e.g., `number`, `string`
88    ///
89    /// Returns `None` for complex types like arrays, tuples, functions, etc.
90    pub fn as_simple_name(&self) -> Option<&str> {
91        match self {
92            TypeAnnotation::Reference(path) => Some(path.as_str()),
93            TypeAnnotation::Basic(name) => Some(name.as_str()),
94            _ => None,
95        }
96    }
97
98    /// Extract the type name string for Basic or Reference variants.
99    /// Handles the `Basic(name) | Reference(path)` pattern uniformly.
100    pub fn as_type_name_str(&self) -> Option<&str> {
101        match self {
102            TypeAnnotation::Basic(name) => Some(name.as_str()),
103            TypeAnnotation::Reference(path) => Some(path.as_str()),
104            _ => None,
105        }
106    }
107
108    /// Convert a type annotation to its full string representation.
109    pub fn to_type_string(&self) -> String {
110        match self {
111            TypeAnnotation::Basic(name) => name.clone(),
112            TypeAnnotation::Reference(path) => path.to_string(),
113            TypeAnnotation::Array(inner) => format!("Array<{}>", inner.to_type_string()),
114            TypeAnnotation::Generic { name, args }
115                if name.as_str() == "Option" && args.len() == 1 =>
116            {
117                format!("{}?", args[0].to_type_string())
118            }
119            TypeAnnotation::Generic { name, args } => {
120                let args_str: Vec<String> = args.iter().map(|a| a.to_type_string()).collect();
121                format!("{}<{}>", name, args_str.join(", "))
122            }
123            TypeAnnotation::Tuple(items) => {
124                let items_str: Vec<String> = items.iter().map(|t| t.to_type_string()).collect();
125                format!("[{}]", items_str.join(", "))
126            }
127            TypeAnnotation::Union(items) => {
128                let items_str: Vec<String> = items.iter().map(|t| t.to_type_string()).collect();
129                items_str.join(" | ")
130            }
131            TypeAnnotation::Void => "void".to_string(),
132            TypeAnnotation::Never => "never".to_string(),
133            TypeAnnotation::Null => "null".to_string(),
134            TypeAnnotation::Undefined => "undefined".to_string(),
135            TypeAnnotation::Object(fields) => {
136                let fields_str: Vec<String> = fields
137                    .iter()
138                    .map(|f| {
139                        let opt = if f.optional { "?" } else { "" };
140                        // Include @alias in the type string so the Python extension
141                        // can use the wire name when looking up dict keys.
142                        let alias = f
143                            .annotations
144                            .iter()
145                            .find(|a| a.name == "alias")
146                            .and_then(|a| a.args.first())
147                            .and_then(|arg| match arg {
148                                super::expressions::Expr::Literal(
149                                    super::literals::Literal::String(s),
150                                    _,
151                                ) => Some(s.as_str()),
152                                _ => None,
153                            });
154                        let alias_str = alias.map(|a| format!("@\"{}\" ", a)).unwrap_or_default();
155                        format!(
156                            "{}{}{}: {}",
157                            alias_str,
158                            f.name,
159                            opt,
160                            f.type_annotation.to_type_string()
161                        )
162                    })
163                    .collect();
164                format!("{{{}}}", fields_str.join(", "))
165            }
166            _ => "any".to_string(),
167        }
168    }
169}
170
171#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
172pub struct ObjectTypeField {
173    pub name: String,
174    pub optional: bool,
175    pub type_annotation: TypeAnnotation,
176    /// Field annotations (e.g. `@alias("wire name")`)
177    #[serde(default)]
178    pub annotations: Vec<Annotation>,
179}
180
181#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
182pub struct FunctionParam {
183    pub name: Option<String>,
184    pub optional: bool,
185    pub type_annotation: TypeAnnotation,
186}
187
188/// A generic parameter on a function, type, trait or impl header.
189///
190/// Two variants:
191/// - [`TypeParam::Type`] — the classic type-level parameter (`<T>`,
192///   `<T: Bound>`, `<T = DefaultType>`).
193/// - [`TypeParam::Const`] — a value-level const generic parameter
194///   (`<const N: int>`, `<const N: int = 4>`). Added in B.2 of the
195///   const-generics track; B.3 binds const args at monomorphization
196///   and B.4 substitutes them into specialized bodies. Until then
197///   most consumers treat const params as a placeholder — the AST
198///   captures the name, declared type, and optional default expression
199///   so later phases can wire them end-to-end.
200#[derive(Debug, Clone, Serialize, Deserialize)]
201pub enum TypeParam {
202    /// A type parameter: `T`, `T: Bound`, `T = DefaultType`.
203    Type {
204        name: String,
205        #[serde(default)]
206        span: Span,
207        #[serde(default)]
208        doc_comment: Option<DocComment>,
209        /// Default type argument: `T = int`
210        default_type: Option<TypeAnnotation>,
211        /// Trait bounds: `T: Comparable + Displayable`
212        #[serde(default)]
213        trait_bounds: Vec<TypePath>,
214    },
215    /// A const generic parameter: `const N: int`, `const N: int = 4`.
216    ///
217    /// Parsed in B.1/B.2. Monomorphization (B.3) and body substitution
218    /// (B.4) are not yet implemented — most compiler passes currently
219    /// treat this variant as a stub / placeholder.
220    Const {
221        name: String,
222        #[serde(default)]
223        span: Span,
224        #[serde(default)]
225        doc_comment: Option<DocComment>,
226        /// Declared type: `const N: int`.
227        ty: TypeAnnotation,
228        /// Optional default value expression: `= 4`.
229        default: Option<super::expressions::Expr>,
230    },
231}
232
233impl TypeParam {
234    /// The parameter's declared name (`T` or `N`), common to both variants.
235    pub fn name(&self) -> &str {
236        match self {
237            TypeParam::Type { name, .. } | TypeParam::Const { name, .. } => name,
238        }
239    }
240
241    /// The parameter's source span, common to both variants.
242    pub fn span(&self) -> &Span {
243        match self {
244            TypeParam::Type { span, .. } | TypeParam::Const { span, .. } => span,
245        }
246    }
247
248    /// The parameter's attached doc comment, common to both variants.
249    pub fn doc_comment(&self) -> Option<&DocComment> {
250        match self {
251            TypeParam::Type { doc_comment, .. } | TypeParam::Const { doc_comment, .. } => {
252                doc_comment.as_ref()
253            }
254        }
255    }
256
257    /// Trait bounds, only present on the `Type` variant.
258    ///
259    /// Returns an empty slice for const generics (they have no bounds).
260    pub fn trait_bounds(&self) -> &[TypePath] {
261        match self {
262            TypeParam::Type { trait_bounds, .. } => trait_bounds.as_slice(),
263            // TODO(B.3): const generics may eventually grow predicates
264            // (e.g. `where N > 0`), but not in this phase.
265            TypeParam::Const { .. } => &[],
266        }
267    }
268
269    /// Default type argument, only meaningful for the `Type` variant.
270    pub fn default_type(&self) -> Option<&TypeAnnotation> {
271        match self {
272            TypeParam::Type { default_type, .. } => default_type.as_ref(),
273            TypeParam::Const { .. } => None,
274        }
275    }
276
277    /// True iff this is a const generic parameter.
278    pub fn is_const(&self) -> bool {
279        matches!(self, TypeParam::Const { .. })
280    }
281
282    /// Convenience constructor for a plain type parameter with no bounds
283    /// or default — the shape that stdlib bootstrap / method desugaring
284    /// most often needs.
285    pub fn simple(name: impl Into<String>) -> Self {
286        TypeParam::Type {
287            name: name.into(),
288            span: Span::DUMMY,
289            doc_comment: None,
290            default_type: None,
291            trait_bounds: Vec::new(),
292        }
293    }
294}
295
296/// A predicate in a where clause: `T: Comparable + Display`
297#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
298pub struct WherePredicate {
299    pub type_name: String,
300    pub bounds: Vec<TypePath>,
301}
302
303impl PartialEq for TypeParam {
304    fn eq(&self, other: &Self) -> bool {
305        match (self, other) {
306            (
307                TypeParam::Type {
308                    name: ln,
309                    doc_comment: ld,
310                    default_type: ldt,
311                    trait_bounds: ltb,
312                    ..
313                },
314                TypeParam::Type {
315                    name: rn,
316                    doc_comment: rd,
317                    default_type: rdt,
318                    trait_bounds: rtb,
319                    ..
320                },
321            ) => ln == rn && ld == rd && ldt == rdt && ltb == rtb,
322            (
323                TypeParam::Const {
324                    name: ln,
325                    doc_comment: ld,
326                    ty: lty,
327                    default: ldef,
328                    ..
329                },
330                TypeParam::Const {
331                    name: rn,
332                    doc_comment: rd,
333                    ty: rty,
334                    default: rdef,
335                    ..
336                },
337            ) => ln == rn && ld == rd && lty == rty && ldef == rdef,
338            _ => false,
339        }
340    }
341}
342
343#[derive(Debug, Clone, Serialize, Deserialize)]
344pub struct TypeAliasDef {
345    pub name: String,
346    #[serde(default)]
347    pub doc_comment: Option<DocComment>,
348    pub type_params: Option<Vec<TypeParam>>,
349    pub type_annotation: TypeAnnotation,
350    /// Meta parameter overrides: type Percent4 = Percent { decimals: 4 }
351    pub meta_param_overrides: Option<std::collections::HashMap<String, super::expressions::Expr>>,
352}
353
354/// A body-less trait member signature: property, method, or index signature.
355///
356/// Used inside `TraitMember::Required(..)` to express the obligation an impl
357/// must satisfy. Trait bodies that include a body (`method`) use
358/// `TraitMember::Default(MethodDef)` instead.
359#[derive(Debug, Clone, Serialize, Deserialize)]
360pub enum TraitMemberSignature {
361    /// Property signature
362    Property {
363        name: String,
364        optional: bool,
365        type_annotation: TypeAnnotation,
366        #[serde(default)]
367        span: Span,
368        #[serde(default)]
369        doc_comment: Option<DocComment>,
370    },
371    /// Method signature
372    Method {
373        name: String,
374        optional: bool,
375        params: Vec<FunctionParam>,
376        return_type: TypeAnnotation,
377        /// Whether this is an async method
378        is_async: bool,
379        #[serde(default)]
380        span: Span,
381        #[serde(default)]
382        doc_comment: Option<DocComment>,
383    },
384    /// Index signature
385    IndexSignature {
386        param_name: String,
387        param_type: String, // "string" or "number"
388        return_type: TypeAnnotation,
389        #[serde(default)]
390        span: Span,
391        #[serde(default)]
392        doc_comment: Option<DocComment>,
393    },
394}
395
396impl TraitMemberSignature {
397    pub fn span(&self) -> Span {
398        match self {
399            TraitMemberSignature::Property { span, .. }
400            | TraitMemberSignature::Method { span, .. }
401            | TraitMemberSignature::IndexSignature { span, .. } => *span,
402        }
403    }
404
405    pub fn doc_comment(&self) -> Option<&DocComment> {
406        match self {
407            TraitMemberSignature::Property { doc_comment, .. }
408            | TraitMemberSignature::Method { doc_comment, .. }
409            | TraitMemberSignature::IndexSignature { doc_comment, .. } => doc_comment.as_ref(),
410        }
411    }
412}
413
414#[derive(Debug, Clone, Serialize, Deserialize)]
415pub struct EnumDef {
416    pub name: String,
417    #[serde(default)]
418    pub doc_comment: Option<DocComment>,
419    pub type_params: Option<Vec<TypeParam>>,
420    pub members: Vec<EnumMember>,
421    /// Annotations applied to the enum (e.g., `@with_label() enum Color { ... }`)
422    #[serde(default)]
423    pub annotations: Vec<super::Annotation>,
424}
425
426#[derive(Debug, Clone, Serialize, Deserialize)]
427pub struct EnumMember {
428    pub name: String,
429    pub kind: EnumMemberKind,
430    #[serde(default)]
431    pub span: Span,
432    #[serde(default)]
433    pub doc_comment: Option<DocComment>,
434}
435
436#[derive(Debug, Clone, Serialize, Deserialize)]
437pub enum EnumMemberKind {
438    /// Unit variant: Variant or Variant = 1
439    Unit { value: Option<EnumValue> },
440    /// Tuple variant: Variant(Type, Type)
441    Tuple(Vec<TypeAnnotation>),
442    /// Struct variant: Variant { field: Type, ... }
443    Struct(Vec<ObjectTypeField>),
444}
445
446#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
447pub enum EnumValue {
448    String(String),
449    Number(f64),
450}
451
452/// A member of a trait definition: either required (signature only) or default (with body)
453#[derive(Debug, Clone, Serialize, Deserialize)]
454pub enum TraitMember {
455    /// Required method — implementors must provide this
456    Required(TraitMemberSignature),
457    /// Default method — used if implementor does not override
458    Default(MethodDef),
459    /// Associated type declaration: `type Item;` or `type Item: Comparable;`
460    AssociatedType {
461        name: String,
462        bounds: Vec<TypeAnnotation>,
463        #[serde(default)]
464        span: Span,
465        #[serde(default)]
466        doc_comment: Option<DocComment>,
467    },
468}
469
470impl TraitMember {
471    pub fn span(&self) -> Span {
472        match self {
473            TraitMember::Required(member) => member.span(),
474            TraitMember::Default(method) => method.span,
475            TraitMember::AssociatedType { span, .. } => *span,
476        }
477    }
478
479    pub fn doc_comment(&self) -> Option<&DocComment> {
480        match self {
481            TraitMember::Required(member) => member.doc_comment(),
482            TraitMember::Default(method) => method.doc_comment.as_ref(),
483            TraitMember::AssociatedType { doc_comment, .. } => doc_comment.as_ref(),
484        }
485    }
486}
487
488/// A concrete binding for an associated type inside an `impl` block:
489/// `type Item = number;`
490#[derive(Debug, Clone, Serialize, Deserialize)]
491pub struct AssociatedTypeBinding {
492    pub name: String,
493    pub concrete_type: TypeAnnotation,
494}
495
496/// Trait definition — supports required signatures and default methods.
497///
498/// ```shape
499/// trait Queryable<T> {
500///     filter(predicate: (T) => bool): Self    // required
501///     method execute() -> Result<Table<T>> {   // default
502///         return Ok(self.filter(|_| true))
503///     }
504/// }
505/// ```
506#[derive(Debug, Clone, Serialize, Deserialize)]
507pub struct TraitDef {
508    pub name: String,
509    #[serde(default)]
510    pub doc_comment: Option<DocComment>,
511    pub type_params: Option<Vec<TypeParam>>,
512    /// Supertrait bounds: `trait Foo: Bar + Baz { ... }`
513    #[serde(default)]
514    pub super_traits: Vec<TypeAnnotation>,
515    pub members: Vec<TraitMember>,
516    /// Annotations applied to the trait (e.g., `@documented("...") trait Foo { ... }`)
517    #[serde(default)]
518    pub annotations: Vec<super::Annotation>,
519    /// `comptime trait Foo { ... }` — parser-only flag (J-CT.0).
520    /// Type-checker validation lands in J-CT.1; comptime evaluator dispatch in
521    /// J-CT.2. Until those land, a `true` value here parses cleanly but has no
522    /// downstream semantic effect.
523    #[serde(default)]
524    pub is_comptime: bool,
525}
526
527/// Impl block — implements a trait for a type
528///
529/// ```shape
530/// impl Queryable<T> for Table<T> {
531///     method filter(predicate) { /* ... */ }
532///     method execute() { Ok(self) }
533/// }
534/// ```
535///
536/// Under the hood, compiles identically to an extend block (UFCS desugaring)
537/// plus trait validation (all required methods present with correct arities).
538#[derive(Debug, Clone, Serialize, Deserialize)]
539pub struct ImplBlock {
540    /// The trait being implemented (e.g., Queryable<T>)
541    pub trait_name: TypeName,
542    /// The type implementing the trait (e.g., Table<T>)
543    pub target_type: TypeName,
544    /// Optional named implementation selector:
545    /// `impl Display for User as JsonDisplay { ... }`
546    pub impl_name: Option<String>,
547    /// Method implementations
548    pub methods: Vec<MethodDef>,
549    /// Associated type bindings: `type Item = number;`
550    pub associated_type_bindings: Vec<AssociatedTypeBinding>,
551    /// Where clause: `where T: Display + Comparable`
552    pub where_clause: Option<Vec<WherePredicate>>,
553    /// `comptime impl Foo for Bar { ... }` — parser-only flag (J-CT.0).
554    /// Type-checker validation lands in J-CT.1; comptime evaluator dispatch in
555    /// J-CT.2. Until those land, a `true` value here parses cleanly but has no
556    /// downstream semantic effect.
557    #[serde(default)]
558    pub is_comptime: bool,
559}
560
561/// Type extension for adding methods to existing types
562#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
563pub struct ExtendStatement {
564    /// The type being extended (e.g., "Vec")
565    pub type_name: TypeName,
566    /// Methods being added to the type
567    pub methods: Vec<MethodDef>,
568}
569
570#[derive(Debug, Clone, Serialize, Deserialize)]
571pub struct MethodDef {
572    /// Method name
573    pub name: String,
574    #[serde(default)]
575    pub span: Span,
576    /// Declaring module path for compiler/runtime provenance checks.
577    ///
578    /// This is injected by the module loader for loaded modules and is not part
579    /// of user-authored source syntax.
580    #[serde(default, skip_serializing_if = "Option::is_none")]
581    pub declaring_module_path: Option<String>,
582    #[serde(default)]
583    pub doc_comment: Option<DocComment>,
584    /// Annotations applied to this method (e.g., `@traced`)
585    #[serde(default)]
586    pub annotations: Vec<super::functions::Annotation>,
587    /// Type parameters for generic methods (e.g., `method map<U>(...)`)
588    #[serde(default)]
589    pub type_params: Option<Vec<TypeParam>>,
590    /// Method parameters
591    pub params: Vec<super::functions::FunctionParameter>,
592    /// Optional when clause for conditional method definitions
593    pub when_clause: Option<Box<super::expressions::Expr>>,
594    /// Optional return type annotation
595    pub return_type: Option<TypeAnnotation>,
596    /// Method body
597    pub body: Vec<super::statements::Statement>,
598    /// Whether this is an async method
599    pub is_async: bool,
600}
601
602impl PartialEq for MethodDef {
603    fn eq(&self, other: &Self) -> bool {
604        self.name == other.name
605            && self.doc_comment == other.doc_comment
606            && self.annotations == other.annotations
607            && self.type_params == other.type_params
608            && self.params == other.params
609            && self.when_clause == other.when_clause
610            && self.return_type == other.return_type
611            && self.body == other.body
612            && self.is_async == other.is_async
613    }
614}
615
616#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
617pub enum TypeName {
618    /// Simple type name (e.g., "Vec", "Table", "foo::Bar")
619    Simple(TypePath),
620    /// Generic type name (e.g., "Table<Row>")
621    Generic {
622        name: TypePath,
623        type_args: Vec<TypeAnnotation>,
624    },
625}
626
627// ============================================================================
628// Struct Type Definitions
629// ============================================================================
630
631/// Struct type definition — pure data with named fields
632///
633/// ```shape
634/// type Point { x: number, y: number }
635/// type DataVec<V, K = Timestamp> { index: Vec<K>, data: Vec<V> }
636/// ```
637#[derive(Debug, Clone, Serialize, Deserialize)]
638pub struct StructTypeDef {
639    pub name: String,
640    #[serde(default)]
641    pub doc_comment: Option<DocComment>,
642    pub type_params: Option<Vec<TypeParam>>,
643    pub fields: Vec<StructField>,
644    /// Inline method definitions inside the type body
645    #[serde(default)]
646    pub methods: Vec<MethodDef>,
647    /// Annotations applied to the struct (e.g., `@derive_debug type Foo { ... }`)
648    pub annotations: Vec<Annotation>,
649    /// Optional native layout metadata for `type C`.
650    #[serde(default)]
651    pub native_layout: Option<NativeLayoutBinding>,
652}
653
654/// Native layout binding metadata for `type C` declarations.
655#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
656pub struct NativeLayoutBinding {
657    /// ABI name (currently `"C"`).
658    pub abi: String,
659}
660
661/// A field in a struct type definition
662///
663/// Comptime fields are type-level constants baked at compile time.
664/// They occupy zero runtime slots (no ValueSlot in TypedObject).
665/// Access resolves to a constant push at compile time.
666///
667/// ```shape
668/// type Currency {
669///     comptime symbol: string = "$",
670///     comptime decimals: number = 2,
671///     amount: number,
672/// }
673/// ```
674#[derive(Debug, Clone, Serialize, Deserialize)]
675pub struct StructField {
676    pub annotations: Vec<Annotation>,
677    pub is_comptime: bool,
678    pub name: String,
679    #[serde(default)]
680    pub span: Span,
681    #[serde(default)]
682    pub doc_comment: Option<DocComment>,
683    pub type_annotation: TypeAnnotation,
684    pub default_value: Option<super::expressions::Expr>,
685}
686
687impl PartialEq for StructField {
688    fn eq(&self, other: &Self) -> bool {
689        self.annotations == other.annotations
690            && self.is_comptime == other.is_comptime
691            && self.name == other.name
692            && self.doc_comment == other.doc_comment
693            && self.type_annotation == other.type_annotation
694            && self.default_value == other.default_value
695    }
696}