correlated condition analysis for simple guards #3185

pyrefly/lib/binding/bindings.rs

@@ -78,6 +78,7 @@ use crate::binding::binding::TypeAliasRefBinding;
 use crate::binding::binding::TypeParameter;
 use crate::binding::expr::Usage;
 use crate::binding::metadata::BindingsMetadata;
+use crate::binding::narrow::AtomicNarrowOp;
 use crate::binding::narrow::NarrowOp;
 use crate::binding::narrow::NarrowOps;
 use crate::binding::scope::Exportable;
@@ -1131,6 +1132,7 @@ impl<'a> BindingsBuilder<'a> {
             FlowStyle::Other
             | FlowStyle::ClassField { .. }
             | FlowStyle::PossiblyUninitialized
+            | FlowStyle::InitializedIfGuardTruthy { .. }
             | FlowStyle::MaybeInitialized(_)
             | FlowStyle::Uninitialized => {
                 self.special_export_from_binding_idx(idx, visited_names, visited_keys)
@@ -1746,6 +1748,11 @@ impl<'a> BindingsBuilder<'a> {
                     Binding::Narrow(initial_idx, Box::new(op.clone()), use_location),
                 );
                 self.scopes.narrow_in_current_flow(name, narrowed_idx);
+                // Correlated-condition analysis: a bare `if <name>:` clears matching
+                // `InitializedIfGuardTruthy` guards.
+                if matches!(op, NarrowOp::Atomic(None, AtomicNarrowOp::IsTruthy)) {
+                    self.scopes.clear_matching_truthy_guard(name.into_key());
+                }
             }
         }
     }

pyrefly/lib/binding/scope.rs

@@ -643,6 +643,12 @@ impl FlowInfo {
                     initial_value: None,
                 } => InitializedInFlow::No,
                 FlowStyle::PossiblyUninitialized => InitializedInFlow::Conditionally,
+                FlowStyle::InitializedIfGuardTruthy {
+                    termination_keys, ..
+                } if termination_keys.is_empty() => InitializedInFlow::Conditionally,
+                FlowStyle::InitializedIfGuardTruthy {
+                    termination_keys, ..
+                } => InitializedInFlow::DeferredCheck(termination_keys.clone()),
                 _ => InitializedInFlow::Yes,
             })
     }
@@ -682,6 +688,15 @@ pub enum FlowStyle {
     ClassDef,
     /// The name is possibly uninitialized (perhaps due to merging branches)
     PossiblyUninitialized,
+    /// The name was defined only on the truthy branch of `if <guard>:` (no else).
+    /// Treated as initialized inside a later `if <guard>:` block, as long as `guard`
+    /// still has `guard_value_idx` (not reassigned). `termination_keys` is the
+    /// `MaybeInitialized` fallback used when the guard doesn't match.
+    InitializedIfGuardTruthy {
+        guard: Name,
+        guard_value_idx: Idx<Key>,
+        termination_keys: Vec<Idx<Key>>,
+    },
     /// The name may or may not be initialized depending on whether certain branches
     /// terminate (have `Never` type). The termination keys are checked at solve time;
     /// if all of them have `Never` type, the name is considered initialized.
@@ -720,15 +735,26 @@ impl FlowStyle {
                 //
                 // Treating it as Other reduces false positives at the expense of
                 // some false negatives.
-                (FlowStyle::Uninitialized | FlowStyle::PossiblyUninitialized, _)
-                | (_, FlowStyle::Uninitialized | FlowStyle::PossiblyUninitialized) => {
-                    match merge_style {
-                        MergeStyle::BoolOp => {
-                            merged = FlowStyle::Other;
-                        }
-                        _ => return FlowStyle::PossiblyUninitialized,
+                //
+                // `InitializedIfGuardTruthy` degrades to `PossiblyUninitialized` on
+                // further merges: we can no longer prove the guard pattern holds.
+                (
+                    FlowStyle::Uninitialized
+                    | FlowStyle::PossiblyUninitialized
+                    | FlowStyle::InitializedIfGuardTruthy { .. },
+                    _,
+                )
+                | (
+                    _,
+                    FlowStyle::Uninitialized
+                    | FlowStyle::PossiblyUninitialized
+                    | FlowStyle::InitializedIfGuardTruthy { .. },
+                ) => match merge_style {
+                    MergeStyle::BoolOp => {
+                        merged = FlowStyle::Other;
                     }
-                }
+                    _ => return FlowStyle::PossiblyUninitialized,
+                },
                 // Two MaybeInitialized: combine termination keys from both branches.
                 // Each branch independently needs its keys to be Never for that path
                 // to be initialized, so we collect all keys.
@@ -792,6 +818,7 @@ impl FlowStyle {
             | FlowStyle::Other => self,
             FlowStyle::Uninitialized
             | FlowStyle::PossiblyUninitialized
+            | FlowStyle::InitializedIfGuardTruthy { .. }
             | FlowStyle::MaybeInitialized { .. } => FlowStyle::Other,
         }
     }
@@ -2059,6 +2086,65 @@ impl Scopes {
         Some(self.current().flow.get_info(name)?.value()?.idx)
     }
 
+    /// Snapshot the names in the current flow.
+    pub fn current_flow_names(&self) -> SmallSet<Name> {
+        self.current().flow.info.keys().cloned().collect()
+    }
+
+    /// For names introduced by the just-finished `if <guard>:` (i.e. not in
+    /// `pre_fork_names`), upgrade `PossiblyUninitialized`/`MaybeInitialized` to
+    /// `InitializedIfGuardTruthy`. Preserves any `MaybeInitialized` termination keys
+    /// as the fallback when the guard doesn't match.
+    pub fn upgrade_to_guarded(
+        &mut self,
+        guard: &Name,
+        guard_value_idx: Idx<Key>,
+        pre_fork_names: &SmallSet<Name>,
+    ) {
+        for (name, info) in self.current_mut().flow.info.iter_mut() {
+            if pre_fork_names.contains(name) {
+                continue;
+            }
+            let Some(value) = info.value_mut() else {
+                continue;
+            };
+            let termination_keys = match &value.style {
+                FlowStyle::PossiblyUninitialized => Vec::new(),
+                FlowStyle::MaybeInitialized(keys) => keys.clone(),
+                _ => continue,
+            };
+            value.style = FlowStyle::InitializedIfGuardTruthy {
+                guard: guard.clone(),
+                guard_value_idx,
+                termination_keys,
+            };
+        }
+    }
+
+    /// Clear `InitializedIfGuardTruthy` entries whose guard matches `narrowed_guard`
+    /// and whose recorded `guard_value_idx` still equals the guard's current flow
+    /// value idx (i.e. the guard hasn't been reassigned).
+    pub fn clear_matching_truthy_guard(&mut self, narrowed_guard: &Name) {
+        let Some(current_idx) = self.current_flow_idx(narrowed_guard) else {
+            return;
+        };
+        for info in self.current_mut().flow.info.values_mut() {
+            let Some(value) = info.value_mut() else {
+                continue;
+            };
+            if let FlowStyle::InitializedIfGuardTruthy {
+                guard,
+                guard_value_idx,
+                ..
+            } = &value.style
+                && guard == narrowed_guard
+                && *guard_value_idx == current_idx
+            {
+                value.style = FlowStyle::Other;
+            }
+        }
+    }
+
     /// Get the flow idx for `name` in the current fork's **base** flow (the flow captured
     /// at `start_fork` time, before any branches were started).
     ///

pyrefly/lib/binding/stmt.rs

@@ -991,6 +991,18 @@ impl<'a> BindingsBuilder<'a> {
                 // are in the base flow and visible after the if-statement. This mirrors the
                 // fix for ternary expressions in expr.rs (Expr::If handling).
                 self.ensure_expr(&mut x.test, &mut Usage::Narrowing(None));
+                // Correlated-condition analysis: for a bare `if <name>:` with no elif/else,
+                // remember the guard so we can mark newly-defined names as
+                // `InitializedIfGuardTruthy` after the merge.
+                let truthy_guard = if x.elif_else_clauses.is_empty()
+                    && let Expr::Name(ref guard_name) = *x.test
+                    && let Some(guard_value_idx) = self.scopes.current_flow_idx(&guard_name.id)
+                {
+                    let pre_fork_names = self.scopes.current_flow_names();
+                    Some((guard_name.id.clone(), guard_value_idx, pre_fork_names))
+                } else {
+                    None
+                };
                 self.start_fork(if_range);
                 // Type narrowing operations that are carried over from one branch to the next. For example, in:
                 //   if x is None:
@@ -1090,6 +1102,11 @@ impl<'a> BindingsBuilder<'a> {
                 } else {
                     self.finish_non_exhaustive_fork(&negated_prev_ops, exhaustive_key);
                 }
+                // Post-merge: apply correlated-condition analysis (see `truthy_guard` above).
+                if let Some((guard, guard_value_idx, pre_fork_names)) = truthy_guard {
+                    self.scopes
+                        .upgrade_to_guarded(&guard, guard_value_idx, &pre_fork_names);
+                }
                 // If we have a statically evaluated test like `sys.version_info`, we should set `is_definitely_unreachable` to false
                 // to reduce false positive unreachable errors, since some code paths can still be hit at runtime
                 if contains_static_test_with_no_else && !is_definitely_unreachable {

pyrefly/lib/test/flow_branching.rs

@@ -1902,6 +1902,42 @@ def f(x: int) -> str:
 "#,
 );
 
+// An `else` branch means the name's uninitialized-ness survives the merge via
+// paths that are not controlled by `<guard>`, so the guarded upgrade is skipped.
+testcase!(
+    test_correlated_if_does_not_cross_else,
+    r#"
+def main(a: bool) -> int:
+    if a:
+        b = 3
+    else:
+        pass
+    if a:
+        return b  # E: `b` may be uninitialized
+    return 9
+"#,
+);
+
+// When the guard branch contains a NoReturn call, the `MaybeInitialized` termination
+// keys are preserved through `InitializedIfGuardTruthy`, so reads in unrelated code
+// still correctly suppress the error via the Never-check.
+testcase!(
+    test_correlated_if_preserves_noreturn_fallback,
+    r#"
+from typing import NoReturn
+
+def raises() -> NoReturn:
+    raise Exception()
+
+def main(a: bool) -> int:
+    if a:
+        b = 3
+    else:
+        raises()
+    return b
+"#,
+);
+
 testcase!(
     test_declared_variable_with_noreturn_else_false_positive,
     r#"
@@ -2407,15 +2443,14 @@ for x in range(10):
 // `if a:`, the variable is guaranteed to be initialized because the same
 // condition guards both the definition and the use.
 testcase!(
-    bug = "false positive: b is always initialized when a is truthy",
     test_guarded_initialization_basic,
     r#"
 def f(a: bool) -> int:
     if a:
         b = 3
     c = 5
     if a:
-        return b  # E: `b` may be uninitialized
+        return b
     return 9
     "#,
 );
@@ -2445,21 +2480,19 @@ def f(a: bool, c: bool) -> int:
 );
 
 testcase!(
-    bug = "false positive: b and c are always initialized when a is truthy",
     test_guarded_initialization_multiple_variables,
     r#"
 def f(a: bool) -> int:
     if a:
         b = 3
         c = 4
     if a:
-        return b + c  # E: `b` may be uninitialized  # E: `c` may be uninitialized
+        return b + c
     return 0
     "#,
 );
 
 testcase!(
-    bug = "false positive: b is always initialized when a is truthy",
     test_guarded_initialization_with_intermediate_statements,
     r#"
 def f(a: bool) -> int:
@@ -2468,7 +2501,7 @@ def f(a: bool) -> int:
     x = 5
     y = x + 1
     if a:
-        return b  # E: `b` may be uninitialized
+        return b
     return 9
     "#,
 );
@@ -2488,14 +2521,13 @@ def f(a: bool) -> int:
 );
 
 testcase!(
-    bug = "false positive: b is always initialized when a is truthy",
     test_guarded_initialization_repeated_use,
     r#"
 def f(a: bool) -> None:
     if a:
         b = 3
     if a:
-        print(b)  # E: `b` may be uninitialized
+        print(b)
     if a:
         print(b)
     "#,