Cache parenthesized expression boundaries in the formatter

crates/ruff_benchmark/benches/formatter.rs

@@ -9,7 +9,7 @@ use ruff_benchmark::{
 };
 use ruff_python_formatter::{PreviewMode, PyFormatOptions, format_module_ast};
 use ruff_python_parser::{Mode, ParseOptions, parse};
-use ruff_python_trivia::CommentRanges;
+use ruff_python_trivia::TriviaRanges;
 
 #[cfg(target_os = "windows")]
 #[global_allocator]
@@ -53,11 +53,11 @@ fn benchmark_formatter(criterion: &mut Criterion) {
                     .expect("Input should be a valid Python code");
 
                 b.iter(|| {
-                    let comment_ranges = CommentRanges::from(parsed.tokens());
+                    let trivia_ranges = TriviaRanges::from(parsed.tokens());
                     let options = PyFormatOptions::from_extension(Path::new(case.name()))
                         .with_preview(PreviewMode::Enabled);
                     let formatted =
-                        format_module_ast(&parsed, &comment_ranges, case.code(), options)
+                        format_module_ast(&parsed, &trivia_ranges, case.code(), options)
                             .expect("Formatting to succeed");
 
                     formatted.print().expect("Printing to succeed")

crates/ruff_python_ast/src/token/tokens.rs

@@ -1,8 +1,9 @@
 use std::{iter::FusedIterator, ops::Deref};
 
 use super::{Token, TokenKind};
-use ruff_python_trivia::CommentRanges;
+use ruff_python_trivia::{CommentRanges, ParenthesizedExpressions, TriviaRanges};
 use ruff_text_size::{Ranged as _, TextRange, TextSize};
+use rustc_hash::FxHashSet;
 
 /// Tokens represents a vector of lexed [`Token`].
 #[derive(Debug, Clone, PartialEq, Eq)]
@@ -279,15 +280,62 @@ impl FusedIterator for TokenAt {}
 impl From<&Tokens> for CommentRanges {
     fn from(tokens: &Tokens) -> Self {
         let mut ranges = vec![];
+
         for token in tokens {
             if token.kind() == TokenKind::Comment {
                 ranges.push(token.range());
             }
         }
+
         CommentRanges::new(ranges)
     }
 }
 
+impl From<&Tokens> for TriviaRanges {
+    fn from(tokens: &Tokens) -> Self {
+        let mut comments = vec![];
+        let mut parenthesized = FxHashSet::default();
+        let mut stack = Vec::<Option<TextSize>>::new();
+        let mut previous_end = None;
+
+        for token in tokens {
+            if token.kind() == TokenKind::Comment {
+                comments.push(token.range());
+            }
+
+            if token.kind().is_trivia() {
+                continue;
+            }
+
+            match token.kind() {
+                TokenKind::Lpar => {
+                    if let Some(start) = stack.last_mut() {
+                        start.get_or_insert(token.start());
+                    }
+                    stack.push(None);
+                }
+                TokenKind::Rpar => {
+                    if let (Some(Some(start)), Some(end)) = (stack.pop(), previous_end) {
+                        parenthesized.insert(TextRange::new(start, end));
+                    }
+                }
+                _ => {
+                    if let Some(start) = stack.last_mut() {
+                        start.get_or_insert(token.start());
+                    }
+                }
+            }
+
+            previous_end = Some(token.end());
+        }
+
+        TriviaRanges::new(
+            CommentRanges::new(comments),
+            ParenthesizedExpressions::new(parenthesized),
+        )
+    }
+}
+
 /// An iterator over the [`Token`]s with context.
 ///
 /// This struct is created by the [`iter_with_context`] method on [`Tokens`]. Refer to its

crates/ruff_python_formatter/src/cli.rs

@@ -8,7 +8,7 @@ use clap::{Parser, ValueEnum};
 use ruff_formatter::SourceCode;
 use ruff_python_ast::{PySourceType, PythonVersion};
 use ruff_python_parser::{ParseOptions, parse};
-use ruff_python_trivia::CommentRanges;
+use ruff_python_trivia::TriviaRanges;
 use ruff_text_size::Ranged;
 
 use crate::comments::collect_comments;
@@ -70,15 +70,15 @@ pub fn format_and_debug_print(source: &str, cli: &Cli, source_path: &Path) -> Re
         .with_target_version(cli.target_version);
 
     let source_code = SourceCode::new(source);
-    let comment_ranges = CommentRanges::from(parsed.tokens());
-    let formatted = format_module_ast(&parsed, &comment_ranges, source, options)
-        .context("Failed to format node")?;
+    let trivia = TriviaRanges::from(parsed.tokens());
+    let formatted =
+        format_module_ast(&parsed, &trivia, source, options).context("Failed to format node")?;
     if cli.print_ir {
         println!("{}", formatted.document().display(source_code));
     }
     if cli.print_comments {
         // Print preceding, following and enclosing nodes
-        let decorated_comments = collect_comments(parsed.syntax(), source_code, &comment_ranges);
+        let decorated_comments = collect_comments(parsed.syntax(), source_code, trivia.comments());
         if !decorated_comments.is_empty() {
             println!("# Comment decoration: Range, Preceding, Following, Enclosing, Comment");
         }

crates/ruff_python_formatter/src/comments/debug.rs

@@ -186,7 +186,7 @@ mod tests {
     use ruff_formatter::SourceCode;
     use ruff_python_ast::{AnyNodeRef, AtomicNodeIndex};
     use ruff_python_ast::{StmtBreak, StmtContinue};
-    use ruff_python_trivia::{CommentLinePosition, CommentRanges};
+    use ruff_python_trivia::CommentLinePosition;
     use ruff_text_size::{TextRange, TextSize};
 
     use crate::comments::map::MultiMap;
@@ -238,8 +238,7 @@ break;
             ),
         );
 
-        let comment_ranges = CommentRanges::default();
-        let comments = Comments::new(comments_map, &comment_ranges);
+        let comments = Comments::new(comments_map);
 
         assert_debug_snapshot!(comments.debug(source_code));
     }

crates/ruff_python_formatter/src/comments/mod.rs

@@ -97,7 +97,7 @@ pub(crate) use format::{
 };
 use ruff_formatter::{SourceCode, SourceCodeSlice};
 use ruff_python_ast::AnyNodeRef;
-use ruff_python_trivia::{CommentLinePosition, CommentRanges, SuppressionKind};
+use ruff_python_trivia::{CommentLinePosition, SuppressionKind, TriviaRanges};
 use ruff_text_size::{Ranged, TextRange};
 pub(crate) use visitor::collect_comments;
 
@@ -219,53 +219,41 @@ pub(crate) struct Comments<'a> {
 }
 
 impl<'a> Comments<'a> {
-    fn new(comments: CommentsMap<'a>, comment_ranges: &'a CommentRanges) -> Self {
+    fn new(comments: CommentsMap<'a>) -> Self {
         Self {
-            data: Rc::new(CommentsData {
-                comments,
-                comment_ranges,
-            }),
+            data: Rc::new(CommentsData { comments }),
         }
     }
 
-    /// Effectively a [`Default`] implementation that works around the lifetimes for tests
+    /// Creates an empty comment map for tests.
     #[cfg(test)]
-    pub(crate) fn from_ranges(comment_ranges: &'a CommentRanges) -> Self {
-        Self {
-            data: Rc::new(CommentsData {
-                comments: CommentsMap::default(),
-                comment_ranges,
-            }),
-        }
-    }
-
-    pub(crate) fn ranges(&self) -> &'a CommentRanges {
-        self.data.comment_ranges
+    pub(crate) fn empty() -> Self {
+        Self::new(CommentsMap::default())
     }
 
     /// Extracts the comments from the AST.
     pub(crate) fn from_ast(
         root: impl Into<AnyNodeRef<'a>>,
         source_code: SourceCode<'a>,
-        comment_ranges: &'a CommentRanges,
+        trivia: &'a TriviaRanges,
     ) -> Self {
         fn collect_comments<'a>(
             root: AnyNodeRef<'a>,
             source_code: SourceCode<'a>,
-            comment_ranges: &'a CommentRanges,
+            trivia: &'a TriviaRanges,
         ) -> Comments<'a> {
-            let map = if comment_ranges.is_empty() {
+            let map = if trivia.comments().is_empty() {
                 CommentsMap::new()
             } else {
-                let mut builder = CommentsMapBuilder::new(source_code.as_str(), comment_ranges);
-                CommentsVisitor::new(source_code, comment_ranges, &mut builder).visit(root);
+                let mut builder = CommentsMapBuilder::new(source_code.as_str(), trivia);
+                CommentsVisitor::new(source_code, trivia.comments(), &mut builder).visit(root);
                 builder.finish()
             };
 
-            Comments::new(map, comment_ranges)
+            Comments::new(map)
         }
 
-        collect_comments(root.into(), source_code, comment_ranges)
+        collect_comments(root.into(), source_code, trivia)
     }
 
     /// Returns `true` if the given `node` has any comments.
@@ -496,9 +484,6 @@ impl LeadingDanglingTrailingComments<'_> {
 #[derive(Debug)]
 struct CommentsData<'a> {
     comments: CommentsMap<'a>,
-
-    /// We need those for backwards lexing
-    comment_ranges: &'a CommentRanges,
 }
 
 pub(crate) fn has_skip_comment(trailing_comments: &[SourceComment], source: &str) -> bool {
@@ -518,13 +503,13 @@ mod tests {
     use ruff_formatter::SourceCode;
     use ruff_python_ast::{Mod, PySourceType};
     use ruff_python_parser::{ParseOptions, Parsed, parse};
-    use ruff_python_trivia::CommentRanges;
+    use ruff_python_trivia::TriviaRanges;
 
     use crate::comments::Comments;
 
     struct CommentsTestCase<'a> {
         parsed: Parsed<Mod>,
-        comment_ranges: CommentRanges,
+        trivia: TriviaRanges,
         source_code: SourceCode<'a>,
     }
 
@@ -534,17 +519,17 @@ mod tests {
             let source_type = PySourceType::Python;
             let parsed = parse(source, ParseOptions::from(source_type))
                 .expect("Expect source to be valid Python");
-            let comment_ranges = CommentRanges::from(parsed.tokens());
+            let trivia = TriviaRanges::from(parsed.tokens());
 
             CommentsTestCase {
                 parsed,
-                comment_ranges,
+                trivia,
                 source_code,
             }
         }
 
         fn to_comments(&self) -> Comments<'_> {
-            Comments::from_ast(self.parsed.syntax(), self.source_code, &self.comment_ranges)
+            Comments::from_ast(self.parsed.syntax(), self.source_code, &self.trivia)
         }
     }