Fix false positive unbound-name for walrus operator in while condition

Summary:
The while-loop condition always evaluates at least once, so any walrus target in the condition is guaranteed to be assigned after the loop. However, `setup_loop` snapshots the current flow as the loop's base *before* `ensure_expr` processes the test, so walrus-defined names only entered the loop flow and were missing from the base. When `teardown_loop` merged, it saw those names missing from the base and marked them `PossiblyUninitialized`.

The fix propagates newly-defined names from the while condition into the loop's base flow after `ensure_expr`, mirroring the existing `propagate_new_flow_entries_to_fork_base` used for elif conditions.

Fixes #3272

Reviewed By: grievejia

Differential Revision: D103246547

fbshipit-source-id: 00560948fcd587cf98764a87e0d0da27e9a4eea4

Author: migeed-z

pyrefly/lib/binding/scope.rs

@@ -525,6 +525,26 @@ impl Flow {
     }
 }
 
+/// Copy new or newly-initialized entries from `source` into `target`.
+/// Entries absent from `target` are inserted; entries present but uninitialized
+/// are overwritten when `source` has an initialized binding.
+fn propagate_new_flow_entries(
+    source: &SmallMap<Name, FlowInfo>,
+    target: &mut SmallMap<Name, FlowInfo>,
+) {
+    for (name, info) in source.iter() {
+        if let Some(existing) = target.get(name) {
+            if matches!(existing.initialized(), InitializedInFlow::No)
+                && !matches!(info.initialized(), InitializedInFlow::No)
+            {
+                target.insert(name.clone(), info.clone());
+            }
+        } else {
+            target.insert(name.clone(), info.clone());
+        }
+    }
+}
+
 /// Bound names can accumulate facet narrows from long assignment chains (e.g. huge
 /// literal dictionaries). Limiting how many consecutive narrows we remember keeps
 /// the flow graph shallow enough to avoid recursive explosions in the solver.
@@ -2142,19 +2162,21 @@ impl Scopes {
             .forks
             .last_mut()
             .expect("propagate_new_flow_entries_to_fork_base called outside of a fork");
-        for (name, info) in scope.flow.info.iter() {
-            if let Some(existing) = fork.base.info.get(name) {
-                // Update the base entry if it is uninitialized but the branch
-                // has an initialized binding (e.g. walrus targeting `x: int`).
-                if matches!(existing.initialized(), InitializedInFlow::No)
-                    && !matches!(info.initialized(), InitializedInFlow::No)
-                {
-                    fork.base.info.insert(name.clone(), info.clone());
-                }
-            } else {
-                fork.base.info.insert(name.clone(), info.clone());
-            }
-        }
+        propagate_new_flow_entries(&scope.flow.info, &mut fork.base.info);
+    }
+
+    /// Copy walrus-defined names from the current flow into the innermost
+    /// loop's base flow. The while-loop condition always evaluates at least
+    /// once, so any walrus target there is guaranteed to be assigned after the
+    /// loop. Without this propagation, `teardown_loop` sees the name only in
+    /// the loop flow and marks it `PossiblyUninitialized`.
+    pub fn propagate_new_flow_entries_to_loop_base(&mut self) {
+        let scope = self.current_mut();
+        let loop_ = scope
+            .loops
+            .last_mut()
+            .expect("propagate_new_flow_entries_to_loop_base called outside of a loop");
+        propagate_new_flow_entries(&scope.flow.info, &mut loop_.base.info);
     }
 
     /// Return the current binding index and flow style for `name`, if it exists

pyrefly/lib/binding/stmt.rs

@@ -963,6 +963,9 @@ impl<'a> BindingsBuilder<'a> {
                 // made in the loop (e.g. if we reassign the test variable).
                 // Typecheck the test condition during solving.
                 self.ensure_expr(&mut x.test, &mut Usage::Narrowing(None));
+                // The while condition always evaluates at least once, so walrus
+                // targets are guaranteed to be assigned after the loop.
+                self.scopes.propagate_new_flow_entries_to_loop_base();
                 let is_while_true = self.sys_info.evaluate_bool(&x.test) == Some(true);
                 let narrow_ops = NarrowOps::from_expr(self, Some(&x.test));
                 self.bind_narrow_ops(

pyrefly/lib/test/flow_branching.rs

@@ -1487,6 +1487,96 @@ def f(a: int) -> int:
     "#,
 );
 
+testcase!(
+    test_walrus_in_while_post_loop,
+    r#"
+from typing import Callable, Any
+
+class Cat:
+    def equals(self, other: Any) -> bool:
+        return False
+
+def main(f: Callable[[], Cat]) -> None:
+    while (a := f()).equals(1):
+        break
+    print(a)
+    "#,
+);
+
+testcase!(
+    test_walrus_in_while_simple,
+    r#"
+def f() -> int:
+    return 1
+
+def main() -> None:
+    while (x := f()) > 0:
+        break
+    print(x)
+    "#,
+);
+
+testcase!(
+    test_walrus_in_while_with_else,
+    r#"
+def f() -> int:
+    return 1
+
+def main() -> None:
+    while (x := f()) > 0:
+        pass
+    else:
+        pass
+    print(x)
+    "#,
+);
+
+testcase!(
+    test_walrus_in_while_pre_declared_uninitialized,
+    r#"
+def f() -> int:
+    return 1
+
+def main() -> None:
+    x: int
+    while (x := f()) > 0:
+        break
+    print(x)
+    "#,
+);
+
+testcase!(
+    bug = "walrus in while overwrites pre-bound type instead of narrowing; yields str | int",
+    test_walrus_in_while_pre_bound_type_precision,
+    r#"
+from typing import assert_type
+
+def f_int() -> int:
+    return 1
+
+def main() -> None:
+    x = ""
+    while (x := f_int()) > 0:
+        break
+    assert_type(x, int)  # E: assert_type(Literal[''] | int, int) failed
+    "#,
+);
+
+testcase!(
+    bug =
+        "BoolOp laxness causes false negative for walrus in while short-circuit context, see #1251",
+    test_walrus_in_while_bool_op,
+    r#"
+def cond() -> bool: ...
+def get() -> int: ...
+
+def main() -> None:
+    while cond() and (x := get()):
+        break
+    print(x)
+    "#,
+);
+
 testcase!(
     test_trycatch_implicit_return,
     r#"