diff --git a/crates/red_knot_python_semantic/resources/mdtest/terminal_statements.md b/crates/red_knot_python_semantic/resources/mdtest/terminal_statements.md index 5888847b84f7d2..6a3427e211c845 100644 --- a/crates/red_knot_python_semantic/resources/mdtest/terminal_statements.md +++ b/crates/red_knot_python_semantic/resources/mdtest/terminal_statements.md @@ -635,5 +635,6 @@ def _(cond: bool): if True: return - reveal_type(x) # revealed: Literal["a"] + # TODO: Literal["a"] + reveal_type(x) # revealed: Literal["a", "b"] ``` diff --git a/crates/red_knot_python_semantic/src/semantic_index/builder.rs b/crates/red_knot_python_semantic/src/semantic_index/builder.rs index 9dac8ef56583d7..ca75b22379ac3a 100644 --- a/crates/red_knot_python_semantic/src/semantic_index/builder.rs +++ b/crates/red_knot_python_semantic/src/semantic_index/builder.rs @@ -413,6 +413,13 @@ impl<'db> SemanticIndexBuilder<'db> { .record_visibility_constraint(ScopedVisibilityConstraintId::AMBIGUOUS); } + /// Simplifies (resets) visibility constraints on all live bindings and declarations that did + /// not see any new definitions since the given snapshot. + fn simplify_visibility_constraints(&mut self, snapshot: FlowSnapshot) { + self.current_use_def_map_mut() + .simplify_visibility_constraints(snapshot); + } + fn push_assignment(&mut self, assignment: CurrentAssignment<'db>) { self.current_assignments.push(assignment); } @@ -1066,6 +1073,8 @@ where for post_clause_state in post_clauses { self.flow_merge(post_clause_state); } + + self.simplify_visibility_constraints(no_branch_taken); } ast::Stmt::While(ast::StmtWhile { test, @@ -1121,7 +1130,7 @@ where // To model this correctly, we need two copies of the while condition constraint, // since the first and later evaluations might produce different results. let post_body = self.flow_snapshot(); - self.flow_restore(pre_loop); + self.flow_restore(pre_loop.clone()); self.record_negated_visibility_constraint(first_vis_constraint_id); self.flow_merge(post_body); self.record_negated_constraint(constraint); @@ -1136,6 +1145,8 @@ where self.record_visibility_constraint_id(body_vis_constraint_id); self.flow_merge(snapshot); } + + self.simplify_visibility_constraints(pre_loop); } ast::Stmt::With(ast::StmtWith { items, @@ -1280,7 +1291,7 @@ where .is_some_and(|case| case.guard.is_none() && case.pattern.is_wildcard()) { post_case_snapshots.push(self.flow_snapshot()); - self.flow_restore(after_subject); + self.flow_restore(after_subject.clone()); for id in &vis_constraints { self.record_negated_visibility_constraint(*id); @@ -1290,6 +1301,8 @@ where for post_clause_state in post_case_snapshots { self.flow_merge(post_clause_state); } + + self.simplify_visibility_constraints(after_subject); } ast::Stmt::Try(ast::StmtTry { body, @@ -1570,12 +1583,13 @@ where self.visit_expr(body); let visibility_constraint = self.record_visibility_constraint(constraint); let post_body = self.flow_snapshot(); - self.flow_restore(pre_if); + self.flow_restore(pre_if.clone()); self.record_negated_constraint(constraint); self.visit_expr(orelse); self.record_negated_visibility_constraint(visibility_constraint); self.flow_merge(post_body); + self.simplify_visibility_constraints(pre_if); } ast::Expr::ListComp( list_comprehension @ ast::ExprListComp { @@ -1632,6 +1646,8 @@ where range: _, op, }) => { + let pre_op = self.flow_snapshot(); + let mut snapshots = vec![]; let mut visibility_constraints = vec![]; @@ -1678,6 +1694,8 @@ where for snapshot in snapshots { self.flow_merge(snapshot); } + + self.simplify_visibility_constraints(pre_op); } _ => { walk_expr(self, expr); diff --git a/crates/red_knot_python_semantic/src/semantic_index/use_def.rs b/crates/red_knot_python_semantic/src/semantic_index/use_def.rs index f8d4ae8f23f6a7..861a4c3132bd47 100644 --- a/crates/red_knot_python_semantic/src/semantic_index/use_def.rs +++ b/crates/red_knot_python_semantic/src/semantic_index/use_def.rs @@ -575,6 +575,40 @@ impl<'db> UseDefMapBuilder<'db> { .add_and_constraint(self.scope_start_visibility, constraint); } + /// This method resets the visibility constraints for all symbols to a previous state + /// *if* there have been no new declarations or bindings since then. Consider the + /// following example: + /// ```py + /// x = 0 + /// y = 0 + /// if test_a: + /// y = 1 + /// elif test_b: + /// y = 2 + /// elif test_c: + /// y = 3 + /// + /// # RESET + /// ``` + /// We build a complex visibility constraint for the `y = 0` binding. We build the same + /// constraint for the `x = 0` binding as well, but at the `RESET` point, we can get rid + /// of it, as the `if`-`elif`-`elif` chain doesn't include any new bindings of `x`. + pub(super) fn simplify_visibility_constraints(&mut self, snapshot: FlowSnapshot) { + debug_assert!(self.symbol_states.len() >= snapshot.symbol_states.len()); + + self.scope_start_visibility = snapshot.scope_start_visibility; + + // Note that this loop terminates when we reach a symbol not present in the snapshot. + // This means we keep visibility constraints for all new symbols, which is intended, + // since these symbols have been introduced in the corresponding branch, which might + // be subject to visibility constraints. We only simplify/reset visibility constraints + // for symbols that have the same bindings and declarations present compared to the + // snapshot. + for (current, snapshot) in self.symbol_states.iter_mut().zip(snapshot.symbol_states) { + current.simplify_visibility_constraints(snapshot); + } + } + pub(super) fn record_declaration( &mut self, symbol: ScopedSymbolId, diff --git a/crates/red_knot_python_semantic/src/semantic_index/use_def/symbol_state.rs b/crates/red_knot_python_semantic/src/semantic_index/use_def/symbol_state.rs index ce9c577a267f94..4c0c72bc6226af 100644 --- a/crates/red_knot_python_semantic/src/semantic_index/use_def/symbol_state.rs +++ b/crates/red_knot_python_semantic/src/semantic_index/use_def/symbol_state.rs @@ -371,6 +371,16 @@ impl SymbolState { .record_visibility_constraint(visibility_constraints, constraint); } + pub(super) fn simplify_visibility_constraints(&mut self, snapshot_state: SymbolState) { + if self.bindings.live_bindings == snapshot_state.bindings.live_bindings { + self.bindings.visibility_constraints = snapshot_state.bindings.visibility_constraints; + } + if self.declarations.live_declarations == snapshot_state.declarations.live_declarations { + self.declarations.visibility_constraints = + snapshot_state.declarations.visibility_constraints; + } + } + /// Record a newly-encountered declaration of this symbol. pub(super) fn record_declaration(&mut self, declaration_id: ScopedDefinitionId) { self.declarations.record_declaration(declaration_id);