Skip to main content

lashlang/
ast.rs

1use compact_str::CompactString;
2use serde::{Deserialize, Serialize};
3use std::fmt;
4
5use crate::lexer::Span;
6
7pub type AstString = CompactString;
8
9#[derive(Clone, Debug, Serialize, Deserialize)]
10pub struct Program {
11    #[serde(default, skip_serializing_if = "Vec::is_empty")]
12    pub declarations: Vec<Declaration>,
13    pub main: Expr,
14    #[serde(default, skip_serializing_if = "Vec::is_empty")]
15    pub declaration_spans: Vec<Span>,
16    #[serde(default, skip_serializing_if = "Vec::is_empty")]
17    pub expression_spans: Vec<Span>,
18}
19
20impl Program {
21    pub fn block(expressions: Vec<Expr>) -> Self {
22        Self {
23            declarations: Vec::new(),
24            main: Expr::Block(expressions),
25            declaration_spans: Vec::new(),
26            expression_spans: Vec::new(),
27        }
28    }
29
30    pub(crate) fn module_with_spans(
31        declarations: Vec<Declaration>,
32        declaration_spans: Vec<Span>,
33        expressions: Vec<Expr>,
34        expression_spans: Vec<Span>,
35    ) -> Self {
36        Self {
37            declarations,
38            main: Expr::Block(expressions),
39            declaration_spans,
40            expression_spans,
41        }
42    }
43
44    pub fn process(&self, name: &str) -> Option<&ProcessDecl> {
45        self.declarations
46            .iter()
47            .find_map(|declaration| match declaration {
48                Declaration::Process(process) if process.name.as_str() == name => Some(process),
49                _ => None,
50            })
51    }
52}
53
54impl PartialEq for Program {
55    fn eq(&self, other: &Self) -> bool {
56        self.declarations == other.declarations && self.main == other.main
57    }
58}
59
60#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
61pub enum Declaration {
62    Type(TypeDecl),
63    Process(ProcessDecl),
64}
65
66#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
67pub struct TypeDecl {
68    pub name: AstString,
69    pub ty: TypeExpr,
70}
71
72#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
73pub struct ProcessDecl {
74    pub name: AstString,
75    pub params: Vec<ProcessParam>,
76    #[serde(default, skip_serializing_if = "Vec::is_empty")]
77    pub signals: Vec<ProcessSignalDecl>,
78    #[serde(default, skip_serializing_if = "Option::is_none")]
79    pub return_ty: Option<TypeExpr>,
80    #[serde(default, skip_serializing_if = "Option::is_none")]
81    pub label: Option<LabelMetadata>,
82    pub body: Expr,
83}
84
85#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
86pub struct ProcessParam {
87    pub name: AstString,
88    pub ty: TypeExpr,
89}
90
91#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
92pub struct ProcessSignalDecl {
93    pub name: AstString,
94    pub ty: TypeExpr,
95}
96
97#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
98pub struct LabelMetadata {
99    pub title: AstString,
100    #[serde(default, skip_serializing_if = "Option::is_none")]
101    pub description: Option<AstString>,
102}
103
104#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
105pub struct AssignTarget {
106    pub root: AstString,
107    #[serde(default, skip_serializing_if = "Vec::is_empty")]
108    pub steps: Vec<AssignPathStep>,
109}
110
111impl AssignTarget {
112    pub fn variable(root: AstString) -> Self {
113        Self {
114            root,
115            steps: Vec::new(),
116        }
117    }
118
119    pub fn is_simple(&self) -> bool {
120        self.steps.is_empty()
121    }
122}
123
124#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
125pub enum AssignPathStep {
126    Field(AstString),
127    Index(Expr),
128}
129
130#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
131pub enum Expr {
132    Block(Vec<Expr>),
133    LabelAnnotated {
134        label: LabelMetadata,
135        expr: Box<Expr>,
136    },
137    Null,
138    Bool(bool),
139    Number(f64),
140    String(AstString),
141    Variable(AstString),
142    List(Vec<Expr>),
143    Record(Vec<(AstString, Expr)>),
144    Assign {
145        target: AssignTarget,
146        expr: Box<Expr>,
147    },
148    If {
149        condition: Box<Expr>,
150        then_block: Box<Expr>,
151        else_block: Box<Expr>,
152    },
153    For {
154        binding: AstString,
155        iterable: Box<Expr>,
156        body: Box<Expr>,
157    },
158    While {
159        condition: Box<Expr>,
160        body: Box<Expr>,
161    },
162    Break,
163    Continue,
164    StartProcess(ProcessStartExpr),
165    ProcessRef {
166        process: AstString,
167    },
168    HostDescriptorConstructor {
169        type_name: AstString,
170        input: Box<Expr>,
171    },
172    ResourceRef(ResourceRefExpr),
173    ReceiverCall {
174        receiver: Box<Expr>,
175        operation: AstString,
176        args: Vec<Expr>,
177    },
178    Await(Box<Expr>),
179    SleepFor(Box<Expr>),
180    SleepUntil(Box<Expr>),
181    WaitSignal {
182        name: AstString,
183    },
184    SignalRun {
185        run: Box<Expr>,
186        name: AstString,
187        payload: Box<Expr>,
188    },
189    ResultUnwrap(Box<Expr>),
190    Cancel(Box<Expr>),
191    Print(Box<Expr>),
192    Submit(Option<Box<Expr>>),
193    Yield(Box<Expr>),
194    Wake(Box<Expr>),
195    Finish(Option<Box<Expr>>),
196    Fail(Box<Expr>),
197    BuiltinCall {
198        name: AstString,
199        args: Vec<Expr>,
200    },
201    Field {
202        target: Box<Expr>,
203        field: AstString,
204    },
205    Index {
206        target: Box<Expr>,
207        index: Box<Expr>,
208    },
209    Unary {
210        op: UnaryOp,
211        expr: Box<Expr>,
212    },
213    Binary {
214        left: Box<Expr>,
215        op: BinaryOp,
216        right: Box<Expr>,
217    },
218    TypeLiteral(Box<TypeExpr>),
219}
220
221impl Expr {
222    /// Yields every direct child expression of `self` in evaluation order.
223    ///
224    /// This is the single structural-traversal primitive: any pass that only
225    /// needs to recurse into the sub-expressions of a node (without caring
226    /// about the node's own kind) can fold over `children()` instead of
227    /// re-spelling the full `match`. Leaf nodes (`Null`, `Bool`, `Number`,
228    /// `String`, `Variable`, `Break`, `Continue`, `WaitSignal`,
229    /// `ResourceRef`, `ProcessRef`, `HostDescriptorConstructor` metadata, and
230    /// `TypeLiteral`) yield nothing.
231    ///
232    /// `Assign` includes any dynamic index expressions in its `target` path
233    /// (in path order) before the assigned value, matching the order in which
234    /// the compiler and linker visit them.
235    pub fn children(&self) -> ExprChildren<'_> {
236        let mut buffer = SmallExprVec::new();
237        match self {
238            Expr::Null
239            | Expr::Bool(_)
240            | Expr::Number(_)
241            | Expr::String(_)
242            | Expr::Variable(_)
243            | Expr::Break
244            | Expr::Continue
245            | Expr::WaitSignal { .. }
246            | Expr::ProcessRef { .. }
247            | Expr::ResourceRef(_)
248            | Expr::TypeLiteral(_) => {}
249            Expr::Block(expressions) | Expr::List(expressions) => {
250                buffer.extend(expressions.iter());
251            }
252            Expr::LabelAnnotated { expr, .. } => buffer.push(expr),
253            Expr::Record(entries) => buffer.extend(entries.iter().map(|(_, value)| value)),
254            Expr::Assign { target, expr } => {
255                for step in &target.steps {
256                    if let AssignPathStep::Index(index) = step {
257                        buffer.push(index);
258                    }
259                }
260                buffer.push(expr);
261            }
262            Expr::If {
263                condition,
264                then_block,
265                else_block,
266            } => {
267                buffer.push(condition);
268                buffer.push(then_block);
269                buffer.push(else_block);
270            }
271            Expr::For { iterable, body, .. } => {
272                buffer.push(iterable);
273                buffer.push(body);
274            }
275            Expr::While { condition, body } => {
276                buffer.push(condition);
277                buffer.push(body);
278            }
279            Expr::StartProcess(start) => buffer.extend(start.args.iter().map(|(_, value)| value)),
280            Expr::HostDescriptorConstructor { input, .. } => buffer.push(input),
281            Expr::ReceiverCall { receiver, args, .. } => {
282                buffer.push(receiver);
283                buffer.extend(args.iter());
284            }
285            Expr::SignalRun { run, payload, .. } => {
286                buffer.push(run);
287                buffer.push(payload);
288            }
289            Expr::Await(expr)
290            | Expr::SleepFor(expr)
291            | Expr::SleepUntil(expr)
292            | Expr::ResultUnwrap(expr)
293            | Expr::Cancel(expr)
294            | Expr::Print(expr)
295            | Expr::Yield(expr)
296            | Expr::Wake(expr)
297            | Expr::Fail(expr)
298            | Expr::Unary { expr, .. } => buffer.push(expr),
299            Expr::Submit(expr) | Expr::Finish(expr) => {
300                if let Some(expr) = expr {
301                    buffer.push(expr);
302                }
303            }
304            Expr::BuiltinCall { args, .. } => buffer.extend(args.iter()),
305            Expr::Field { target, .. } => buffer.push(target),
306            Expr::Index { target, index } => {
307                buffer.push(target);
308                buffer.push(index);
309            }
310            Expr::Binary { left, right, .. } => {
311                buffer.push(left);
312                buffer.push(right);
313            }
314        }
315        ExprChildren {
316            buffer,
317            position: 0,
318        }
319    }
320}
321
322type SmallExprVec<'expr> = smallvec::SmallVec<[&'expr Expr; 3]>;
323
324/// Iterator over the direct child expressions yielded by [`Expr::children`].
325pub struct ExprChildren<'expr> {
326    buffer: SmallExprVec<'expr>,
327    position: usize,
328}
329
330impl<'expr> Iterator for ExprChildren<'expr> {
331    type Item = &'expr Expr;
332
333    fn next(&mut self) -> Option<Self::Item> {
334        let item = self.buffer.get(self.position).copied();
335        if item.is_some() {
336            self.position += 1;
337        }
338        item
339    }
340
341    fn size_hint(&self) -> (usize, Option<usize>) {
342        let remaining = self.buffer.len() - self.position;
343        (remaining, Some(remaining))
344    }
345}
346
347impl ExactSizeIterator for ExprChildren<'_> {}
348
349pub trait ExprVisitor {
350    fn visit_expr(&mut self, expr: &Expr) {
351        walk_expr(self, expr);
352    }
353}
354
355pub fn walk_expr<V>(visitor: &mut V, expr: &Expr)
356where
357    V: ExprVisitor + ?Sized,
358{
359    for child in expr.children() {
360        visitor.visit_expr(child);
361    }
362}
363
364pub trait ExprFolder {
365    fn fold_expr(&mut self, expr: Expr) -> Expr {
366        fold_expr_children(self, expr)
367    }
368}
369
370pub fn fold_expr_children<F>(folder: &mut F, expr: Expr) -> Expr
371where
372    F: ExprFolder + ?Sized,
373{
374    match expr {
375        Expr::Block(expressions) => Expr::Block(
376            expressions
377                .into_iter()
378                .map(|expr| folder.fold_expr(expr))
379                .collect(),
380        ),
381        Expr::LabelAnnotated { label, expr } => Expr::LabelAnnotated {
382            label,
383            expr: Box::new(folder.fold_expr(*expr)),
384        },
385        Expr::List(items) => Expr::List(
386            items
387                .into_iter()
388                .map(|expr| folder.fold_expr(expr))
389                .collect(),
390        ),
391        Expr::Record(entries) => Expr::Record(
392            entries
393                .into_iter()
394                .map(|(name, value)| (name, folder.fold_expr(value)))
395                .collect(),
396        ),
397        Expr::Assign { target, expr } => Expr::Assign {
398            target: fold_assign_target(folder, target),
399            expr: Box::new(folder.fold_expr(*expr)),
400        },
401        Expr::If {
402            condition,
403            then_block,
404            else_block,
405        } => Expr::If {
406            condition: Box::new(folder.fold_expr(*condition)),
407            then_block: Box::new(folder.fold_expr(*then_block)),
408            else_block: Box::new(folder.fold_expr(*else_block)),
409        },
410        Expr::For {
411            binding,
412            iterable,
413            body,
414        } => Expr::For {
415            binding,
416            iterable: Box::new(folder.fold_expr(*iterable)),
417            body: Box::new(folder.fold_expr(*body)),
418        },
419        Expr::While { condition, body } => Expr::While {
420            condition: Box::new(folder.fold_expr(*condition)),
421            body: Box::new(folder.fold_expr(*body)),
422        },
423        Expr::StartProcess(mut start) => {
424            start.args = start
425                .args
426                .into_iter()
427                .map(|(name, value)| (name, folder.fold_expr(value)))
428                .collect();
429            Expr::StartProcess(start)
430        }
431        Expr::ProcessRef { process } => Expr::ProcessRef { process },
432        Expr::HostDescriptorConstructor { type_name, input } => Expr::HostDescriptorConstructor {
433            type_name,
434            input: Box::new(folder.fold_expr(*input)),
435        },
436        Expr::ReceiverCall {
437            receiver,
438            operation,
439            args,
440        } => Expr::ReceiverCall {
441            receiver: Box::new(folder.fold_expr(*receiver)),
442            operation,
443            args: args
444                .into_iter()
445                .map(|expr| folder.fold_expr(expr))
446                .collect(),
447        },
448        Expr::Await(expr) => Expr::Await(Box::new(folder.fold_expr(*expr))),
449        Expr::SleepFor(expr) => Expr::SleepFor(Box::new(folder.fold_expr(*expr))),
450        Expr::SleepUntil(expr) => Expr::SleepUntil(Box::new(folder.fold_expr(*expr))),
451        Expr::SignalRun { run, name, payload } => Expr::SignalRun {
452            run: Box::new(folder.fold_expr(*run)),
453            name,
454            payload: Box::new(folder.fold_expr(*payload)),
455        },
456        Expr::ResultUnwrap(expr) => Expr::ResultUnwrap(Box::new(folder.fold_expr(*expr))),
457        Expr::Cancel(expr) => Expr::Cancel(Box::new(folder.fold_expr(*expr))),
458        Expr::Print(expr) => Expr::Print(Box::new(folder.fold_expr(*expr))),
459        Expr::Submit(expr) => Expr::Submit(expr.map(|expr| Box::new(folder.fold_expr(*expr)))),
460        Expr::Yield(expr) => Expr::Yield(Box::new(folder.fold_expr(*expr))),
461        Expr::Wake(expr) => Expr::Wake(Box::new(folder.fold_expr(*expr))),
462        Expr::Finish(expr) => Expr::Finish(expr.map(|expr| Box::new(folder.fold_expr(*expr)))),
463        Expr::Fail(expr) => Expr::Fail(Box::new(folder.fold_expr(*expr))),
464        Expr::BuiltinCall { name, args } => Expr::BuiltinCall {
465            name,
466            args: args
467                .into_iter()
468                .map(|expr| folder.fold_expr(expr))
469                .collect(),
470        },
471        Expr::Field { target, field } => Expr::Field {
472            target: Box::new(folder.fold_expr(*target)),
473            field,
474        },
475        Expr::Index { target, index } => Expr::Index {
476            target: Box::new(folder.fold_expr(*target)),
477            index: Box::new(folder.fold_expr(*index)),
478        },
479        Expr::Unary { op, expr } => Expr::Unary {
480            op,
481            expr: Box::new(folder.fold_expr(*expr)),
482        },
483        Expr::Binary { left, op, right } => Expr::Binary {
484            left: Box::new(folder.fold_expr(*left)),
485            op,
486            right: Box::new(folder.fold_expr(*right)),
487        },
488        leaf @ (Expr::Null
489        | Expr::Bool(_)
490        | Expr::Number(_)
491        | Expr::String(_)
492        | Expr::Variable(_)
493        | Expr::Break
494        | Expr::Continue
495        | Expr::ResourceRef(_)
496        | Expr::WaitSignal { .. }
497        | Expr::TypeLiteral(_)) => leaf,
498    }
499}
500
501fn fold_assign_target<F>(folder: &mut F, target: AssignTarget) -> AssignTarget
502where
503    F: ExprFolder + ?Sized,
504{
505    AssignTarget {
506        root: target.root,
507        steps: target
508            .steps
509            .into_iter()
510            .map(|step| match step {
511                AssignPathStep::Field(field) => AssignPathStep::Field(field),
512                AssignPathStep::Index(index) => AssignPathStep::Index(folder.fold_expr(index)),
513            })
514            .collect(),
515    }
516}
517
518#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
519pub enum TypeExpr {
520    Any,
521    Str,
522    Int,
523    Float,
524    Bool,
525    Dict,
526    /// The literal `null` type; usually only useful as part of a
527    /// `Union` (e.g. `str | null` for a nullable field).
528    Null,
529    Enum(Vec<AstString>),
530    List(Box<TypeExpr>),
531    Object(Vec<TypeField>),
532    Ref(AstString),
533    Process {
534        input: Box<TypeExpr>,
535        output: Box<TypeExpr>,
536        input_count: usize,
537    },
538    TriggerHandle(Box<TypeExpr>),
539    /// Union of alternative type shapes, e.g. `str | int | null`.
540    /// Always has two or more variants; single-variant parses collapse
541    /// to the underlying `TypeExpr` in the parser.
542    Union(Vec<TypeExpr>),
543}
544
545pub fn format_type_expr(ty: &TypeExpr) -> String {
546    match ty {
547        TypeExpr::Any => "any".to_string(),
548        TypeExpr::Str => "str".to_string(),
549        TypeExpr::Int => "int".to_string(),
550        TypeExpr::Float => "float".to_string(),
551        TypeExpr::Bool => "bool".to_string(),
552        TypeExpr::Dict => "dict".to_string(),
553        TypeExpr::Null => "null".to_string(),
554        TypeExpr::Enum(values) => format!(
555            "enum[{}]",
556            values
557                .iter()
558                .map(|value| format!("\"{value}\""))
559                .collect::<Vec<_>>()
560                .join(", ")
561        ),
562        TypeExpr::List(item) => format!("list[{}]", format_type_expr(item)),
563        TypeExpr::Object(fields) => {
564            let fields = fields
565                .iter()
566                .map(|field| {
567                    let optional = if field.optional { "?" } else { "" };
568                    format!(
569                        "{}: {}{}",
570                        field.name,
571                        format_type_expr(&field.ty),
572                        optional
573                    )
574                })
575                .collect::<Vec<_>>()
576                .join(", ");
577            format!("{{ {fields} }}")
578        }
579        TypeExpr::Ref(name) => name.to_string(),
580        TypeExpr::Process { input, output, .. } => {
581            format!(
582                "Process<{}, {}>",
583                format_type_expr(input),
584                format_type_expr(output)
585            )
586        }
587        TypeExpr::TriggerHandle(event) => {
588            format!("TriggerHandle<{}>", format_type_expr(event))
589        }
590        TypeExpr::Union(items) => items
591            .iter()
592            .map(format_type_expr)
593            .collect::<Vec<_>>()
594            .join(" | "),
595    }
596}
597
598impl fmt::Display for TypeExpr {
599    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
600        f.write_str(&format_type_expr(self))
601    }
602}
603
604#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
605pub struct TypeField {
606    pub name: AstString,
607    pub ty: TypeExpr,
608    pub optional: bool,
609}
610
611#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
612pub struct ProcessStartExpr {
613    pub process: AstString,
614    pub args: Vec<(AstString, Expr)>,
615}
616
617#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
618pub struct ResourceRefExpr {
619    #[serde(default, skip_serializing_if = "Vec::is_empty")]
620    pub path: Vec<AstString>,
621    pub resource_type: AstString,
622    pub alias: AstString,
623}
624
625impl ResourceRefExpr {
626    pub fn unresolved(path: Vec<AstString>) -> Self {
627        Self {
628            path,
629            resource_type: AstString::default(),
630            alias: AstString::default(),
631        }
632    }
633
634    pub fn resolved(
635        path: Vec<AstString>,
636        resource_type: impl Into<AstString>,
637        alias: impl Into<AstString>,
638    ) -> Self {
639        Self {
640            path,
641            resource_type: resource_type.into(),
642            alias: alias.into(),
643        }
644    }
645
646    pub fn path_string(&self) -> String {
647        if self.path.is_empty() {
648            format!("{}.{}", self.resource_type, self.alias)
649        } else {
650            self.path
651                .iter()
652                .map(AstString::as_str)
653                .collect::<Vec<_>>()
654                .join(".")
655        }
656    }
657}
658
659#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
660pub enum UnaryOp {
661    Negate,
662    Not,
663}
664
665#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
666pub enum BinaryOp {
667    Add,
668    Subtract,
669    Multiply,
670    Divide,
671    Modulo,
672    Equal,
673    NotEqual,
674    Less,
675    LessEqual,
676    Greater,
677    GreaterEqual,
678    And,
679    Or,
680}
681
682#[cfg(test)]
683mod tests {
684    use super::*;
685
686    #[test]
687    fn type_expr_formatting_covers_nested_shapes() {
688        let ty = TypeExpr::Object(vec![
689            TypeField {
690                name: "status".into(),
691                ty: TypeExpr::Enum(vec!["ok".into(), "err".into()]),
692                optional: false,
693            },
694            TypeField {
695                name: "tags".into(),
696                ty: TypeExpr::List(Box::new(TypeExpr::Str)),
697                optional: true,
698            },
699            TypeField {
700                name: "owner".into(),
701                ty: TypeExpr::Ref("User".into()),
702                optional: false,
703            },
704            TypeField {
705                name: "value".into(),
706                ty: TypeExpr::Union(vec![TypeExpr::Int, TypeExpr::Null]),
707                optional: false,
708            },
709        ]);
710
711        assert_eq!(
712            format_type_expr(&ty),
713            r#"{ status: enum["ok", "err"], tags: list[str]?, owner: User, value: int | null }"#
714        );
715        assert_eq!(ty.to_string(), format_type_expr(&ty));
716    }
717
718    fn var(name: &str) -> Expr {
719        Expr::Variable(name.into())
720    }
721
722    fn child_vars(expr: &Expr) -> Vec<String> {
723        expr.children()
724            .map(|child| match child {
725                Expr::Variable(name) => name.to_string(),
726                other => format!("{other:?}"),
727            })
728            .collect()
729    }
730
731    #[test]
732    fn children_yields_leaves_as_empty() {
733        for leaf in [
734            Expr::Null,
735            Expr::Bool(true),
736            Expr::Number(1.0),
737            Expr::String("s".into()),
738            var("x"),
739            Expr::Break,
740            Expr::Continue,
741            Expr::WaitSignal {
742                name: "ready".into(),
743            },
744            Expr::TypeLiteral(Box::new(TypeExpr::Str)),
745        ] {
746            let children: Vec<_> = leaf.children().collect();
747            assert!(children.is_empty(), "{leaf:?} should have no children");
748        }
749    }
750
751    #[test]
752    fn children_yields_composite_subexpressions_in_order() {
753        let block = Expr::Block(vec![var("a"), var("b"), var("c")]);
754        assert_eq!(child_vars(&block), ["a", "b", "c"]);
755
756        let record = Expr::Record(vec![("k1".into(), var("v1")), ("k2".into(), var("v2"))]);
757        assert_eq!(child_vars(&record), ["v1", "v2"]);
758
759        let if_expr = Expr::If {
760            condition: Box::new(var("cond")),
761            then_block: Box::new(var("then")),
762            else_block: Box::new(var("else")),
763        };
764        assert_eq!(child_vars(&if_expr), ["cond", "then", "else"]);
765
766        let while_expr = Expr::While {
767            condition: Box::new(var("cond")),
768            body: Box::new(var("body")),
769        };
770        assert_eq!(child_vars(&while_expr), ["cond", "body"]);
771
772        let receiver = Expr::ReceiverCall {
773            receiver: Box::new(var("recv")),
774            operation: "op".into(),
775            args: vec![var("arg0"), var("arg1")],
776        };
777        assert_eq!(child_vars(&receiver), ["recv", "arg0", "arg1"]);
778
779        let binary = Expr::Binary {
780            left: Box::new(var("left")),
781            op: BinaryOp::Add,
782            right: Box::new(var("right")),
783        };
784        assert_eq!(child_vars(&binary), ["left", "right"]);
785    }
786
787    #[test]
788    fn children_yields_assign_index_steps_before_value() {
789        let assign = Expr::Assign {
790            target: AssignTarget {
791                root: "root".into(),
792                steps: vec![
793                    AssignPathStep::Field("field".into()),
794                    AssignPathStep::Index(var("idx")),
795                ],
796            },
797            expr: Box::new(var("value")),
798        };
799        // Field steps contribute no child expressions; the dynamic index is
800        // yielded before the assigned value.
801        assert_eq!(child_vars(&assign), ["idx", "value"]);
802    }
803
804    #[test]
805    fn children_handles_optional_finish_and_submit() {
806        assert!(Expr::Submit(None).children().next().is_none());
807        assert!(Expr::Finish(None).children().next().is_none());
808        assert_eq!(
809            child_vars(&Expr::Submit(Some(Box::new(var("done"))))),
810            ["done"]
811        );
812    }
813
814    #[test]
815    fn children_size_hint_is_exact() {
816        let block = Expr::Block(vec![var("a"), var("b"), var("c"), var("d")]);
817        let iter = block.children();
818        assert_eq!(iter.len(), 4);
819        assert_eq!(iter.size_hint(), (4, Some(4)));
820    }
821
822    #[test]
823    fn visitor_walks_descendants_through_single_child_boundary() {
824        struct VariableCollector(Vec<String>);
825
826        impl ExprVisitor for VariableCollector {
827            fn visit_expr(&mut self, expr: &Expr) {
828                if let Expr::Variable(name) = expr {
829                    self.0.push(name.to_string());
830                }
831                walk_expr(self, expr);
832            }
833        }
834
835        let expr = Expr::While {
836            condition: Box::new(var("ready")),
837            body: Box::new(Expr::Block(vec![
838                Expr::Assign {
839                    target: AssignTarget {
840                        root: "items".into(),
841                        steps: vec![AssignPathStep::Index(var("idx"))],
842                    },
843                    expr: Box::new(var("value")),
844                },
845                Expr::Submit(Some(Box::new(var("done")))),
846            ])),
847        };
848
849        let mut collector = VariableCollector(Vec::new());
850        collector.visit_expr(&expr);
851
852        assert_eq!(collector.0, ["ready", "idx", "value", "done"]);
853    }
854
855    #[test]
856    fn folder_reconstructs_owned_expr_trees() {
857        struct RenameVariables;
858
859        impl ExprFolder for RenameVariables {
860            fn fold_expr(&mut self, expr: Expr) -> Expr {
861                match expr {
862                    Expr::Variable(name) => Expr::Variable(format!("renamed_{name}").into()),
863                    other => fold_expr_children(self, other),
864                }
865            }
866        }
867
868        let expr = Expr::Assign {
869            target: AssignTarget {
870                root: "items".into(),
871                steps: vec![AssignPathStep::Index(var("idx"))],
872            },
873            expr: Box::new(Expr::List(vec![var("first"), var("second")])),
874        };
875
876        let mut folder = RenameVariables;
877        let folded = folder.fold_expr(expr);
878
879        let Expr::Assign { target, expr } = folded else {
880            panic!("expected assign");
881        };
882        assert!(matches!(
883            target.steps.as_slice(),
884            [AssignPathStep::Index(Expr::Variable(name))] if name.as_str() == "renamed_idx"
885        ));
886        let Expr::List(items) = *expr else {
887            panic!("expected list");
888        };
889        assert_eq!(items, vec![var("renamed_first"), var("renamed_second")]);
890    }
891}