Remove impossible states from version-comparison representation (#4696)

crates/ruff/src/rules/pyupgrade/rules/outdated_version_block.rs

@@ -32,64 +32,99 @@ impl AlwaysAutofixableViolation for OutdatedVersionBlock {
     }
 }
 
+/// The metadata for a version-comparison block.
 #[derive(Debug)]
 struct BlockMetadata {
-    /// The first non-whitespace token in the block.
+    /// The first `if` or `elif` token in the block, used to signal the start of the
-    starter: Tok,
+    /// version-comparison block.
-    /// The location of the first `elif` token, if any.
+    leading_token: StartToken,
-    elif: Option<TextSize>,
+    /// The first `elif` or `else` token following the start token, if any, used to signal the end
-    /// The location of the `else` token, if any.
+    /// of the version-comparison block.
-    else_: Option<TextSize>,
+    trailing_token: Option<EndToken>,
 }
 
-impl BlockMetadata {
+/// The set of tokens that can start a block, i.e., the first token in an `if` statement.
-    const fn new(starter: Tok, elif: Option<TextSize>, else_: Option<TextSize>) -> Self {
+#[derive(Debug)]
-        Self {
+enum StartTok {
-            starter,
+    If,
-            elif,
+    Elif,
-            else_,
+}
+
+impl StartTok {
+    fn from_tok(tok: &Tok) -> Option<Self> {
+        match tok {
+            Tok::If => Some(Self::If),
+            Tok::Elif => Some(Self::Elif),
+            _ => None,
         }
     }
 }
 
+#[derive(Debug)]
+struct StartToken {
+    tok: StartTok,
+    range: TextRange,
+}
+
+/// The set of tokens that can end a block, i.e., the first token in the subsequent `elif` or `else`
+/// branch that follows an `if` or `elif` statement.
+#[derive(Debug)]
+enum EndTok {
+    Elif,
+    Else,
+}
+
+impl EndTok {
+    fn from_tok(tok: &Tok) -> Option<Self> {
+        match tok {
+            Tok::Elif => Some(Self::Elif),
+            Tok::Else => Some(Self::Else),
+            _ => None,
+        }
+    }
+}
+
+#[derive(Debug)]
+struct EndToken {
+    tok: EndTok,
+    range: TextRange,
+}
+
 fn metadata<T>(locator: &Locator, located: &T, body: &[Stmt]) -> Option<BlockMetadata>
 where
     T: Ranged,
 {
     indentation(locator, located)?;
 
-    let mut starter = None;
+    let mut iter = lexer::lex_starts_at(
-    let mut elif = None;
-    let mut else_ = None;
-
-    let body_end = body.last().map(Ranged::end).unwrap();
-    for (tok, range) in lexer::lex_starts_at(
         locator.slice(located.range()),
         Mode::Module,
         located.start(),
     )
-    .flatten()
+    .flatten();
-    .filter(|(tok, _)| matches!(tok, Tok::If | Tok::Elif | Tok::Else))
+
-    {
+    // First the leading `if` or `elif` token.
-        if starter.is_none() {
+    let (tok, range) = iter.next()?;
-            starter = Some(tok.clone());
+    let leading_token = StartToken {
-        } else if range.start() < body_end {
+        tok: StartTok::from_tok(&tok)?,
-            // Continue until the end of the `if` body, thus ensuring that we don't
+        range,
-            // accidentally pick up an `else` or `elif` in the nested block.
+    };
-            continue;
+
-        // We only care about the first `elif` or `else` following the `if`.
+    // Skip any tokens until we reach the end of the `if` body.
-        } else if matches!(tok, Tok::Elif) {
+    let body_end = body.last()?.range().end();
-            elif = Some(range.start());
+
-            break;
+    // Find the trailing `elif` or `else` token, if any.
-        } else if matches!(tok, Tok::Else) {
+    let trailing_token = iter
-            else_ = Some(range.start());
+        .skip_while(|(_, range)| range.start() < body_end)
-            break;
+        .find_map(|(tok, range)| EndTok::from_tok(&tok).map(|tok| EndToken { tok, range }));
-        }
+
-    }
+    Some(BlockMetadata {
-    Some(BlockMetadata::new(starter.unwrap(), elif, else_))
+        leading_token,
+        trailing_token,
+    })
 }
 
-/// Converts a `BigInt` to a `u32`, if the number is negative, it will return 0
+/// Converts a `BigInt` to a `u32`. If the number is negative, it will return 0.
 fn bigint_to_u32(number: &BigInt) -> u32 {
     let the_number = number.to_u32_digits();
     match the_number.0 {
@@ -147,11 +182,11 @@ fn compare_version(if_version: &[u32], py_version: PythonVersion, or_equal: bool
 fn fix_py2_block(
     checker: &mut Checker,
     stmt: &Stmt,
-    body: &[Stmt],
     orelse: &[Stmt],
     block: &BlockMetadata,
 ) -> Option<Fix> {
-    if orelse.is_empty() {
+    let leading_token = &block.leading_token;
+    let Some(trailing_token) = &block.trailing_token else {
         // Delete the entire statement. If this is an `elif`, know it's the only child
         // of its parent, so avoid passing in the parent at all. Otherwise,
         // `delete_stmt` will erroneously include a `pass`.
@@ -160,7 +195,7 @@ fn fix_py2_block(
         let defined_in = checker.semantic_model().stmt_parent();
         return match delete_stmt(
             defined_by,
-            if block.starter == Tok::If {
+            if matches!(block.leading_token.tok, StartTok::If) {
                 defined_in
             } else {
                 None
@@ -180,13 +215,13 @@ fn fix_py2_block(
                 None
             }
         };
-    }
+    };
 
+    match (&leading_token.tok, &trailing_token.tok) {
         // If we only have an `if` and an `else`, dedent the `else` block.
-    if block.starter == Tok::If && block.elif.is_none() {
+        (StartTok::If, EndTok::Else) => {
-        let start = orelse.first().unwrap();
+            let start = orelse.first()?;
-        let end = orelse.last().unwrap();
+            let end = orelse.last()?;
-
             if indentation(checker.locator, start).is_none() {
                 // Inline `else` block (e.g., `else: x = 1`).
                 #[allow(deprecated)]
@@ -217,22 +252,27 @@ fn fix_py2_block(
                         ))
                     })
             }
-    } else {
-        let mut end_location = orelse.last().unwrap().start();
-        if block.starter == Tok::If && block.elif.is_some() {
-            // Turn the `elif` into an `if`.
-            end_location = block.elif.unwrap() + TextSize::from(2);
-        } else if block.starter == Tok::Elif {
-            if let Some(elif) = block.elif {
-                end_location = elif;
-            } else if let Some(else_) = block.else_ {
-                end_location = else_;
-            } else {
-                end_location = body.last().unwrap().end();
-            }
         }
+        (StartTok::If, EndTok::Elif) => {
+            // If we have an `if` and an `elif`, turn the `elif` into an `if`.
+            let start_location = leading_token.range.start();
+            let end_location = trailing_token.range.start() + TextSize::from(2);
             #[allow(deprecated)]
-        Some(Fix::unspecified(Edit::deletion(stmt.start(), end_location)))
+            Some(Fix::unspecified(Edit::deletion(
+                start_location,
+                end_location,
+            )))
+        }
+        (StartTok::Elif, _) => {
+            // If we have an `elif`, delete up to the `else` or the end of the statement.
+            let start_location = leading_token.range.start();
+            let end_location = trailing_token.range.start();
+            #[allow(deprecated)]
+            Some(Fix::unspecified(Edit::deletion(
+                start_location,
+                end_location,
+            )))
+        }
     }
 }
 
@@ -244,13 +284,12 @@ fn fix_py3_block(
     body: &[Stmt],
     block: &BlockMetadata,
 ) -> Option<Fix> {
-    match block.starter {
+    match block.leading_token.tok {
-        Tok::If => {
+        StartTok::If => {
             // If the first statement is an if, use the body of this statement, and ignore
             // the rest.
-            let start = body.first().unwrap();
+            let start = body.first()?;
-            let end = body.last().unwrap();
+            let end = body.last()?;
-
             if indentation(checker.locator, start).is_none() {
                 // Inline `if` block (e.g., `if ...: x = 1`).
                 #[allow(deprecated)]
@@ -282,10 +321,10 @@ fn fix_py3_block(
                     })
             }
         }
-        Tok::Elif => {
+        StartTok::Elif => {
             // Replace the `elif` with an `else, preserve the body of the elif, and remove
             // the rest.
-            let end = body.last().unwrap();
+            let end = body.last()?;
             let text = checker.locator.slice(TextRange::new(test.end(), end.end()));
             #[allow(deprecated)]
             Some(Fix::unspecified(Edit::range_replacement(
@@ -293,7 +332,6 @@ fn fix_py3_block(
                 stmt.range(),
             )))
         }
-        _ => None,
     }
 }
 
@@ -336,9 +374,7 @@ pub(crate) fn outdated_version_block(
                         let mut diagnostic = Diagnostic::new(OutdatedVersionBlock, stmt.range());
                         if checker.patch(diagnostic.kind.rule()) {
                             if let Some(block) = metadata(checker.locator, stmt, body) {
-                                if let Some(fix) =
+                                if let Some(fix) = fix_py2_block(checker, stmt, orelse, &block) {
-                                    fix_py2_block(checker, stmt, body, orelse, &block)
-                                {
                                     diagnostic.set_fix(fix);
                                 }
                             }
@@ -369,7 +405,7 @@ pub(crate) fn outdated_version_block(
                     let mut diagnostic = Diagnostic::new(OutdatedVersionBlock, stmt.range());
                     if checker.patch(diagnostic.kind.rule()) {
                         if let Some(block) = metadata(checker.locator, stmt, body) {
-                            if let Some(fix) = fix_py2_block(checker, stmt, body, orelse, &block) {
+                            if let Some(fix) = fix_py2_block(checker, stmt, orelse, &block) {
                                 diagnostic.set_fix(fix);
                             }
                         }