Use single lookup for leading, dangling, and trailing comments (#6589)

crates/ruff_python_formatter/src/expression/expr_attribute.rs

@@ -2,7 +2,7 @@ use ruff_formatter::{write, FormatRuleWithOptions};
 use ruff_python_ast::node::AnyNodeRef;
 use ruff_python_ast::{Constant, Expr, ExprAttribute, ExprConstant};
 
-use crate::comments::{leading_comments, trailing_comments};
+use crate::comments::{leading_comments, trailing_comments, SourceComment};
 use crate::expression::parentheses::{
     is_expression_parenthesized, NeedsParentheses, OptionalParentheses, Parentheses,
 };
@@ -150,7 +150,7 @@ impl FormatNodeRule<ExprAttribute> for FormatExprAttribute {
 
     fn fmt_dangling_comments(
         &self,
-        _node: &ExprAttribute,
+        _dangling_comments: &[SourceComment],
         _f: &mut PyFormatter,
     ) -> FormatResult<()> {
         // handle in `fmt_fields`

crates/ruff_python_formatter/src/expression/expr_bin_op.rs

@@ -9,7 +9,7 @@ use smallvec::SmallVec;
 use ruff_formatter::{format_args, write, FormatOwnedWithRule, FormatRefWithRule};
 use ruff_python_ast::node::{AnyNodeRef, AstNode};
 
-use crate::comments::{trailing_comments, trailing_node_comments};
+use crate::comments::{trailing_comments, trailing_node_comments, SourceComment};
 use crate::expression::expr_constant::ExprConstantLayout;
 use crate::expression::parentheses::{
     in_parentheses_only_group, in_parentheses_only_soft_line_break,
@@ -147,7 +147,11 @@ impl FormatNodeRule<ExprBinOp> for FormatExprBinOp {
         }
     }
 
-    fn fmt_dangling_comments(&self, _node: &ExprBinOp, _f: &mut PyFormatter) -> FormatResult<()> {
+    fn fmt_dangling_comments(
+        &self,
+        _dangling_comments: &[SourceComment],
+        _f: &mut PyFormatter,
+    ) -> FormatResult<()> {
         // Handled inside of `fmt_fields`
         Ok(())
     }

crates/ruff_python_formatter/src/expression/expr_bool_op.rs

@@ -1,7 +1,7 @@
 use crate::comments::leading_comments;
 use crate::expression::parentheses::{
     in_parentheses_only_group, in_parentheses_only_soft_line_break_or_space, NeedsParentheses,
-    OptionalParentheses, Parentheses,
+    OptionalParentheses,
 };
 use crate::prelude::*;
 use ruff_formatter::{write, FormatOwnedWithRule, FormatRefWithRule, FormatRuleWithOptions};
@@ -12,20 +12,20 @@ use super::parentheses::is_expression_parenthesized;
 
 #[derive(Default)]
 pub struct FormatExprBoolOp {
-    parentheses: Option<Parentheses>,
-    chained: bool,
+    layout: BoolOpLayout,
 }
 
-pub struct BoolOpLayout {
-    pub(crate) parentheses: Option<Parentheses>,
-    pub(crate) chained: bool,
+#[derive(Default, Copy, Clone)]
+pub enum BoolOpLayout {
+    #[default]
+    Default,
+    Chained,
 }
 
 impl FormatRuleWithOptions<ExprBoolOp, PyFormatContext<'_>> for FormatExprBoolOp {
     type Options = BoolOpLayout;
     fn with_options(mut self, options: Self::Options) -> Self {
-        self.parentheses = options.parentheses;
-        self.chained = options.chained;
+        self.layout = options;
         self
     }
 }
@@ -68,7 +68,7 @@ impl FormatNodeRule<ExprBoolOp> for FormatExprBoolOp {
             Ok(())
         });
 
-        if self.chained {
+        if matches!(self.layout, BoolOpLayout::Chained) {
             // Chained boolean operations should not be given a new group
             inner.fmt(f)
         } else {
@@ -101,13 +101,7 @@ impl Format<PyFormatContext<'_>> for FormatValue<'_> {
                 ) =>
             {
                 // Mark chained boolean operations e.g. `x and y or z` and avoid creating a new group
-                write!(
-                    f,
-                    [bool_op.format().with_options(BoolOpLayout {
-                        parentheses: None,
-                        chained: true,
-                    })]
-                )
+                write!(f, [bool_op.format().with_options(BoolOpLayout::Chained)])
             }
             _ => write!(f, [in_parentheses_only_group(&self.value.format())]),
         }

crates/ruff_python_formatter/src/expression/expr_compare.rs

@@ -1,27 +1,16 @@
-use crate::comments::leading_comments;
+use crate::comments::{leading_comments, SourceComment};
 use crate::expression::parentheses::{
     in_parentheses_only_group, in_parentheses_only_soft_line_break_or_space, NeedsParentheses,
-    OptionalParentheses, Parentheses,
+    OptionalParentheses,
 };
 use crate::prelude::*;
 use crate::FormatNodeRule;
-use ruff_formatter::{write, FormatOwnedWithRule, FormatRefWithRule, FormatRuleWithOptions};
+use ruff_formatter::{write, FormatOwnedWithRule, FormatRefWithRule};
 use ruff_python_ast::node::AnyNodeRef;
 use ruff_python_ast::{CmpOp, ExprCompare};
 
 #[derive(Default)]
-pub struct FormatExprCompare {
-    parentheses: Option<Parentheses>,
-}
-
-impl FormatRuleWithOptions<ExprCompare, PyFormatContext<'_>> for FormatExprCompare {
-    type Options = Option<Parentheses>;
-
-    fn with_options(mut self, options: Self::Options) -> Self {
-        self.parentheses = options;
-        self
-    }
-}
+pub struct FormatExprCompare;
 
 impl FormatNodeRule<ExprCompare> for FormatExprCompare {
     fn fmt_fields(&self, item: &ExprCompare, f: &mut PyFormatter) -> FormatResult<()> {
@@ -70,7 +59,11 @@ impl FormatNodeRule<ExprCompare> for FormatExprCompare {
         in_parentheses_only_group(&inner).fmt(f)
     }
 
-    fn fmt_dangling_comments(&self, _node: &ExprCompare, _f: &mut PyFormatter) -> FormatResult<()> {
+    fn fmt_dangling_comments(
+        &self,
+        _dangling_comments: &[SourceComment],
+        _f: &mut PyFormatter,
+    ) -> FormatResult<()> {
         // Node can not have dangling comments
         Ok(())
     }

crates/ruff_python_formatter/src/expression/expr_constant.rs

@@ -1,3 +1,4 @@
+use crate::comments::SourceComment;
 use ruff_formatter::FormatRuleWithOptions;
 use ruff_python_ast::node::AnyNodeRef;
 use ruff_python_ast::{Constant, ExprConstant, Ranged};
@@ -65,7 +66,7 @@ impl FormatNodeRule<ExprConstant> for FormatExprConstant {
 
     fn fmt_dangling_comments(
         &self,
-        _node: &ExprConstant,
+        _dangling_comments: &[SourceComment],
         _f: &mut PyFormatter,
     ) -> FormatResult<()> {
         Ok(())

crates/ruff_python_formatter/src/expression/expr_dict.rs

@@ -4,7 +4,7 @@ use ruff_python_ast::Ranged;
 use ruff_python_ast::{Expr, ExprDict};
 use ruff_text_size::TextRange;
 
-use crate::comments::leading_comments;
+use crate::comments::{leading_comments, SourceComment};
 use crate::expression::parentheses::{
     empty_parenthesized, parenthesized, NeedsParentheses, OptionalParentheses,
 };
@@ -89,7 +89,11 @@ impl FormatNodeRule<ExprDict> for FormatExprDict {
             .fmt(f)
     }
 
-    fn fmt_dangling_comments(&self, _node: &ExprDict, _f: &mut PyFormatter) -> FormatResult<()> {
+    fn fmt_dangling_comments(
+        &self,
+        _dangling_comments: &[SourceComment],
+        _f: &mut PyFormatter,
+    ) -> FormatResult<()> {
         // Handled by `fmt_fields`
         Ok(())
     }

crates/ruff_python_formatter/src/expression/expr_generator_exp.rs

@@ -2,7 +2,7 @@ use ruff_formatter::{format_args, write, Buffer, FormatResult, FormatRuleWithOpt
 use ruff_python_ast::node::AnyNodeRef;
 use ruff_python_ast::ExprGeneratorExp;
 
-use crate::comments::leading_comments;
+use crate::comments::{leading_comments, SourceComment};
 use crate::context::PyFormatContext;
 use crate::expression::parentheses::{parenthesized, NeedsParentheses, OptionalParentheses};
 use crate::prelude::*;
@@ -81,7 +81,7 @@ impl FormatNodeRule<ExprGeneratorExp> for FormatExprGeneratorExp {
 
     fn fmt_dangling_comments(
         &self,
-        _node: &ExprGeneratorExp,
+        _dangling_comments: &[SourceComment],
         _f: &mut PyFormatter,
     ) -> FormatResult<()> {
         // Handled as part of `fmt_fields`

crates/ruff_python_formatter/src/expression/expr_lambda.rs

@@ -1,4 +1,4 @@
-use crate::comments::dangling_node_comments;
+use crate::comments::{dangling_node_comments, SourceComment};
 use crate::context::PyFormatContext;
 use crate::expression::parentheses::{NeedsParentheses, OptionalParentheses};
 use crate::other::parameters::ParametersParentheses;
@@ -55,7 +55,11 @@ impl FormatNodeRule<ExprLambda> for FormatExprLambda {
         )
     }
 
-    fn fmt_dangling_comments(&self, _node: &ExprLambda, _f: &mut PyFormatter) -> FormatResult<()> {
+    fn fmt_dangling_comments(
+        &self,
+        _dangling_comments: &[SourceComment],
+        _f: &mut PyFormatter,
+    ) -> FormatResult<()> {
         // Override. Dangling comments are handled in `fmt_fields`.
         Ok(())
     }

crates/ruff_python_formatter/src/expression/expr_list.rs

@@ -1,3 +1,4 @@
+use crate::comments::SourceComment;
 use ruff_formatter::prelude::format_with;
 use ruff_python_ast::node::AnyNodeRef;
 use ruff_python_ast::{ExprList, Ranged};
@@ -37,7 +38,11 @@ impl FormatNodeRule<ExprList> for FormatExprList {
             .fmt(f)
     }
 
-    fn fmt_dangling_comments(&self, _node: &ExprList, _f: &mut PyFormatter) -> FormatResult<()> {
+    fn fmt_dangling_comments(
+        &self,
+        _dangling_comments: &[SourceComment],
+        _f: &mut PyFormatter,
+    ) -> FormatResult<()> {
         // Handled as part of `fmt_fields`
         Ok(())
     }

crates/ruff_python_formatter/src/expression/expr_list_comp.rs

@@ -2,6 +2,7 @@ use ruff_formatter::{format_args, write, FormatResult};
 use ruff_python_ast::node::AnyNodeRef;
 use ruff_python_ast::ExprListComp;
 
+use crate::comments::SourceComment;
 use crate::context::PyFormatContext;
 use crate::expression::parentheses::{parenthesized, NeedsParentheses, OptionalParentheses};
 use crate::prelude::*;
@@ -45,7 +46,7 @@ impl FormatNodeRule<ExprListComp> for FormatExprListComp {
 
     fn fmt_dangling_comments(
         &self,
-        _node: &ExprListComp,
+        _dangling_comments: &[SourceComment],
         _f: &mut PyFormatter,
     ) -> FormatResult<()> {
         // Handled as part of `fmt_fields`

crates/ruff_python_formatter/src/expression/expr_name.rs

@@ -1,3 +1,4 @@
+use crate::comments::SourceComment;
 use crate::expression::parentheses::{NeedsParentheses, OptionalParentheses};
 use crate::prelude::*;
 use crate::FormatNodeRule;
@@ -23,7 +24,11 @@ impl FormatNodeRule<ExprName> for FormatExprName {
         write!(f, [source_text_slice(*range, ContainsNewlines::No)])
     }
 
-    fn fmt_dangling_comments(&self, _node: &ExprName, _f: &mut PyFormatter) -> FormatResult<()> {
+    fn fmt_dangling_comments(
+        &self,
+        _dangling_comments: &[SourceComment],
+        _f: &mut PyFormatter,
+    ) -> FormatResult<()> {
         // Node cannot have dangling comments
         Ok(())
     }

crates/ruff_python_formatter/src/expression/expr_set.rs

@@ -1,3 +1,4 @@
+use crate::comments::SourceComment;
 use ruff_python_ast::node::AnyNodeRef;
 use ruff_python_ast::{ExprSet, Ranged};
 
@@ -28,7 +29,11 @@ impl FormatNodeRule<ExprSet> for FormatExprSet {
             .fmt(f)
     }
 
-    fn fmt_dangling_comments(&self, _node: &ExprSet, _f: &mut PyFormatter) -> FormatResult<()> {
+    fn fmt_dangling_comments(
+        &self,
+        _dangling_comments: &[SourceComment],
+        _f: &mut PyFormatter,
+    ) -> FormatResult<()> {
         // Handled as part of `fmt_fields`
         Ok(())
     }

crates/ruff_python_formatter/src/expression/expr_set_comp.rs

@@ -2,6 +2,7 @@ use ruff_formatter::{format_args, write, Buffer, FormatResult};
 use ruff_python_ast::node::AnyNodeRef;
 use ruff_python_ast::ExprSetComp;
 
+use crate::comments::SourceComment;
 use crate::context::PyFormatContext;
 use crate::expression::parentheses::{parenthesized, NeedsParentheses, OptionalParentheses};
 use crate::prelude::*;
@@ -43,7 +44,11 @@ impl FormatNodeRule<ExprSetComp> for FormatExprSetComp {
         )
     }
 
-    fn fmt_dangling_comments(&self, _node: &ExprSetComp, _f: &mut PyFormatter) -> FormatResult<()> {
+    fn fmt_dangling_comments(
+        &self,
+        _dangling_comments: &[SourceComment],
+        _f: &mut PyFormatter,
+    ) -> FormatResult<()> {
         // Handled as part of `fmt_fields`
         Ok(())
     }

crates/ruff_python_formatter/src/expression/expr_starred.rs

@@ -1,5 +1,6 @@
 use ruff_python_ast::ExprStarred;
 
+use crate::comments::SourceComment;
 use ruff_formatter::write;
 use ruff_python_ast::node::AnyNodeRef;
 
@@ -22,9 +23,11 @@ impl FormatNodeRule<ExprStarred> for FormatExprStarred {
         write!(f, [text("*"), value.format()])
     }
 
-    fn fmt_dangling_comments(&self, node: &ExprStarred, f: &mut PyFormatter) -> FormatResult<()> {
-        debug_assert_eq!(f.context().comments().dangling_comments(node), []);
-
+    fn fmt_dangling_comments(
+        &self,
+        _dangling_comments: &[SourceComment],
+        _f: &mut PyFormatter,
+    ) -> FormatResult<()> {
         Ok(())
     }
 }

crates/ruff_python_formatter/src/expression/expr_subscript.rs

@@ -2,7 +2,7 @@ use ruff_formatter::{format_args, write, FormatRuleWithOptions};
 use ruff_python_ast::node::{AnyNodeRef, AstNode};
 use ruff_python_ast::{Expr, ExprSubscript};
 
-use crate::comments::trailing_comments;
+use crate::comments::{trailing_comments, SourceComment};
 use crate::context::PyFormatContext;
 use crate::context::{NodeLevel, WithNodeLevel};
 use crate::expression::expr_tuple::TupleParentheses;
@@ -81,7 +81,7 @@ impl FormatNodeRule<ExprSubscript> for FormatExprSubscript {
 
     fn fmt_dangling_comments(
         &self,
-        _node: &ExprSubscript,
+        _dangling_comments: &[SourceComment],
         _f: &mut PyFormatter,
     ) -> FormatResult<()> {
         // Handled inside of `fmt_fields`

crates/ruff_python_formatter/src/expression/expr_tuple.rs

@@ -5,6 +5,7 @@ use ruff_python_ast::{Expr, Ranged};
 use ruff_text_size::TextRange;
 
 use crate::builders::parenthesize_if_expands;
+use crate::comments::SourceComment;
 use crate::expression::parentheses::{
     empty_parenthesized, parenthesized, NeedsParentheses, OptionalParentheses,
 };
@@ -162,7 +163,11 @@ impl FormatNodeRule<ExprTuple> for FormatExprTuple {
         }
     }
 
-    fn fmt_dangling_comments(&self, _node: &ExprTuple, _f: &mut PyFormatter) -> FormatResult<()> {
+    fn fmt_dangling_comments(
+        &self,
+        _dangling_comments: &[SourceComment],
+        _f: &mut PyFormatter,
+    ) -> FormatResult<()> {
         // Handled in `fmt_fields`
         Ok(())
     }

crates/ruff_python_formatter/src/expression/mod.rs

@@ -16,8 +16,6 @@ use crate::expression::parentheses::{
 };
 use crate::prelude::*;
 
-use self::expr_bool_op::BoolOpLayout;
-
 pub(crate) mod expr_attribute;
 pub(crate) mod expr_await;
 pub(crate) mod expr_bin_op;
@@ -69,13 +67,7 @@ impl FormatRule<Expr, PyFormatContext<'_>> for FormatExpr {
         let parentheses = self.parentheses;
 
         let format_expr = format_with(|f| match expression {
-            Expr::BoolOp(expr) => expr
-                .format()
-                .with_options(BoolOpLayout {
-                    parentheses: Some(parentheses),
-                    chained: false,
-                })
-                .fmt(f),
+            Expr::BoolOp(expr) => expr.format().fmt(f),
             Expr::NamedExpr(expr) => expr.format().fmt(f),
             Expr::BinOp(expr) => expr.format().fmt(f),
             Expr::UnaryOp(expr) => expr.format().fmt(f),
@@ -90,7 +82,7 @@ impl FormatRule<Expr, PyFormatContext<'_>> for FormatExpr {
             Expr::Await(expr) => expr.format().fmt(f),
             Expr::Yield(expr) => expr.format().fmt(f),
             Expr::YieldFrom(expr) => expr.format().fmt(f),
-            Expr::Compare(expr) => expr.format().with_options(Some(parentheses)).fmt(f),
+            Expr::Compare(expr) => expr.format().fmt(f),
             Expr::Call(expr) => expr.format().fmt(f),
             Expr::FormattedValue(expr) => expr.format().fmt(f),
             Expr::FString(expr) => expr.format().fmt(f),