From d249801bf2bcd8afa30b1ffc1c39a5c65982d28a Mon Sep 17 00:00:00 2001 From: David Peter Date: Wed, 18 Dec 2024 11:32:56 +0100 Subject: [PATCH] Implement if-expressions --- .../mdtest/statically_known_branches.md | 59 ++++++++++++++++--- .../src/semantic_index/builder.rs | 31 +++++----- .../src/semantic_index/use_def.rs | 11 +--- .../semantic_index/use_def/symbol_state.rs | 38 ++++++------ .../src/visibility_constraints.rs | 6 +- 5 files changed, 88 insertions(+), 57 deletions(-) diff --git a/crates/red_knot_python_semantic/resources/mdtest/statically_known_branches.md b/crates/red_knot_python_semantic/resources/mdtest/statically_known_branches.md index 8cefccc342..1a588d220d 100644 --- a/crates/red_knot_python_semantic/resources/mdtest/statically_known_branches.md +++ b/crates/red_knot_python_semantic/resources/mdtest/statically_known_branches.md @@ -41,8 +41,6 @@ if typing.TYPE_CHECKING: def f(s: SomeType) -> None: ... ``` -The rest of this document contains tests for various cases where this feature can be used. - ## Common use cases This section makes sure that we can handle all commonly encountered patterns of static conditions. @@ -164,6 +162,9 @@ other ## If statements +The rest of this document contains tests for various control flow elements. This section tests `if` +statements. + ### Always false #### If @@ -349,7 +350,7 @@ else: reveal_type(x) # revealed: Literal[2, 3, 4] ``` -#### `elif` without `else` branch +#### `elif` without `else` branch, always true ```py def flag() -> bool: ... @@ -364,6 +365,21 @@ elif True: reveal_type(x) # revealed: Literal[2, 3] ``` +#### `elif` without `else` branch, always false + +```py +def flag() -> bool: ... + +x = 1 + +if flag(): + x = 2 +elif False: + x = 3 + +reveal_type(x) # revealed: Literal[1, 2] +``` + ### Nested conditionals #### `if True` inside `if True` @@ -769,22 +785,31 @@ reveal_type(x) # revealed: Literal[3, 4] ## If expressions -See also: tests in [expression/if.md](expression/if.md). +Note that the result type of an `if`-expression can be precisely inferred if the condition is +statically known. This is a plain type inference feature that does not need support for statically +known branches. The tests for this feature are in [expression/if.md](expression/if.md). + +The tests here make sure that we also handle assignment expressions inside `if`-expressions +correctly. + +### Type inference ### Always true ```py -x = 1 if True else 2 +x = (y := 1) if True else (y := 2) reveal_type(x) # revealed: Literal[1] +reveal_type(y) # revealed: Literal[1] ``` ### Always false ```py -x = 1 if False else 2 +x = (y := 1) if False else (y := 2) reveal_type(x) # revealed: Literal[2] +reveal_type(y) # revealed: Literal[2] ``` ## Boolean expressions @@ -795,6 +820,10 @@ reveal_type(x) # revealed: Literal[2] (x := 1) or (x := 2) reveal_type(x) # revealed: Literal[1] + +(y := 1) or (y := 2) or (y := 3) or (y := 4) + +reveal_type(y) # revealed: Literal[1] ``` ### Always true, `and` @@ -803,22 +832,34 @@ reveal_type(x) # revealed: Literal[1] (x := 1) and (x := 2) reveal_type(x) # revealed: Literal[2] + +(y := 1) and (y := 2) and (y := 3) and (y := 4) + +reveal_type(y) # revealed: Literal[4] ``` ### Always false, `or` ```py -(x := 0) or (x := 2) +(x := 0) or (x := 1) -reveal_type(x) # revealed: Literal[2] +reveal_type(x) # revealed: Literal[1] + +(y := 0) or (y := 0) or (y := 1) or (y := 2) + +reveal_type(y) # revealed: Literal[1] ``` ### Always false, `and` ```py -(x := 0) and (x := 2) +(x := 0) and (x := 1) reveal_type(x) # revealed: Literal[0] + +(y := 0) and (y := 1) and (y := 2) and (y := 3) + +reveal_type(y) # revealed: Literal[0] ``` ## While loops 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 9fb3fae453..acf0fcae96 100644 --- a/crates/red_knot_python_semantic/src/semantic_index/builder.rs +++ b/crates/red_knot_python_semantic/src/semantic_index/builder.rs @@ -296,14 +296,6 @@ impl<'db> SemanticIndexBuilder<'db> { self.current_use_def_map_mut().record_constraint(constraint) } - fn add_visibility_constraint( - &mut self, - constraint: ScopedConstraintId, - ) -> ScopedVisibilityConstraintId { - self.current_use_def_map_mut() - .add_visibility_constraint(VisibilityConstraint::VisibleIf(constraint)) - } - fn record_visibility_constraint( &mut self, constraint: ScopedConstraintId, @@ -1325,14 +1317,17 @@ where }) => { self.visit_expr(test); let pre_if = self.flow_snapshot(); - let (_, constraint) = self.record_expression_constraint(test); + let (constraint_id, constraint) = self.record_expression_constraint(test); self.visit_expr(body); + let visibility_constraint = self.record_visibility_constraint(constraint_id); 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.reset_visibility_constraints(pre_if); } ast::Expr::ListComp( list_comprehension @ ast::ExprListComp { @@ -1392,12 +1387,13 @@ where let pre_op = self.flow_snapshot(); let mut snapshots = vec![]; - let mut constraints = vec![]; + let mut visibility_constraints = vec![]; + let mut last_constraint = None; for (index, value) in values.iter().enumerate() { self.visit_expr(value); - if let Some(last_constraint_id) = constraints.last() { - self.record_visibility_constraint(*last_constraint_id); + if let Some(id) = last_constraint { + visibility_constraints.push(self.record_visibility_constraint(id)); } // Snapshot is taken after visiting the expression but before adding the constraint. snapshots.push(self.flow_snapshot()); @@ -1408,7 +1404,7 @@ where BoolOp::And => self.record_constraint(constraint), BoolOp::Or => self.record_negated_constraint(constraint), }; - constraints.push(id); + last_constraint = Some(id); } } @@ -1416,9 +1412,10 @@ where let first = snapshots.next().expect("at least one value"); self.flow_restore(first); - for id in constraints { - let vid = self.add_visibility_constraint(id); - self.record_negated_visibility_constraint(vid); + // debug_assert(constraints.len() == snapshots.len()); + + for id in visibility_constraints { + self.record_negated_visibility_constraint(id); } for snapshot in snapshots { 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 bd5c24c5c5..78045de6b9 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 @@ -472,7 +472,7 @@ impl<'db> UseDefMapBuilder<'db> { all_definitions: IndexVec::from_iter([None]), all_constraints: IndexVec::new(), visibility_constraints: VisibilityConstraints::new(), - unbound_visibility: ScopedVisibilityConstraintId::ALWAYS_VISIBLE, + unbound_visibility: ScopedVisibilityConstraintId::ALWAYS_TRUE, bindings_by_use: IndexVec::new(), definitions_by_definition: FxHashMap::default(), symbol_states: IndexVec::new(), @@ -504,18 +504,11 @@ impl<'db> UseDefMapBuilder<'db> { constraint_id } - pub(super) fn add_visibility_constraint( - &mut self, - constraint: VisibilityConstraint, - ) -> ScopedVisibilityConstraintId { - self.visibility_constraints.add(constraint) - } - pub(super) fn record_visibility_constraint( &mut self, constraint: VisibilityConstraint, ) -> ScopedVisibilityConstraintId { - let new_constraint_id = self.add_visibility_constraint(constraint); + let new_constraint_id = self.visibility_constraints.add(constraint); for state in &mut self.symbol_states { state.record_visibility_constraint(&mut self.visibility_constraints, new_constraint_id); } 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 e0f2c224ba..f473c78942 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 @@ -98,7 +98,7 @@ type ConstraintsIntoIterator = smallvec::IntoIter; pub(crate) struct ScopedVisibilityConstraintId; impl ScopedVisibilityConstraintId { - pub(crate) const ALWAYS_VISIBLE: ScopedVisibilityConstraintId = + pub(crate) const ALWAYS_TRUE: ScopedVisibilityConstraintId = ScopedVisibilityConstraintId::from_u32(0); } @@ -135,7 +135,7 @@ impl SymbolDeclarations { self.visibility_constraints = VisibilityConstraintPerBinding::with_capacity(1); self.visibility_constraints - .push(ScopedVisibilityConstraintId::ALWAYS_VISIBLE); + .push(ScopedVisibilityConstraintId::ALWAYS_TRUE); } /// Add given visibility constraint to all live bindings. @@ -199,7 +199,7 @@ impl SymbolBindings { self.visibility_constraints = VisibilityConstraintPerBinding::with_capacity(1); self.visibility_constraints - .push(ScopedVisibilityConstraintId::ALWAYS_VISIBLE); + .push(ScopedVisibilityConstraintId::ALWAYS_TRUE); } /// Add given constraint to all live bindings. @@ -655,7 +655,7 @@ mod tests { fn unbound() { let constraints = AllConstraints::new(); let visibility_constraints = VisibilityConstraints::new(); - let sym = SymbolState::undefined(ScopedVisibilityConstraintId::ALWAYS_VISIBLE); + let sym = SymbolState::undefined(ScopedVisibilityConstraintId::ALWAYS_TRUE); assert_bindings(&constraints, &visibility_constraints, &sym, &["unbound<>"]); } @@ -664,7 +664,7 @@ mod tests { fn with() { let constraints = AllConstraints::new(); let visibility_constraints = VisibilityConstraints::new(); - let mut sym = SymbolState::undefined(ScopedVisibilityConstraintId::ALWAYS_VISIBLE); + let mut sym = SymbolState::undefined(ScopedVisibilityConstraintId::ALWAYS_TRUE); sym.record_binding(ScopedDefinitionId::from_u32(1)); assert_bindings(&constraints, &visibility_constraints, &sym, &["1<>"]); @@ -674,7 +674,7 @@ mod tests { fn record_constraint() { let constraints = AllConstraints::new(); let visibility_constraints = VisibilityConstraints::new(); - let mut sym = SymbolState::undefined(ScopedVisibilityConstraintId::ALWAYS_VISIBLE); + let mut sym = SymbolState::undefined(ScopedVisibilityConstraintId::ALWAYS_TRUE); sym.record_binding(ScopedDefinitionId::from_u32(1)); sym.record_constraint(ScopedConstraintId::from_u32(0)); @@ -687,11 +687,11 @@ mod tests { let mut visibility_constraints = VisibilityConstraints::new(); // merging the same definition with the same constraint keeps the constraint - let mut sym1a = SymbolState::undefined(ScopedVisibilityConstraintId::ALWAYS_VISIBLE); + let mut sym1a = SymbolState::undefined(ScopedVisibilityConstraintId::ALWAYS_TRUE); sym1a.record_binding(ScopedDefinitionId::from_u32(1)); sym1a.record_constraint(ScopedConstraintId::from_u32(0)); - let mut sym1b = SymbolState::undefined(ScopedVisibilityConstraintId::ALWAYS_VISIBLE); + let mut sym1b = SymbolState::undefined(ScopedVisibilityConstraintId::ALWAYS_TRUE); sym1b.record_binding(ScopedDefinitionId::from_u32(1)); sym1b.record_constraint(ScopedConstraintId::from_u32(0)); @@ -700,11 +700,11 @@ mod tests { assert_bindings(&constraints, &visibility_constraints, &sym1, &["1<0>"]); // merging the same definition with differing constraints drops all constraints - let mut sym2a = SymbolState::undefined(ScopedVisibilityConstraintId::ALWAYS_VISIBLE); + let mut sym2a = SymbolState::undefined(ScopedVisibilityConstraintId::ALWAYS_TRUE); sym2a.record_binding(ScopedDefinitionId::from_u32(2)); sym2a.record_constraint(ScopedConstraintId::from_u32(1)); - let mut sym1b = SymbolState::undefined(ScopedVisibilityConstraintId::ALWAYS_VISIBLE); + let mut sym1b = SymbolState::undefined(ScopedVisibilityConstraintId::ALWAYS_TRUE); sym1b.record_binding(ScopedDefinitionId::from_u32(2)); sym1b.record_constraint(ScopedConstraintId::from_u32(2)); @@ -713,11 +713,11 @@ mod tests { assert_bindings(&constraints, &visibility_constraints, &sym2, &["2<>"]); // merging a constrained definition with unbound keeps both - let mut sym3a = SymbolState::undefined(ScopedVisibilityConstraintId::ALWAYS_VISIBLE); + let mut sym3a = SymbolState::undefined(ScopedVisibilityConstraintId::ALWAYS_TRUE); sym3a.record_binding(ScopedDefinitionId::from_u32(3)); sym3a.record_constraint(ScopedConstraintId::from_u32(3)); - let sym2b = SymbolState::undefined(ScopedVisibilityConstraintId::ALWAYS_VISIBLE); + let sym2b = SymbolState::undefined(ScopedVisibilityConstraintId::ALWAYS_TRUE); sym3a.merge(sym2b, &mut visibility_constraints); let sym3 = sym3a; @@ -743,7 +743,7 @@ mod tests { fn no_declaration() { let constraints = AllConstraints::new(); let visibility_constraints = VisibilityConstraints::new(); - let sym = SymbolState::undefined(ScopedVisibilityConstraintId::ALWAYS_VISIBLE); + let sym = SymbolState::undefined(ScopedVisibilityConstraintId::ALWAYS_TRUE); assert_declarations(&constraints, &visibility_constraints, &sym, &["undeclared"]); } @@ -752,7 +752,7 @@ mod tests { fn record_declaration() { let constraints = AllConstraints::new(); let visibility_constraints = VisibilityConstraints::new(); - let mut sym = SymbolState::undefined(ScopedVisibilityConstraintId::ALWAYS_VISIBLE); + let mut sym = SymbolState::undefined(ScopedVisibilityConstraintId::ALWAYS_TRUE); sym.record_declaration(ScopedDefinitionId::from_u32(1)); assert_declarations(&constraints, &visibility_constraints, &sym, &["1"]); @@ -762,7 +762,7 @@ mod tests { fn record_declaration_override() { let constraints = AllConstraints::new(); let visibility_constraints = VisibilityConstraints::new(); - let mut sym = SymbolState::undefined(ScopedVisibilityConstraintId::ALWAYS_VISIBLE); + let mut sym = SymbolState::undefined(ScopedVisibilityConstraintId::ALWAYS_TRUE); sym.record_declaration(ScopedDefinitionId::from_u32(1)); sym.record_declaration(ScopedDefinitionId::from_u32(2)); @@ -773,10 +773,10 @@ mod tests { fn record_declaration_merge() { let constraints = AllConstraints::new(); let mut visibility_constraints = VisibilityConstraints::new(); - let mut sym = SymbolState::undefined(ScopedVisibilityConstraintId::ALWAYS_VISIBLE); + let mut sym = SymbolState::undefined(ScopedVisibilityConstraintId::ALWAYS_TRUE); sym.record_declaration(ScopedDefinitionId::from_u32(1)); - let mut sym2 = SymbolState::undefined(ScopedVisibilityConstraintId::ALWAYS_VISIBLE); + let mut sym2 = SymbolState::undefined(ScopedVisibilityConstraintId::ALWAYS_TRUE); sym2.record_declaration(ScopedDefinitionId::from_u32(2)); sym.merge(sym2, &mut visibility_constraints); @@ -788,10 +788,10 @@ mod tests { fn record_declaration_merge_partial_undeclared() { let constraints = AllConstraints::new(); let mut visibility_constraints = VisibilityConstraints::new(); - let mut sym = SymbolState::undefined(ScopedVisibilityConstraintId::ALWAYS_VISIBLE); + let mut sym = SymbolState::undefined(ScopedVisibilityConstraintId::ALWAYS_TRUE); sym.record_declaration(ScopedDefinitionId::from_u32(1)); - let sym2 = SymbolState::undefined(ScopedVisibilityConstraintId::ALWAYS_VISIBLE); + let sym2 = SymbolState::undefined(ScopedVisibilityConstraintId::ALWAYS_TRUE); sym.merge(sym2, &mut visibility_constraints); diff --git a/crates/red_knot_python_semantic/src/visibility_constraints.rs b/crates/red_knot_python_semantic/src/visibility_constraints.rs index 157e746333..fa953ec6a0 100644 --- a/crates/red_knot_python_semantic/src/visibility_constraints.rs +++ b/crates/red_knot_python_semantic/src/visibility_constraints.rs @@ -43,10 +43,10 @@ impl VisibilityConstraints { ) -> ScopedVisibilityConstraintId { match (&self.constraints[a], &self.constraints[b]) { (_, VisibilityConstraint::VisibleIfNot(id)) if a == *id => { - ScopedVisibilityConstraintId::ALWAYS_VISIBLE + ScopedVisibilityConstraintId::ALWAYS_TRUE } (VisibilityConstraint::VisibleIfNot(id), _) if *id == b => { - ScopedVisibilityConstraintId::ALWAYS_VISIBLE + ScopedVisibilityConstraintId::ALWAYS_TRUE } _ => self.add(VisibilityConstraint::Merged(a, b)), } @@ -57,7 +57,7 @@ impl VisibilityConstraints { a: ScopedVisibilityConstraintId, b: ScopedVisibilityConstraintId, ) -> ScopedVisibilityConstraintId { - if a == ScopedVisibilityConstraintId::ALWAYS_VISIBLE { + if a == ScopedVisibilityConstraintId::ALWAYS_TRUE { b } else { self.add(VisibilityConstraint::Sequence(a, b))