[red-knot] Add control flow for try/except blocks

crates/red_knot_python_semantic/src/semantic_index/builder.rs

@@ -31,6 +31,9 @@ use super::definition::{
     DefinitionCategory, ExceptHandlerDefinitionNodeRef, MatchPatternDefinitionNodeRef,
     WithItemDefinitionNodeRef,
 };
+use except_handlers::{TryBlockContexts, TryBlockContextsRefMut, TryBlockContextsStack};
+
+mod except_handlers;
 
 pub(super) struct SemanticIndexBuilder<'db> {
     // Builder state
@@ -44,6 +47,8 @@ pub(super) struct SemanticIndexBuilder<'db> {
     current_match_case: Option<CurrentMatchCase<'db>>,
     /// Flow states at each `break` in the current loop.
     loop_break_states: Vec<FlowSnapshot>,
+    /// Per-scope contexts regarding nested `try`/`except` statements
+    try_block_contexts_stack: TryBlockContextsStack,
 
     // Semantic Index fields
     scopes: IndexVec<FileScopeId, Scope>,
@@ -67,6 +72,7 @@ impl<'db> SemanticIndexBuilder<'db> {
             current_assignment: None,
             current_match_case: None,
             loop_break_states: vec![],
+            try_block_contexts_stack: TryBlockContextsStack::default(),
 
             scopes: IndexVec::new(),
             symbol_tables: IndexVec::new(),
@@ -92,6 +98,16 @@ impl<'db> SemanticIndexBuilder<'db> {
             .expect("Always to have a root scope")
     }
 
+    fn try_block_contexts(&self) -> &TryBlockContexts {
+        self.try_block_contexts_stack.current_try_block_context()
+    }
+
+    fn try_block_contexts_mut(&self) -> TryBlockContextsRefMut {
+        self.try_block_contexts_stack
+            .current_try_block_context()
+            .borrow_mut()
+    }
+
     fn push_scope(&mut self, node: NodeWithScopeRef) {
         let parent = self.current_scope();
         self.push_scope_with_parent(node, Some(parent));
@@ -105,6 +121,7 @@ impl<'db> SemanticIndexBuilder<'db> {
             kind: node.scope_kind(),
             descendents: children_start..children_start,
         };
+        self.try_block_contexts_stack.push_context();
 
         let file_scope_id = self.scopes.push(scope);
         self.symbol_tables.push(SymbolTableBuilder::new());
@@ -134,6 +151,7 @@ impl<'db> SemanticIndexBuilder<'db> {
         let children_end = self.scopes.next_index();
         let scope = &mut self.scopes[id];
         scope.descendents = scope.descendents.start..children_end;
+        self.try_block_contexts_stack.pop_context();
         id
     }
 
@@ -222,6 +240,10 @@ impl<'db> SemanticIndexBuilder<'db> {
             DefinitionCategory::Binding => use_def.record_binding(symbol, definition),
         }
 
+        if let Some(current_try_block) = self.try_block_contexts_mut().current_try_block() {
+            current_try_block.record_definition(self);
+        }
+
         definition
     }
 
@@ -738,9 +760,55 @@ where
                 is_star,
                 range: _,
             }) => {
+                // Save the state prior to visiting any of the `try` block.
+                //
+                // Potentially none of the `try` block could have been executed prior to executing
+                // the `except` block(s), so we will merge this state with all of the intermediate
+                // states during the `try` block before visiting the `except` blocks.
+                let pre_try_block_state = self.flow_snapshot();
+                self.try_block_contexts().push_try_block();
+
+                // Visit the `try` block!
+                //
+                // If `visiting_try_block == true`, this informs `.add_definition()` to record the state
+                // in `self.try_block_definition_states` after each new definition is recorded.
                 self.visit_body(body);
 
-                for except_handler in handlers {
+                // Save the state immediately *after* visiting the `try` block
+                // but *before* we prepare for visiting the `except` block(s).
+                //
+                // We will revert to this state prior to visiting the the `else` block,
+                // as there necessarily must have been 0 `except` blocks executed
+                // if we hit the `else` block.
+                let post_try_block_state = self.flow_snapshot();
+
+                // Prepare for visiting the `except` block(s)
+
+                let visiting_try_block_states = self
+                    .try_block_contexts()
+                    .pop_try_block()
+                    .expect("A `try` block should have been pushed to the stack prior to visiting the `body` field");
+                if let Some(parent_try_block) =
+                    self.try_block_contexts().borrow_mut().current_try_block()
+                {
+                    parent_try_block.enter_nested_try_stmt();
+                }
+                self.flow_restore(pre_try_block_state);
+                for state in visiting_try_block_states.snapshots() {
+                    self.flow_merge(state.clone());
+                }
+                let pre_except_state = self.flow_snapshot();
+
+                // TODO: we currently detect a `try`/`except` with a bare `except:` branch as being exhaustive,
+                // but not a `try`/`except` with `except BaseException:`.
+                // They should be treated the same way. --Alex.
+                let mut exhaustive_except_branches = false;
+
+                let mut post_except_states = vec![];
+                let num_handlers = handlers.len();
+
+                // Visit the `except` blocks!
+                for (i, except_handler) in handlers.iter().enumerate() {
                     let ast::ExceptHandler::ExceptHandler(except_handler) = except_handler;
                     let ast::ExceptHandlerExceptHandler {
                         name: symbol_name,
@@ -751,6 +819,8 @@ where
 
                     if let Some(handled_exceptions) = handled_exceptions {
                         self.visit_expr(handled_exceptions);
+                    } else {
+                        exhaustive_except_branches = true;
                     }
 
                     // If `handled_exceptions` above was `None`, it's something like `except as e:`,
@@ -769,10 +839,56 @@ where
                     }
 
                     self.visit_body(handler_body);
+
+                    // Each `except` block is mutually exclusive with all other `except` blocks.
+                    post_except_states.push(self.flow_snapshot());
+
+                    // It's unnecessary to do the `self.flow_restore()` call for the final except handler,
+                    // as we'll immediately call `self.flow_restore()` to a different state
+                    // as soon as this loop over the handlers terminates.
+                    if i < (num_handlers - 1) {
+                        self.flow_restore(pre_except_state.clone());
+                    }
                 }
 
+                // If we get to the `else` block, we know that 0 of the `except` blocks can have been executed,
+                // and the entire `try` block must have been executed:
+                self.flow_restore(post_try_block_state);
                 self.visit_body(orelse);
+
+                // Visiting the `finally` block!
+                //
+                // The state here could be that we visited exactly one `except` block,
+                // or that we visited the `else` block, **or** (unless the `except` branches
+                // are exhaustive) that we visited none of them!
+                // (An exception **could** have been raised that wasn't handled by any
+                // of the except handlers! But it would still be caught by the `finally`,
+                // which **always** runs.)
+                //
+                // N.B. Not all `try`/`except` clauses have `else` and/or `finally`!
+                // (In fact, some of them have `finally` but not `except`!)
+                // That's actually an irrelevant detail for us here, though, because
+                // a nonexistent `else` or `finally` block is just represented
+                // as an empty `Vec` in our AST, and visiting an empty `Vec` will not
+                // change the flow state at all.
+                for state in post_except_states {
+                    self.flow_merge(state);
+                }
+                if !exhaustive_except_branches {
+                    for state in visiting_try_block_states.into_snapshots() {
+                        self.flow_merge(state);
+                    }
+                }
+
                 self.visit_body(finalbody);
+
+                // Account for the fact that we could be visiting a nested `try` block,
+                // in which case the outer `try` block must be kept informed about all the possible
+                // between-definition states we could have encountered in the inner `try` block(!)
+                if let Some(current_try_block) = self.try_block_contexts_mut().current_try_block() {
+                    current_try_block.exit_nested_try_stmt();
+                    current_try_block.record_definition(self);
+                }
             }
             _ => {
                 walk_stmt(self, stmt);

crates/red_knot_python_semantic/src/semantic_index/builder/except_handlers.rs

@@ -0,0 +1,88 @@
+use std::cell::{RefCell, RefMut};
+
+use super::SemanticIndexBuilder;
+use crate::semantic_index::use_def::FlowSnapshot;
+
+#[derive(Debug, Default)]
+pub(super) struct TryBlockContextsStack(Vec<TryBlockContexts>);
+
+impl TryBlockContextsStack {
+    pub(super) fn push_context(&mut self) {
+        self.0.push(TryBlockContexts::default());
+    }
+
+    pub(super) fn current_try_block_context(&self) -> &TryBlockContexts {
+        self.0
+            .last()
+            .expect("There should always be at least one `TryBlockContexts` on the stack")
+    }
+
+    pub(super) fn pop_context(&mut self) {
+        let popped_context = self.0.pop();
+        assert_ne!(
+            popped_context, None,
+            "pop_context() should never be called on an empty stack \
+(this indicates an unbalanced `push_context()`/`pop_context()` pair of calls)"
+        );
+    }
+}
+
+#[derive(Debug, Default, PartialEq, Eq)]
+pub(super) struct TryBlockContexts(RefCell<Vec<TryBlockContext>>);
+
+impl TryBlockContexts {
+    pub(super) fn push_try_block(&self) {
+        self.0.borrow_mut().push(TryBlockContext::default());
+    }
+
+    pub(super) fn pop_try_block(&self) -> Option<TryBlockContext> {
+        self.0.borrow_mut().pop()
+    }
+
+    pub(super) fn borrow_mut(&self) -> TryBlockContextsRefMut {
+        TryBlockContextsRefMut(self.0.borrow_mut())
+    }
+}
+
+#[derive(Debug)]
+pub(super) struct TryBlockContextsRefMut<'a>(RefMut<'a, Vec<TryBlockContext>>);
+
+impl<'a> TryBlockContextsRefMut<'a> {
+    pub(super) fn current_try_block(&mut self) -> Option<&mut TryBlockContext> {
+        self.0.last_mut()
+    }
+}
+
+#[derive(Debug, Default, PartialEq, Eq)]
+pub(super) struct TryBlockContext {
+    snapshots: Vec<FlowSnapshot>,
+    visiting_nested_try_stmt: bool,
+}
+
+impl TryBlockContext {
+    pub(super) fn record_definition(&mut self, builder: &SemanticIndexBuilder) {
+        // The `if !self.visiting_nested_try_stmt` check here isn't necessary for correctness.
+        // It's a worthwhile optimisation, however: if we're visiting a nested `try/except/else/finally`
+        // block, we'll push a single snapshot on completion of visiting that block,
+        // meaning that pushing intermediate snapshots during visitation of that block
+        if !self.visiting_nested_try_stmt {
+            self.snapshots.push(builder.flow_snapshot());
+        }
+    }
+
+    pub(super) fn enter_nested_try_stmt(&mut self) {
+        self.visiting_nested_try_stmt = true;
+    }
+
+    pub(super) fn exit_nested_try_stmt(&mut self) {
+        self.visiting_nested_try_stmt = false;
+    }
+
+    pub(super) fn snapshots(&self) -> &[FlowSnapshot] {
+        &self.snapshots
+    }
+
+    pub(super) fn into_snapshots(self) -> Vec<FlowSnapshot> {
+        self.snapshots
+    }
+}

crates/red_knot_python_semantic/src/semantic_index/use_def.rs

@@ -426,7 +426,7 @@ impl<'map, 'db> Iterator for DeclarationsIterator<'map, 'db> {
 impl std::iter::FusedIterator for DeclarationsIterator<'_, '_> {}
 
 /// A snapshot of the definitions and constraints state at a particular point in control flow.
-#[derive(Clone, Debug)]
+#[derive(Clone, Debug, PartialEq, Eq)]
 pub(super) struct FlowSnapshot {
     symbol_states: IndexVec<ScopedSymbolId, SymbolState>,
 }