Skip to main content

php_ast/
visitor.rs

1use std::ops::ControlFlow;
2
3use crate::ast::*;
4
5/// Visitor trait for immutable AST traversal.
6///
7/// All methods return `ControlFlow<()>`:
8/// - `ControlFlow::Continue(())` — keep walking.
9/// - `ControlFlow::Break(())` — stop the entire traversal immediately.
10///
11/// Default implementations recursively walk child nodes, so implementors
12/// only need to override the node types they care about.
13///
14/// To **skip** a subtree, override the method and return `Continue(())`
15/// without calling the corresponding `walk_*` function.
16///
17/// # Example
18///
19/// ```
20/// use php_ast::visitor::{Visitor, walk_expr};
21/// use php_ast::ast::*;
22/// use std::ops::ControlFlow;
23///
24/// struct VarCounter { count: usize }
25///
26/// impl<'arena, 'src> Visitor<'arena, 'src> for VarCounter {
27///     fn visit_expr(&mut self, expr: &Expr<'arena, 'src>) -> ControlFlow<()> {
28///         if matches!(&expr.kind, ExprKind::Variable(_)) {
29///             self.count += 1;
30///         }
31///         walk_expr(self, expr)
32///     }
33/// }
34/// ```
35pub trait Visitor<'arena, 'src> {
36    fn visit_program(&mut self, program: &Program<'arena, 'src>) -> ControlFlow<()> {
37        walk_program(self, program)
38    }
39
40    fn visit_stmt(&mut self, stmt: &Stmt<'arena, 'src>) -> ControlFlow<()> {
41        walk_stmt(self, stmt)
42    }
43
44    fn visit_block(&mut self, block: &Block<'arena, 'src>) -> ControlFlow<()> {
45        walk_block(self, block)
46    }
47
48    fn visit_expr(&mut self, expr: &Expr<'arena, 'src>) -> ControlFlow<()> {
49        walk_expr(self, expr)
50    }
51
52    fn visit_param(&mut self, param: &Param<'arena, 'src>) -> ControlFlow<()> {
53        walk_param(self, param)
54    }
55
56    fn visit_arg(&mut self, arg: &Arg<'arena, 'src>) -> ControlFlow<()> {
57        walk_arg(self, arg)
58    }
59
60    fn visit_class_member(&mut self, member: &ClassMember<'arena, 'src>) -> ControlFlow<()> {
61        walk_class_member(self, member)
62    }
63
64    fn visit_enum_member(&mut self, member: &EnumMember<'arena, 'src>) -> ControlFlow<()> {
65        walk_enum_member(self, member)
66    }
67
68    fn visit_property_hook(&mut self, hook: &PropertyHook<'arena, 'src>) -> ControlFlow<()> {
69        walk_property_hook(self, hook)
70    }
71
72    fn visit_type_hint(&mut self, type_hint: &TypeHint<'arena, 'src>) -> ControlFlow<()> {
73        walk_type_hint(self, type_hint)
74    }
75
76    fn visit_attribute(&mut self, attribute: &Attribute<'arena, 'src>) -> ControlFlow<()> {
77        walk_attribute(self, attribute)
78    }
79
80    fn visit_catch_clause(&mut self, catch: &CatchClause<'arena, 'src>) -> ControlFlow<()> {
81        walk_catch_clause(self, catch)
82    }
83
84    fn visit_match_arm(&mut self, arm: &MatchArm<'arena, 'src>) -> ControlFlow<()> {
85        walk_match_arm(self, arm)
86    }
87
88    fn visit_closure_use_var(&mut self, _var: &ClosureUseVar<'src>) -> ControlFlow<()> {
89        ControlFlow::Continue(())
90    }
91
92    fn visit_trait_use(&mut self, trait_use: &TraitUseDecl<'arena, 'src>) -> ControlFlow<()> {
93        walk_trait_use(self, trait_use)
94    }
95
96    fn visit_trait_adaptation(
97        &mut self,
98        _adaptation: &TraitAdaptation<'arena, 'src>,
99    ) -> ControlFlow<()> {
100        ControlFlow::Continue(())
101    }
102
103    fn visit_name(&mut self, _name: &Name<'arena, 'src>) -> ControlFlow<()> {
104        ControlFlow::Continue(())
105    }
106
107    /// Called for each comment when driven by [`walk_comments`].
108    ///
109    /// Comments live in [`ParseResult::comments`](php_rs_parser::ParseResult::comments)
110    /// separately from AST nodes. Use [`walk_comments`] to drive this hook.
111    fn visit_comment(&mut self, _comment: &Comment<'src>) -> ControlFlow<()> {
112        ControlFlow::Continue(())
113    }
114}
115
116// =============================================================================
117// Walk functions
118// =============================================================================
119
120/// Calls [`Visitor::visit_name`] on `name`.
121///
122/// The default [`Visitor::visit_name`] is a leaf — it does nothing and returns
123/// `Continue(())`. Override it if you need to inspect name nodes.
124pub fn walk_name<'arena, 'src, V: Visitor<'arena, 'src> + ?Sized>(
125    visitor: &mut V,
126    name: &Name<'arena, 'src>,
127) -> ControlFlow<()> {
128    visitor.visit_name(name)
129}
130
131/// Calls [`Visitor::visit_comment`] for each comment in `comments`.
132///
133/// Comments are stored separately from the AST in
134/// `ParseResult::comments` — this function is the bridge.
135/// Call it alongside [`walk_program`] to cover the full source:
136///
137/// ```ignore
138/// walk_comments(&mut visitor, &result.comments)?;
139/// visitor.visit_program(&result.program)?;
140/// ```
141///
142/// Order (before vs. after the program walk) is up to the caller.
143pub fn walk_comments<'arena, 'src, V: Visitor<'arena, 'src> + ?Sized>(
144    visitor: &mut V,
145    comments: &[Comment<'src>],
146) -> ControlFlow<()> {
147    for comment in comments {
148        visitor.visit_comment(comment)?;
149    }
150    ControlFlow::Continue(())
151}
152
153/// Visits every top-level statement in `program` by calling [`Visitor::visit_stmt`].
154///
155/// This is the entry point for a full-file traversal. Call it from
156/// [`Visitor::visit_program`] (which is the default) or drive it directly:
157///
158/// ```
159/// # use php_ast::visitor::{Visitor, walk_program};
160/// # use php_ast::ast::*;
161/// # use std::ops::ControlFlow;
162/// # struct V;
163/// # impl<'a, 'b> Visitor<'a, 'b> for V {}
164/// # fn example<'a, 'b>(v: &mut V, program: &Program<'a, 'b>) {
165/// walk_program(v, program);
166/// # }
167/// ```
168pub fn walk_program<'arena, 'src, V: Visitor<'arena, 'src> + ?Sized>(
169    visitor: &mut V,
170    program: &Program<'arena, 'src>,
171) -> ControlFlow<()> {
172    for stmt in program.stmts.iter() {
173        visitor.visit_stmt(stmt)?;
174    }
175    ControlFlow::Continue(())
176}
177
178/// Visits each statement in a block.
179pub fn walk_block<'arena, 'src, V: Visitor<'arena, 'src> + ?Sized>(
180    visitor: &mut V,
181    block: &Block<'arena, 'src>,
182) -> ControlFlow<()> {
183    for stmt in block.stmts.iter() {
184        visitor.visit_stmt(stmt)?;
185    }
186    ControlFlow::Continue(())
187}
188
189/// Dispatches `stmt` to the appropriate child visitors based on its [`StmtKind`].
190///
191/// Call this from [`Visitor::visit_stmt`] to recurse into a statement's children.
192/// Omit the call to skip the subtree entirely.
193pub fn walk_stmt<'arena, 'src, V: Visitor<'arena, 'src> + ?Sized>(
194    visitor: &mut V,
195    stmt: &Stmt<'arena, 'src>,
196) -> ControlFlow<()> {
197    match &stmt.kind {
198        StmtKind::Expression(expr) => {
199            visitor.visit_expr(expr)?;
200        }
201        StmtKind::Echo(exprs) => {
202            for expr in exprs.iter() {
203                visitor.visit_expr(expr)?;
204            }
205        }
206        StmtKind::Return(expr) => {
207            if let Some(expr) = expr {
208                visitor.visit_expr(expr)?;
209            }
210        }
211        StmtKind::Block(block) => {
212            visitor.visit_block(block)?;
213        }
214        StmtKind::If(if_stmt) => {
215            visitor.visit_expr(&if_stmt.condition)?;
216            visitor.visit_stmt(if_stmt.then_branch)?;
217            for elseif in if_stmt.elseif_branches.iter() {
218                visitor.visit_expr(&elseif.condition)?;
219                visitor.visit_stmt(&elseif.body)?;
220            }
221            if let Some(else_branch) = &if_stmt.else_branch {
222                visitor.visit_stmt(else_branch)?;
223            }
224        }
225        StmtKind::While(while_stmt) => {
226            visitor.visit_expr(&while_stmt.condition)?;
227            visitor.visit_stmt(while_stmt.body)?;
228        }
229        StmtKind::For(for_stmt) => {
230            for expr in for_stmt.init.iter() {
231                visitor.visit_expr(expr)?;
232            }
233            for expr in for_stmt.condition.iter() {
234                visitor.visit_expr(expr)?;
235            }
236            for expr in for_stmt.update.iter() {
237                visitor.visit_expr(expr)?;
238            }
239            visitor.visit_stmt(for_stmt.body)?;
240        }
241        StmtKind::Foreach(foreach_stmt) => {
242            visitor.visit_expr(&foreach_stmt.expr)?;
243            if let Some(key) = &foreach_stmt.key {
244                visitor.visit_expr(key)?;
245            }
246            visitor.visit_expr(&foreach_stmt.value)?;
247            visitor.visit_stmt(foreach_stmt.body)?;
248        }
249        StmtKind::DoWhile(do_while) => {
250            visitor.visit_stmt(do_while.body)?;
251            visitor.visit_expr(&do_while.condition)?;
252        }
253        StmtKind::Function(func) => {
254            walk_function_like(visitor, &func.attributes, &func.params, &func.return_type)?;
255            visitor.visit_block(func.body)?;
256        }
257        StmtKind::Break(expr) | StmtKind::Continue(expr) => {
258            if let Some(expr) = expr {
259                visitor.visit_expr(expr)?;
260            }
261        }
262        StmtKind::Switch(switch_stmt) => {
263            visitor.visit_expr(&switch_stmt.expr)?;
264            for case in switch_stmt.body.cases.iter() {
265                if let Some(value) = &case.value {
266                    visitor.visit_expr(value)?;
267                }
268                for stmt in case.body.iter() {
269                    visitor.visit_stmt(stmt)?;
270                }
271            }
272        }
273        StmtKind::Throw(expr) => {
274            visitor.visit_expr(expr)?;
275        }
276        StmtKind::TryCatch(tc) => {
277            visitor.visit_block(tc.body)?;
278            for catch in tc.catches.iter() {
279                visitor.visit_catch_clause(catch)?;
280            }
281            if let Some(finally) = tc.finally {
282                visitor.visit_block(finally)?;
283            }
284        }
285        StmtKind::Declare(decl) => {
286            for (_, expr) in decl.directives.iter() {
287                visitor.visit_expr(expr)?;
288            }
289            if let Some(body) = decl.body {
290                visitor.visit_stmt(body)?;
291            }
292        }
293        StmtKind::Unset(exprs) | StmtKind::Global(exprs) => {
294            for expr in exprs.iter() {
295                visitor.visit_expr(expr)?;
296            }
297        }
298        StmtKind::Class(class) => {
299            walk_attributes(visitor, &class.attributes)?;
300            if let Some(extends) = &class.extends {
301                visitor.visit_name(extends)?;
302            }
303            for name in class.implements.iter() {
304                visitor.visit_name(name)?;
305            }
306            for member in class.body.members.iter() {
307                visitor.visit_class_member(member)?;
308            }
309        }
310        StmtKind::Interface(iface) => {
311            walk_attributes(visitor, &iface.attributes)?;
312            for name in iface.extends.iter() {
313                visitor.visit_name(name)?;
314            }
315            for member in iface.body.members.iter() {
316                visitor.visit_class_member(member)?;
317            }
318        }
319        StmtKind::Trait(trait_decl) => {
320            walk_attributes(visitor, &trait_decl.attributes)?;
321            for member in trait_decl.body.members.iter() {
322                visitor.visit_class_member(member)?;
323            }
324        }
325        StmtKind::Enum(enum_decl) => {
326            walk_attributes(visitor, &enum_decl.attributes)?;
327            if let Some(scalar_type) = &enum_decl.scalar_type {
328                visitor.visit_name(scalar_type)?;
329            }
330            for name in enum_decl.implements.iter() {
331                visitor.visit_name(name)?;
332            }
333            for member in enum_decl.body.members.iter() {
334                visitor.visit_enum_member(member)?;
335            }
336        }
337        StmtKind::Namespace(ns) => {
338            if let NamespaceBody::Braced(block) = &ns.body {
339                visitor.visit_block(block)?;
340            }
341        }
342        StmtKind::Const(items) => {
343            for item in items.iter() {
344                walk_attributes(visitor, &item.attributes)?;
345                visitor.visit_expr(&item.value)?;
346            }
347        }
348        StmtKind::StaticVar(vars) => {
349            for var in vars.iter() {
350                if let Some(default) = &var.default {
351                    visitor.visit_expr(default)?;
352                }
353            }
354        }
355        StmtKind::Use(decl) => {
356            for item in decl.uses.iter() {
357                visitor.visit_name(&item.name)?;
358            }
359        }
360        StmtKind::Goto(_)
361        | StmtKind::Label(_)
362        | StmtKind::Nop
363        | StmtKind::InlineHtml(_)
364        | StmtKind::HaltCompiler(_)
365        | StmtKind::Error => {}
366    }
367    ControlFlow::Continue(())
368}
369
370/// Dispatches `expr` to the appropriate child visitors based on its [`ExprKind`].
371///
372/// Call this from [`Visitor::visit_expr`] to recurse into an expression's children.
373/// Omit the call to skip the subtree entirely.
374pub fn walk_expr<'arena, 'src, V: Visitor<'arena, 'src> + ?Sized>(
375    visitor: &mut V,
376    expr: &Expr<'arena, 'src>,
377) -> ControlFlow<()> {
378    match &expr.kind {
379        ExprKind::Assign(assign) => {
380            visitor.visit_expr(assign.target)?;
381            visitor.visit_expr(assign.value)?;
382        }
383        ExprKind::Binary(binary) => {
384            visitor.visit_expr(binary.left)?;
385            visitor.visit_expr(binary.right)?;
386        }
387        ExprKind::UnaryPrefix(unary) => {
388            visitor.visit_expr(unary.operand)?;
389        }
390        ExprKind::UnaryPostfix(unary) => {
391            visitor.visit_expr(unary.operand)?;
392        }
393        ExprKind::Ternary(ternary) => {
394            visitor.visit_expr(ternary.condition)?;
395            if let Some(then_expr) = &ternary.then_expr {
396                visitor.visit_expr(then_expr)?;
397            }
398            visitor.visit_expr(ternary.else_expr)?;
399        }
400        ExprKind::NullCoalesce(nc) => {
401            visitor.visit_expr(nc.left)?;
402            visitor.visit_expr(nc.right)?;
403        }
404        ExprKind::FunctionCall(call) => {
405            visitor.visit_expr(call.name)?;
406            for arg in call.args.iter() {
407                visitor.visit_arg(arg)?;
408            }
409        }
410        ExprKind::Array(elements) => {
411            for elem in elements.iter() {
412                if let Some(key) = &elem.key {
413                    visitor.visit_expr(key)?;
414                }
415                visitor.visit_expr(&elem.value)?;
416            }
417        }
418        ExprKind::ArrayAccess(access) => {
419            visitor.visit_expr(access.array)?;
420            if let Some(index) = &access.index {
421                visitor.visit_expr(index)?;
422            }
423        }
424        ExprKind::Print(expr) => {
425            visitor.visit_expr(expr)?;
426        }
427        ExprKind::Parenthesized(expr) => {
428            visitor.visit_expr(expr)?;
429        }
430        ExprKind::Cast(_, expr) => {
431            visitor.visit_expr(expr)?;
432        }
433        ExprKind::ErrorSuppress(expr) => {
434            visitor.visit_expr(expr)?;
435        }
436        ExprKind::Isset(exprs) => {
437            for expr in exprs.iter() {
438                visitor.visit_expr(expr)?;
439            }
440        }
441        ExprKind::Empty(expr) => {
442            visitor.visit_expr(expr)?;
443        }
444        ExprKind::Include(_, expr) => {
445            visitor.visit_expr(expr)?;
446        }
447        ExprKind::Eval(expr) => {
448            visitor.visit_expr(expr)?;
449        }
450        ExprKind::Exit(expr) => {
451            if let Some(expr) = expr {
452                visitor.visit_expr(expr)?;
453            }
454        }
455        ExprKind::Clone(expr) => {
456            visitor.visit_expr(expr)?;
457        }
458        ExprKind::CloneWith(object, overrides) => {
459            visitor.visit_expr(object)?;
460            visitor.visit_expr(overrides)?;
461        }
462        ExprKind::New(new_expr) => {
463            visitor.visit_expr(new_expr.class)?;
464            for arg in new_expr.args.iter() {
465                visitor.visit_arg(arg)?;
466            }
467        }
468        ExprKind::PropertyAccess(access) | ExprKind::NullsafePropertyAccess(access) => {
469            visitor.visit_expr(access.object)?;
470            visitor.visit_expr(access.property)?;
471        }
472        ExprKind::MethodCall(call) | ExprKind::NullsafeMethodCall(call) => {
473            visitor.visit_expr(call.object)?;
474            visitor.visit_expr(call.method)?;
475            for arg in call.args.iter() {
476                visitor.visit_arg(arg)?;
477            }
478        }
479        ExprKind::StaticPropertyAccess(access) | ExprKind::ClassConstAccess(access) => {
480            visitor.visit_expr(access.class)?;
481            visitor.visit_expr(access.member)?;
482        }
483        ExprKind::ClassConstAccessDynamic { class, member }
484        | ExprKind::StaticPropertyAccessDynamic { class, member } => {
485            visitor.visit_expr(class)?;
486            visitor.visit_expr(member)?;
487        }
488        ExprKind::StaticMethodCall(call) => {
489            visitor.visit_expr(call.class)?;
490            visitor.visit_expr(call.method)?;
491            for arg in call.args.iter() {
492                visitor.visit_arg(arg)?;
493            }
494        }
495        ExprKind::StaticDynMethodCall(call) => {
496            visitor.visit_expr(call.class)?;
497            visitor.visit_expr(call.method)?;
498            for arg in call.args.iter() {
499                visitor.visit_arg(arg)?;
500            }
501        }
502        ExprKind::Closure(closure) => {
503            walk_function_like(
504                visitor,
505                &closure.attributes,
506                &closure.params,
507                &closure.return_type,
508            )?;
509            for use_var in closure.use_vars.iter() {
510                visitor.visit_closure_use_var(use_var)?;
511            }
512            visitor.visit_block(closure.body)?;
513        }
514        ExprKind::ArrowFunction(arrow) => {
515            walk_function_like(
516                visitor,
517                &arrow.attributes,
518                &arrow.params,
519                &arrow.return_type,
520            )?;
521            visitor.visit_expr(arrow.body)?;
522        }
523        ExprKind::Match(match_expr) => {
524            visitor.visit_expr(match_expr.subject)?;
525            for arm in match_expr.arms.iter() {
526                visitor.visit_match_arm(arm)?;
527            }
528        }
529        ExprKind::ThrowExpr(expr) => {
530            visitor.visit_expr(expr)?;
531        }
532        ExprKind::Yield(yield_expr) => {
533            if let Some(key) = &yield_expr.key {
534                visitor.visit_expr(key)?;
535            }
536            if let Some(value) = &yield_expr.value {
537                visitor.visit_expr(value)?;
538            }
539        }
540        ExprKind::AnonymousClass(class) => {
541            walk_attributes(visitor, &class.attributes)?;
542            for member in class.body.members.iter() {
543                visitor.visit_class_member(member)?;
544            }
545        }
546        ExprKind::InterpolatedString(parts)
547        | ExprKind::Heredoc { parts, .. }
548        | ExprKind::ShellExec(parts) => {
549            for part in parts.iter() {
550                if let StringPart::Expr(e) = part {
551                    visitor.visit_expr(e)?;
552                }
553            }
554        }
555        ExprKind::VariableVariable(inner) => {
556            visitor.visit_expr(inner)?;
557        }
558        ExprKind::CallableCreate(cc) => match &cc.kind {
559            CallableCreateKind::Function(name) => visitor.visit_expr(name)?,
560            CallableCreateKind::Method { object, method }
561            | CallableCreateKind::NullsafeMethod { object, method } => {
562                visitor.visit_expr(object)?;
563                visitor.visit_expr(method)?;
564            }
565            CallableCreateKind::StaticMethod { class, method } => {
566                visitor.visit_expr(class)?;
567                visitor.visit_expr(method)?;
568            }
569        },
570        ExprKind::Int(_)
571        | ExprKind::Float(_)
572        | ExprKind::String(_)
573        | ExprKind::Bool(_)
574        | ExprKind::Null
575        | ExprKind::Omit
576        | ExprKind::Variable(_)
577        | ExprKind::Identifier(_)
578        | ExprKind::MagicConst(_)
579        | ExprKind::Nowdoc { .. }
580        | ExprKind::Error => {}
581    }
582    ControlFlow::Continue(())
583}
584
585/// Visits a function/method parameter's attributes, type hint, default expression, and property hooks.
586pub fn walk_param<'arena, 'src, V: Visitor<'arena, 'src> + ?Sized>(
587    visitor: &mut V,
588    param: &Param<'arena, 'src>,
589) -> ControlFlow<()> {
590    walk_attributes(visitor, &param.attributes)?;
591    if let Some(type_hint) = &param.type_hint {
592        visitor.visit_type_hint(type_hint)?;
593    }
594    if let Some(default) = &param.default {
595        visitor.visit_expr(default)?;
596    }
597    for hook in param.hooks.iter() {
598        visitor.visit_property_hook(hook)?;
599    }
600    ControlFlow::Continue(())
601}
602
603/// Visits the value expression of a call argument.
604pub fn walk_arg<'arena, 'src, V: Visitor<'arena, 'src> + ?Sized>(
605    visitor: &mut V,
606    arg: &Arg<'arena, 'src>,
607) -> ControlFlow<()> {
608    visitor.visit_expr(&arg.value)
609}
610
611/// Dispatches a class member (property, method, constant, or trait use) to its child visitors.
612pub fn walk_class_member<'arena, 'src, V: Visitor<'arena, 'src> + ?Sized>(
613    visitor: &mut V,
614    member: &ClassMember<'arena, 'src>,
615) -> ControlFlow<()> {
616    match &member.kind {
617        ClassMemberKind::Property(prop) => {
618            walk_property_decl(visitor, prop)?;
619        }
620        ClassMemberKind::Method(method) => {
621            walk_method_decl(visitor, method)?;
622        }
623        ClassMemberKind::ClassConst(cc) => {
624            walk_class_const_decl(visitor, cc)?;
625        }
626        ClassMemberKind::TraitUse(trait_use) => {
627            visitor.visit_trait_use(trait_use)?;
628        }
629    }
630    ControlFlow::Continue(())
631}
632
633/// Visits a property hook's attributes, parameters, and body statements or expression.
634pub fn walk_property_hook<'arena, 'src, V: Visitor<'arena, 'src> + ?Sized>(
635    visitor: &mut V,
636    hook: &PropertyHook<'arena, 'src>,
637) -> ControlFlow<()> {
638    walk_attributes(visitor, &hook.attributes)?;
639    for param in hook.params.iter() {
640        visitor.visit_param(param)?;
641    }
642    match &hook.body {
643        PropertyHookBody::Block(block) => {
644            visitor.visit_block(block)?;
645        }
646        PropertyHookBody::Expression(expr) => {
647            visitor.visit_expr(expr)?;
648        }
649        PropertyHookBody::Abstract => {}
650    }
651    ControlFlow::Continue(())
652}
653
654/// Dispatches an enum member (case, method, constant, or trait use) to its child visitors.
655pub fn walk_enum_member<'arena, 'src, V: Visitor<'arena, 'src> + ?Sized>(
656    visitor: &mut V,
657    member: &EnumMember<'arena, 'src>,
658) -> ControlFlow<()> {
659    match &member.kind {
660        EnumMemberKind::Case(case) => {
661            walk_attributes(visitor, &case.attributes)?;
662            if let Some(value) = &case.value {
663                visitor.visit_expr(value)?;
664            }
665        }
666        EnumMemberKind::Method(method) => {
667            walk_method_decl(visitor, method)?;
668        }
669        EnumMemberKind::ClassConst(cc) => {
670            walk_class_const_decl(visitor, cc)?;
671        }
672        EnumMemberKind::TraitUse(trait_use) => {
673            visitor.visit_trait_use(trait_use)?;
674        }
675    }
676    ControlFlow::Continue(())
677}
678
679/// Visits the inner types of a type hint (recursing into nullable, union, and intersection).
680pub fn walk_type_hint<'arena, 'src, V: Visitor<'arena, 'src> + ?Sized>(
681    visitor: &mut V,
682    type_hint: &TypeHint<'arena, 'src>,
683) -> ControlFlow<()> {
684    match &type_hint.kind {
685        TypeHintKind::Nullable(inner) => {
686            visitor.visit_type_hint(inner)?;
687        }
688        TypeHintKind::Union(types) | TypeHintKind::Intersection(types) => {
689            for ty in types.iter() {
690                visitor.visit_type_hint(ty)?;
691            }
692        }
693        TypeHintKind::Named(name) => {
694            visitor.visit_name(name)?;
695        }
696        TypeHintKind::Keyword(_, _) => {}
697    }
698    ControlFlow::Continue(())
699}
700
701/// Visits an attribute's name and argument expressions.
702pub fn walk_attribute<'arena, 'src, V: Visitor<'arena, 'src> + ?Sized>(
703    visitor: &mut V,
704    attribute: &Attribute<'arena, 'src>,
705) -> ControlFlow<()> {
706    visitor.visit_name(&attribute.name)?;
707    for arg in attribute.args.iter() {
708        visitor.visit_arg(arg)?;
709    }
710    ControlFlow::Continue(())
711}
712
713/// Visits a catch clause's caught type names and body statements.
714pub fn walk_catch_clause<'arena, 'src, V: Visitor<'arena, 'src> + ?Sized>(
715    visitor: &mut V,
716    catch: &CatchClause<'arena, 'src>,
717) -> ControlFlow<()> {
718    for ty in catch.types.iter() {
719        visitor.visit_name(ty)?;
720    }
721    visitor.visit_block(catch.body)?;
722    ControlFlow::Continue(())
723}
724
725/// Visits a match arm's condition expressions (if any) and body expression.
726pub fn walk_match_arm<'arena, 'src, V: Visitor<'arena, 'src> + ?Sized>(
727    visitor: &mut V,
728    arm: &MatchArm<'arena, 'src>,
729) -> ControlFlow<()> {
730    if let Some(conditions) = &arm.conditions {
731        for cond in conditions.iter() {
732            visitor.visit_expr(cond)?;
733        }
734    }
735    visitor.visit_expr(&arm.body)
736}
737
738/// Visits a trait use declaration's trait names and adaptations (`insteadof`, `as`).
739pub fn walk_trait_use<'arena, 'src, V: Visitor<'arena, 'src> + ?Sized>(
740    visitor: &mut V,
741    trait_use: &TraitUseDecl<'arena, 'src>,
742) -> ControlFlow<()> {
743    for name in trait_use.traits.iter() {
744        visitor.visit_name(name)?;
745    }
746    for adaptation in trait_use.adaptations.iter() {
747        visitor.visit_trait_adaptation(adaptation)?;
748    }
749    ControlFlow::Continue(())
750}
751
752// =============================================================================
753// Internal helpers — shared walking logic to avoid duplication
754// =============================================================================
755
756/// Walks the common parts of any function-like construct:
757/// attributes → params → optional return type.
758fn walk_function_like<'arena, 'src, V: Visitor<'arena, 'src> + ?Sized>(
759    visitor: &mut V,
760    attributes: &[Attribute<'arena, 'src>],
761    params: &[Param<'arena, 'src>],
762    return_type: &Option<TypeHint<'arena, 'src>>,
763) -> ControlFlow<()> {
764    walk_attributes(visitor, attributes)?;
765    for param in params.iter() {
766        visitor.visit_param(param)?;
767    }
768    if let Some(ret) = return_type {
769        visitor.visit_type_hint(ret)?;
770    }
771    ControlFlow::Continue(())
772}
773
774/// Walks a method declaration (shared by ClassMember and EnumMember).
775fn walk_method_decl<'arena, 'src, V: Visitor<'arena, 'src> + ?Sized>(
776    visitor: &mut V,
777    method: &MethodDecl<'arena, 'src>,
778) -> ControlFlow<()> {
779    walk_function_like(
780        visitor,
781        &method.attributes,
782        &method.params,
783        &method.return_type,
784    )?;
785    if let Some(body) = method.body {
786        visitor.visit_block(body)?;
787    }
788    ControlFlow::Continue(())
789}
790
791/// Walks a class constant declaration (shared by ClassMember and EnumMember).
792fn walk_class_const_decl<'arena, 'src, V: Visitor<'arena, 'src> + ?Sized>(
793    visitor: &mut V,
794    cc: &ClassConstDecl<'arena, 'src>,
795) -> ControlFlow<()> {
796    walk_attributes(visitor, &cc.attributes)?;
797    if let Some(type_hint) = &cc.type_hint {
798        visitor.visit_type_hint(type_hint)?;
799    }
800    visitor.visit_expr(&cc.value)
801}
802
803/// Walks a property declaration.
804fn walk_property_decl<'arena, 'src, V: Visitor<'arena, 'src> + ?Sized>(
805    visitor: &mut V,
806    prop: &PropertyDecl<'arena, 'src>,
807) -> ControlFlow<()> {
808    walk_attributes(visitor, &prop.attributes)?;
809    if let Some(type_hint) = &prop.type_hint {
810        visitor.visit_type_hint(type_hint)?;
811    }
812    if let Some(default) = &prop.default {
813        visitor.visit_expr(default)?;
814    }
815    for hook in prop.hooks.iter() {
816        visitor.visit_property_hook(hook)?;
817    }
818    ControlFlow::Continue(())
819}
820
821fn walk_attributes<'arena, 'src, V: Visitor<'arena, 'src> + ?Sized>(
822    visitor: &mut V,
823    attributes: &[Attribute<'arena, 'src>],
824) -> ControlFlow<()> {
825    for attr in attributes.iter() {
826        visitor.visit_attribute(attr)?;
827    }
828    ControlFlow::Continue(())
829}
830
831// =============================================================================
832// ScopeVisitor — scope-aware traversal
833// =============================================================================
834
835/// Lexical scope context passed to each [`ScopeVisitor`] method.
836///
837/// Represents the immediately enclosing namespace, class-like definition,
838/// and named function or method at the point a node is visited.
839/// All fields are `None` when the node is at the global top level.
840///
841/// **Namespace** is set when inside a braced or simple `namespace` declaration.
842/// **`class_name`** is set inside `class`, `interface`, `trait`, and `enum`
843/// declarations; it is `None` for anonymous classes.
844/// **`function_name`** is set inside named functions and methods; it is `None`
845/// inside closures and arrow functions.
846#[derive(Debug, Clone, Copy, Default)]
847pub struct Scope<'src> {
848    /// Current namespace, or `None` for the global namespace.
849    ///
850    /// This is a borrowed slice of the original source string, so copying or
851    /// cloning the scope is always allocation-free.
852    pub namespace: Option<&'src str>,
853    /// Name of the immediately enclosing class-like declaration, or `None`.
854    pub class_name: Option<&'src str>,
855    /// Name of the immediately enclosing named function or method, or `None`.
856    pub function_name: Option<&'src str>,
857}
858
859/// A scope-aware variant of [`Visitor`].
860///
861/// Every visit method receives a [`Scope`] describing the lexical context at
862/// that node — the current namespace, enclosing class-like declaration, and
863/// enclosing named function or method.  All methods have no-op default
864/// implementations, so implementors override only what they need.
865///
866/// Drive traversal with [`ScopeWalker`], which maintains the scope
867/// automatically and calls these methods with the current context.
868///
869/// # Example
870///
871/// ```
872/// use php_ast::visitor::{ScopeVisitor, ScopeWalker, Scope};
873/// use php_ast::ast::*;
874/// use std::ops::ControlFlow;
875///
876/// struct MethodCollector { methods: Vec<String> }
877///
878/// impl<'arena, 'src> ScopeVisitor<'arena, 'src> for MethodCollector {
879///     fn visit_class_member(
880///         &mut self,
881///         member: &ClassMember<'arena, 'src>,
882///         scope: &Scope<'src>,
883///     ) -> ControlFlow<()> {
884///         if let ClassMemberKind::Method(m) = &member.kind {
885///             self.methods.push(format!(
886///                 "{}::{}",
887///                 scope.class_name.unwrap_or("<anon>"),
888///                 m.name
889///             ));
890///         }
891///         ControlFlow::Continue(())
892///     }
893/// }
894/// ```
895pub trait ScopeVisitor<'arena, 'src> {
896    fn visit_program(
897        &mut self,
898        _program: &Program<'arena, 'src>,
899        _scope: &Scope<'src>,
900    ) -> ControlFlow<()> {
901        ControlFlow::Continue(())
902    }
903    fn visit_stmt(&mut self, _stmt: &Stmt<'arena, 'src>, _scope: &Scope<'src>) -> ControlFlow<()> {
904        ControlFlow::Continue(())
905    }
906    fn visit_expr(&mut self, _expr: &Expr<'arena, 'src>, _scope: &Scope<'src>) -> ControlFlow<()> {
907        ControlFlow::Continue(())
908    }
909    fn visit_param(
910        &mut self,
911        _param: &Param<'arena, 'src>,
912        _scope: &Scope<'src>,
913    ) -> ControlFlow<()> {
914        ControlFlow::Continue(())
915    }
916    fn visit_arg(&mut self, _arg: &Arg<'arena, 'src>, _scope: &Scope<'src>) -> ControlFlow<()> {
917        ControlFlow::Continue(())
918    }
919    fn visit_class_member(
920        &mut self,
921        _member: &ClassMember<'arena, 'src>,
922        _scope: &Scope<'src>,
923    ) -> ControlFlow<()> {
924        ControlFlow::Continue(())
925    }
926    fn visit_enum_member(
927        &mut self,
928        _member: &EnumMember<'arena, 'src>,
929        _scope: &Scope<'src>,
930    ) -> ControlFlow<()> {
931        ControlFlow::Continue(())
932    }
933    fn visit_property_hook(
934        &mut self,
935        _hook: &PropertyHook<'arena, 'src>,
936        _scope: &Scope<'src>,
937    ) -> ControlFlow<()> {
938        ControlFlow::Continue(())
939    }
940    fn visit_type_hint(
941        &mut self,
942        _type_hint: &TypeHint<'arena, 'src>,
943        _scope: &Scope<'src>,
944    ) -> ControlFlow<()> {
945        ControlFlow::Continue(())
946    }
947    fn visit_attribute(
948        &mut self,
949        _attribute: &Attribute<'arena, 'src>,
950        _scope: &Scope<'src>,
951    ) -> ControlFlow<()> {
952        ControlFlow::Continue(())
953    }
954    fn visit_catch_clause(
955        &mut self,
956        _catch: &CatchClause<'arena, 'src>,
957        _scope: &Scope<'src>,
958    ) -> ControlFlow<()> {
959        ControlFlow::Continue(())
960    }
961    fn visit_match_arm(
962        &mut self,
963        _arm: &MatchArm<'arena, 'src>,
964        _scope: &Scope<'src>,
965    ) -> ControlFlow<()> {
966        ControlFlow::Continue(())
967    }
968    fn visit_closure_use_var(
969        &mut self,
970        _var: &ClosureUseVar<'src>,
971        _scope: &Scope<'src>,
972    ) -> ControlFlow<()> {
973        ControlFlow::Continue(())
974    }
975
976    fn visit_trait_use(
977        &mut self,
978        _trait_use: &TraitUseDecl<'arena, 'src>,
979        _scope: &Scope<'src>,
980    ) -> ControlFlow<()> {
981        ControlFlow::Continue(())
982    }
983
984    fn visit_trait_adaptation(
985        &mut self,
986        _adaptation: &TraitAdaptation<'arena, 'src>,
987        _scope: &Scope<'src>,
988    ) -> ControlFlow<()> {
989        ControlFlow::Continue(())
990    }
991
992    fn visit_comment(&mut self, _comment: &Comment<'src>, _scope: &Scope<'src>) -> ControlFlow<()> {
993        ControlFlow::Continue(())
994    }
995}
996
997/// Drives a [`ScopeVisitor`] over an AST, maintaining [`Scope`] automatically.
998///
999/// `ScopeWalker` wraps a [`ScopeVisitor`] and tracks the lexical scope as it
1000/// descends the tree, updating scope before visiting children and restoring it
1001/// on exit from scope-defining nodes (functions, classes, namespaces).
1002///
1003/// # Usage
1004///
1005/// ```no_run
1006/// # use php_ast::visitor::{ScopeWalker, ScopeVisitor, Scope};
1007/// # use php_ast::ast::*;
1008/// # use std::ops::ControlFlow;
1009/// # struct MyVisitor;
1010/// # impl<'a, 'b> ScopeVisitor<'a, 'b> for MyVisitor {}
1011/// # fn parse<'a, 'b>(_: &'a bumpalo::Bump, _: &'b str) -> Program<'a, 'b> { unimplemented!() }
1012/// let arena = bumpalo::Bump::new();
1013/// let src = "<?php class Foo { public function bar() {} }";
1014/// let program = parse(&arena, src);
1015/// let mut walker = ScopeWalker::new(src, MyVisitor);
1016/// walker.walk(&program);
1017/// let _my_visitor = walker.into_inner();
1018/// ```
1019pub struct ScopeWalker<'src, V> {
1020    inner: V,
1021    scope: Scope<'src>,
1022    src: &'src str,
1023}
1024
1025impl<'src, V> ScopeWalker<'src, V> {
1026    /// Creates a new `ScopeWalker` wrapping `inner`.
1027    ///
1028    /// `src` must be the same source string that was passed to the parser that
1029    /// produced the [`Program`] you will walk.  It is used to derive
1030    /// zero-allocation [`Scope::namespace`] slices for qualified namespace
1031    /// names (e.g. `Foo\Bar`).
1032    pub fn new(src: &'src str, inner: V) -> Self {
1033        Self {
1034            inner,
1035            scope: Scope::default(),
1036            src,
1037        }
1038    }
1039
1040    /// Consumes the walker and returns the inner visitor.
1041    pub fn into_inner(self) -> V {
1042        self.inner
1043    }
1044
1045    /// Returns a reference to the inner visitor.
1046    pub fn inner(&self) -> &V {
1047        &self.inner
1048    }
1049
1050    /// Returns a mutable reference to the inner visitor.
1051    pub fn inner_mut(&mut self) -> &mut V {
1052        &mut self.inner
1053    }
1054}
1055
1056impl<'arena, 'src, V: ScopeVisitor<'arena, 'src>> ScopeWalker<'src, V> {
1057    /// Walks `program`, calling [`ScopeVisitor`] methods with scope context.
1058    pub fn walk(&mut self, program: &Program<'arena, 'src>) -> ControlFlow<()> {
1059        self.visit_program(program)
1060    }
1061}
1062
1063impl<'arena, 'src, V: ScopeVisitor<'arena, 'src>> Visitor<'arena, 'src> for ScopeWalker<'src, V> {
1064    fn visit_program(&mut self, program: &Program<'arena, 'src>) -> ControlFlow<()> {
1065        self.inner.visit_program(program, &self.scope)?;
1066        walk_program(self, program)
1067    }
1068
1069    fn visit_stmt(&mut self, stmt: &Stmt<'arena, 'src>) -> ControlFlow<()> {
1070        self.inner.visit_stmt(stmt, &self.scope)?;
1071        match &stmt.kind {
1072            StmtKind::Function(func) => {
1073                let prev_fn = std::mem::replace(&mut self.scope.function_name, func.name.as_str());
1074                walk_stmt(self, stmt)?;
1075                self.scope.function_name = prev_fn;
1076            }
1077            StmtKind::Class(class) => {
1078                let prev_class = self.scope.class_name;
1079                let prev_fn = self.scope.function_name.take();
1080                self.scope.class_name = class.name.and_then(|n| n.as_str());
1081                walk_stmt(self, stmt)?;
1082                self.scope.class_name = prev_class;
1083                self.scope.function_name = prev_fn;
1084            }
1085            StmtKind::Interface(iface) => {
1086                let prev_class = std::mem::replace(&mut self.scope.class_name, iface.name.as_str());
1087                let prev_fn = self.scope.function_name.take();
1088                walk_stmt(self, stmt)?;
1089                self.scope.class_name = prev_class;
1090                self.scope.function_name = prev_fn;
1091            }
1092            StmtKind::Trait(trait_decl) => {
1093                let prev_class =
1094                    std::mem::replace(&mut self.scope.class_name, trait_decl.name.as_str());
1095                let prev_fn = self.scope.function_name.take();
1096                walk_stmt(self, stmt)?;
1097                self.scope.class_name = prev_class;
1098                self.scope.function_name = prev_fn;
1099            }
1100            StmtKind::Enum(enum_decl) => {
1101                let prev_class =
1102                    std::mem::replace(&mut self.scope.class_name, enum_decl.name.as_str());
1103                let prev_fn = self.scope.function_name.take();
1104                walk_stmt(self, stmt)?;
1105                self.scope.class_name = prev_class;
1106                self.scope.function_name = prev_fn;
1107            }
1108            StmtKind::Namespace(ns) => {
1109                let ns_str = ns.name.as_ref().map(|n| n.src_repr(self.src));
1110                match ns.body {
1111                    NamespaceBody::Braced(_) => {
1112                        let prev_ns = self.scope.namespace;
1113                        let prev_class = self.scope.class_name.take();
1114                        let prev_fn = self.scope.function_name.take();
1115                        self.scope.namespace = ns_str;
1116                        walk_stmt(self, stmt)?;
1117                        self.scope.namespace = prev_ns;
1118                        self.scope.class_name = prev_class;
1119                        self.scope.function_name = prev_fn;
1120                    }
1121                    NamespaceBody::Simple => {
1122                        // Simple namespace: update scope and leave it set for
1123                        // the remainder of the file (no push/pop).
1124                        self.scope.namespace = ns_str;
1125                        self.scope.class_name = None;
1126                        self.scope.function_name = None;
1127                    }
1128                }
1129            }
1130            _ => {
1131                walk_stmt(self, stmt)?;
1132            }
1133        }
1134        ControlFlow::Continue(())
1135    }
1136
1137    fn visit_expr(&mut self, expr: &Expr<'arena, 'src>) -> ControlFlow<()> {
1138        self.inner.visit_expr(expr, &self.scope)?;
1139        match &expr.kind {
1140            ExprKind::Closure(_) | ExprKind::ArrowFunction(_) => {
1141                let prev_fn = self.scope.function_name.take();
1142                walk_expr(self, expr)?;
1143                self.scope.function_name = prev_fn;
1144            }
1145            ExprKind::AnonymousClass(_) => {
1146                let prev_class = self.scope.class_name.take();
1147                let prev_fn = self.scope.function_name.take();
1148                walk_expr(self, expr)?;
1149                self.scope.class_name = prev_class;
1150                self.scope.function_name = prev_fn;
1151            }
1152            _ => {
1153                walk_expr(self, expr)?;
1154            }
1155        }
1156        ControlFlow::Continue(())
1157    }
1158
1159    fn visit_class_member(&mut self, member: &ClassMember<'arena, 'src>) -> ControlFlow<()> {
1160        self.inner.visit_class_member(member, &self.scope)?;
1161        if let ClassMemberKind::Method(method) = &member.kind {
1162            let prev_fn = std::mem::replace(&mut self.scope.function_name, method.name.as_str());
1163            walk_class_member(self, member)?;
1164            self.scope.function_name = prev_fn;
1165        } else {
1166            walk_class_member(self, member)?;
1167        }
1168        ControlFlow::Continue(())
1169    }
1170
1171    fn visit_enum_member(&mut self, member: &EnumMember<'arena, 'src>) -> ControlFlow<()> {
1172        self.inner.visit_enum_member(member, &self.scope)?;
1173        if let EnumMemberKind::Method(method) = &member.kind {
1174            let prev_fn = std::mem::replace(&mut self.scope.function_name, method.name.as_str());
1175            walk_enum_member(self, member)?;
1176            self.scope.function_name = prev_fn;
1177        } else {
1178            walk_enum_member(self, member)?;
1179        }
1180        ControlFlow::Continue(())
1181    }
1182
1183    fn visit_param(&mut self, param: &Param<'arena, 'src>) -> ControlFlow<()> {
1184        self.inner.visit_param(param, &self.scope)?;
1185        walk_param(self, param)
1186    }
1187
1188    fn visit_arg(&mut self, arg: &Arg<'arena, 'src>) -> ControlFlow<()> {
1189        self.inner.visit_arg(arg, &self.scope)?;
1190        walk_arg(self, arg)
1191    }
1192
1193    fn visit_property_hook(&mut self, hook: &PropertyHook<'arena, 'src>) -> ControlFlow<()> {
1194        self.inner.visit_property_hook(hook, &self.scope)?;
1195        walk_property_hook(self, hook)
1196    }
1197
1198    fn visit_type_hint(&mut self, type_hint: &TypeHint<'arena, 'src>) -> ControlFlow<()> {
1199        self.inner.visit_type_hint(type_hint, &self.scope)?;
1200        walk_type_hint(self, type_hint)
1201    }
1202
1203    fn visit_attribute(&mut self, attribute: &Attribute<'arena, 'src>) -> ControlFlow<()> {
1204        self.inner.visit_attribute(attribute, &self.scope)?;
1205        walk_attribute(self, attribute)
1206    }
1207
1208    fn visit_catch_clause(&mut self, catch: &CatchClause<'arena, 'src>) -> ControlFlow<()> {
1209        self.inner.visit_catch_clause(catch, &self.scope)?;
1210        walk_catch_clause(self, catch)
1211    }
1212
1213    fn visit_match_arm(&mut self, arm: &MatchArm<'arena, 'src>) -> ControlFlow<()> {
1214        self.inner.visit_match_arm(arm, &self.scope)?;
1215        walk_match_arm(self, arm)
1216    }
1217
1218    fn visit_closure_use_var(&mut self, var: &ClosureUseVar<'src>) -> ControlFlow<()> {
1219        self.inner.visit_closure_use_var(var, &self.scope)
1220    }
1221
1222    fn visit_trait_use(&mut self, trait_use: &TraitUseDecl<'arena, 'src>) -> ControlFlow<()> {
1223        self.inner.visit_trait_use(trait_use, &self.scope)?;
1224        walk_trait_use(self, trait_use)
1225    }
1226
1227    fn visit_trait_adaptation(
1228        &mut self,
1229        adaptation: &TraitAdaptation<'arena, 'src>,
1230    ) -> ControlFlow<()> {
1231        self.inner.visit_trait_adaptation(adaptation, &self.scope)
1232    }
1233
1234    fn visit_comment(&mut self, comment: &Comment<'src>) -> ControlFlow<()> {
1235        self.inner.visit_comment(comment, &self.scope)
1236    }
1237}
1238
1239#[cfg(test)]
1240mod tests {
1241    use super::*;
1242    use crate::Span;
1243    // =========================================================================
1244    // Unit tests with hand-built ASTs
1245    // =========================================================================
1246
1247    struct VarCounter {
1248        count: usize,
1249    }
1250
1251    impl<'arena, 'src> Visitor<'arena, 'src> for VarCounter {
1252        fn visit_expr(&mut self, expr: &Expr<'arena, 'src>) -> ControlFlow<()> {
1253            if matches!(&expr.kind, ExprKind::Variable(_)) {
1254                self.count += 1;
1255            }
1256            walk_expr(self, expr)
1257        }
1258    }
1259
1260    #[test]
1261    fn counts_variables() {
1262        let arena = bumpalo::Bump::new();
1263        let var_x = arena.alloc(Expr {
1264            kind: ExprKind::Variable(NameStr::__src("x")),
1265            span: Span::DUMMY,
1266        });
1267        let var_y = arena.alloc(Expr {
1268            kind: ExprKind::Variable(NameStr::__src("y")),
1269            span: Span::DUMMY,
1270        });
1271        let var_z = arena.alloc(Expr {
1272            kind: ExprKind::Variable(NameStr::__src("z")),
1273            span: Span::DUMMY,
1274        });
1275        let binary = arena.alloc(Expr {
1276            kind: ExprKind::Binary(BinaryExpr {
1277                left: var_y,
1278                op: BinaryOp::Add,
1279                right: var_z,
1280            }),
1281            span: Span::DUMMY,
1282        });
1283        let assign = arena.alloc(Expr {
1284            kind: ExprKind::Assign(AssignExpr {
1285                target: var_x,
1286                op: AssignOp::Assign,
1287                value: binary,
1288                by_ref: false,
1289            }),
1290            span: Span::DUMMY,
1291        });
1292        let mut stmts = ArenaVec::new_in(&arena);
1293        stmts.push(Stmt {
1294            kind: StmtKind::Expression(assign),
1295            span: Span::DUMMY,
1296            doc_comment: None,
1297        });
1298        let program = Program {
1299            stmts,
1300            span: Span::DUMMY,
1301        };
1302
1303        let mut v = VarCounter { count: 0 };
1304        let _ = v.visit_program(&program);
1305        assert_eq!(v.count, 3);
1306    }
1307
1308    #[test]
1309    fn early_termination() {
1310        let arena = bumpalo::Bump::new();
1311        let var_a = arena.alloc(Expr {
1312            kind: ExprKind::Variable(NameStr::__src("a")),
1313            span: Span::DUMMY,
1314        });
1315        let var_b = arena.alloc(Expr {
1316            kind: ExprKind::Variable(NameStr::__src("b")),
1317            span: Span::DUMMY,
1318        });
1319        let binary = arena.alloc(Expr {
1320            kind: ExprKind::Binary(BinaryExpr {
1321                left: var_a,
1322                op: BinaryOp::Add,
1323                right: var_b,
1324            }),
1325            span: Span::DUMMY,
1326        });
1327        let mut stmts = ArenaVec::new_in(&arena);
1328        stmts.push(Stmt {
1329            kind: StmtKind::Expression(binary),
1330            span: Span::DUMMY,
1331            doc_comment: None,
1332        });
1333        let program = Program {
1334            stmts,
1335            span: Span::DUMMY,
1336        };
1337
1338        struct FindFirst {
1339            found: Option<String>,
1340        }
1341        impl<'arena, 'src> Visitor<'arena, 'src> for FindFirst {
1342            fn visit_expr(&mut self, expr: &Expr<'arena, 'src>) -> ControlFlow<()> {
1343                if let ExprKind::Variable(name) = &expr.kind {
1344                    self.found = Some(name.to_string());
1345                    return ControlFlow::Break(());
1346                }
1347                walk_expr(self, expr)
1348            }
1349        }
1350
1351        let mut finder = FindFirst { found: None };
1352        let result = finder.visit_program(&program);
1353        assert!(result.is_break());
1354        assert_eq!(finder.found.as_deref(), Some("a"));
1355    }
1356
1357    #[test]
1358    fn skip_subtree() {
1359        let arena = bumpalo::Bump::new();
1360        // 1 + 2; function foo() { 3 + 4; }
1361        let one = arena.alloc(Expr {
1362            kind: ExprKind::Int(1),
1363            span: Span::DUMMY,
1364        });
1365        let two = arena.alloc(Expr {
1366            kind: ExprKind::Int(2),
1367            span: Span::DUMMY,
1368        });
1369        let top = arena.alloc(Expr {
1370            kind: ExprKind::Binary(BinaryExpr {
1371                left: one,
1372                op: BinaryOp::Add,
1373                right: two,
1374            }),
1375            span: Span::DUMMY,
1376        });
1377        let three = arena.alloc(Expr {
1378            kind: ExprKind::Int(3),
1379            span: Span::DUMMY,
1380        });
1381        let four = arena.alloc(Expr {
1382            kind: ExprKind::Int(4),
1383            span: Span::DUMMY,
1384        });
1385        let inner = arena.alloc(Expr {
1386            kind: ExprKind::Binary(BinaryExpr {
1387                left: three,
1388                op: BinaryOp::Add,
1389                right: four,
1390            }),
1391            span: Span::DUMMY,
1392        });
1393        let mut func_body_stmts = ArenaVec::new_in(&arena);
1394        func_body_stmts.push(Stmt {
1395            kind: StmtKind::Expression(inner),
1396            span: Span::DUMMY,
1397            doc_comment: None,
1398        });
1399        let func_body = arena.alloc(Block {
1400            stmts: func_body_stmts,
1401            span: Span::DUMMY,
1402        });
1403        let func = arena.alloc(FunctionDecl {
1404            name: Ident::name("foo"),
1405            params: ArenaVec::new_in(&arena),
1406            body: func_body,
1407            return_type: None,
1408            by_ref: false,
1409            attributes: ArenaVec::new_in(&arena),
1410            doc_comment: None,
1411        });
1412        let mut stmts = ArenaVec::new_in(&arena);
1413        stmts.push(Stmt {
1414            kind: StmtKind::Expression(top),
1415            span: Span::DUMMY,
1416            doc_comment: None,
1417        });
1418        stmts.push(Stmt {
1419            kind: StmtKind::Function(func),
1420            span: Span::DUMMY,
1421            doc_comment: None,
1422        });
1423        let program = Program {
1424            stmts,
1425            span: Span::DUMMY,
1426        };
1427
1428        struct SkipFunctions {
1429            expr_count: usize,
1430        }
1431        impl<'arena, 'src> Visitor<'arena, 'src> for SkipFunctions {
1432            fn visit_expr(&mut self, expr: &Expr<'arena, 'src>) -> ControlFlow<()> {
1433                self.expr_count += 1;
1434                walk_expr(self, expr)
1435            }
1436            fn visit_stmt(&mut self, stmt: &Stmt<'arena, 'src>) -> ControlFlow<()> {
1437                if matches!(&stmt.kind, StmtKind::Function(_)) {
1438                    return ControlFlow::Continue(());
1439                }
1440                walk_stmt(self, stmt)
1441            }
1442        }
1443
1444        let mut v = SkipFunctions { expr_count: 0 };
1445        let _ = v.visit_program(&program);
1446        // Only top-level: binary(1, 2) = 3 exprs
1447        assert_eq!(v.expr_count, 3);
1448    }
1449}