Improve performance of loading range suppressions

crates/ruff_linter/src/linter.rs

@@ -404,7 +404,8 @@ pub fn add_noqa_to_path(
     );
 
     // Parse range suppression comments
-    let suppressions = Suppressions::from_tokens(settings, locator.contents(), parsed.tokens());
+    let suppressions =
+        Suppressions::from_tokens(settings, locator.contents(), parsed.tokens(), &indexer);
 
     // Generate diagnostics, ignoring any existing `noqa` directives.
     let diagnostics = check_path(
@@ -470,7 +471,8 @@ pub fn lint_only(
     );
 
     // Parse range suppression comments
-    let suppressions = Suppressions::from_tokens(settings, locator.contents(), parsed.tokens());
+    let suppressions =
+        Suppressions::from_tokens(settings, locator.contents(), parsed.tokens(), &indexer);
 
     // Generate diagnostics.
     let diagnostics = check_path(
@@ -579,7 +581,8 @@ pub fn lint_fix<'a>(
         );
 
         // Parse range suppression comments
-        let suppressions = Suppressions::from_tokens(settings, locator.contents(), parsed.tokens());
+        let suppressions =
+            Suppressions::from_tokens(settings, locator.contents(), parsed.tokens(), &indexer);
 
         // Generate diagnostics.
         let diagnostics = check_path(
@@ -961,7 +964,8 @@ mod tests {
             &locator,
             &indexer,
         );
-        let suppressions = Suppressions::from_tokens(settings, locator.contents(), parsed.tokens());
+        let suppressions =
+            Suppressions::from_tokens(settings, locator.contents(), parsed.tokens(), &indexer);
         let mut diagnostics = check_path(
             path,
             None,

crates/ruff_linter/src/rules/pyflakes/mod.rs

@@ -957,7 +957,7 @@ mod tests {
             &indexer,
         );
         let suppressions =
-            Suppressions::from_tokens(&settings, locator.contents(), parsed.tokens());
+            Suppressions::from_tokens(&settings, locator.contents(), parsed.tokens(), &indexer);
         let mut messages = check_path(
             Path::new("<filename>"),
             None,

crates/ruff_linter/src/suppression.rs

@@ -2,14 +2,15 @@ use compact_str::CompactString;
 use core::fmt;
 use ruff_db::diagnostic::Diagnostic;
 use ruff_diagnostics::{Edit, Fix};
-use ruff_python_ast::token::{TokenKind, Tokens};
+use ruff_python_ast::token::{Token, TokenKind, Tokens};
 use ruff_python_ast::whitespace::indentation;
-use rustc_hash::FxHashSet;
+use ruff_python_index::Indexer;
+use rustc_hash::{FxHashMap, FxHashSet};
 use std::cell::Cell;
 use std::{error::Error, fmt::Formatter};
 use thiserror::Error;
 
-use ruff_python_trivia::Cursor;
+use ruff_python_trivia::{CommentRanges, Cursor};
 use ruff_text_size::{Ranged, TextLen, TextRange, TextSize, TextSlice};
 use smallvec::{SmallVec, smallvec};
 
@@ -125,10 +126,17 @@ pub struct Suppressions {
 }
 
 impl Suppressions {
-    pub fn from_tokens(settings: &LinterSettings, source: &str, tokens: &Tokens) -> Suppressions {
+    pub fn from_tokens(
+        settings: &LinterSettings,
+        source: &str,
+        tokens: &Tokens,
+        indexer: &Indexer,
+    ) -> Suppressions {
         if is_range_suppressions_enabled(settings) {
             let builder = SuppressionsBuilder::new(source);
-            builder.load_from_tokens(tokens)
+            builder.load_from_tokens_indexed(tokens, indexer)
+            // let builder = SuppressionsBuilder::new(source);
+            // builder.load_from_tokens(tokens)
         } else {
             Suppressions::default()
         }
@@ -303,7 +311,7 @@ impl Suppressions {
     }
 }
 
-#[derive(Default)]
+#[derive(Debug, Default)]
 pub(crate) struct SuppressionsBuilder<'a> {
     source: &'a str,
 
@@ -312,6 +320,7 @@ pub(crate) struct SuppressionsBuilder<'a> {
     errors: Vec<ParseError>,
 
     pending: Vec<PendingSuppressionComment<'a>>,
+    pending_by_indent: FxHashMap<TextRange, Vec<SuppressionComment>>,
 }
 
 impl<'a> SuppressionsBuilder<'a> {
@@ -322,6 +331,89 @@ impl<'a> SuppressionsBuilder<'a> {
         }
     }
 
+    pub(crate) fn load_from_tokens_indexed(
+        mut self,
+        tokens: &Tokens,
+        indexer: &Indexer,
+    ) -> Suppressions {
+        let global_indent = TextRange::empty(0.into());
+
+        dbg!(indexer.block_ranges());
+        'outer: for comment_range in indexer.comment_ranges() {
+            dbg!(comment_range);
+            dbg!(self.source.slice(&comment_range));
+            let mut parser = SuppressionParser::new(self.source, comment_range);
+            match parser.parse_comment() {
+                Ok(comment) => {
+                    let Some(indent) = indentation(self.source, &comment_range) else {
+                        // trailing suppressions are not supported
+                        self.invalid.push(InvalidSuppression {
+                            kind: InvalidSuppressionKind::Trailing,
+                            comment,
+                        });
+                        continue;
+                    };
+                    let comment_indent = indent.text_len();
+
+                    let token_index = match tokens.binary_search_by_start(comment_range.start()) {
+                        Ok(index) => index,
+                        Err(index) => index,
+                    };
+                    let precedes_dedent = tokens[token_index..]
+                        .iter()
+                        .find(|t| !t.kind().is_trivia())
+                        .is_some_and(|token| matches!(token.kind(), TokenKind::Dedent));
+
+                    let block_ranges = indexer.block_ranges().containing(&comment_range);
+                    dbg!(&block_ranges);
+
+                    // no blocks, global scope
+                    if block_ranges.is_empty() && comment_indent == 0.into() {
+                        self.pending_by_indent
+                            .entry(global_indent)
+                            .or_default()
+                            .push(comment);
+                        continue 'outer;
+                    }
+
+                    for block in indexer
+                        .block_ranges()
+                        .containing(&comment_range)
+                        .iter()
+                        .rev()
+                    {
+                        let block_indent = block.indent.len();
+
+                        if comment_indent == block_indent {
+                            self.pending_by_indent
+                                .entry(block.indent)
+                                .or_default()
+                                .push(comment);
+                            continue 'outer;
+                        } else if comment_indent < block_indent && precedes_dedent {
+                            continue;
+                        }
+                        break;
+                    }
+
+                    // weirdly indented? ¯\_(ツ)_/¯
+                    self.invalid.push(InvalidSuppression {
+                        kind: InvalidSuppressionKind::Indentation,
+                        comment,
+                    });
+                }
+                Err(ParseError {
+                    kind: ParseErrorKind::NotASuppression,
+                    ..
+                }) => {}
+                Err(error) => self.errors.push(error),
+            }
+        }
+
+        dbg!(self);
+
+        Suppressions::default()
+    }
     pub(crate) fn load_from_tokens(mut self, tokens: &Tokens) -> Suppressions {
         let default_indent = "";
         let mut indents: Vec<&str> = vec![];
@@ -642,6 +734,7 @@ mod tests {
 
     use insta::assert_debug_snapshot;
     use itertools::Itertools;
+    use ruff_python_index::Indexer;
     use ruff_python_parser::{Mode, ParseOptions, parse};
     use ruff_text_size::{TextRange, TextSize};
     use similar::DiffableStr;
@@ -1568,10 +1661,12 @@ def bar():
         /// Parse all suppressions and errors in a module for testing
         fn debug(source: &'_ str) -> DebugSuppressions<'_> {
             let parsed = parse(source, ParseOptions::from(Mode::Module)).unwrap();
+            let indexer = Indexer::from_tokens(parsed.tokens(), source);
             let suppressions = Suppressions::from_tokens(
                 &LinterSettings::default().with_preview_mode(),
                 source,
                 parsed.tokens(),
+                &indexer,
             );
             DebugSuppressions {
                 source,

crates/ruff_linter/src/test.rs

@@ -235,7 +235,8 @@ pub(crate) fn test_contents<'a>(
         &locator,
         &indexer,
     );
-    let suppressions = Suppressions::from_tokens(settings, locator.contents(), parsed.tokens());
+    let suppressions =
+        Suppressions::from_tokens(settings, locator.contents(), parsed.tokens(), &indexer);
     let messages = check_path(
         path,
         path.parent()
@@ -303,7 +304,7 @@ pub(crate) fn test_contents<'a>(
             );
 
             let suppressions =
-                Suppressions::from_tokens(settings, locator.contents(), parsed.tokens());
+                Suppressions::from_tokens(settings, locator.contents(), parsed.tokens(), &indexer);
             let fixed_messages = check_path(
                 path,
                 None,

crates/ruff_python_index/src/indexer.rs

@@ -4,7 +4,7 @@
 use ruff_python_ast::Stmt;
 use ruff_python_ast::token::{TokenKind, Tokens};
 use ruff_python_trivia::{
-    CommentRanges, has_leading_content, has_trailing_content, is_python_whitespace,
+    BlockRanges, CommentRanges, has_leading_content, has_trailing_content, is_python_whitespace,
 };
 use ruff_source_file::LineRanges;
 use ruff_text_size::{Ranged, TextRange, TextSize};
@@ -26,6 +26,9 @@ pub struct Indexer {
 
     /// The range of all comments in the source document.
     comment_ranges: CommentRanges,
+
+    /// The ranges of all indent/dedent blocks in the source document.
+    block_ranges: BlockRanges,
 }
 
 impl Indexer {
@@ -36,6 +39,8 @@ impl Indexer {
         let mut multiline_ranges_builder = MultilineRangesBuilder::default();
         let mut continuation_lines = Vec::new();
         let mut comment_ranges = Vec::new();
+        let mut indent_ranges = Vec::new();
+        let mut dedent_ranges = Vec::new();
 
         // Token, end
         let mut prev_end = TextSize::default();
@@ -76,6 +81,12 @@ impl Indexer {
                 TokenKind::Comment => {
                     comment_ranges.push(token.range());
                 }
+                TokenKind::Indent => {
+                    indent_ranges.push(token.range());
+                }
+                TokenKind::Dedent => {
+                    dedent_ranges.push(token.range());
+                }
                 _ => {}
             }
 
@@ -87,6 +98,7 @@ impl Indexer {
             interpolated_string_ranges: interpolated_string_ranges_builder.finish(),
             multiline_ranges: multiline_ranges_builder.finish(),
             comment_ranges: CommentRanges::new(comment_ranges),
+            block_ranges: BlockRanges::new(indent_ranges, dedent_ranges),
         }
     }
 
@@ -95,6 +107,11 @@ impl Indexer {
         &self.comment_ranges
     }
 
+    /// Returns the block indent/dedent ranges.
+    pub const fn block_ranges(&self) -> &BlockRanges {
+        &self.block_ranges
+    }
+
     /// Returns the byte offset ranges of interpolated strings.
     pub const fn interpolated_string_ranges(&self) -> &InterpolatedStringRanges {
         &self.interpolated_string_ranges

crates/ruff_python_trivia/src/block_ranges.rs

@@ -0,0 +1,48 @@
+use ruff_text_size::TextRange;
+
+/// Stores the ranges of indents and dedents sorted by [`TextRange::start`] in increasing order.
+#[derive(Clone, Debug, Default)]
+pub struct BlockRanges {
+    raw: Vec<BlockRange>,
+}
+
+#[derive(Clone, Copy, Debug, Default)]
+pub struct BlockRange {
+    pub indent: TextRange,
+    pub dedent: TextRange,
+}
+
+impl BlockRanges {
+    pub fn new(indent_ranges: Vec<TextRange>, dedent_ranges: Vec<TextRange>) -> Self {
+        let mut index = 0;
+        let mut stack = Vec::new();
+        let mut blocks = Vec::new();
+
+        for dedent in &dedent_ranges {
+            while index < indent_ranges.len() && indent_ranges[index].end() < dedent.start() {
+                stack.push(indent_ranges[index]);
+                index += 1;
+            }
+
+            if let Some(indent) = stack.pop() {
+                blocks.push(BlockRange {
+                    indent,
+                    dedent: *dedent,
+                });
+            }
+        }
+
+        blocks.sort_by_key(|b| b.indent.start());
+
+        Self { raw: blocks }
+    }
+
+    pub fn containing(&self, range: &TextRange) -> Vec<&BlockRange> {
+        self.raw
+            .iter()
+            .filter(|block| {
+                block.indent.start() <= range.start() && block.dedent.end() > range.end()
+            })
+            .collect()
+    }
+}

crates/ruff_python_trivia/src/lib.rs

@@ -1,3 +1,4 @@
+mod block_ranges;
 mod comment_ranges;
 mod comments;
 mod cursor;
@@ -6,6 +7,7 @@ pub mod textwrap;
 mod tokenizer;
 mod whitespace;
 
+pub use block_ranges::BlockRanges;
 pub use comment_ranges::CommentRanges;
 pub use comments::*;
 pub use cursor::*;

crates/ruff_server/src/lint.rs

@@ -120,8 +120,12 @@ pub(crate) fn check(
     let directives = extract_directives(parsed.tokens(), Flags::all(), &locator, &indexer);
 
     // Parse range suppression comments
-    let suppressions =
-        Suppressions::from_tokens(&settings.linter, locator.contents(), parsed.tokens());
+    let suppressions = Suppressions::from_tokens(
+        &settings.linter,
+        locator.contents(),
+        parsed.tokens(),
+        &indexer,
+    );
 
     // Generate checks.
     let diagnostics = check_path(

crates/ruff_wasm/src/lib.rs

@@ -213,8 +213,12 @@ impl Workspace {
             &indexer,
         );
 
-        let suppressions =
-            Suppressions::from_tokens(&self.settings.linter, locator.contents(), parsed.tokens());
+        let suppressions = Suppressions::from_tokens(
+            &self.settings.linter,
+            locator.contents(),
+            parsed.tokens(),
+            &indexer,
+        );
 
         // Generate checks.
         let diagnostics = check_path(