Support suppression comments on the start or end of the range

Author: MichaReiser

crates/red_knot_python_semantic/resources/mdtest/suppressions/type-ignore.md

@@ -10,9 +10,8 @@ a = 4 + test  # type: ignore
 
 ## Multiline ranges
 
-A diagnostic with a multiline range can be suppressed
-by a comment on the same line as the diagnostic's start or end.
-This is the same behavior as Mypy's.
+A diagnostic with a multiline range can be suppressed by a comment on the same line as the
+diagnostic's start or end. This is the same behavior as Mypy's.
 
 ```py
 # fmt: off
@@ -31,16 +30,16 @@ y = (
 )
 ```
 
-This is different from Pyright where Pyright uses the intersection between
-the diagnostic range and any suppression comment. This can be problematic for nested expressions
-because a suppression in a child expression now suppresses errors in the outer expression.
+Pyright diverges from this behavior and instead applies a suppression if its range intersects with
+the diagnostic range. This can be problematic for nested expressions because a suppression in a
+child expression now suppresses errors in the outer expression.
 
-For example, the `type: ignore` comment in this example
-suppresses the error of adding `2` to `"test"` and
-adding `"other"` to the result of the cast.
+For example, the `type: ignore` comment in this example suppresses the error of adding `2` to
+`"test"` and adding `"other"` to the result of the cast.
 
 ```py path=nested.py
-from typing import cast 
+# fmt: off
+from typing import cast
 
 y = (
     cast(int, "test" +
@@ -87,8 +86,8 @@ a = test \
 
 ## Codes
 
-Mypy supports `type: ignore[code]`. Red Knot doesn't understand mypy's
-rule names. Therefore, ignore the codes and suppress all errors.
+Mypy supports `type: ignore[code]`. Red Knot doesn't understand mypy's rule names. Therefore, ignore
+the codes and suppress all errors.
 
 ```py
 a = test  # type: ignore[name-defined]

crates/red_knot_python_semantic/src/suppression.rs

@@ -1,7 +1,6 @@
 use ruff_python_parser::TokenKind;
 use ruff_source_file::LineRanges;
 use ruff_text_size::{Ranged, TextRange};
-use std::ops::RangeBounds;
 
 use ruff_db::{files::File, parsed::parsed_module, source::source_text};
 
@@ -12,6 +11,8 @@ pub(crate) fn suppressions(db: &dyn Db, file: File) -> Suppressions {
     let source = source_text(db.upcast(), file);
     let parsed = parsed_module(db.upcast(), file);
 
+    // TODO: Support `type: ignore` comments at the
+    //   [start of the file](https://typing.readthedocs.io/en/latest/spec/directives.html#type-ignore-comments).
     let mut suppressions = Vec::default();
     let mut line_start = source.bom_start_offset();
 
@@ -50,7 +51,9 @@ pub(crate) struct Suppressions {
 impl Suppressions {
     /// Finds a suppression for the specified lint.
     ///
-    /// Returns the first matching suppression if more than one suppression apply for the current line.
+    /// Returns the first matching suppression if more than one suppression apply to `range` and `id`.
+    ///
+    /// Returns `None` if the lint isn't suppressed.
     pub(crate) fn find_suppression(&self, range: TextRange, _id: LintId) -> Option<&Suppression> {
         // TODO(micha):
         //   * Test if the suppression suppresses the passed lint
@@ -63,48 +66,22 @@ impl Suppressions {
     /// start or end offset. This means the suppression is on the same line
     /// as the diagnostic's start or end.
     fn for_range(&self, range: TextRange) -> impl Iterator<Item = &Suppression> + '_ {
-        // First find the index of the suppression that starts closest to the range's start.
-        let start_offset = self
+        // First find the index of the suppression comment that ends right before the range
+        // starts. This allows us to skip suppressions that are not relevant for the range.
+        let end_offset = self
             .suppressions
             .binary_search_by_key(&range.start(), |suppression| {
-                suppression.suppressed_range.start()
+                suppression.suppressed_range.end()
             })
             .unwrap_or_else(|index| index);
 
-        // Search backward for suppressions that start before the range's start
-        // but overlap with `range`.
-        //
-        // ```python
-        // y = (
-        //     4 / 0  # type: ignore
-        //     ^----^ range
-        // ^--- suppression range --^
-        // )
-        // ```
-        self.suppressions[..start_offset]
+        // From here, search the remaining suppression comments for one that
+        // contains the range's start or end offset. Stop the search
+        // as soon as the suppression's range and the range no longer overlap.
+        self.suppressions[end_offset..]
             .iter()
-            .rev()
+            // Stop searching if the suppression starts after the range we're looking for.
             .take_while(move |suppression| range.end() >= suppression.suppressed_range.start())
-            .chain(
-                // Search forward for suppressions that start at or after the range's start
-                // but overlap with `range`.
-                //
-                // ```python
-                // y = (
-                //     4 /
-                //     ^--- range start
-                // ^------ suppression start
-                //     0  # type: ignore
-                //      ^- range end    ^---suppression end
-                // )
-                // ```
-                //
-                self.suppressions[start_offset..]
-                    .iter()
-                    .take_while(move |suppression| {
-                        range.end() >= suppression.suppressed_range.start()
-                    }),
-            )
             .filter(move |suppression| {
                 // Don't use intersect to avoid that suppressions on inner-expression
                 // ignore errors for outer expressions