[ty] improve complex TDD-based narrowing

crates/ruff_benchmark/benches/ty.rs

@@ -658,6 +658,70 @@ class E(Enum):
     });
 }
 
+/// Benchmark for narrowing a large union type through multiple match statements.
+///
+/// This is extracted from egglog-python's `pretty.py`, where a ~30-class union type
+/// (`AllDecls`) is narrowed by exhaustive match statements.
+///
+/// Sample code structure:
+/// ```python
+/// from __future__ import annotations
+/// from dataclasses import dataclass
+///
+/// @dataclass
+/// class C0:
+///     value: int
+/// ...
+///
+/// AllDecls = C0 | C1 | ...
+///
+/// def process(decl: AllDecls) -> None:
+///     match decl:
+///         case C0(): pass
+///         ...
+///         case _: pass
+/// ```
+fn benchmark_large_union_narrowing(criterion: &mut Criterion) {
+    const NUM_CLASSES: usize = 30;
+    const NUM_MATCH_BRANCHES: usize = 29;
+
+    setup_rayon();
+
+    let mut code =
+        "from __future__ import annotations\nfrom dataclasses import dataclass\n\n".to_string();
+
+    for i in 0..NUM_CLASSES {
+        writeln!(&mut code, "@dataclass\nclass C{i}:\n    value: int\n").ok();
+    }
+
+    code.push_str("AllDecls = ");
+    for i in 0..NUM_CLASSES {
+        if i > 0 {
+            code.push_str(" | ");
+        }
+        write!(&mut code, "C{i}").ok();
+    }
+    code.push_str("\n\n");
+
+    code.push_str("def process(decl: AllDecls) -> None:\n    match decl:\n");
+    for i in 0..NUM_MATCH_BRANCHES {
+        writeln!(&mut code, "        case C{i}():\n            pass",).ok();
+    }
+    code.push_str("        case _:\n            pass\n\n");
+
+    criterion.bench_function("ty_micro[large_union_narrowing]", |b| {
+        b.iter_batched_ref(
+            || setup_micro_case(&code),
+            |case| {
+                let Case { db, .. } = case;
+                let result = db.check();
+                assert_eq!(result.len(), 0);
+            },
+            BatchSize::SmallInput,
+        );
+    });
+}
+
 struct ProjectBenchmark<'a> {
     project: InstalledProject<'a>,
     fs: MemoryFileSystem,
@@ -820,6 +884,7 @@ criterion_group!(
     benchmark_many_enum_members,
     benchmark_many_enum_members_2,
     benchmark_very_large_tuple,
+    benchmark_large_union_narrowing,
 );
 criterion_group!(project, anyio, attrs, hydra, datetype);
 criterion_main!(check_file, micro, project);

crates/ty_python_semantic/resources/mdtest/binary/integers.md

@@ -1,5 +1,11 @@
 # Binary operations on integers
 
+> Developer's note: This is mainly a test for the behavior of the type inferer. The constant
+> evaluator (`resolve_to_literal`) of `SemanticIndexBuilder` is implemented separately from the type
+> inferer, so if you modify the contents of this file or the type inferer, please also modify the
+> implementation of `resolve_to_literal` and the unit tests (semantic_index/tests/const_eval\_\*) at
+> the same time.
+
 ## Basic Arithmetic
 
 ```py

crates/ty_python_semantic/resources/mdtest/loops/while_loop.md

@@ -247,21 +247,45 @@ Here the loop condition forces `x` to be `False` at loop exit, because there is
 def random() -> bool:
     return True
 
-x = random()
-reveal_type(x)  # revealed: bool
-while x:
-    pass
-reveal_type(x)  # revealed: Literal[False]
+def _(x: bool):
+    while x:
+        pass
+    reveal_type(x)  # revealed: Literal[False]
 ```
 
 However, we can't narrow `x` like this when there's a `break` in the loop:
 
 ```py
-x = random()
-while x:
-    if random():
+def _(x: bool):
+    while x:
+        if random():
+            break
+    reveal_type(x)  # revealed: bool
+
+def _(x: bool):
+    while x:
+        pass
+    reveal_type(x)  # revealed: Literal[False]
+
+    x = random()
+    while x:
+        if random():
+            break
+    reveal_type(x)  # revealed: bool
+
+def _(y: int | None):
+    x = 1
+    while True:
+        if x == 0:
+            break
+
+        if y is None:
+            y = 0
+            continue
+
         break
-reveal_type(x)  # revealed: bool
+
+    reveal_type(y)  # revealed: int
 ```
 
 ### Non-static loop conditions

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

@@ -254,7 +254,7 @@ def _(x: Literal["foo", b"bar"] | int):
             pass
         case b"bar" if reveal_type(x):  # revealed: Literal[b"bar"] | int
             pass
-        case _ if reveal_type(x):  # revealed: int | Literal["foo", b"bar"]
+        case _ if reveal_type(x):  # revealed: Literal["foo", b"bar"] | int
             pass
 ```
 

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

@@ -180,27 +180,56 @@ def _(x: int | None):
 ```
 
 ```py
+from typing import Final
+
 def _(x: int | None):
     if 1 + 1 == 2:
         if x is None:
             return
         reveal_type(x)  # revealed: int
 
-    # TODO: should be `int` (the else-branch of `1 + 1 == 2` is unreachable)
-    reveal_type(x)  # revealed: int | None
+    reveal_type(x)  # revealed: int
+
+# non-constant but always-true condition
+needs_inference: Final = True
+
+def _(x: int | None):
+    if needs_inference:
+        if x is None:
+            return
+        reveal_type(x)  # revealed: int
+
+    reveal_type(x)  # revealed: int
 ```
 
 This also works when the always-true condition is nested inside a narrowing branch:
 
 ```py
+from typing import Literal
+
 def _(x: int | None):
     if x is None:
         if 1 + 1 == 2:
             return
 
-    # TODO: should be `int` (the inner always-true branch makes the outer
-    # if-branch terminal)
-    reveal_type(x)  # revealed: int | None
+    reveal_type(x)  # revealed: int
+
+def _(x: int | None):
+    if x is None:
+        if needs_inference:
+            return
+
+    reveal_type(x)  # revealed: int
+
+def always_true(val: object) -> Literal[True]:
+    return True
+
+def _(x: int | None):
+    if x is None:
+        if always_true(x):
+            return
+
+    reveal_type(x)  # revealed: int
 ```
 
 ## Narrowing from `assert` should not affect reassigned variables

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

@@ -31,14 +31,14 @@ else:
     reveal_type(x)  # revealed: Never
 
 if x or not x:
-    reveal_type(x)  # revealed: Literal[0, -1, "", "foo", b"", b"bar"] | bool | None | tuple[()]
+    reveal_type(x)  # revealed: Literal[-1, 0, "foo", "", b"bar", b""] | bool | None | tuple[()]
 else:
     reveal_type(x)  # revealed: Never
 
 if not (x or not x):
     reveal_type(x)  # revealed: Never
 else:
-    reveal_type(x)  # revealed: Literal[0, -1, "", "foo", b"", b"bar"] | bool | None | tuple[()]
+    reveal_type(x)  # revealed: Literal[-1, 0, "foo", "", b"bar", b""] | bool | None | tuple[()]
 
 if (isinstance(x, int) or isinstance(x, str)) and x:
     reveal_type(x)  # revealed: Literal[-1, True, "foo"]

crates/ty_python_semantic/src/semantic_index.rs

@@ -925,6 +925,50 @@ mod tests {
             .collect()
     }
 
+    /// A function to test how the constant evaluator of `SemanticIndexBuilder` evaluates an expression
+    /// (the evaluation should match that of `TypeInferenceBuilder`).
+    /// For example, for the input `x = 1\nif cond: x = 2\nx`, if `cond` evaluates to `AlwaysTrue`, it returns `vec![2]`,
+    /// if it evaluates to `AlwaysFalse`, it returns `vec![1]`, ​​if it evaluates to `Ambiguous`, it returns `vec![1, 2]`.
+    fn reachable_bindings_for_terminal_use(content: &str) -> Vec<i64> {
+        let TestCase { db, file } = test_case(content);
+        let scope = global_scope(&db, file);
+        let module = parsed_module(&db, file).load(&db);
+        let ast = module.syntax();
+
+        let terminal_expr = ast
+            .body
+            .last()
+            .and_then(ast::Stmt::as_expr_stmt)
+            .map(|stmt| stmt.value.as_ref())
+            .expect("expected terminal expression statement");
+        let terminal_name = terminal_expr
+            .as_name_expr()
+            .expect("terminal expression should be a name");
+
+        let use_id = terminal_name.scoped_use_id(&db, scope);
+        let use_def = use_def_map(&db, scope);
+
+        use_def
+            .bindings_at_use(use_id)
+            .filter_map(|binding_with_constraints| {
+                let definition = binding_with_constraints.binding.definition()?;
+                let DefinitionKind::Assignment(assignment) = definition.kind(&db) else {
+                    return None;
+                };
+
+                let ast::Expr::NumberLiteral(ast::ExprNumberLiteral {
+                    value: ast::Number::Int(value),
+                    ..
+                }) = assignment.value(&module)
+                else {
+                    return None;
+                };
+
+                value.as_i64()
+            })
+            .collect::<Vec<_>>()
+    }
+
     #[test]
     fn empty() {
         let TestCase { db, file } = test_case("");
@@ -1590,6 +1634,71 @@ class C[T]:
         assert_eq!(*num, 1);
     }
 
+    #[test]
+    fn const_eval_lshift_overflow_is_ambiguous() {
+        let values = reachable_bindings_for_terminal_use(
+            "
+x = 1
+if 1 << 63:
+    x = 2
+x
+",
+        );
+        assert_eq!(values, vec![1, 2]);
+    }
+
+    #[test]
+    fn const_eval_lshift_zero_short_circuit() {
+        let values = reachable_bindings_for_terminal_use(
+            "
+x = 1
+if 0 << 4000000000000000000:
+    x = 2
+x
+",
+        );
+        assert_eq!(values, vec![1]);
+    }
+
+    #[test]
+    fn const_eval_rshift_large_positive() {
+        let values = reachable_bindings_for_terminal_use(
+            "
+x = 1
+if 1 >> 5000000000:
+    x = 2
+x
+",
+        );
+        assert_eq!(values, vec![1]);
+    }
+
+    #[test]
+    fn const_eval_rshift_large_negative_operand() {
+        let values = reachable_bindings_for_terminal_use(
+            "
+x = 1
+if (-1) >> 5000000000:
+    x = 2
+x
+",
+        );
+        assert_eq!(values, vec![2]);
+    }
+
+    #[test]
+    fn const_eval_negative_lshift_is_ambiguous() {
+        let values = reachable_bindings_for_terminal_use(
+            "
+x = 1
+if 42 << -3:
+    x = 2
+x
+",
+        );
+        assert_eq!(values, vec![1, 2]);
+    }
+
     #[test]
     fn expression_scope() {
         let TestCase { db, file } = test_case("x = 1;\ndef test():\n  y = 4");