Add enter and leave_node methods to Preoder visitor (#6422)

crates/ruff_python_ast/src/visitor/preorder.rs

@@ -1,3 +1,4 @@
+use crate::node::AnyNodeRef;
 use crate::{
     self as ast, Alias, Arguments, BoolOp, CmpOp, Comprehension, Constant, Decorator,
     ElifElseClause, ExceptHandler, Expr, Keyword, MatchCase, Mod, Operator, Parameter,
@@ -7,64 +8,89 @@ use crate::{
 
 /// Visitor that traverses all nodes recursively in pre-order.
 pub trait PreorderVisitor<'a> {
+    #[allow(clippy::inline_always)]
+    #[inline(always)]
+    fn enter_node(&mut self, _node: AnyNodeRef<'a>) -> TraversalSignal {
+        TraversalSignal::Traverse
+    }
+
+    #[inline(always)]
+    fn leave_node(&mut self, _node: AnyNodeRef<'a>) {}
+
+    #[inline]
     fn visit_mod(&mut self, module: &'a Mod) {
         walk_module(self, module);
     }
 
+    #[inline]
     fn visit_stmt(&mut self, stmt: &'a Stmt) {
         walk_stmt(self, stmt);
     }
 
+    #[inline]
     fn visit_annotation(&mut self, expr: &'a Expr) {
         walk_annotation(self, expr);
     }
 
+    #[inline]
     fn visit_expr(&mut self, expr: &'a Expr) {
         walk_expr(self, expr);
     }
 
+    #[inline]
     fn visit_decorator(&mut self, decorator: &'a Decorator) {
         walk_decorator(self, decorator);
     }
 
+    #[inline]
     fn visit_constant(&mut self, _constant: &'a Constant) {}
 
+    #[inline]
     fn visit_bool_op(&mut self, bool_op: &'a BoolOp) {
         walk_bool_op(self, bool_op);
     }
 
+    #[inline]
     fn visit_operator(&mut self, operator: &'a Operator) {
         walk_operator(self, operator);
     }
 
+    #[inline]
     fn visit_unary_op(&mut self, unary_op: &'a UnaryOp) {
         walk_unary_op(self, unary_op);
     }
 
+    #[inline]
     fn visit_cmp_op(&mut self, cmp_op: &'a CmpOp) {
         walk_cmp_op(self, cmp_op);
     }
 
+    #[inline]
     fn visit_comprehension(&mut self, comprehension: &'a Comprehension) {
         walk_comprehension(self, comprehension);
     }
 
+    #[inline]
     fn visit_except_handler(&mut self, except_handler: &'a ExceptHandler) {
         walk_except_handler(self, except_handler);
     }
 
+    #[inline]
     fn visit_format_spec(&mut self, format_spec: &'a Expr) {
         walk_format_spec(self, format_spec);
     }
 
+    #[inline]
     fn visit_arguments(&mut self, arguments: &'a Arguments) {
         walk_arguments(self, arguments);
     }
 
+    #[inline]
     fn visit_parameters(&mut self, parameters: &'a Parameters) {
         walk_parameters(self, parameters);
     }
 
+    #[inline]
     fn visit_parameter(&mut self, arg: &'a Parameter) {
         walk_parameter(self, arg);
     }
@@ -73,38 +99,47 @@ pub trait PreorderVisitor<'a> {
         walk_parameter_with_default(self, parameter_with_default);
     }
 
+    #[inline]
     fn visit_keyword(&mut self, keyword: &'a Keyword) {
         walk_keyword(self, keyword);
     }
 
+    #[inline]
     fn visit_alias(&mut self, alias: &'a Alias) {
         walk_alias(self, alias);
     }
 
+    #[inline]
     fn visit_with_item(&mut self, with_item: &'a WithItem) {
         walk_with_item(self, with_item);
     }
 
+    #[inline]
     fn visit_type_params(&mut self, type_params: &'a TypeParams) {
         walk_type_params(self, type_params);
     }
 
+    #[inline]
     fn visit_type_param(&mut self, type_param: &'a TypeParam) {
         walk_type_param(self, type_param);
     }
 
+    #[inline]
     fn visit_match_case(&mut self, match_case: &'a MatchCase) {
         walk_match_case(self, match_case);
     }
 
+    #[inline]
     fn visit_pattern(&mut self, pattern: &'a Pattern) {
         walk_pattern(self, pattern);
     }
 
+    #[inline]
     fn visit_body(&mut self, body: &'a [Stmt]) {
         walk_body(self, body);
     }
 
+    #[inline]
     fn visit_elif_else_clause(&mut self, elif_else_clause: &'a ElifElseClause) {
         walk_elif_else_clause(self, elif_else_clause);
     }
@@ -114,12 +149,17 @@ pub fn walk_module<'a, V>(visitor: &mut V, module: &'a Mod)
 where
     V: PreorderVisitor<'a> + ?Sized,
 {
+    let node = AnyNodeRef::from(module);
+    if visitor.enter_node(node).is_traverse() {
         match module {
             Mod::Module(ast::ModModule { body, range: _ }) => {
                 visitor.visit_body(body);
             }
             Mod::Expression(ast::ModExpression { body, range: _ }) => visitor.visit_expr(body),
         }
+    }
+
+    visitor.leave_node(node);
 }
 
 pub fn walk_body<'a, V>(visitor: &mut V, body: &'a [Stmt])
@@ -135,6 +175,9 @@ pub fn walk_stmt<'a, V>(visitor: &mut V, stmt: &'a Stmt)
 where
     V: PreorderVisitor<'a> + ?Sized,
 {
+    let node = AnyNodeRef::from(stmt);
+
+    if visitor.enter_node(node).is_traverse() {
         match stmt {
             Stmt::Expr(ast::StmtExpr { value, range: _ }) => visitor.visit_expr(value),
 
@@ -377,23 +420,50 @@ where
             | Stmt::Nonlocal(_)
             | Stmt::LineMagic(_) => {}
         }
+    }
+
+    visitor.leave_node(node);
+}
+
+#[derive(Copy, Clone, Eq, PartialEq, Debug)]
+pub enum TraversalSignal {
+    Traverse,
+    Skip,
+}
+
+impl TraversalSignal {
+    const fn is_traverse(self) -> bool {
+        matches!(self, TraversalSignal::Traverse)
+    }
 }
 
 pub fn walk_annotation<'a, V: PreorderVisitor<'a> + ?Sized>(visitor: &mut V, expr: &'a Expr) {
+    let node = AnyNodeRef::from(expr);
+    if visitor.enter_node(node).is_traverse() {
         visitor.visit_expr(expr);
+    }
+
+    visitor.leave_node(node);
 }
 
 pub fn walk_decorator<'a, V>(visitor: &mut V, decorator: &'a Decorator)
 where
     V: PreorderVisitor<'a> + ?Sized,
 {
+    let node = AnyNodeRef::from(decorator);
+    if visitor.enter_node(node).is_traverse() {
         visitor.visit_expr(&decorator.expression);
+    }
+
+    visitor.leave_node(node);
 }
 
 pub fn walk_expr<'a, V>(visitor: &mut V, expr: &'a Expr)
 where
     V: PreorderVisitor<'a> + ?Sized,
 {
+    let node = AnyNodeRef::from(expr);
+    if visitor.enter_node(node).is_traverse() {
         match expr {
             Expr::BoolOp(ast::ExprBoolOp {
                 op,
@@ -651,34 +721,49 @@ where
             }
             Expr::LineMagic(_) => (),
         }
+    }
+
+    visitor.leave_node(node);
 }
 
 pub fn walk_comprehension<'a, V>(visitor: &mut V, comprehension: &'a Comprehension)
 where
     V: PreorderVisitor<'a> + ?Sized,
 {
+    let node = AnyNodeRef::from(comprehension);
+    if visitor.enter_node(node).is_traverse() {
         visitor.visit_expr(&comprehension.target);
         visitor.visit_expr(&comprehension.iter);
 
         for expr in &comprehension.ifs {
             visitor.visit_expr(expr);
         }
+    }
+
+    visitor.leave_node(node);
 }
 
 pub fn walk_elif_else_clause<'a, V>(visitor: &mut V, elif_else_clause: &'a ElifElseClause)
 where
     V: PreorderVisitor<'a> + ?Sized,
 {
+    let node = AnyNodeRef::from(elif_else_clause);
+    if visitor.enter_node(node).is_traverse() {
         if let Some(test) = &elif_else_clause.test {
             visitor.visit_expr(test);
         }
         visitor.visit_body(&elif_else_clause.body);
+    }
+
+    visitor.leave_node(node);
 }
 
 pub fn walk_except_handler<'a, V>(visitor: &mut V, except_handler: &'a ExceptHandler)
 where
     V: PreorderVisitor<'a> + ?Sized,
 {
+    let node = AnyNodeRef::from(except_handler);
+    if visitor.enter_node(node).is_traverse() {
         match except_handler {
             ExceptHandler::ExceptHandler(ast::ExceptHandlerExceptHandler {
                 range: _,
@@ -692,19 +777,28 @@ where
                 visitor.visit_body(body);
             }
         }
+    }
+    visitor.leave_node(node);
 }
 
 pub fn walk_format_spec<'a, V: PreorderVisitor<'a> + ?Sized>(
     visitor: &mut V,
     format_spec: &'a Expr,
 ) {
+    let node = AnyNodeRef::from(format_spec);
+    if visitor.enter_node(node).is_traverse() {
         visitor.visit_expr(format_spec);
+    }
+
+    visitor.leave_node(node);
 }
 
 pub fn walk_arguments<'a, V>(visitor: &mut V, arguments: &'a Arguments)
 where
     V: PreorderVisitor<'a> + ?Sized,
 {
+    let node = AnyNodeRef::from(arguments);
+    if visitor.enter_node(node).is_traverse() {
         for arg in &arguments.args {
             visitor.visit_expr(arg);
         }
@@ -712,12 +806,17 @@ where
         for keyword in &arguments.keywords {
             visitor.visit_keyword(keyword);
         }
+    }
+
+    visitor.leave_node(node);
 }
 
 pub fn walk_parameters<'a, V>(visitor: &mut V, parameters: &'a Parameters)
 where
     V: PreorderVisitor<'a> + ?Sized,
 {
+    let node = AnyNodeRef::from(parameters);
+    if visitor.enter_node(node).is_traverse() {
         for arg in parameters.posonlyargs.iter().chain(&parameters.args) {
             visitor.visit_parameter_with_default(arg);
         }
@@ -733,15 +832,23 @@ where
         if let Some(arg) = &parameters.kwarg {
             visitor.visit_parameter(arg);
         }
+    }
+
+    visitor.leave_node(node);
 }
 
 pub fn walk_parameter<'a, V>(visitor: &mut V, parameter: &'a Parameter)
 where
     V: PreorderVisitor<'a> + ?Sized,
 {
+    let node = AnyNodeRef::from(parameter);
+
+    if visitor.enter_node(node).is_traverse() {
         if let Some(expr) = &parameter.annotation {
             visitor.visit_annotation(expr);
         }
+    }
+    visitor.leave_node(node);
 }
 
 pub fn walk_parameter_with_default<'a, V>(
@@ -750,10 +857,15 @@ pub fn walk_parameter_with_default<'a, V>(
 ) where
     V: PreorderVisitor<'a> + ?Sized,
 {
+    let node = AnyNodeRef::from(parameter_with_default);
+    if visitor.enter_node(node).is_traverse() {
         visitor.visit_parameter(&parameter_with_default.parameter);
         if let Some(expr) = &parameter_with_default.default {
             visitor.visit_expr(expr);
         }
+    }
+
+    visitor.leave_node(node);
 }
 
 #[inline]
@@ -761,33 +873,48 @@ pub fn walk_keyword<'a, V>(visitor: &mut V, keyword: &'a Keyword)
 where
     V: PreorderVisitor<'a> + ?Sized,
 {
+    let node = AnyNodeRef::from(keyword);
+
+    if visitor.enter_node(node).is_traverse() {
         visitor.visit_expr(&keyword.value);
+    }
+    visitor.leave_node(node);
 }
 
 pub fn walk_with_item<'a, V>(visitor: &mut V, with_item: &'a WithItem)
 where
     V: PreorderVisitor<'a> + ?Sized,
 {
+    let node = AnyNodeRef::from(with_item);
+    if visitor.enter_node(node).is_traverse() {
         visitor.visit_expr(&with_item.context_expr);
 
         if let Some(expr) = &with_item.optional_vars {
             visitor.visit_expr(expr);
         }
+    }
+    visitor.leave_node(node);
 }
 
 pub fn walk_type_params<'a, V>(visitor: &mut V, type_params: &'a TypeParams)
 where
     V: PreorderVisitor<'a> + ?Sized,
 {
+    let node = AnyNodeRef::from(type_params);
+    if visitor.enter_node(node).is_traverse() {
         for type_param in &type_params.type_params {
             visitor.visit_type_param(type_param);
         }
+    }
+    visitor.leave_node(node);
 }
 
 pub fn walk_type_param<'a, V>(visitor: &mut V, type_param: &'a TypeParam)
 where
     V: PreorderVisitor<'a> + ?Sized,
 {
+    let node = AnyNodeRef::from(type_param);
+    if visitor.enter_node(node).is_traverse() {
         match type_param {
             TypeParam::TypeVar(TypeParamTypeVar {
                 bound,
@@ -800,23 +927,31 @@ where
             }
             TypeParam::TypeVarTuple(_) | TypeParam::ParamSpec(_) => {}
         }
+    }
+    visitor.leave_node(node);
 }
 
 pub fn walk_match_case<'a, V>(visitor: &mut V, match_case: &'a MatchCase)
 where
     V: PreorderVisitor<'a> + ?Sized,
 {
+    let node = AnyNodeRef::from(match_case);
+    if visitor.enter_node(node).is_traverse() {
         visitor.visit_pattern(&match_case.pattern);
         if let Some(expr) = &match_case.guard {
             visitor.visit_expr(expr);
         }
         visitor.visit_body(&match_case.body);
+    }
+    visitor.leave_node(node);
 }
 
 pub fn walk_pattern<'a, V>(visitor: &mut V, pattern: &'a Pattern)
 where
     V: PreorderVisitor<'a> + ?Sized,
 {
+    let node = AnyNodeRef::from(pattern);
+    if visitor.enter_node(node).is_traverse() {
         match pattern {
             Pattern::MatchValue(ast::PatternMatchValue { value, range: _ }) => {
                 visitor.visit_expr(value);
@@ -879,6 +1014,8 @@ where
                 }
             }
         }
+    }
+    visitor.leave_node(node);
 }
 
 pub fn walk_bool_op<'a, V>(_visitor: &mut V, _bool_op: &'a BoolOp)
@@ -909,8 +1046,11 @@ where
 }
 
 #[inline]
-pub fn walk_alias<'a, V>(_visitor: &mut V, _alias: &'a Alias)
+pub fn walk_alias<'a, V>(visitor: &mut V, alias: &'a Alias)
 where
     V: PreorderVisitor<'a> + ?Sized,
 {
+    let node = AnyNodeRef::from(alias);
+    visitor.enter_node(node);
+    visitor.leave_node(node);
 }

crates/ruff_python_formatter/src/comments/visitor.rs

@@ -1,10 +1,6 @@
 use std::iter::Peekable;
 
-use ruff_python_ast::{
-    Alias, Arguments, Comprehension, Decorator, ElifElseClause, ExceptHandler, Expr, Keyword,
-    MatchCase, Mod, Parameter, ParameterWithDefault, Parameters, Pattern, Ranged, Stmt, TypeParam,
-    TypeParams, WithItem,
-};
+use ruff_python_ast::{Mod, Ranged, Stmt};
 use ruff_text_size::{TextRange, TextSize};
 
 use ruff_formatter::{SourceCode, SourceCodeSlice};
@@ -48,14 +44,22 @@ impl<'a> CommentsVisitor<'a> {
         self.finish()
     }
 
-    fn start_node<N>(&mut self, node: N) -> TraversalSignal
-    where
-        N: Into<AnyNodeRef<'a>>,
-    {
-        self.start_node_impl(node.into())
+    // Try to skip the subtree if
+    // * there are no comments
+    // * if the next comment comes after this node (meaning, this nodes subtree contains no comments)
+    fn can_skip(&mut self, node_end: TextSize) -> bool {
+        self.comment_ranges
+            .peek()
+            .map_or(true, |next_comment| next_comment.start() >= node_end)
     }
 
-    fn start_node_impl(&mut self, node: AnyNodeRef<'a>) -> TraversalSignal {
+    fn finish(self) -> CommentsMap<'a> {
+        self.builder.finish()
+    }
+}
+
+impl<'ast> PreorderVisitor<'ast> for CommentsVisitor<'ast> {
+    fn enter_node(&mut self, node: AnyNodeRef<'ast>) -> TraversalSignal {
         let node_range = node.range();
 
         let enclosing_node = self.parents.last().copied().unwrap_or(node);
@@ -95,23 +99,7 @@ impl<'a> CommentsVisitor<'a> {
         }
     }
 
-    // Try to skip the subtree if
-    // * there are no comments
-    // * if the next comment comes after this node (meaning, this nodes subtree contains no comments)
-    fn can_skip(&mut self, node_end: TextSize) -> bool {
-        self.comment_ranges
-            .peek()
-            .map_or(true, |next_comment| next_comment.start() >= node_end)
-    }
-
-    fn finish_node<N>(&mut self, node: N)
-    where
-        N: Into<AnyNodeRef<'a>>,
-    {
-        self.finish_node_impl(node.into());
-    }
-
-    fn finish_node_impl(&mut self, node: AnyNodeRef<'a>) {
+    fn leave_node(&mut self, node: AnyNodeRef<'ast>) {
         // We are leaving this node, pop it from the parent stack.
         self.parents.pop();
 
@@ -146,19 +134,6 @@ impl<'a> CommentsVisitor<'a> {
         self.preceding_node = Some(node);
     }
 
-    fn finish(self) -> CommentsMap<'a> {
-        self.builder.finish()
-    }
-}
-
-impl<'ast> PreorderVisitor<'ast> for CommentsVisitor<'ast> {
-    fn visit_mod(&mut self, module: &'ast Mod) {
-        if self.start_node(module).is_traverse() {
-            walk_module(self, module);
-        }
-        self.finish_node(module);
-    }
-
     fn visit_body(&mut self, body: &'ast [Stmt]) {
         match body {
             [] => {
@@ -178,140 +153,6 @@ impl<'ast> PreorderVisitor<'ast> for CommentsVisitor<'ast> {
             }
         }
     }
-
-    fn visit_stmt(&mut self, stmt: &'ast Stmt) {
-        if self.start_node(stmt).is_traverse() {
-            walk_stmt(self, stmt);
-        }
-        self.finish_node(stmt);
-    }
-
-    fn visit_annotation(&mut self, expr: &'ast Expr) {
-        if self.start_node(expr).is_traverse() {
-            walk_expr(self, expr);
-        }
-        self.finish_node(expr);
-    }
-
-    fn visit_decorator(&mut self, decorator: &'ast Decorator) {
-        if self.start_node(decorator).is_traverse() {
-            walk_decorator(self, decorator);
-        }
-        self.finish_node(decorator);
-    }
-
-    fn visit_expr(&mut self, expr: &'ast Expr) {
-        if self.start_node(expr).is_traverse() {
-            walk_expr(self, expr);
-        }
-        self.finish_node(expr);
-    }
-
-    fn visit_comprehension(&mut self, comprehension: &'ast Comprehension) {
-        if self.start_node(comprehension).is_traverse() {
-            walk_comprehension(self, comprehension);
-        }
-        self.finish_node(comprehension);
-    }
-
-    fn visit_except_handler(&mut self, except_handler: &'ast ExceptHandler) {
-        if self.start_node(except_handler).is_traverse() {
-            walk_except_handler(self, except_handler);
-        }
-        self.finish_node(except_handler);
-    }
-
-    fn visit_format_spec(&mut self, format_spec: &'ast Expr) {
-        if self.start_node(format_spec).is_traverse() {
-            walk_expr(self, format_spec);
-        }
-        self.finish_node(format_spec);
-    }
-
-    fn visit_arguments(&mut self, arguments: &'ast Arguments) {
-        if self.start_node(arguments).is_traverse() {
-            walk_arguments(self, arguments);
-        }
-        self.finish_node(arguments);
-    }
-
-    fn visit_parameters(&mut self, parameters: &'ast Parameters) {
-        if self.start_node(parameters).is_traverse() {
-            walk_parameters(self, parameters);
-        }
-        self.finish_node(parameters);
-    }
-
-    fn visit_parameter(&mut self, arg: &'ast Parameter) {
-        if self.start_node(arg).is_traverse() {
-            walk_parameter(self, arg);
-        }
-        self.finish_node(arg);
-    }
-
-    fn visit_parameter_with_default(&mut self, parameter_with_default: &'ast ParameterWithDefault) {
-        if self.start_node(parameter_with_default).is_traverse() {
-            walk_parameter_with_default(self, parameter_with_default);
-        }
-        self.finish_node(parameter_with_default);
-    }
-
-    fn visit_keyword(&mut self, keyword: &'ast Keyword) {
-        if self.start_node(keyword).is_traverse() {
-            walk_keyword(self, keyword);
-        }
-        self.finish_node(keyword);
-    }
-
-    fn visit_alias(&mut self, alias: &'ast Alias) {
-        if self.start_node(alias).is_traverse() {
-            walk_alias(self, alias);
-        }
-        self.finish_node(alias);
-    }
-
-    fn visit_with_item(&mut self, with_item: &'ast WithItem) {
-        if self.start_node(with_item).is_traverse() {
-            walk_with_item(self, with_item);
-        }
-
-        self.finish_node(with_item);
-    }
-
-    fn visit_match_case(&mut self, match_case: &'ast MatchCase) {
-        if self.start_node(match_case).is_traverse() {
-            walk_match_case(self, match_case);
-        }
-        self.finish_node(match_case);
-    }
-
-    fn visit_pattern(&mut self, pattern: &'ast Pattern) {
-        if self.start_node(pattern).is_traverse() {
-            walk_pattern(self, pattern);
-        }
-        self.finish_node(pattern);
-    }
-
-    fn visit_elif_else_clause(&mut self, elif_else_clause: &'ast ElifElseClause) {
-        if self.start_node(elif_else_clause).is_traverse() {
-            walk_elif_else_clause(self, elif_else_clause);
-        }
-        self.finish_node(elif_else_clause);
-    }
-
-    fn visit_type_params(&mut self, type_params: &'ast TypeParams) {
-        if self.start_node(type_params).is_traverse() {
-            walk_type_params(self, type_params);
-        }
-        self.finish_node(type_params);
-    }
-
-    fn visit_type_param(&mut self, type_param: &'ast TypeParam) {
-        if self.start_node(type_param).is_traverse() {
-            walk_type_param(self, type_param);
-        }
-        self.finish_node(type_param);
-    }
 }
 
 fn text_position(comment_range: TextRange, source_code: SourceCode) -> CommentLinePosition {
@@ -663,18 +504,6 @@ impl<'a> CommentPlacement<'a> {
     }
 }
 
-#[derive(Copy, Clone, Eq, PartialEq, Debug)]
-enum TraversalSignal {
-    Traverse,
-    Skip,
-}
-
-impl TraversalSignal {
-    const fn is_traverse(self) -> bool {
-        matches!(self, TraversalSignal::Traverse)
-    }
-}
-
 #[derive(Clone, Debug, Default)]
 struct CommentsBuilder<'a> {
     comments: CommentsMap<'a>,