From d7ce548893132cddd677be537b4fb8b4f2cd0ef6 Mon Sep 17 00:00:00 2001
From: David Peter <sharkdp@users.noreply.github.com>
Date: Fri, 13 Dec 2024 07:40:14 +0100
Subject: [PATCH] [red-knot] Add narrowing for 'while' loops (#14947)

## Summary

Add type narrowing for `while` loops and corresponding `else` branches.

closes #14861

## Test Plan

New Markdown tests.
---
 .../resources/mdtest/narrow/while.md          | 58 +++++++++++++++++++
 .../src/semantic_index/builder.rs             |  2 +
 .../src/types/infer.rs                        |  2 +-
 3 files changed, 61 insertions(+), 1 deletion(-)
 create mode 100644 crates/red_knot_python_semantic/resources/mdtest/narrow/while.md

diff --git a/crates/red_knot_python_semantic/resources/mdtest/narrow/while.md b/crates/red_knot_python_semantic/resources/mdtest/narrow/while.md
new file mode 100644
index 0000000000..ac91a576c9
--- /dev/null
+++ b/crates/red_knot_python_semantic/resources/mdtest/narrow/while.md
@@ -0,0 +1,58 @@
+# Narrowing in `while` loops
+
+We only make sure that narrowing works for `while` loops in general, we do not exhaustively test all
+narrowing forms here, as they are covered in other tests.
+
+Note how type narrowing works subtly different from `if` ... `else`, because the negated constraint
+is retained after the loop.
+
+## Basic `while` loop
+
+```py
+def next_item() -> int | None: ...
+
+x = next_item()
+
+while x is not None:
+    reveal_type(x)  # revealed: int
+    x = next_item()
+
+reveal_type(x)  # revealed: None
+```
+
+## `while` loop with `else`
+
+```py
+def next_item() -> int | None: ...
+
+x = next_item()
+
+while x is not None:
+    reveal_type(x)  # revealed: int
+    x = next_item()
+else:
+    reveal_type(x)  # revealed: None
+
+reveal_type(x)  # revealed: None
+```
+
+## Nested `while` loops
+
+```py
+from typing import Literal
+
+def next_item() -> Literal[1, 2, 3]: ...
+
+x = next_item()
+
+while x != 1:
+    reveal_type(x)  # revealed: Literal[2, 3]
+
+    while x != 2:
+        # TODO: this should be Literal[1, 3]; Literal[3] is only correct
+        # in the first loop iteration
+        reveal_type(x)  # revealed: Literal[3]
+        x = next_item()
+
+    x = next_item()
+```
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 c75af214d8..27e657aba9 100644
--- a/crates/red_knot_python_semantic/src/semantic_index/builder.rs
+++ b/crates/red_knot_python_semantic/src/semantic_index/builder.rs
@@ -833,6 +833,7 @@ where
                 self.visit_expr(test);
 
                 let pre_loop = self.flow_snapshot();
+                let constraint = self.record_expression_constraint(test);
 
                 // Save aside any break states from an outer loop
                 let saved_break_states = std::mem::take(&mut self.loop_break_states);
@@ -852,6 +853,7 @@ where
                 // We may execute the `else` clause without ever executing the body, so merge in
                 // the pre-loop state before visiting `else`.
                 self.flow_merge(pre_loop);
+                self.record_negated_constraint(constraint);
                 self.visit_body(orelse);
 
                 // Breaking out of a while loop bypasses the `else` clause, so merge in the break
diff --git a/crates/red_knot_python_semantic/src/types/infer.rs b/crates/red_knot_python_semantic/src/types/infer.rs
index 6b4393f588..1b732ce531 100644
--- a/crates/red_knot_python_semantic/src/types/infer.rs
+++ b/crates/red_knot_python_semantic/src/types/infer.rs
@@ -2092,7 +2092,7 @@ impl<'db> TypeInferenceBuilder<'db> {
             orelse,
         } = while_statement;
 
-        self.infer_expression(test);
+        self.infer_standalone_expression(test);
         self.infer_body(body);
         self.infer_body(orelse);
     }