see if text ranges work instead

crates/ruff_python_formatter/src/comments/mod.rs

@@ -99,7 +99,6 @@ use ruff_formatter::{SourceCode, SourceCodeSlice};
 use ruff_python_ast::AnyNodeRef;
 use ruff_python_trivia::{CommentLinePosition, CommentRanges, SuppressionKind};
 use ruff_text_size::{Ranged, TextRange};
-use rustc_hash::FxHashSet;
 pub(crate) use visitor::collect_comments;
 
 use crate::comments::debug::{DebugComment, DebugComments};
@@ -512,21 +511,51 @@ pub(crate) fn has_skip_comment(trailing_comments: &[SourceComment], source: &str
     })
 }
 
-pub(crate) struct SuppressedNodes<'a>(FxHashSet<NodeRefEqualityKey<'a>>);
+pub(crate) struct SuppressedNodeRanges(Vec<TextRange>);
 
-impl<'a> SuppressedNodes<'a> {
+impl<'a> SuppressedNodeRanges {
     pub(crate) fn from_comments(comments: &Comments<'a>, source: &'a str) -> Self {
         let map = &comments.clone().data.comments;
-        Self(
-            map.keys()
-                .copied()
-                .filter(|key| has_skip_comment(map.trailing(key), source))
-                .collect(),
-        )
+
+        let mut ranges = map
+            .keys()
+            .copied()
+            .filter_map(|key| suppressed_range(key.node(), map.trailing(&key), source))
+            .collect::<Vec<_>>();
+
+        ranges.sort_by_key(ruff_text_size::Ranged::start);
+
+        Self(ranges)
     }
 
     pub(crate) fn contains(&self, node: AnyNodeRef<'_>) -> bool {
-        self.0.contains(&NodeRefEqualityKey::from_ref(node))
+        let target = node.range();
+        self.0
+            .binary_search_by(|range| {
+                if range.eq(&target) {
+                    std::cmp::Ordering::Equal
+                } else if range.start() < target.start() {
+                    std::cmp::Ordering::Less
+                } else {
+                    std::cmp::Ordering::Greater
+                }
+            })
+            .is_ok()
+    }
+}
+
+fn suppressed_range<'a>(
+    node: AnyNodeRef<'a>,
+    trailing_comments: &[SourceComment],
+    source: &'a str,
+) -> Option<TextRange> {
+    if has_skip_comment(trailing_comments, source) {
+        let end = node.end();
+        let start = node.start();
+
+        Some(TextRange::new(start, end))
+    } else {
+        None
     }
 }
 

crates/ruff_python_formatter/src/context.rs

@@ -7,14 +7,14 @@ use ruff_python_ast::str::Quote;
 use ruff_python_parser::Tokens;
 
 use crate::PyFormatOptions;
-use crate::comments::{Comments, SuppressedNodes};
+use crate::comments::{Comments, SuppressedNodeRanges};
 use crate::other::interpolated_string::InterpolatedStringContext;
 
 pub struct PyFormatContext<'a> {
     options: PyFormatOptions,
     contents: &'a str,
     comments: Comments<'a>,
-    suppressed_nodes: SuppressedNodes<'a>,
+    suppressed_nodes: SuppressedNodeRanges,
     tokens: &'a Tokens,
     node_level: NodeLevel,
     indent_level: IndentLevel,
@@ -36,7 +36,7 @@ impl<'a> PyFormatContext<'a> {
         options: PyFormatOptions,
         contents: &'a str,
         comments: Comments<'a>,
-        suppressed_nodes: SuppressedNodes<'a>,
+        suppressed_nodes: SuppressedNodeRanges,
         tokens: &'a Tokens,
     ) -> Self {
         Self {

crates/ruff_python_formatter/src/lib.rs

@@ -14,7 +14,8 @@ use ruff_python_trivia::CommentRanges;
 use ruff_text_size::{Ranged, TextRange};
 
 use crate::comments::{
-    Comments, SourceComment, SuppressedNodes, has_skip_comment, leading_comments, trailing_comments,
+    Comments, SourceComment, SuppressedNodeRanges, has_skip_comment, leading_comments,
+    trailing_comments,
 };
 pub use crate::context::PyFormatContext;
 pub use crate::db::Db;
@@ -174,7 +175,7 @@ where
 {
     let source_code = SourceCode::new(source);
     let comments = Comments::from_ast(parsed.syntax(), source_code, comment_ranges);
-    let suppressed_nodes = SuppressedNodes::from_comments(&comments, source);
+    let suppressed_nodes = SuppressedNodeRanges::from_comments(&comments, source);
 
     let formatted = format!(
         PyFormatContext::new(options, source, comments, suppressed_nodes, parsed.tokens()),

crates/ruff_python_formatter/src/range.rs

@@ -12,7 +12,7 @@ use ruff_python_trivia::{
 };
 use ruff_text_size::{Ranged, TextLen, TextRange, TextSize};
 
-use crate::comments::{Comments, SuppressedNodes};
+use crate::comments::{Comments, SuppressedNodeRanges};
 use crate::context::{IndentLevel, NodeLevel};
 use crate::prelude::*;
 use crate::statement::suite::DocstringStmt;
@@ -77,7 +77,7 @@ pub fn format_range(
     let source_code = SourceCode::new(source);
     let comment_ranges = CommentRanges::from(parsed.tokens());
     let comments = Comments::from_ast(parsed.syntax(), source_code, &comment_ranges);
-    let suppressed_nodes = SuppressedNodes::from_comments(&comments, source);
+    let suppressed_nodes = SuppressedNodeRanges::from_comments(&comments, source);
 
     let mut context = PyFormatContext::new(
         options.with_source_map_generation(SourceMapGeneration::Enabled),

crates/ruff_python_formatter/src/statement/suite.rs

@@ -945,7 +945,7 @@ mod tests {
     use ruff_python_trivia::CommentRanges;
 
     use crate::PyFormatOptions;
-    use crate::comments::{Comments, SuppressedNodes};
+    use crate::comments::{Comments, SuppressedNodeRanges};
     use crate::prelude::*;
     use crate::statement::suite::SuiteKind;
 
@@ -974,7 +974,7 @@ def trailing_func():
         let parsed = parse_module(source).unwrap();
         let comment_ranges = CommentRanges::from(parsed.tokens());
         let comments = Comments::from_ranges(&comment_ranges);
-        let suppressed_nodes = SuppressedNodes::from_comments(&comments, source);
+        let suppressed_nodes = SuppressedNodeRanges::from_comments(&comments, source);
 
         let context = PyFormatContext::new(
             PyFormatOptions::default(),

crates/ruff_python_formatter/src/string/docstring.rs

@@ -20,7 +20,7 @@ use {
 };
 
 use super::NormalizedString;
-use crate::comments::SuppressedNodes;
+use crate::comments::SuppressedNodeRanges;
 use crate::preview::is_no_chaperone_for_escaped_quote_in_triple_quoted_docstring_enabled;
 use crate::string::StringQuotes;
 use crate::{DocstringCodeLineWidth, FormatModuleError, prelude::*};
@@ -1586,7 +1586,7 @@ fn docstring_format_source(
     let comment_ranges = CommentRanges::from(parsed.tokens());
     let source_code = ruff_formatter::SourceCode::new(source);
     let comments = crate::Comments::from_ast(parsed.syntax(), source_code, &comment_ranges);
-    let suppressed_nodes = SuppressedNodes::from_comments(&comments, source_code.as_str());
+    let suppressed_nodes = SuppressedNodeRanges::from_comments(&comments, source_code.as_str());
 
     let ctx = PyFormatContext::new(options, source, comments, suppressed_nodes, parsed.tokens())
         .in_docstring(docstring_quote_style);