[ty] fix too-many-cycle panics when inferring literal type loop variables

crates/ty_python_semantic/resources/corpus/cyclic_loop_var.py

@@ -0,0 +1,5 @@
+# Regression test for https://github.com/astral-sh/ty/issues/3011
+
+x = None
+for _ in range(10):
+    x = 0 if x is None else x + 1

crates/ty_python_semantic/src/types.rs

@@ -7290,3 +7290,9 @@ impl EvaluationMode {
 #[cfg(not(debug_assertions))]
 #[cfg(target_pointer_width = "64")]
 static_assertions::assert_eq_size!(Type, [u8; 16]);
+
+// Make sure that `LiteralValueTypeInner` stays at 12 bytes.
+// The `LiteralFlags` byte must fit in the discriminant's padding.
+#[cfg(not(debug_assertions))]
+#[cfg(target_pointer_width = "64")]
+static_assertions::assert_eq_size!(literal::LiteralValueType, [u8; 12]);

crates/ty_python_semantic/src/types/literal.rs

@@ -1,29 +1,72 @@
+use bitflags::bitflags;
 use compact_str::CompactString;
 use ruff_python_ast::name::Name;
 
 use crate::Db;
+use crate::types::set_theoretic::RecursivelyDefined;
 use crate::types::{ClassLiteral, KnownClass, Type};
 
 /// A literal value. See [`LiteralValueTypeKind`] for details.
 #[derive(Copy, Clone, Debug, PartialEq, Eq, Hash, salsa::Update, get_size2::GetSize)]
 pub struct LiteralValueType<'db>(LiteralValueTypeInner<'db>);
 
-// This enum effectively contains two variants, `Promotable(LiteralValueKind)` and `Unpromotable(LiteralValueKind)`,
-// but flattened to reduce the size of the type.
+/// Each variant carries a [`LiteralFlags`] byte alongside its payload.
+/// Because the flags byte fits into the padding between the enum discriminant
+/// (1 byte) and the 4-byte-aligned payload, the overall size of this enum
+/// stays at 12 bytes, the same as the original flagless representation.
 #[derive(Copy, Clone, Debug, PartialEq, Eq, Hash, salsa::Update, get_size2::GetSize)]
 enum LiteralValueTypeInner<'db> {
-    PromotableInt(IntLiteralType),
-    PromotableBool(bool),
-    PromotableString(StringLiteralType<'db>),
-    PromotableEnum(EnumLiteralType<'db>),
-    PromotableBytes(BytesLiteralType<'db>),
-    PromotableLiteralString,
-    UnpromotableInt(IntLiteralType),
-    UnpromotableBool(bool),
-    UnpromotableString(StringLiteralType<'db>),
-    UnpromotableEnum(EnumLiteralType<'db>),
-    UnpromotableBytes(BytesLiteralType<'db>),
-    UnpromotableLiteralString,
+    Int(IntLiteralType, LiteralFlags),
+    Bool(bool, LiteralFlags),
+    String(StringLiteralType<'db>, LiteralFlags),
+    Enum(EnumLiteralType<'db>, LiteralFlags),
+    Bytes(BytesLiteralType<'db>, LiteralFlags),
+    LiteralString(LiteralFlags),
+}
+
+bitflags! {
+    /// Bit-packed flags for promotability and recursive-definition status.
+    ///
+    /// Stored in each [`LiteralValueTypeInner`] variant, fitting into the
+    /// discriminant's padding so that the enum size is unchanged.
+    #[derive(Copy, Clone, Debug, PartialEq, Eq, Hash, salsa::Update)]
+    struct LiteralFlags: u8 {
+        const PROMOTABLE = 1 << 0;
+        const RECURSIVELY_DEFINED = 1 << 1;
+    }
+}
+
+impl get_size2::GetSize for LiteralFlags {}
+
+impl LiteralFlags {
+    fn new(promotable: bool, recursively_defined: RecursivelyDefined) -> Self {
+        let mut flags = Self::empty();
+        flags.set(Self::PROMOTABLE, promotable);
+        flags.set(Self::RECURSIVELY_DEFINED, recursively_defined.is_yes());
+        flags
+    }
+
+    const fn is_promotable(self) -> bool {
+        self.intersects(Self::PROMOTABLE)
+    }
+
+    const fn recursively_defined(self) -> RecursivelyDefined {
+        if self.intersects(Self::RECURSIVELY_DEFINED) {
+            RecursivelyDefined::Yes
+        } else {
+            RecursivelyDefined::No
+        }
+    }
+
+    fn with_promotable(mut self, promotable: bool) -> Self {
+        self.set(Self::PROMOTABLE, promotable);
+        self
+    }
+
+    fn with_recursively_defined(mut self, value: RecursivelyDefined) -> Self {
+        self.set(Self::RECURSIVELY_DEFINED, value.is_yes());
+        self
+    }
 }
 
 #[derive(Copy, Clone, Debug, PartialEq, Eq, Hash, salsa::Update, get_size2::GetSize)]
@@ -55,104 +98,85 @@ impl<'db> LiteralValueType<'db> {
         }
     }
 
+    fn flags(self) -> LiteralFlags {
+        match self.0 {
+            LiteralValueTypeInner::Int(_, f)
+            | LiteralValueTypeInner::Bool(_, f)
+            | LiteralValueTypeInner::String(_, f)
+            | LiteralValueTypeInner::Enum(_, f)
+            | LiteralValueTypeInner::Bytes(_, f)
+            | LiteralValueTypeInner::LiteralString(f) => f,
+        }
+    }
+
+    fn map_flags(self, func: impl FnOnce(LiteralFlags) -> LiteralFlags) -> Self {
+        Self(match self.0 {
+            LiteralValueTypeInner::Int(v, f) => LiteralValueTypeInner::Int(v, func(f)),
+            LiteralValueTypeInner::Bool(v, f) => LiteralValueTypeInner::Bool(v, func(f)),
+            LiteralValueTypeInner::String(v, f) => LiteralValueTypeInner::String(v, func(f)),
+            LiteralValueTypeInner::Enum(v, f) => LiteralValueTypeInner::Enum(v, func(f)),
+            LiteralValueTypeInner::Bytes(v, f) => LiteralValueTypeInner::Bytes(v, func(f)),
+            LiteralValueTypeInner::LiteralString(f) => {
+                LiteralValueTypeInner::LiteralString(func(f))
+            }
+        })
+    }
+
+    pub(crate) fn with_recursively_defined(self, value: RecursivelyDefined) -> Self {
+        self.map_flags(|f| f.with_recursively_defined(value))
+    }
+
+    pub(crate) fn recursively_defined(self) -> RecursivelyDefined {
+        self.flags().recursively_defined()
+    }
+
     /// Creates a literal value that may be promoted.
     pub(crate) fn promotable(kind: impl Into<LiteralValueTypeKind<'db>>) -> LiteralValueType<'db> {
-        let repr = match kind.into() {
-            LiteralValueTypeKind::Int(int) => LiteralValueTypeInner::PromotableInt(int),
-            LiteralValueTypeKind::Bool(bool) => LiteralValueTypeInner::PromotableBool(bool),
-            LiteralValueTypeKind::String(string) => LiteralValueTypeInner::PromotableString(string),
-            LiteralValueTypeKind::Enum(e) => LiteralValueTypeInner::PromotableEnum(e),
-            LiteralValueTypeKind::Bytes(bytes) => LiteralValueTypeInner::PromotableBytes(bytes),
-            LiteralValueTypeKind::LiteralString => LiteralValueTypeInner::PromotableLiteralString,
-        };
-
-        Self(repr)
+        let flags = LiteralFlags::new(true, RecursivelyDefined::No);
+        Self(match kind.into() {
+            LiteralValueTypeKind::Int(v) => LiteralValueTypeInner::Int(v, flags),
+            LiteralValueTypeKind::Bool(v) => LiteralValueTypeInner::Bool(v, flags),
+            LiteralValueTypeKind::String(v) => LiteralValueTypeInner::String(v, flags),
+            LiteralValueTypeKind::Enum(v) => LiteralValueTypeInner::Enum(v, flags),
+            LiteralValueTypeKind::Bytes(v) => LiteralValueTypeInner::Bytes(v, flags),
+            LiteralValueTypeKind::LiteralString => LiteralValueTypeInner::LiteralString(flags),
+        })
     }
 
     /// Creates a literal value that should not be promoted.
     pub(crate) fn unpromotable(
         kind: impl Into<LiteralValueTypeKind<'db>>,
     ) -> LiteralValueType<'db> {
-        let repr = match kind.into() {
-            LiteralValueTypeKind::Int(int) => LiteralValueTypeInner::UnpromotableInt(int),
-            LiteralValueTypeKind::Bool(bool) => LiteralValueTypeInner::UnpromotableBool(bool),
-            LiteralValueTypeKind::String(string) => {
-                LiteralValueTypeInner::UnpromotableString(string)
-            }
-            LiteralValueTypeKind::Enum(e) => LiteralValueTypeInner::UnpromotableEnum(e),
-            LiteralValueTypeKind::Bytes(bytes) => LiteralValueTypeInner::UnpromotableBytes(bytes),
-            LiteralValueTypeKind::LiteralString => LiteralValueTypeInner::UnpromotableLiteralString,
-        };
-
-        Self(repr)
+        let flags = LiteralFlags::new(false, RecursivelyDefined::No);
+        Self(match kind.into() {
+            LiteralValueTypeKind::Int(v) => LiteralValueTypeInner::Int(v, flags),
+            LiteralValueTypeKind::Bool(v) => LiteralValueTypeInner::Bool(v, flags),
+            LiteralValueTypeKind::String(v) => LiteralValueTypeInner::String(v, flags),
+            LiteralValueTypeKind::Enum(v) => LiteralValueTypeInner::Enum(v, flags),
+            LiteralValueTypeKind::Bytes(v) => LiteralValueTypeInner::Bytes(v, flags),
+            LiteralValueTypeKind::LiteralString => LiteralValueTypeInner::LiteralString(flags),
+        })
     }
 
     /// Returns the unpromotable form of this literal value.
     #[must_use]
     pub(crate) fn to_unpromotable(self) -> Self {
-        let repr = match self.0 {
-            LiteralValueTypeInner::PromotableInt(int) => {
-                LiteralValueTypeInner::UnpromotableInt(int)
-            }
-            LiteralValueTypeInner::PromotableBool(bool) => {
-                LiteralValueTypeInner::UnpromotableBool(bool)
-            }
-            LiteralValueTypeInner::PromotableString(string) => {
-                LiteralValueTypeInner::UnpromotableString(string)
-            }
-            LiteralValueTypeInner::PromotableEnum(e) => LiteralValueTypeInner::UnpromotableEnum(e),
-            LiteralValueTypeInner::PromotableBytes(bytes) => {
-                LiteralValueTypeInner::UnpromotableBytes(bytes)
-            }
-            LiteralValueTypeInner::PromotableLiteralString => {
-                LiteralValueTypeInner::UnpromotableLiteralString
-            }
-            LiteralValueTypeInner::UnpromotableInt(_)
-            | LiteralValueTypeInner::UnpromotableBool(_)
-            | LiteralValueTypeInner::UnpromotableString(_)
-            | LiteralValueTypeInner::UnpromotableEnum(_)
-            | LiteralValueTypeInner::UnpromotableBytes(_)
-            | LiteralValueTypeInner::UnpromotableLiteralString => self.0,
-        };
-
-        Self(repr)
+        self.map_flags(|f| f.with_promotable(false))
     }
 
     /// Returns `true` if this literal value should be eagerly promoted to its instance type.
     pub(crate) fn is_promotable(self) -> bool {
-        match self.0 {
-            LiteralValueTypeInner::PromotableInt(_)
-            | LiteralValueTypeInner::PromotableBool(_)
-            | LiteralValueTypeInner::PromotableString(_)
-            | LiteralValueTypeInner::PromotableEnum(_)
-            | LiteralValueTypeInner::PromotableBytes(_)
-            | LiteralValueTypeInner::PromotableLiteralString => true,
-
-            LiteralValueTypeInner::UnpromotableInt(_)
-            | LiteralValueTypeInner::UnpromotableBool(_)
-            | LiteralValueTypeInner::UnpromotableString(_)
-            | LiteralValueTypeInner::UnpromotableEnum(_)
-            | LiteralValueTypeInner::UnpromotableBytes(_)
-            | LiteralValueTypeInner::UnpromotableLiteralString => false,
-        }
+        self.flags().is_promotable()
     }
 
     pub(crate) fn kind(self) -> LiteralValueTypeKind<'db> {
         match self.0 {
-            LiteralValueTypeInner::UnpromotableInt(int)
-            | LiteralValueTypeInner::PromotableInt(int) => LiteralValueTypeKind::Int(int),
-            LiteralValueTypeInner::UnpromotableBool(bool)
-            | LiteralValueTypeInner::PromotableBool(bool) => LiteralValueTypeKind::Bool(bool),
-            LiteralValueTypeInner::UnpromotableString(string)
-            | LiteralValueTypeInner::PromotableString(string) => {
-                LiteralValueTypeKind::String(string)
-            }
-            LiteralValueTypeInner::UnpromotableEnum(e)
-            | LiteralValueTypeInner::PromotableEnum(e) => LiteralValueTypeKind::Enum(e),
-            LiteralValueTypeInner::UnpromotableBytes(bytes)
-            | LiteralValueTypeInner::PromotableBytes(bytes) => LiteralValueTypeKind::Bytes(bytes),
-            LiteralValueTypeInner::UnpromotableLiteralString
-            | LiteralValueTypeInner::PromotableLiteralString => LiteralValueTypeKind::LiteralString,
+            LiteralValueTypeInner::Int(v, _) => LiteralValueTypeKind::Int(v),
+            LiteralValueTypeInner::Bool(v, _) => LiteralValueTypeKind::Bool(v),
+            LiteralValueTypeInner::String(v, _) => LiteralValueTypeKind::String(v),
+            LiteralValueTypeInner::Enum(v, _) => LiteralValueTypeKind::Enum(v),
+            LiteralValueTypeInner::Bytes(v, _) => LiteralValueTypeKind::Bytes(v),
+            LiteralValueTypeInner::LiteralString(_) => LiteralValueTypeKind::LiteralString,
         }
     }
 

crates/ty_python_semantic/src/types/set_theoretic.rs

@@ -841,7 +841,7 @@ pub enum RecursivelyDefined {
 }
 
 impl RecursivelyDefined {
-    const fn is_yes(self) -> bool {
+    pub(crate) const fn is_yes(self) -> bool {
         matches!(self, RecursivelyDefined::Yes)
     }
 

crates/ty_python_semantic/src/types/set_theoretic/builder.rs

@@ -462,6 +462,8 @@ impl<'db> UnionBuilder<'db> {
                 }
             }
             Type::LiteralValue(literal) => {
+                self.recursively_defined =
+                    self.recursively_defined.or(literal.recursively_defined());
                 match literal.kind() {
                     // If adding a string literal, look for an existing `UnionElement::StringLiterals` to
                     // add it to, or an existing element that is a super-type of string literals, which
@@ -825,22 +827,34 @@ impl<'db> UnionBuilder<'db> {
             match element {
                 UnionElement::IntLiterals(literals) => {
                     types.extend(literals.into_iter().map(|(literal, promotable)| {
-                        Type::from(LiteralValueType::new(literal, promotable))
+                        Type::from(
+                            LiteralValueType::new(literal, promotable)
+                                .with_recursively_defined(self.recursively_defined),
+                        )
                     }));
                 }
                 UnionElement::StringLiterals(literals) => {
                     types.extend(literals.into_iter().map(|(literal, promotable)| {
-                        Type::from(LiteralValueType::new(literal, promotable))
+                        Type::from(
+                            LiteralValueType::new(literal, promotable)
+                                .with_recursively_defined(self.recursively_defined),
+                        )
                     }));
                 }
                 UnionElement::BytesLiterals(literals) => {
                     types.extend(literals.into_iter().map(|(literal, promotable)| {
-                        Type::from(LiteralValueType::new(literal, promotable))
+                        Type::from(
+                            LiteralValueType::new(literal, promotable)
+                                .with_recursively_defined(self.recursively_defined),
+                        )
                     }));
                 }
                 UnionElement::EnumLiterals { literals, .. } => {
                     types.extend(literals.into_iter().map(|(literal, promotable)| {
-                        Type::from(LiteralValueType::new(literal, promotable))
+                        Type::from(
+                            LiteralValueType::new(literal, promotable)
+                                .with_recursively_defined(self.recursively_defined),
+                        )
                     }));
                 }
                 UnionElement::Type(ty) => types.push(ty),