From 8cbd433a3130c692ea8f78894fdfacd6e4a50c27 Mon Sep 17 00:00:00 2001
From: Dhruv Manilawala <dhruvmanila@gmail.com>
Date: Wed, 14 May 2025 02:57:48 +0530
Subject: [PATCH] [ty] Add cycle handling for unpacking targets (#18078)

## Summary

This PR adds cycle handling for `infer_unpack_types` based on the
analysis in astral-sh/ty#364.

Fixes: astral-sh/ty#364

## Test Plan

Add a cycle handling test for unpacking in `cycle.md`
---
 .../resources/mdtest/cycle.md                 | 18 +++++++++++
 crates/ty_python_semantic/src/types/infer.rs  | 15 +++++++++-
 .../ty_python_semantic/src/types/unpacker.rs  | 30 +++++++++++++++++--
 3 files changed, 60 insertions(+), 3 deletions(-)

diff --git a/crates/ty_python_semantic/resources/mdtest/cycle.md b/crates/ty_python_semantic/resources/mdtest/cycle.md
index dea0a124e5..0bd3b5b2a6 100644
--- a/crates/ty_python_semantic/resources/mdtest/cycle.md
+++ b/crates/ty_python_semantic/resources/mdtest/cycle.md
@@ -13,3 +13,21 @@ def f(x: f):
 
 reveal_type(f)  # revealed: def f(x: Unknown) -> Unknown
 ```
+
+## Unpacking
+
+See: <https://github.com/astral-sh/ty/issues/364>
+
+```py
+class Point:
+    def __init__(self, x: int = 0, y: int = 0) -> None:
+        self.x = x
+        self.y = y
+
+    def replace_with(self, other: "Point") -> None:
+        self.x, self.y = other.x, other.y
+
+p = Point()
+reveal_type(p.x)  # revealed: Unknown | int
+reveal_type(p.y)  # revealed: Unknown | int
+```
diff --git a/crates/ty_python_semantic/src/types/infer.rs b/crates/ty_python_semantic/src/types/infer.rs
index 5b7533d714..be64625dc2 100644
--- a/crates/ty_python_semantic/src/types/infer.rs
+++ b/crates/ty_python_semantic/src/types/infer.rs
@@ -307,7 +307,7 @@ fn single_expression_cycle_initial<'db>(
 /// involved in an unpacking operation. It returns a result-like object that can be used to get the
 /// type of the variables involved in this unpacking along with any violations that are detected
 /// during this unpacking.
-#[salsa::tracked(returns(ref))]
+#[salsa::tracked(returns(ref), cycle_fn=unpack_cycle_recover, cycle_initial=unpack_cycle_initial)]
 pub(super) fn infer_unpack_types<'db>(db: &'db dyn Db, unpack: Unpack<'db>) -> UnpackResult<'db> {
     let file = unpack.file(db);
     let _span =
@@ -318,6 +318,19 @@ pub(super) fn infer_unpack_types<'db>(db: &'db dyn Db, unpack: Unpack<'db>) -> U
     unpacker.finish()
 }
 
+fn unpack_cycle_recover<'db>(
+    _db: &'db dyn Db,
+    _value: &UnpackResult<'db>,
+    _count: u32,
+    _unpack: Unpack<'db>,
+) -> salsa::CycleRecoveryAction<UnpackResult<'db>> {
+    salsa::CycleRecoveryAction::Iterate
+}
+
+fn unpack_cycle_initial<'db>(_db: &'db dyn Db, _unpack: Unpack<'db>) -> UnpackResult<'db> {
+    UnpackResult::cycle_fallback(Type::Never)
+}
+
 /// Returns the type of the nearest enclosing class for the given scope.
 ///
 /// This function walks up the ancestor scopes starting from the given scope,
diff --git a/crates/ty_python_semantic/src/types/unpacker.rs b/crates/ty_python_semantic/src/types/unpacker.rs
index d3c3026f4a..5a3ffbc65e 100644
--- a/crates/ty_python_semantic/src/types/unpacker.rs
+++ b/crates/ty_python_semantic/src/types/unpacker.rs
@@ -287,6 +287,7 @@ impl<'db> Unpacker<'db> {
         UnpackResult {
             diagnostics: self.context.finish(),
             targets: self.targets,
+            cycle_fallback_type: None,
         }
     }
 }
@@ -295,21 +296,46 @@ impl<'db> Unpacker<'db> {
 pub(crate) struct UnpackResult<'db> {
     targets: FxHashMap<ScopedExpressionId, Type<'db>>,
     diagnostics: TypeCheckDiagnostics,
+
+    /// The fallback type for missing expressions.
+    ///
+    /// This is used only when constructing a cycle-recovery `UnpackResult`.
+    cycle_fallback_type: Option<Type<'db>>,
 }
 
 impl<'db> UnpackResult<'db> {
     /// Returns the inferred type for a given sub-expression of the left-hand side target
     /// of an unpacking assignment.
     ///
-    /// Panics if a scoped expression ID is passed in that does not correspond to a sub-
+    /// # Panics
+    ///
+    /// May panic if a scoped expression ID is passed in that does not correspond to a sub-
     /// expression of the target.
     #[track_caller]
     pub(crate) fn expression_type(&self, expr_id: ScopedExpressionId) -> Type<'db> {
-        self.targets[&expr_id]
+        self.try_expression_type(expr_id).expect(
+            "expression should belong to this `UnpackResult` and \
+            `Unpacker` should have inferred a type for it",
+        )
+    }
+
+    pub(crate) fn try_expression_type(&self, expr_id: ScopedExpressionId) -> Option<Type<'db>> {
+        self.targets
+            .get(&expr_id)
+            .copied()
+            .or(self.cycle_fallback_type)
     }
 
     /// Returns the diagnostics in this unpacking assignment.
     pub(crate) fn diagnostics(&self) -> &TypeCheckDiagnostics {
         &self.diagnostics
     }
+
+    pub(crate) fn cycle_fallback(cycle_fallback_type: Type<'db>) -> Self {
+        Self {
+            targets: FxHashMap::default(),
+            diagnostics: TypeCheckDiagnostics::default(),
+            cycle_fallback_type: Some(cycle_fallback_type),
+        }
+    }
 }