Remove suite body tracking from SemanticModel (#5848)

## Summary

The `SemanticModel` currently stores the "body" of a given `Suite`,
along with the current statement index. This is used to support "next
sibling" queries, but we only use this in exactly one place -- the rule
that simplifies constructs like this to `any` or `all`:

```python
for x in y:
    if x == 0:
        return True
return False
```

Instead of tracking the state, we can just do a (slightly more
expensive) traversal, by finding the node within its parent and
returning the next node in the body.

Note that we'll only have to do this extremely rarely -- namely, for
functions that contain something like:

```python
for x in y:
    if x == 0:
        return True
```

crates/ruff_python_ast/src/lib.rs

@@ -15,6 +15,7 @@ pub mod statement_visitor;
 pub mod stmt_if;
 pub mod str;
 pub mod token_kind;
+pub mod traversal;
 pub mod types;
 pub mod typing;
 pub mod visitor;

crates/ruff_python_ast/src/traversal.rs

@@ -0,0 +1,113 @@
+//! Utilities for manually traversing a Python AST.
+use rustpython_ast::{ExceptHandler, Stmt, Suite};
+use rustpython_parser::ast;
+
+/// Given a [`Stmt`] and its parent, return the [`Suite`] that contains the [`Stmt`].
+pub fn suite<'a>(stmt: &'a Stmt, parent: &'a Stmt) -> Option<&'a Suite> {
+    match parent {
+        Stmt::FunctionDef(ast::StmtFunctionDef { body, .. }) => Some(body),
+        Stmt::AsyncFunctionDef(ast::StmtAsyncFunctionDef { body, .. }) => Some(body),
+        Stmt::ClassDef(ast::StmtClassDef { body, .. }) => Some(body),
+        Stmt::For(ast::StmtFor { body, orelse, .. }) => {
+            if body.contains(stmt) {
+                Some(body)
+            } else if orelse.contains(stmt) {
+                Some(orelse)
+            } else {
+                None
+            }
+        }
+        Stmt::AsyncFor(ast::StmtAsyncFor { body, orelse, .. }) => {
+            if body.contains(stmt) {
+                Some(body)
+            } else if orelse.contains(stmt) {
+                Some(orelse)
+            } else {
+                None
+            }
+        }
+        Stmt::While(ast::StmtWhile { body, orelse, .. }) => {
+            if body.contains(stmt) {
+                Some(body)
+            } else if orelse.contains(stmt) {
+                Some(orelse)
+            } else {
+                None
+            }
+        }
+        Stmt::If(ast::StmtIf {
+            body,
+            elif_else_clauses,
+            ..
+        }) => {
+            if body.contains(stmt) {
+                Some(body)
+            } else {
+                elif_else_clauses
+                    .iter()
+                    .map(|elif_else_clause| &elif_else_clause.body)
+                    .find(|body| body.contains(stmt))
+            }
+        }
+        Stmt::With(ast::StmtWith { body, .. }) => Some(body),
+        Stmt::AsyncWith(ast::StmtAsyncWith { body, .. }) => Some(body),
+        Stmt::Match(ast::StmtMatch { cases, .. }) => cases
+            .iter()
+            .map(|case| &case.body)
+            .find(|body| body.contains(stmt)),
+        Stmt::Try(ast::StmtTry {
+            body,
+            handlers,
+            orelse,
+            finalbody,
+            ..
+        }) => {
+            if body.contains(stmt) {
+                Some(body)
+            } else if orelse.contains(stmt) {
+                Some(orelse)
+            } else if finalbody.contains(stmt) {
+                Some(finalbody)
+            } else {
+                handlers
+                    .iter()
+                    .filter_map(ExceptHandler::as_except_handler)
+                    .map(|handler| &handler.body)
+                    .find(|body| body.contains(stmt))
+            }
+        }
+        Stmt::TryStar(ast::StmtTryStar {
+            body,
+            handlers,
+            orelse,
+            finalbody,
+            ..
+        }) => {
+            if body.contains(stmt) {
+                Some(body)
+            } else if orelse.contains(stmt) {
+                Some(orelse)
+            } else if finalbody.contains(stmt) {
+                Some(finalbody)
+            } else {
+                handlers
+                    .iter()
+                    .filter_map(ExceptHandler::as_except_handler)
+                    .map(|handler| &handler.body)
+                    .find(|body| body.contains(stmt))
+            }
+        }
+        _ => None,
+    }
+}
+
+/// Given a [`Stmt`] and its containing [`Suite`], return the next [`Stmt`] in the [`Suite`].
+pub fn next_sibling<'a>(stmt: &'a Stmt, suite: &'a Suite) -> Option<&'a Stmt> {
+    let mut iter = suite.iter();
+    while let Some(sibling) = iter.next() {
+        if sibling == stmt {
+            return iter.next();
+        }
+    }
+    None
+}