Skip to main content

shape_ast/parser/
extensions.rs

1//! AST parser for language extensions
2
3use crate::ast::{
4    AnnotationDef, AnnotationHandler, AnnotationHandlerParam, AnnotationHandlerType,
5    FunctionParameter,
6};
7use crate::error::{Result, ShapeError};
8use crate::parser::Rule;
9use crate::parser::expressions::control_flow::parse_block_expr;
10use crate::parser::pair_span;
11use crate::parser::types::parse_type_annotation;
12use pest::iterators::Pair;
13
14/// Parse an annotation definition with lifecycle handlers
15///
16/// Grammar:
17/// ```pest
18/// annotation_def = {
19///     "annotation" ~ ident ~ "(" ~ annotation_def_params? ~ ")" ~ "{" ~ annotation_body ~ "}"
20/// }
21/// annotation_body = { annotation_handler* }
22/// annotation_handler = {
23///     annotation_handler_name ~ "(" ~ annotation_handler_params? ~ ")" ~ return_type? ~ block_expr
24/// }
25/// ```
26pub fn parse_annotation_def(pair: Pair<Rule>) -> Result<AnnotationDef> {
27    let span = pair_span(&pair);
28    let mut inner = pair.into_inner();
29
30    let name_pair = inner.next().ok_or_else(|| ShapeError::ParseError {
31        message: "Missing annotation name".to_string(),
32        location: None,
33    })?;
34    let name = name_pair.as_str().to_string();
35    let name_span = pair_span(&name_pair);
36
37    let mut params: Vec<FunctionParameter> = Vec::new();
38    let mut handlers = Vec::new();
39    let mut allowed_targets: Option<Vec<crate::ast::AnnotationTargetKind>> = None;
40
41    for part in inner {
42        match part.as_rule() {
43            Rule::annotation_def_params => {
44                for param_pair in part.into_inner() {
45                    if param_pair.as_rule() == Rule::ident {
46                        let pattern = crate::ast::DestructurePattern::Identifier(
47                            param_pair.as_str().to_string(),
48                            pair_span(&param_pair),
49                        );
50                        params.push(FunctionParameter {
51                            pattern,
52                            is_const: false,
53                            is_reference: false,
54                            is_mut_reference: false,
55                            type_annotation: None,
56                            default_value: None,
57                        });
58                    }
59                }
60            }
61            Rule::annotation_body => {
62                for body_item in part.into_inner() {
63                    let item = if body_item.as_rule() == Rule::annotation_body_item {
64                        body_item
65                            .into_inner()
66                            .next()
67                            .ok_or_else(|| ShapeError::ParseError {
68                                message: "empty annotation body item".to_string(),
69                                location: None,
70                            })?
71                    } else {
72                        body_item
73                    };
74
75                    match item.as_rule() {
76                        Rule::annotation_handler => {
77                            handlers.push(parse_annotation_handler(item)?);
78                        }
79                        Rule::annotation_targets_decl => {
80                            let mut targets = Vec::new();
81                            for target_pair in item.into_inner() {
82                                if target_pair.as_rule() != Rule::annotation_target_kind {
83                                    continue;
84                                }
85                                let kind = match target_pair.as_str() {
86                                    "function" => crate::ast::AnnotationTargetKind::Function,
87                                    "type" => crate::ast::AnnotationTargetKind::Type,
88                                    "module" => crate::ast::AnnotationTargetKind::Module,
89                                    "expression" => crate::ast::AnnotationTargetKind::Expression,
90                                    "block" => crate::ast::AnnotationTargetKind::Block,
91                                    "await_expr" => crate::ast::AnnotationTargetKind::AwaitExpr,
92                                    "binding" => crate::ast::AnnotationTargetKind::Binding,
93                                    other => {
94                                        return Err(ShapeError::ParseError {
95                                            message: format!(
96                                                "unknown annotation target kind '{}'",
97                                                other
98                                            ),
99                                            location: None,
100                                        });
101                                    }
102                                };
103                                targets.push(kind);
104                            }
105                            allowed_targets = Some(targets);
106                        }
107                        _ => {}
108                    }
109                }
110            }
111            _ => {}
112        }
113    }
114
115    Ok(AnnotationDef {
116        name,
117        name_span,
118        params,
119        allowed_targets,
120        handlers,
121        span,
122    })
123}
124
125/// Parse a single annotation lifecycle handler
126fn parse_annotation_handler(pair: Pair<Rule>) -> Result<AnnotationHandler> {
127    let span = pair_span(&pair);
128    let mut inner = pair.into_inner();
129
130    // Parse handler type (on_define, before, after, metadata, comptime pre/post)
131    let handler_name_pair = inner.next().ok_or_else(|| ShapeError::ParseError {
132        message: "Missing annotation handler name".to_string(),
133        location: None,
134    })?;
135
136    let handler_kind = handler_name_pair
137        .as_str()
138        .split_whitespace()
139        .collect::<Vec<_>>()
140        .join(" ");
141
142    let handler_type = match handler_kind.as_str() {
143        "on_define" => AnnotationHandlerType::OnDefine,
144        "before" => AnnotationHandlerType::Before,
145        "after" => AnnotationHandlerType::After,
146        "metadata" => AnnotationHandlerType::Metadata,
147        "comptime pre" => AnnotationHandlerType::ComptimePre,
148        "comptime post" => AnnotationHandlerType::ComptimePost,
149        other => {
150            return Err(ShapeError::ParseError {
151                message: format!(
152                    "unknown annotation handler type '{}'. Expected `on_define`, `before`, `after`, `metadata`, `comptime pre`, or `comptime post`",
153                    other
154                ),
155                location: None,
156            });
157        }
158    };
159
160    // Parse handler parameters
161    let mut params = Vec::new();
162    let mut return_type = None;
163    let mut body = None;
164
165    for part in inner {
166        match part.as_rule() {
167            Rule::annotation_handler_params => {
168                for param_pair in part.into_inner() {
169                    if param_pair.as_rule() == Rule::annotation_handler_param {
170                        let raw = param_pair.as_str().trim();
171                        let (name, is_variadic) = if let Some(rest) = raw.strip_prefix("...") {
172                            (rest.trim().to_string(), true)
173                        } else {
174                            (raw.to_string(), false)
175                        };
176                        if !name.is_empty() {
177                            params.push(AnnotationHandlerParam { name, is_variadic });
178                        }
179                    }
180                }
181            }
182            Rule::return_type => {
183                // return_type = "->" ~ type_annotation
184                if let Some(type_pair) = part.into_inner().next() {
185                    return_type = Some(parse_type_annotation(type_pair)?);
186                }
187            }
188            Rule::block_expr => {
189                body = Some(parse_block_expr(part)?);
190            }
191            _ => {}
192        }
193    }
194
195    let body = body.ok_or_else(|| ShapeError::ParseError {
196        message: "Missing annotation handler body".to_string(),
197        location: None,
198    })?;
199
200    Ok(AnnotationHandler {
201        handler_type,
202        params,
203        return_type,
204        body,
205        span,
206    })
207}
208
209pub fn parse_extend_statement(pair: Pair<Rule>) -> Result<crate::ast::ExtendStatement> {
210    use crate::ast::types::TypeName;
211
212    let mut inner = pair.into_inner();
213
214    // Parse type_name: ident ~ ("<" ~ type_annotation ~ ">")?
215    let type_name_pair = inner.next().ok_or_else(|| ShapeError::ParseError {
216        message: "Missing type name in extend statement".to_string(),
217        location: None,
218    })?;
219
220    let type_name = {
221        let mut tn_inner = type_name_pair.into_inner();
222        let name_pair = tn_inner.next().ok_or_else(|| ShapeError::ParseError {
223            message: "Missing type name identifier".to_string(),
224            location: None,
225        })?;
226        let name = name_pair.as_str().to_string();
227        let type_args: Vec<_> = tn_inner
228            .map(|p| parse_type_annotation(p))
229            .collect::<Result<Vec<_>>>()?;
230        if type_args.is_empty() {
231            TypeName::Simple(name)
232        } else {
233            TypeName::Generic { name, type_args }
234        }
235    };
236
237    // Parse method_def*
238    let mut methods = Vec::new();
239    for method_pair in inner {
240        if method_pair.as_rule() == Rule::method_def {
241            methods.push(super::types::parse_method_def_shared(method_pair)?);
242        }
243    }
244
245    Ok(crate::ast::ExtendStatement { type_name, methods })
246}
247
248/// Parse an impl block: `impl TraitName for TypeName [as ImplName] { method_def* }`
249///
250/// Reuses the same method_def parsing as extend blocks.
251pub fn parse_impl_block(pair: Pair<Rule>) -> Result<crate::ast::ImplBlock> {
252    let mut inner = pair.into_inner();
253
254    // First type_name is the trait being implemented
255    let trait_name_pair = inner.next().ok_or_else(|| ShapeError::ParseError {
256        message: "Missing trait name in impl block".to_string(),
257        location: None,
258    })?;
259    let trait_name = parse_type_name(trait_name_pair)?;
260
261    // Second type_name is the target type
262    let target_type_pair = inner.next().ok_or_else(|| ShapeError::ParseError {
263        message: "Missing target type in impl block".to_string(),
264        location: None,
265    })?;
266    let target_type = parse_type_name(target_type_pair)?;
267
268    // Parse optional impl_name, where_clause and impl_member*
269    let mut impl_name = None;
270    let mut methods = Vec::new();
271    let mut associated_type_bindings = Vec::new();
272    let mut where_clause = None;
273    for member_pair in inner {
274        if member_pair.as_rule() == Rule::impl_name {
275            let name_pair =
276                member_pair
277                    .into_inner()
278                    .next()
279                    .ok_or_else(|| ShapeError::ParseError {
280                        message: "Missing impl name after 'as'".to_string(),
281                        location: None,
282                    })?;
283            impl_name = Some(name_pair.as_str().to_string());
284            continue;
285        }
286        if member_pair.as_rule() == Rule::where_clause {
287            where_clause = Some(super::functions::parse_where_clause(member_pair)?);
288            continue;
289        }
290        if member_pair.as_rule() != Rule::impl_member {
291            continue;
292        }
293        for child in member_pair.into_inner() {
294            match child.as_rule() {
295                Rule::associated_type_binding => {
296                    let binding = parse_associated_type_binding(child)?;
297                    associated_type_bindings.push(binding);
298                }
299                Rule::method_def => {
300                    methods.push(super::types::parse_method_def_shared(child)?);
301                }
302                _ => {}
303            }
304        }
305    }
306
307    Ok(crate::ast::ImplBlock {
308        trait_name,
309        target_type,
310        impl_name,
311        methods,
312        associated_type_bindings,
313        where_clause,
314    })
315}
316
317/// Parse `type Item = number;`
318fn parse_associated_type_binding(
319    pair: Pair<Rule>,
320) -> Result<crate::ast::types::AssociatedTypeBinding> {
321    let mut inner = pair.into_inner();
322
323    let name_pair = inner.next().ok_or_else(|| ShapeError::ParseError {
324        message: "expected associated type name in binding".to_string(),
325        location: None,
326    })?;
327    let name = name_pair.as_str().to_string();
328
329    let type_pair = inner.next().ok_or_else(|| ShapeError::ParseError {
330        message: format!("expected type annotation for associated type '{}'", name),
331        location: None,
332    })?;
333    let concrete_type = super::types::parse_type_annotation(type_pair)?;
334
335    Ok(crate::ast::types::AssociatedTypeBinding {
336        name,
337        concrete_type,
338    })
339}
340
341/// Parse a type_name from the grammar (ident with optional generic args)
342fn parse_type_name(pair: Pair<Rule>) -> Result<crate::ast::types::TypeName> {
343    use crate::ast::types::TypeName;
344
345    let mut tn_inner = pair.into_inner();
346    let name_pair = tn_inner.next().ok_or_else(|| ShapeError::ParseError {
347        message: "Missing type name identifier".to_string(),
348        location: None,
349    })?;
350    let name = name_pair.as_str().to_string();
351    let type_args: Vec<_> = tn_inner
352        .map(|p| super::types::parse_type_annotation(p))
353        .collect::<Result<Vec<_>>>()?;
354    if type_args.is_empty() {
355        Ok(TypeName::Simple(name))
356    } else {
357        Ok(TypeName::Generic { name, type_args })
358    }
359}
360
361pub fn parse_optimize_statement(_pair: Pair<Rule>) -> Result<crate::ast::OptimizeStatement> {
362    Err(ShapeError::ParseError {
363        message: "parse_optimize_statement not yet implemented".to_string(),
364        location: None,
365    })
366}