[ty] Eagerly simplify 'True' and 'False' constraints

Author: sharkdp

crates/ty_python_semantic/resources/mdtest/narrow/while.md

@@ -59,3 +59,17 @@ while x != 1:
 
     x = next_item()
 ```
+
+## With `break` statements
+
+```py
+def next_item() -> int | None:
+    return 1
+
+while True:
+    x = next_item()
+    if x is not None:
+        break
+
+reveal_type(x)  # revealed: int
+```

crates/ty_python_semantic/src/semantic_index/builder.rs

@@ -35,8 +35,8 @@ use crate::semantic_index::place::{
     PlaceExprWithFlags, PlaceTableBuilder, Scope, ScopeId, ScopeKind, ScopedPlaceId,
 };
 use crate::semantic_index::predicate::{
-    PatternPredicate, PatternPredicateKind, Predicate, PredicateNode, ScopedPredicateId,
-    StarImportPlaceholderPredicate,
+    PatternPredicate, PatternPredicateKind, Predicate, PredicateInner, PredicateNode,
+    ScopedPredicateId, StarImportPlaceholderPredicate,
 };
 use crate::semantic_index::re_exports::exported_names;
 use crate::semantic_index::reachability_constraints::{
@@ -543,9 +543,14 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
 
     fn build_predicate(&mut self, predicate_node: &ast::Expr) -> Predicate<'db> {
         let expression = self.add_standalone_expression(predicate_node);
-        Predicate {
-            node: PredicateNode::Expression(expression),
-            is_positive: true,
+
+        if let Some(boolean_literal) = predicate_node.as_boolean_literal_expr() {
+            Predicate::Always(boolean_literal.value)
+        } else {
+            Predicate::Inner(PredicateInner {
+                node: PredicateNode::Expression(expression),
+                is_positive: true,
+            })
         }
     }
 
@@ -617,6 +622,7 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
         let reachability_constraint = self
             .current_reachability_constraints_mut()
             .add_atom(predicate_id);
+
         self.current_use_def_map_mut()
             .record_reachability_constraint(reachability_constraint);
         reachability_constraint
@@ -705,10 +711,10 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
             guard,
             countme::Count::default(),
         );
-        let predicate = Predicate {
+        let predicate = Predicate::Inner(PredicateInner {
             node: PredicateNode::Pattern(pattern_predicate),
             is_positive: true,
-        };
+        });
         self.record_narrowing_constraint(predicate);
         predicate
     }
@@ -1666,10 +1672,10 @@ impl<'ast> Visitor<'ast> for SemanticIndexBuilder<'_, 'ast> {
                         self.record_ambiguous_reachability();
                         self.visit_expr(guard);
                         let post_guard_eval = self.flow_snapshot();
-                        let predicate = Predicate {
+                        let predicate = Predicate::Inner(PredicateInner {
                             node: PredicateNode::Expression(guard_expr),
                             is_positive: true,
-                        };
+                        });
                         self.record_negated_narrowing_constraint(predicate);
                         let match_success_guard_failure = self.flow_snapshot();
                         self.flow_restore(post_guard_eval);

crates/ty_python_semantic/src/semantic_index/predicate.rs

@@ -21,19 +21,29 @@ use crate::semantic_index::place::{FileScopeId, ScopeId, ScopedPlaceId};
 #[derive(Ord, PartialOrd, get_size2::GetSize)]
 pub(crate) struct ScopedPredicateId;
 
+impl ScopedPredicateId {
+    /// A special ID that is used for an "always true" predicate.
+    pub(crate) const ALWAYS_TRUE: ScopedPredicateId =
+        ScopedPredicateId(std::num::NonZero::new(0xffff_ffff).unwrap());
+
+    /// A special ID that is used for an "always false" predicate.
+    pub(crate) const ALWAYS_FALSE: ScopedPredicateId =
+        ScopedPredicateId(std::num::NonZero::new(0xffff_fffe).unwrap());
+}
+
 // A collection of predicates for a given scope.
-pub(crate) type Predicates<'db> = IndexVec<ScopedPredicateId, Predicate<'db>>;
+pub(crate) type Predicates<'db> = IndexVec<ScopedPredicateId, PredicateInner<'db>>;
 
 #[derive(Debug, Default)]
 pub(crate) struct PredicatesBuilder<'db> {
-    predicates: IndexVec<ScopedPredicateId, Predicate<'db>>,
+    predicates: IndexVec<ScopedPredicateId, PredicateInner<'db>>,
 }
 
 impl<'db> PredicatesBuilder<'db> {
     /// Adds a predicate. Note that we do not deduplicate predicates. If you add a `Predicate`
     /// more than once, you will get distinct `ScopedPredicateId`s for each one. (This lets you
     /// model predicates that might evaluate to different values at different points of execution.)
-    pub(crate) fn add_predicate(&mut self, predicate: Predicate<'db>) -> ScopedPredicateId {
+    pub(crate) fn add_predicate(&mut self, predicate: PredicateInner<'db>) -> ScopedPredicateId {
         self.predicates.push(predicate)
     }
 
@@ -44,16 +54,27 @@ impl<'db> PredicatesBuilder<'db> {
 }
 
 #[derive(Clone, Copy, Debug, Hash, PartialEq, Eq, salsa::Update, get_size2::GetSize)]
-pub(crate) struct Predicate<'db> {
+pub(crate) struct PredicateInner<'db> {
     pub(crate) node: PredicateNode<'db>,
     pub(crate) is_positive: bool,
 }
 
+#[derive(Clone, Copy, Debug, Hash, PartialEq, Eq, salsa::Update, get_size2::GetSize)]
+pub(crate) enum Predicate<'db> {
+    Always(bool),
+    Inner(PredicateInner<'db>),
+}
+
 impl Predicate<'_> {
     pub(crate) fn negated(self) -> Self {
-        Self {
-            node: self.node,
-            is_positive: !self.is_positive,
+        match self {
+            Predicate::Always(value) => Predicate::Always(!value),
+            Predicate::Inner(PredicateInner { node, is_positive }) => {
+                Predicate::Inner(PredicateInner {
+                    node,
+                    is_positive: !is_positive,
+                })
+            }
         }
     }
 }
@@ -171,9 +192,9 @@ impl<'db> StarImportPlaceholderPredicate<'db> {
 
 impl<'db> From<StarImportPlaceholderPredicate<'db>> for Predicate<'db> {
     fn from(predicate: StarImportPlaceholderPredicate<'db>) -> Self {
-        Predicate {
+        Predicate::Inner(PredicateInner {
             node: PredicateNode::StarImportPlaceholder(predicate),
             is_positive: true,
-        }
+        })
     }
 }

crates/ty_python_semantic/src/semantic_index/reachability_constraints.rs

@@ -204,7 +204,8 @@ use crate::place::{RequiresExplicitReExport, imported_symbol};
 use crate::semantic_index::expression::Expression;
 use crate::semantic_index::place_table;
 use crate::semantic_index::predicate::{
-    PatternPredicate, PatternPredicateKind, Predicate, PredicateNode, Predicates, ScopedPredicateId,
+    PatternPredicate, PatternPredicateKind, PredicateInner, PredicateNode, Predicates,
+    ScopedPredicateId,
 };
 use crate::types::{Truthiness, Type, infer_expression_type};
 
@@ -388,12 +389,18 @@ impl ReachabilityConstraintsBuilder {
         &mut self,
         predicate: ScopedPredicateId,
     ) -> ScopedReachabilityConstraintId {
-        self.add_interior(InteriorNode {
-            atom: predicate,
-            if_true: ALWAYS_TRUE,
-            if_ambiguous: AMBIGUOUS,
-            if_false: ALWAYS_FALSE,
-        })
+        if predicate == ScopedPredicateId::ALWAYS_FALSE {
+            ScopedReachabilityConstraintId::ALWAYS_FALSE
+        } else if predicate == ScopedPredicateId::ALWAYS_TRUE {
+            ScopedReachabilityConstraintId::ALWAYS_TRUE
+        } else {
+            self.add_interior(InteriorNode {
+                atom: predicate,
+                if_true: ALWAYS_TRUE,
+                if_ambiguous: AMBIGUOUS,
+                if_false: ALWAYS_FALSE,
+            })
+        }
     }
 
     /// Adds a new reachability constraint that is the ternary NOT of an existing one.
@@ -672,7 +679,7 @@ impl ReachabilityConstraints {
         }
     }
 
-    fn analyze_single(db: &dyn Db, predicate: &Predicate) -> Truthiness {
+    fn analyze_single(db: &dyn Db, predicate: &PredicateInner) -> Truthiness {
         match predicate.node {
             PredicateNode::Expression(test_expr) => {
                 let ty = infer_expression_type(db, test_expr);

crates/ty_python_semantic/src/semantic_index/use_def.rs

@@ -247,7 +247,8 @@ use crate::semantic_index::place::{
     FileScopeId, PlaceExpr, PlaceExprWithFlags, ScopeKind, ScopedPlaceId,
 };
 use crate::semantic_index::predicate::{
-    Predicate, Predicates, PredicatesBuilder, ScopedPredicateId, StarImportPlaceholderPredicate,
+    Predicate, PredicateInner, Predicates, PredicatesBuilder, ScopedPredicateId,
+    StarImportPlaceholderPredicate,
 };
 use crate::semantic_index::reachability_constraints::{
     ReachabilityConstraints, ReachabilityConstraintsBuilder, ScopedReachabilityConstraintId,
@@ -610,7 +611,7 @@ pub(crate) struct ConstraintsIterator<'map, 'db> {
 }
 
 impl<'db> Iterator for ConstraintsIterator<'_, 'db> {
-    type Item = Predicate<'db>;
+    type Item = PredicateInner<'db>;
 
     fn next(&mut self) -> Option<Self::Item> {
         self.constraint_ids
@@ -806,10 +807,21 @@ impl<'db> UseDefMapBuilder<'db> {
     }
 
     pub(super) fn add_predicate(&mut self, predicate: Predicate<'db>) -> ScopedPredicateId {
-        self.predicates.add_predicate(predicate)
+        match predicate {
+            Predicate::Inner(predicate) => self.predicates.add_predicate(predicate),
+            Predicate::Always(true) => ScopedPredicateId::ALWAYS_TRUE,
+            Predicate::Always(false) => ScopedPredicateId::ALWAYS_FALSE,
+        }
     }
 
     pub(super) fn record_narrowing_constraint(&mut self, predicate: ScopedPredicateId) {
+        if predicate == ScopedPredicateId::ALWAYS_TRUE
+            || predicate == ScopedPredicateId::ALWAYS_FALSE
+        {
+            // No need to record a narrowing constraint for `True` or `False`.
+            return;
+        }
+
         let narrowing_constraint = predicate.into();
         for state in &mut self.place_states {
             state

crates/ty_python_semantic/src/types/narrow.rs

@@ -4,7 +4,7 @@ use crate::semantic_index::expression::Expression;
 use crate::semantic_index::place::{PlaceExpr, PlaceTable, ScopeId, ScopedPlaceId};
 use crate::semantic_index::place_table;
 use crate::semantic_index::predicate::{
-    PatternPredicate, PatternPredicateKind, Predicate, PredicateNode,
+    PatternPredicate, PatternPredicateKind, PredicateInner, PredicateNode,
 };
 use crate::types::function::KnownFunction;
 use crate::types::infer::infer_same_file_expression_type;
@@ -42,7 +42,7 @@ use super::UnionType;
 /// constraint is applied to that symbol, so we'd just return `None`.
 pub(crate) fn infer_narrowing_constraint<'db>(
     db: &'db dyn Db,
-    predicate: Predicate<'db>,
+    predicate: PredicateInner<'db>,
     place: ScopedPlaceId,
 ) -> Option<Type<'db>> {
     let constraints = match predicate.node {