From 45f9fca228ddd4b347ef916e2bc555a612fa7cca Mon Sep 17 00:00:00 2001
From: konsti <konstin@mailbox.org>
Date: Thu, 7 Sep 2023 20:08:28 +0200
Subject: [PATCH] Reuse locator in formatter comments (#7227)

**Summary** The comment visitor used to rebuild the locator for every
comment. Instead, we now keep the locator on the builder. Follow-up to
#6813.

**Test Plan** No formatting changes.
---
 .../ruff_python_formatter/src/comments/mod.rs |  3 ++-
 .../src/comments/visitor.rs                   | 23 +++++++++++--------
 2 files changed, 16 insertions(+), 10 deletions(-)

diff --git a/crates/ruff_python_formatter/src/comments/mod.rs b/crates/ruff_python_formatter/src/comments/mod.rs
index 8a5566d4a6..cbc0ac01df 100644
--- a/crates/ruff_python_formatter/src/comments/mod.rs
+++ b/crates/ruff_python_formatter/src/comments/mod.rs
@@ -104,6 +104,7 @@ use ruff_python_ast::node::AnyNodeRef;
 use ruff_python_ast::visitor::preorder::{PreorderVisitor, TraversalSignal};
 use ruff_python_index::CommentRanges;
 use ruff_python_trivia::PythonWhitespace;
+use ruff_source_file::Locator;
 
 use crate::comments::debug::{DebugComment, DebugComments};
 use crate::comments::map::{LeadingDanglingTrailing, MultiMap};
@@ -325,7 +326,7 @@ impl<'a> Comments<'a> {
         let map = if comment_ranges.is_empty() {
             CommentsMap::new()
         } else {
-            let mut builder = CommentsMapBuilder::default();
+            let mut builder = CommentsMapBuilder::new(Locator::new(source_code.as_str()));
             CommentsVisitor::new(source_code, comment_ranges, &mut builder).visit(root);
             builder.finish()
         };
diff --git a/crates/ruff_python_formatter/src/comments/visitor.rs b/crates/ruff_python_formatter/src/comments/visitor.rs
index f2cfc2408a..c291e635e9 100644
--- a/crates/ruff_python_formatter/src/comments/visitor.rs
+++ b/crates/ruff_python_formatter/src/comments/visitor.rs
@@ -91,8 +91,7 @@ impl<'ast> PreorderVisitor<'ast> for CommentsVisitor<'ast, '_> {
                 slice: self.source_code.slice(*comment_range),
             };
 
-            self.builder
-                .push_comment(comment, &Locator::new(self.source_code.as_str()));
+            self.builder.push_comment(comment);
             self.comment_ranges.next();
         }
 
@@ -132,8 +131,7 @@ impl<'ast> PreorderVisitor<'ast> for CommentsVisitor<'ast, '_> {
                 slice: self.source_code.slice(*comment_range),
             };
 
-            self.builder
-                .push_comment(comment, &Locator::new(self.source_code.as_str()));
+            self.builder.push_comment(comment);
 
             self.comment_ranges.next();
         }
@@ -518,7 +516,7 @@ impl<'a> CommentPlacement<'a> {
 }
 
 pub(super) trait PushComment<'a> {
-    fn push_comment(&mut self, placement: DecoratedComment<'a>, locator: &Locator);
+    fn push_comment(&mut self, placement: DecoratedComment<'a>);
 }
 
 /// A storage for the [`CommentsVisitor`] that just pushes the decorated comments to a [`Vec`] for
@@ -529,21 +527,21 @@ struct CommentsVecBuilder<'a> {
 }
 
 impl<'a> PushComment<'a> for CommentsVecBuilder<'a> {
-    fn push_comment(&mut self, placement: DecoratedComment<'a>, _: &Locator) {
+    fn push_comment(&mut self, placement: DecoratedComment<'a>) {
         self.comments.push(placement);
     }
 }
 
 /// A storage for the [`CommentsVisitor`] that fixes the placement and stores the comments in a
 /// [`CommentsMap`].
-#[derive(Clone, Debug, Default)]
 pub(super) struct CommentsMapBuilder<'a> {
     comments: CommentsMap<'a>,
+    locator: Locator<'a>,
 }
 
 impl<'a> PushComment<'a> for CommentsMapBuilder<'a> {
-    fn push_comment(&mut self, comment: DecoratedComment<'a>, locator: &Locator) {
-        let placement = place_comment(comment, locator);
+    fn push_comment(&mut self, placement: DecoratedComment<'a>) {
+        let placement = place_comment(placement, &self.locator);
         match placement {
             CommentPlacement::Leading { node, comment } => {
                 self.push_leading_comment(node, comment);
@@ -605,6 +603,13 @@ impl<'a> PushComment<'a> for CommentsMapBuilder<'a> {
 }
 
 impl<'a> CommentsMapBuilder<'a> {
+    pub(crate) fn new(locator: Locator<'a>) -> Self {
+        Self {
+            comments: CommentsMap::default(),
+            locator,
+        }
+    }
+
     pub(crate) fn finish(self) -> CommentsMap<'a> {
         self.comments
     }