//! We have Salsa queries for inferring types at four different granularities: scope-level,
//! statement-level, definition-level, and expression-level, plus lightweight queries for
//! focused subregions like function decorators.
//!
//! Scope-level inference is for when we are actually checking a file, and need to check types for
//! everything in that file's scopes, or give a linter access to types of arbitrary expressions
//! (via the [`HasType`](crate::semantic_model::HasType) trait).
//!
//! Statement-level inference is needed in only a few cases. Some expressions need access to
//! type context from a parent expression, e.g., the type of lambda parameter depends on the
//! annotated type of the lambda. Statements are the minimal unit of code that can be inferred
//! without external type context, and so statement-level inference allows us to resolve type
//! context without relying on scope-level inference cycles.
//!
//! Definition-level inference allows us to look up the types of places in other scopes (e.g. for
//! imports) with the minimum inference necessary, so that if we're looking up one place from a
//! very large module, we can avoid a bunch of unnecessary work. Definition-level inference also
//! allows us to handle import cycles without getting into a cycle of scope-level inference
//! queries.
//!
//! The expression-level inference query is needed in only a few cases. Since some assignments can
//! have multiple targets (via `x = y = z` or unpacking `(x, y) = z`, they can be associated with
//! multiple definitions (one per assigned place). In order to avoid inferring the type of the
//! right-hand side once per definition, we infer it as a standalone query, so its result will be
//! cached by Salsa. We also need the expression-level query for inferring types in type guard
//! expressions (e.g. the test clause of an `if` statement.)
//!
//! Inferring types at any of the four region granularities returns a [`ExpressionInference`],
//! [`DefinitionInference`], [`StatementInference`], or [`ScopeInference`], which hold the types for
//! every expression within the inferred region. Some inference types also expose the type of every
//! definition within the inferred region.
//!
//! Some type expressions can require deferred evaluation. This includes all type expressions in
//! stub files, or annotation expressions in modules with `from __future__ import annotations`, or
//! stringified annotations. We have a fourth Salsa query for inferring the deferred types
//! associated with a particular definition. Scope-level inference infers deferred types for all
//! definitions once the rest of the types in the scope have been inferred.
//!
//! Many of our type inference Salsa queries implement cycle recovery via fixed-point iteration. In
//! general, they initiate fixed-point iteration by returning an `Inference` type that returns
//! the `Divergent` type for all expressions, bindings, and declarations, and then they continue iterating
//! the query cycle until a fixed-point is reached. Salsa has a built-in fixed limit on the number
//! of iterations, so if we fail to converge, Salsa will eventually panic. (This should of course
//! be considered a bug.)

use itertools::Either;
use ruff_db::parsed::parsed_module;
use ruff_python_ast as ast;
use ruff_text_size::{Ranged, TextRange};
use rustc_hash::FxHashMap;
use salsa;
use salsa::plumbing::AsId;
use std::borrow::Cow;
pub(super) use ty_python_core::frozen::{FrozenMap, FrozenSet, FrozenValueMap};

use crate::types::diagnostic::TypeCheckDiagnostics;
use crate::types::function::{FunctionDecorators, FunctionType};
use crate::types::generics::Specialization;
use crate::types::unpacker::{UnpackResult, Unpacker};
use crate::types::{
    ClassLiteral, KnownClass, StaticClassLiteral, Type, TypeAndQualifiers, TypeQualifiers,
};
use crate::{Db, FxIndexSet};

use builder::TypeInferenceBuilder;
pub(super) use comparisons::UnsupportedComparisonError;
use ty_python_core::definition::{Definition, DefinitionKind};
use ty_python_core::expression::Expression;
use ty_python_core::scope::ScopeId;
use ty_python_core::statement::StatementInner;
use ty_python_core::unpack::Unpack;
use ty_python_core::{ExpressionNodeKey, SemanticIndex, Statement, semantic_index};

mod builder;
mod comparisons;
#[cfg(test)]
mod tests;

bitflags::bitflags! {
    /// Metadata for expressions inferred as type expressions.
    #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)]
    pub(crate) struct TypeExpressionFlags: u8 {
        /// The expression is syntactically an `Unpack[...]` type expression.
        const UNPACK = 1 << 0;
    }
}

impl get_size2::GetSize for TypeExpressionFlags {}

#[derive(Debug, Clone, Copy, Eq, PartialEq, get_size2::GetSize, salsa::SalsaValue)]
struct TypeAndRange<'db> {
    ty: Type<'db>,
    range: TextRange,
}

type CollectionUseConstraints<'db> = FxHashMap<Definition<'db>, FxIndexSet<Type<'db>>>;

/// Extends the current collection-use constraints with those from the previous cycle iteration.
///
/// Constraints for a collection can appear and disappear while dependent inference results are
/// still changing. Retaining previous constraints makes this part of inference monotonic.
fn extend_collection_use_constraints<'db>(
    current: &mut CollectionUseConstraints<'db>,
    previous: &CollectionUseConstraints<'db>,
) {
    #[expect(
        clippy::iter_over_hash_type,
        reason = "constraints for distinct collection definitions are merged independently"
    )]
    for (collection_def, constraints) in previous {
        current
            .entry(*collection_def)
            .or_default()
            .extend(constraints);
    }
}

/// Infer all types for a [`Definition`] (including sub-expressions).
/// Use when resolving a place use or public type of a place.
#[salsa::tracked(
    returns(ref),
    cycle_initial=|db, id, definition: Definition<'db>| {
        DefinitionInference::cycle_initial(db, definition, Type::divergent(id))
    },
    cycle_fn=|db, cycle, previous: &DefinitionInference<'db>, inference: DefinitionInference<'db>, definition| {
        inference.cycle_normalized(db, previous, cycle, definition)
    },
    heap_size=ruff_memory_usage::heap_size
)]
pub(crate) fn infer_definition_types<'db>(
    db: &'db dyn Db,
    definition: Definition<'db>,
) -> DefinitionInference<'db> {
    let file = definition.file(db);
    let module = parsed_module(db, file).load(db);
    let _span = tracing::trace_span!(
        "infer_definition_types",
        range = ?definition.kind(db).target_range(&module),
        ?file
    )
    .entered();

    let index = semantic_index(db, file);

    TypeInferenceBuilder::new(db, InferenceRegion::Definition(definition), index, &module)
        .finish_definition(definition)
}

/// Returns `true` if the definition refers to a dictionary-key binding that should be discarded.
///
/// For example, inference synthesizes an `x["a"] = "bad"` binding for:
/// ```python
/// def f(x: dict[str, int]):
///     x = {"a": "bad"}  # invalid-assignment
/// ```
/// Since the enclosing assignment was rejected, place resolution must ignore that binding and fall
/// back to the declared value type of `x`.
pub(crate) fn is_discarded_dict_key_assignment<'db>(
    db: &'db dyn Db,
    definition: Definition<'db>,
) -> bool {
    let DefinitionKind::DictKeyAssignment(dict_key_assignment) = definition.kind(db) else {
        return false;
    };

    infer_definition_types(db, dict_key_assignment.assignment()).discards_dict_key_assignments()
}

/// Infer decorator expression types for a function definition.
///
/// This is a lightweight query that avoids the cycle risk of calling
/// `infer_definition_types` when we need to check decorators while
/// already inside definition inference (e.g. checking `Self` in a
/// `@staticmethod`).
#[salsa::tracked(
    returns(ref),
    cycle_initial=|_, _, _| FunctionDecoratorInference::default(),
    heap_size=ruff_memory_usage::heap_size
)]
pub(crate) fn function_known_decorators<'db>(
    db: &'db dyn Db,
    definition: Definition<'db>,
) -> FunctionDecoratorInference<'db> {
    let file = definition.file(db);
    let module = parsed_module(db, file).load(db);
    let index = semantic_index(db, file);

    TypeInferenceBuilder::new(
        db,
        InferenceRegion::FunctionDecorators(definition),
        index,
        &module,
    )
    .finish_function_decorator_inference()
}

pub(crate) fn function_known_decorator_flags<'db>(
    db: &'db dyn Db,
    definition: Definition<'db>,
) -> FunctionDecorators {
    function_known_decorators(db, definition).known_decorators()
}

/// A compact inference result for function decorators.
///
/// Unlike [`DefinitionInference`], this stores only decorator expression types and
/// diagnostics, plus the expression-side state that needs to be merged back into
/// function-definition inference.
#[derive(Debug, Eq, PartialEq, Default, get_size2::GetSize, salsa::SalsaValue)]
pub(crate) struct FunctionDecoratorInference<'db> {
    expression_types: FrozenMap<ExpressionNodeKey, Type<'db>>,
    bindings: Box<[(Definition<'db>, Type<'db>)]>,
    called_functions: Box<[FunctionType<'db>]>,
    known_decorators: FunctionDecorators,
    diagnostics: TypeCheckDiagnostics,
}

impl<'db> FunctionDecoratorInference<'db> {
    pub(crate) fn expression_type(
        &self,
        expression: impl Into<ExpressionNodeKey>,
    ) -> Option<Type<'db>> {
        self.expression_types.get(&expression.into()).copied()
    }

    pub(crate) fn expression_types(
        &self,
    ) -> impl ExactSizeIterator<Item = (ExpressionNodeKey, Type<'db>)> + '_ {
        self.expression_types.iter().copied()
    }

    pub(crate) fn bindings(
        &self,
    ) -> impl ExactSizeIterator<Item = (Definition<'db>, Type<'db>)> + '_ {
        self.bindings.iter().copied()
    }

    pub(crate) fn called_functions(&self) -> &[FunctionType<'db>] {
        &self.called_functions
    }

    pub(crate) fn known_decorators(&self) -> FunctionDecorators {
        self.known_decorators
    }

    pub(crate) fn diagnostics(&self) -> &TypeCheckDiagnostics {
        &self.diagnostics
    }
}

/// Infer types for all deferred type expressions in a [`Definition`].
///
/// Deferred expressions are type expressions (annotations, base classes, aliases...) in a stub
/// file, or in a file with `from __future__ import annotations`, or stringified annotations.
#[salsa::tracked(
    returns(ref),
    cycle_initial=|db, id, definition: Definition<'db>| {
        DefinitionInference::cycle_initial(db, definition, Type::divergent(id))
    },
    cycle_fn=|db, cycle, previous: &DefinitionInference<'db>, inference: DefinitionInference<'db>, definition| {
        inference.cycle_normalized(db, previous, cycle, definition)
    },
    heap_size=ruff_memory_usage::heap_size
)]
pub(crate) fn infer_deferred_types<'db>(
    db: &'db dyn Db,
    definition: Definition<'db>,
) -> DefinitionInference<'db> {
    let file = definition.file(db);
    let module = parsed_module(db, file).load(db);
    let _span = tracing::trace_span!(
        "infer_deferred_types",
        definition = ?definition.as_id(),
        range = ?definition.kind(db).target_range(&module),
        ?file
    )
    .entered();

    let index = semantic_index(db, file);

    TypeInferenceBuilder::new(db, InferenceRegion::Deferred(definition), index, &module)
        .finish_definition(definition)
}

/// Infer all types for a [`ScopeId`], including all definitions and expressions in that scope.
/// Use when checking a scope, or needing to provide a type for an arbitrary expression in the
/// scope.
///
/// Unlike [`infer_scope_types`], this function does not take a type context, as it may infer
/// the parent scope to obtain the necessary type context by which to infer the inner scope.
pub(crate) fn infer_complete_scope_types<'db>(
    db: &'db dyn Db,
    scope: ScopeId<'db>,
) -> &'db ScopeInference<'db> {
    // Scopes that may require type context are inferred during the inference of
    // their outer scope.
    if scope.accepts_type_context(db) {
        let file = scope.file(db);
        let index = semantic_index(db, file);

        if let Some(parent_scope) = index.parent_scope_id(scope.file_scope_id(db)) {
            // Note that nested lambdas or comprehensions may require recursing until we reach
            // an outer scope that is independent of any type context.
            return infer_complete_scope_types(db, parent_scope.to_scope_id(db, file));
        }
    }

    infer_scope_types_impl(db, InferScope::new(db, scope, TypeContext::default()))
}

/// Infer all types for a [`ScopeId`], including all definitions and expressions in that scope.
/// Use when checking a scope, or needing to provide a type for an arbitrary expression in the
/// scope.
///
/// Note that you should generally use [`infer_complete_scope_types`] instead of this method,
/// unless you have already obtained the necessary type context while inferring the parent scope.
/// Inferring a nested scope independently without type context can lead to incorrect inferred
/// types or diagnostics.
pub(crate) fn infer_scope_types<'db>(
    db: &'db dyn Db,
    scope: ScopeId<'db>,
    tcx: TypeContext<'db>,
) -> &'db ScopeInference<'db> {
    infer_scope_types_impl(db, InferScope::new(db, scope, tcx))
}

#[salsa::tracked(
    returns(ref),
    cycle_initial=|_, id, _| ScopeInference::cycle_initial(Type::divergent(id)),
    cycle_fn=|db, cycle, previous: &ScopeInference<'db>, inference: ScopeInference<'db>, _| {
        inference.cycle_normalized(db, previous, cycle)
    },
    heap_size=ruff_memory_usage::heap_size
)]
pub(crate) fn infer_scope_types_impl<'db>(
    db: &'db dyn Db,
    input: InferScope<'db>,
) -> ScopeInference<'db> {
    let (scope, tcx) = input.into_inner(db);
    let file = scope.file(db);
    let _span = tracing::trace_span!("infer_scope_types", scope=?scope.as_id(), ?file).entered();

    let module = parsed_module(db, file).load(db);

    // Using the index here is fine because the code below depends on the AST anyway.
    // The isolation of the query is by the return inferred types.
    let index = semantic_index(db, file);

    TypeInferenceBuilder::new(db, InferenceRegion::Scope(scope, tcx), index, &module).finish_scope()
}

/// Infer all types for an [`Expression`] (including sub-expressions).
/// Use rarely; only for cases where we'd otherwise risk double-inferring an expression: RHS of an
/// assignment, which might be unpacking/multi-target and thus part of multiple definitions, or a
/// type narrowing guard expression (e.g. if statement test node).
pub(crate) fn infer_expression_types<'db>(
    db: &'db dyn Db,
    expression: Expression<'db>,
    tcx: TypeContext<'db>,
) -> &'db ExpressionInference<'db> {
    infer_expression_types_impl(db, InferExpression::new(db, expression, tcx))
}

#[salsa::tracked(
    returns(ref),
    cycle_initial=expression_cycle_initial,
    cycle_fn=|db, cycle, previous: &ExpressionInference<'db>, inference: ExpressionInference<'db>, _| {
        inference.cycle_normalized(db, previous, cycle)
    },
    heap_size=ruff_memory_usage::heap_size
)]
pub(super) fn infer_expression_types_impl<'db>(
    db: &'db dyn Db,
    input: InferExpression<'db>,
) -> ExpressionInference<'db> {
    let (expression, tcx) = input.into_inner(db);

    let file = expression.file(db);
    let module = parsed_module(db, file).load(db);
    let _span = tracing::trace_span!(
        "infer_expression_types",
        expression = ?expression.as_id(),
        range = ?expression.node_ref(db).node(&module).range(),
        ?file
    )
    .entered();

    let index = semantic_index(db, file);

    TypeInferenceBuilder::new(
        db,
        InferenceRegion::Expression(expression, tcx),
        index,
        &module,
    )
    .finish_expression()
}

fn expression_cycle_initial<'db>(
    db: &'db dyn Db,
    id: salsa::Id,
    input: InferExpression<'db>,
) -> ExpressionInference<'db> {
    let (expression, _) = input.into_inner(db);
    let cycle_recovery = Type::divergent(id);
    ExpressionInference::cycle_initial(expression.scope(db), cycle_recovery)
}

/// Infers the type of an `expression` that is guaranteed to be in the same file as the calling query.
///
/// This is a small helper around [`infer_expression_types()`] to reduce the boilerplate.
/// Use [`infer_expression_type()`] if it isn't guaranteed that `expression` is in the same file to
/// avoid cross-file query dependencies.
pub(crate) fn infer_same_file_expression_type<'db>(
    db: &'db dyn Db,
    expression: Expression<'db>,
    tcx: TypeContext<'db>,
) -> Type<'db> {
    let inference = infer_expression_types(db, expression, tcx);
    inference.expression_type(expression.node_ref(db))
}

/// Infers the type of an expression where the expression might come from another file.
///
/// Use this over [`infer_expression_types`] if the expression might come from another file than the
/// enclosing query to avoid cross-file query dependencies.
///
/// Use [`infer_same_file_expression_type`] if it is guaranteed that  `expression` is in the same
/// to avoid unnecessary salsa ingredients. This is normally the case inside the `TypeInferenceBuilder`.
pub(crate) fn infer_expression_type<'db>(
    db: &'db dyn Db,
    expression: Expression<'db>,
    tcx: TypeContext<'db>,
) -> Type<'db> {
    infer_expression_type_impl(db, InferExpression::new(db, expression, tcx))
}

#[salsa::tracked(
    returns(copy),
    cycle_initial=|_, id, _| Type::divergent(id),
    cycle_fn=|db, cycle, previous: &Type<'db>, result: Type<'db>, _| {
        result.cycle_normalized(db, *previous, cycle)
    },
    heap_size=ruff_memory_usage::heap_size
)]
fn infer_expression_type_impl<'db>(db: &'db dyn Db, input: InferExpression<'db>) -> Type<'db> {
    let (expression, _) = input.into_inner(db);

    // It's okay to call the "same file" version here because we're inside a salsa query.
    let inference = infer_expression_types_impl(db, input);

    inference.expression_type(expression.node_ref(db))
}

/// Infer all types for a [`Statement`].
///
/// This is useful when you want to infer a sub-expression with its natural type context, as
/// statements are the minimal unit of code that can be inferred without external type context.
pub(super) fn infer_statement_types<'db>(
    db: &'db dyn Db,
    statement: Statement<'db>,
) -> StatementInference<'db> {
    match statement {
        Statement::Expression(expression) => StatementInference::Expression(
            infer_expression_types(db, expression, TypeContext::default()),
        ),
        Statement::Definition(definition) => {
            StatementInference::Definition(definition, infer_definition_types(db, definition))
        }
        Statement::Other(statement) => {
            StatementInference::Other(infer_statement_types_impl(db, statement))
        }
    }
}

#[salsa::tracked(
    returns(ref),
    cycle_initial=|db, id, statement: StatementInner<'db>| {
        StatementInferenceInner::cycle_initial(statement.scope(db), Type::divergent(id))
    },
    cycle_fn=|db, cycle, previous: &StatementInferenceInner<'db>, inference: StatementInferenceInner<'db>, _| {
        inference.cycle_normalized(db, previous, cycle)
    },
    heap_size=ruff_memory_usage::heap_size
)]
fn infer_statement_types_impl<'db>(
    db: &'db dyn Db,
    statement: StatementInner<'db>,
) -> StatementInferenceInner<'db> {
    let file = statement.file(db);
    let module = parsed_module(db, file).load(db);
    let _span = tracing::trace_span!(
        "infer_statement_types",
        statement = ?statement.as_id(),
        range = ?statement.node_ref(db).node(&module).range(),
        ?file
    )
    .entered();

    let index = semantic_index(db, file);

    TypeInferenceBuilder::new(db, InferenceRegion::Statement(statement), index, &module)
        .finish_statement()
}

/// An `Expression` with an optional `TypeContext`.
///
/// This is a Salsa supertype used as the input to `infer_expression_types` to avoid
/// interning an `ExpressionWithContext` unnecessarily when no type context is provided.
#[derive(Debug, Clone, Copy, Eq, Hash, PartialEq, salsa::Supertype)]
pub(super) enum InferExpression<'db> {
    Bare(Expression<'db>),
    WithContext(ExpressionWithContext<'db>),
}

#[salsa::interned(debug, heap_size=ruff_memory_usage::heap_size)]
pub(super) struct ExpressionWithContext<'db> {
    #[returns(copy)]
    expression: Expression<'db>,
    #[returns(copy)]
    tcx: TypeContext<'db>,
}

impl<'db> InferExpression<'db> {
    pub(super) fn new(
        db: &'db dyn Db,
        expression: Expression<'db>,
        tcx: TypeContext<'db>,
    ) -> InferExpression<'db> {
        if tcx.annotation.is_some() {
            InferExpression::WithContext(ExpressionWithContext::new(db, expression, tcx))
        } else {
            InferExpression::Bare(expression)
        }
    }

    fn into_inner(self, db: &'db dyn Db) -> (Expression<'db>, TypeContext<'db>) {
        match self {
            InferExpression::Bare(expression) => (expression, TypeContext::default()),
            InferExpression::WithContext(expression_with_context) => (
                expression_with_context.expression(db),
                expression_with_context.tcx(db),
            ),
        }
    }
}

/// A `ScopeId` with an optional `TypeContext`.
#[derive(Debug, Clone, Copy, Eq, Hash, PartialEq, salsa::Supertype)]
pub(super) enum InferScope<'db> {
    Bare(ScopeId<'db>),
    WithContext(ScopeWithContext<'db>),
}

#[salsa::interned(debug, heap_size=ruff_memory_usage::heap_size)]
pub(super) struct ScopeWithContext<'db> {
    #[returns(copy)]
    scope: ScopeId<'db>,
    #[returns(copy)]
    tcx: TypeContext<'db>,
}

impl<'db> InferScope<'db> {
    pub(super) fn new(
        db: &'db dyn Db,
        scope: ScopeId<'db>,
        tcx: TypeContext<'db>,
    ) -> InferScope<'db> {
        if tcx.annotation.is_some() {
            InferScope::WithContext(ScopeWithContext::new(db, scope, tcx))
        } else {
            InferScope::Bare(scope)
        }
    }

    fn into_inner(self, db: &'db dyn Db) -> (ScopeId<'db>, TypeContext<'db>) {
        match self {
            InferScope::Bare(scope) => (scope, TypeContext::default()),
            InferScope::WithContext(scope_with_context) => {
                (scope_with_context.scope(db), scope_with_context.tcx(db))
            }
        }
    }
}

/// The type context for a given expression, namely the type annotation
/// in an annotated assignment.
///
/// Knowing the outer type context when inferring an expression can enable
/// more precise inference results, aka "bidirectional type inference".
#[derive(
    Default, Copy, Clone, Debug, PartialEq, Eq, Hash, get_size2::GetSize, salsa::SalsaValue,
)]
pub(crate) struct TypeContext<'db> {
    pub(crate) annotation: Option<Type<'db>>,
}

impl<'db> TypeContext<'db> {
    pub(crate) fn new(annotation: Option<Type<'db>>) -> Self {
        Self { annotation }
    }

    /// If the type annotation is a specialized instance of the given `KnownClass`, returns the
    /// specialization.
    fn known_specialization(
        &self,
        db: &'db dyn Db,
        known_class: KnownClass,
    ) -> Option<Specialization<'db>> {
        self.annotation
            .and_then(|ty| ty.known_specialization(db, known_class))
    }

    pub(crate) fn map(self, f: impl FnOnce(Type<'db>) -> Type<'db>) -> Self {
        Self {
            annotation: self.annotation.map(f),
        }
    }

    pub(crate) fn is_typealias(&self) -> bool {
        self.annotation
            .is_some_and(|ty| ty.is_typealias_special_form())
    }

    /// If the type annotation is a union, returns the target elements that it can be narrowed to.
    pub(crate) fn narrow_targets(&self, db: &'db dyn Db) -> Option<Cow<'db, [Type<'db>]>> {
        let union = self.annotation?.as_union_like(db)?;

        let targets = if union.has_aliases(db) {
            let expanded = union.expand_aliases(db);
            if let Some(union) = expanded.as_union_like(db) {
                Cow::Borrowed(union.elements(db))
            } else {
                Cow::Owned(vec![expanded])
            }
        } else {
            Cow::Borrowed(union.elements(db))
        };

        // TODO: We could theoretically attempt to narrow to every element of
        // the power set of this union. However, this leads to an exponential
        // explosion of inference attempts, and is rarely needed in practice.
        Some(targets)
    }
}

impl<'db> From<Type<'db>> for TypeContext<'db> {
    fn from(annotation: Type<'db>) -> Self {
        Self::new(Some(annotation))
    }
}

/// Infer the types for an [`Unpack`] operation.
///
/// This infers the expression type and performs structural match against the target expression
/// 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),
    cycle_initial=|_, id, _| UnpackResult::cycle_initial(Type::divergent(id)),
    cycle_fn=|db, cycle, previous: &UnpackResult<'db>, result: UnpackResult<'db>, _| {
        result.cycle_normalized(db, previous, cycle)
    },
    heap_size=ruff_memory_usage::heap_size
)]
pub(super) fn infer_unpack_types<'db>(db: &'db dyn Db, unpack: Unpack<'db>) -> UnpackResult<'db> {
    let file = unpack.file(db);
    let module = parsed_module(db, file).load(db);
    let _span = tracing::trace_span!("infer_unpack_types", range=?unpack.range(db, &module), ?file)
        .entered();

    let mut unpacker = Unpacker::new(db, unpack.target_scope(db), &module);
    unpacker.unpack(unpack.target(db, &module), unpack.value(db));
    unpacker.finish()
}

/// Returns the type of the nearest enclosing class for the given scope.
///
/// This function walks up the ancestor scopes starting from the given scope,
/// and finds the closest class definition. This is different to the behaviour of
/// [`TypeInferenceBuilder::class_context_of_current_method`], which will only return
/// `Some(class)` if either the immediate parent scope is a class OR the immediate parent
/// scope is a type-parameters scope and the grandparent scope is a class.
///
/// Returns `None` if no enclosing class is found.
pub(crate) fn nearest_enclosing_class<'db>(
    db: &'db dyn Db,
    semantic: &SemanticIndex<'db>,
    scope: ScopeId,
) -> Option<StaticClassLiteral<'db>> {
    semantic
        .ancestor_scopes(scope.file_scope_id(db))
        .find_map(|(_, ancestor_scope)| {
            let class = ancestor_scope.node().as_class()?;
            let definition = semantic.expect_single_definition(class);
            original_class_type(db, definition).and_then(ClassLiteral::as_static)
        })
}

/// Return the original class literal for a class definition.
///
/// For decorated classes, this is the class object before applying decorators. The public
/// binding may be replaced by a class decorator's return type, but class-body inference still
/// needs the original class object for implicit `self`/`cls`, `Self`, and dataclass logic.
///
/// For example, the public binding for `C` may be the `int` returned by `replace`, but the class
/// body and nested definitions still need the original class object:
/// ```python
/// def replace(cls: type[object]) -> int:
///     return 1
///
/// @replace
/// class C:
///     def method(self) -> None: ...
/// ```
pub(crate) fn original_class_type<'db>(
    db: &'db dyn Db,
    definition: Definition<'db>,
) -> Option<ClassLiteral<'db>> {
    let inference = infer_definition_types(db, definition);
    inference
        .undecorated_type()
        .unwrap_or_else(|| inference.binding_type(definition))
        .as_class_literal()
}

/// Returns the type of the nearest enclosing function for the given scope.
///
/// This function walks up the ancestor scopes starting from the given scope,
/// and finds the closest (non-lambda) function definition.
///
/// Returns `None` if no enclosing function is found.
pub(crate) fn nearest_enclosing_function<'db>(
    db: &'db dyn Db,
    semantic: &SemanticIndex<'db>,
    scope: ScopeId,
) -> Option<FunctionType<'db>> {
    semantic
        .ancestor_scopes(scope.file_scope_id(db))
        .find_map(|(_, ancestor_scope)| {
            let func = ancestor_scope.node().as_function()?;
            let definition = semantic.expect_single_definition(func);
            infer_definition_types(db, definition).function_type(definition)
        })
}

/// A region within which we can infer types.
#[derive(Copy, Clone, Debug)]
pub(crate) enum InferenceRegion<'db> {
    // infer types for a [`Statement`].
    Statement(StatementInner<'db>),
    /// infer types for a standalone [`Expression`]
    Expression(Expression<'db>, TypeContext<'db>),
    /// infer types for a [`Definition`]
    Definition(Definition<'db>),
    /// infer types for the decorators on a function [`Definition`]
    FunctionDecorators(Definition<'db>),
    /// infer deferred types for a [`Definition`]
    Deferred(Definition<'db>),
    /// infer types for an entire [`ScopeId`]
    Scope(ScopeId<'db>, TypeContext<'db>),
}

impl<'db> InferenceRegion<'db> {
    fn scope(self, db: &'db dyn Db) -> ScopeId<'db> {
        match self {
            InferenceRegion::Statement(statement) => statement.scope(db),
            InferenceRegion::Expression(expression, _) => expression.scope(db),
            InferenceRegion::Definition(definition)
            | InferenceRegion::FunctionDecorators(definition)
            | InferenceRegion::Deferred(definition) => definition.scope(db),
            InferenceRegion::Scope(scope, _) => scope,
        }
    }
}

/// The inferred types for a scope region.
#[derive(Debug, Eq, PartialEq, get_size2::GetSize, salsa::SalsaValue)]
pub(crate) struct ScopeInference<'db> {
    /// The types of every expression in this region.
    expressions: FrozenValueMap<ExpressionNodeKey, Type<'db>>,

    /// The extra data that is only present for few inference regions.
    extra: Option<Box<ScopeInferenceExtra<'db>>>,
}

#[derive(Debug, Eq, PartialEq, get_size2::GetSize, Default, salsa::SalsaValue)]
struct ScopeInferenceExtra<'db> {
    /// String annotations found in this region
    string_annotations: FrozenSet<ExpressionNodeKey>,

    /// Type qualifiers (`Required`, `NotRequired`, etc.) for annotation expressions.
    qualifiers: FrozenMap<ExpressionNodeKey, TypeQualifiers>,

    /// Expected types for expression nodes tracked for IDE completion.
    expected_types: FrozenMap<ExpressionNodeKey, Type<'db>>,

    /// Metadata for type expressions in this region.
    type_expression_flags: FrozenMap<ExpressionNodeKey, TypeExpressionFlags>,

    /// The constraints on any collection initializers that are accessed in this region.
    collection_use_constraints: CollectionUseConstraints<'db>,

    /// The fallback type for missing expressions/bindings/declarations or recursive type inference.
    cycle_recovery: Option<Type<'db>>,

    /// The diagnostics for this region.
    diagnostics: TypeCheckDiagnostics,
}

impl<'db> ScopeInference<'db> {
    fn cycle_initial(cycle_recovery: Type<'db>) -> Self {
        Self {
            extra: Some(Box::new(ScopeInferenceExtra {
                cycle_recovery: Some(cycle_recovery),
                ..ScopeInferenceExtra::default()
            })),
            expressions: FrozenValueMap::default(),
        }
    }

    fn cycle_normalized(
        mut self,
        db: &'db dyn Db,
        previous_inference: &ScopeInference<'db>,
        cycle: &salsa::Cycle,
    ) -> ScopeInference<'db> {
        self.expressions.map_values(|expr, ty| {
            ty.cycle_normalized(db, previous_inference.expression_type(expr), cycle)
        });

        if cycle.iteration() > crate::TAINTED_CYCLES
            && let Some(previous_extra) = previous_inference.extra.as_deref()
            && !previous_extra.collection_use_constraints.is_empty()
        {
            let extra = self.extra.get_or_insert_default();
            extend_collection_use_constraints(
                &mut extra.collection_use_constraints,
                &previous_extra.collection_use_constraints,
            );
        }

        self
    }

    pub(crate) fn diagnostics(&self) -> Option<&TypeCheckDiagnostics> {
        self.extra.as_deref().map(|extra| &extra.diagnostics)
    }

    pub(crate) fn expression_type(&self, expression: impl Into<ExpressionNodeKey>) -> Type<'db> {
        self.try_expression_type(expression)
            .unwrap_or_else(Type::unknown)
    }

    pub(crate) fn try_expression_type(
        &self,
        expression: impl Into<ExpressionNodeKey>,
    ) -> Option<Type<'db>> {
        self.expressions
            .get(&expression.into())
            .copied()
            .or_else(|| self.fallback_type())
    }

    /// Get qualifiers for an annotation expression.
    pub(crate) fn qualifiers(&self, expression: impl Into<ExpressionNodeKey>) -> TypeQualifiers {
        self.extra
            .as_deref()
            .and_then(|extra| extra.qualifiers.get(&expression.into()).copied())
            .unwrap_or_default()
    }

    pub(crate) fn try_expected_type(
        &self,
        expression: impl Into<ExpressionNodeKey>,
    ) -> Option<Type<'db>> {
        self.extra
            .as_deref()
            .and_then(|extra| extra.expected_types.get(&expression.into()).copied())
    }

    fn fallback_type(&self) -> Option<Type<'db>> {
        self.extra.as_ref().and_then(|extra| extra.cycle_recovery)
    }

    /// Returns whether the given expression is a string annotation
    /// (the string in `x: "int | None"`).
    pub(crate) fn is_string_annotation(&self, expression: impl Into<ExpressionNodeKey>) -> bool {
        let Some(extra) = &self.extra else {
            return false;
        };

        extra.string_annotations.contains(&expression.into())
    }

    /// Get metadata for a type expression.
    pub(crate) fn type_expression_flags(
        &self,
        expression: impl Into<ExpressionNodeKey>,
    ) -> TypeExpressionFlags {
        self.extra
            .as_ref()
            .and_then(|extra| extra.type_expression_flags.get(&expression.into()).copied())
            .unwrap_or_default()
    }
}

/// The result of inferring a declaration recorded by the semantic index.
#[derive(Debug, Clone, Copy, Eq, PartialEq, get_size2::GetSize)]
pub(crate) enum InferredDeclaration<'db> {
    /// A valid declaration with an inferred declared type.
    Declared(TypeAndQualifiers<'db>),
    /// An invalid declaration that should not participate in declaration resolution.
    ///
    /// For example, `obj.attr: int` when `obj` is not a valid implicit receiver, or
    /// `items[0]: int`, which is never a valid annotation target.
    Rejected,
}

impl<'db> InferredDeclaration<'db> {
    pub(crate) fn declared(self) -> Option<TypeAndQualifiers<'db>> {
        match self {
            InferredDeclaration::Declared(declared) => Some(declared),
            InferredDeclaration::Rejected => None,
        }
    }
}

/// The inferred types for a definition region.
#[derive(Debug, Eq, PartialEq, get_size2::GetSize, salsa::SalsaValue)]
pub(crate) struct DefinitionInference<'db> {
    /// The types of every expression in this region.
    expressions: FrozenMap<ExpressionNodeKey, Type<'db>>,

    /// The scope this region is part of.
    #[cfg(debug_assertions)]
    scope: ScopeId<'db>,

    /// The binding and declaration types inferred in this region.
    types: DefinitionTypes<'db>,

    /// The extra data that is only present for few inference regions.
    extra: Option<Box<DefinitionInferenceExtra<'db>>>,
}

/// The binding and declaration types inferred for a definition region.
///
/// Almost all regions contain a single binding, optionally paired with a declaration for the same
/// definition and type. Keep those cases inline instead of retaining separate allocations.
#[derive(Debug, Default, Eq, PartialEq, get_size2::GetSize, salsa::SalsaValue)]
enum DefinitionTypes<'db> {
    #[default]
    Empty,
    Binding(Type<'db>),
    Declaration(TypeAndQualifiers<'db>),
    BindingAndDeclaration(TypeAndQualifiers<'db>),
    Other(Box<OtherDefinitionTypes<'db>>),
}

type DefinitionBinding<'db> = (Definition<'db>, Type<'db>);
type DefinitionDeclaration<'db> = (Definition<'db>, TypeAndQualifiers<'db>);

#[derive(Debug, Eq, PartialEq, get_size2::GetSize, salsa::SalsaValue)]
struct OtherDefinitionTypes<'db> {
    bindings: Box<[DefinitionBinding<'db>]>,
    declarations: Box<[DefinitionDeclaration<'db>]>,
}

impl<'db> DefinitionTypes<'db> {
    fn from_parts(
        owner: Definition<'db>,
        bindings: Vec<DefinitionBinding<'db>>,
        declarations: Vec<DefinitionDeclaration<'db>>,
    ) -> Self {
        match (&*bindings, &*declarations) {
            ([], []) => Self::Empty,
            ([(definition, ty)], []) if *definition == owner => Self::Binding(*ty),
            ([], [(definition, ty)]) if *definition == owner => Self::Declaration(*ty),
            ([(binding, binding_ty)], [(declaration, declaration_ty)])
                if *binding == owner
                    && binding == declaration
                    && *binding_ty == declaration_ty.inner_type() =>
            {
                Self::BindingAndDeclaration(*declaration_ty)
            }
            _ => Self::Other(Box::new(OtherDefinitionTypes {
                bindings: bindings.into_boxed_slice(),
                declarations: declarations.into_boxed_slice(),
            })),
        }
    }

    fn binding_type(
        &self,
        owner: Definition<'db>,
        definition: Definition<'db>,
    ) -> Option<Type<'db>> {
        self.bindings(owner)
            .find_map(|(candidate, ty)| (candidate == definition).then_some(ty))
    }

    fn declaration_type(
        &self,
        owner: Definition<'db>,
        definition: Definition<'db>,
    ) -> Option<TypeAndQualifiers<'db>> {
        self.declarations(owner)
            .find_map(|(candidate, ty)| (candidate == definition).then_some(ty))
    }

    fn normalize_binding(
        db: &'db dyn Db,
        previous: &DefinitionTypes<'db>,
        cycle: &salsa::Cycle,
        owner: Definition<'db>,
        definition: Definition<'db>,
        ty: Type<'db>,
    ) -> Type<'db> {
        if let Some(previous_ty) = previous.binding_type(owner, definition) {
            ty.cycle_normalized(db, previous_ty, cycle)
        } else {
            ty.recursive_type_normalized(db, cycle)
        }
    }

    fn normalize_declaration(
        db: &'db dyn Db,
        previous: &DefinitionTypes<'db>,
        cycle: &salsa::Cycle,
        owner: Definition<'db>,
        definition: Definition<'db>,
        ty: TypeAndQualifiers<'db>,
    ) -> TypeAndQualifiers<'db> {
        if let Some(previous_ty) = previous.declaration_type(owner, definition) {
            ty.map_type(|inner| inner.cycle_normalized(db, previous_ty.inner_type(), cycle))
        } else {
            ty.map_type(|inner| inner.recursive_type_normalized(db, cycle))
        }
    }

    fn cycle_normalized(
        self,
        db: &'db dyn Db,
        previous: &DefinitionTypes<'db>,
        cycle: &salsa::Cycle,
        owner: Definition<'db>,
    ) -> Self {
        match self {
            Self::Empty => Self::Empty,
            Self::Binding(ty) => Self::Binding(Self::normalize_binding(
                db, previous, cycle, owner, owner, ty,
            )),
            Self::Declaration(ty) => Self::Declaration(Self::normalize_declaration(
                db, previous, cycle, owner, owner, ty,
            )),
            Self::BindingAndDeclaration(declaration_ty) => {
                let binding_ty = Self::normalize_binding(
                    db,
                    previous,
                    cycle,
                    owner,
                    owner,
                    declaration_ty.inner_type(),
                );
                let declaration_ty =
                    Self::normalize_declaration(db, previous, cycle, owner, owner, declaration_ty);

                if binding_ty == declaration_ty.inner_type() {
                    Self::BindingAndDeclaration(declaration_ty)
                } else {
                    Self::Other(Box::new(OtherDefinitionTypes {
                        bindings: Box::new([(owner, binding_ty)]),
                        declarations: Box::new([(owner, declaration_ty)]),
                    }))
                }
            }
            Self::Other(mut other) => {
                for (definition, ty) in &mut other.bindings {
                    *ty = Self::normalize_binding(db, previous, cycle, owner, *definition, *ty);
                }
                for (definition, ty) in &mut other.declarations {
                    *ty = Self::normalize_declaration(db, previous, cycle, owner, *definition, *ty);
                }

                match (&*other.bindings, &*other.declarations) {
                    ([(binding, binding_ty)], [(declaration, declaration_ty)])
                        if *binding == owner
                            && binding == declaration
                            && *binding_ty == declaration_ty.inner_type() =>
                    {
                        Self::BindingAndDeclaration(*declaration_ty)
                    }
                    _ => Self::Other(other),
                }
            }
        }
    }

    fn bindings(
        &self,
        owner: Definition<'db>,
    ) -> impl ExactSizeIterator<Item = (Definition<'db>, Type<'db>)> {
        match self {
            Self::Binding(ty) => Either::Left(std::iter::once((owner, *ty))),
            Self::BindingAndDeclaration(declaration) => {
                Either::Left(std::iter::once((owner, declaration.inner_type())))
            }
            Self::Other(other) => Either::Right(other.bindings.iter().copied()),
            Self::Empty | Self::Declaration(..) => Either::Right([].iter().copied()),
        }
    }

    fn declarations(
        &self,
        owner: Definition<'db>,
    ) -> impl ExactSizeIterator<Item = (Definition<'db>, TypeAndQualifiers<'db>)> {
        match self {
            Self::Declaration(declaration) | Self::BindingAndDeclaration(declaration) => {
                Either::Left(std::iter::once((owner, *declaration)))
            }
            Self::Other(other) => Either::Right(other.declarations.iter().copied()),
            Self::Empty | Self::Binding(..) => Either::Right([].iter().copied()),
        }
    }

    fn declaration_types(&self) -> impl ExactSizeIterator<Item = TypeAndQualifiers<'db>> + '_ {
        match self {
            Self::Declaration(declaration) | Self::BindingAndDeclaration(declaration) => {
                Either::Left(Either::Left(std::iter::once(*declaration)))
            }
            Self::Other(other) => {
                Either::Left(Either::Right(other.declarations.iter().map(|(_, ty)| *ty)))
            }
            Self::Empty | Self::Binding(..) => Either::Right(std::iter::empty()),
        }
    }
}

/// Compact representations for common combinations of extra definition inference data.
/// `Other` stores uncommon combinations that require multiple fields.
#[derive(Debug, Eq, PartialEq, get_size2::GetSize, salsa::SalsaValue)]
enum DefinitionInferenceExtra<'db> {
    /// Type qualifiers are the only extra data for most annotated definitions.
    Qualifiers(FrozenMap<ExpressionNodeKey, TypeQualifiers>),

    /// Deferred definitions are the only extra data for most definitions with annotations.
    Deferred(Box<[Definition<'db>]>),

    Diagnostics(Box<TypeCheckDiagnostics>),

    Undecorated(Box<Type<'db>>),

    /// Deferred definitions and the undecorated type are a common two-field combination.
    DeferredAndUndecorated(Box<DeferredAndUndecorated<'db>>),

    CalledFunctions(Box<[FunctionType<'db>]>),

    ExpectedTypes(FrozenMap<ExpressionNodeKey, Type<'db>>),

    StringAnnotations(FrozenSet<ExpressionNodeKey>),

    DiscardsDictKeyAssignments,

    Other(Box<OtherDefinitionInferenceExtra<'db>>),
}

#[derive(Debug, Eq, PartialEq, get_size2::GetSize, salsa::SalsaValue)]
struct DeferredAndUndecorated<'db> {
    deferred: Box<[Definition<'db>]>,
    undecorated_type: Type<'db>,
}

#[derive(Debug, Eq, PartialEq, get_size2::GetSize, Default, salsa::SalsaValue)]
struct OtherDefinitionInferenceExtra<'db> {
    /// String annotations found in this region
    string_annotations: FrozenSet<ExpressionNodeKey>,

    /// Expected types for expression nodes tracked for IDE completion.
    expected_types: FrozenMap<ExpressionNodeKey, Type<'db>>,

    /// Functions called while inferring this definition.
    called_functions: Box<[FunctionType<'db>]>,

    /// Metadata for type expressions in this region.
    type_expression_flags: FrozenMap<ExpressionNodeKey, TypeExpressionFlags>,

    /// The constraints on any collection initializers that are accessed in this region.
    collection_use_constraints: CollectionUseConstraints<'db>,

    /// The fallback type for missing expressions/bindings/declarations or recursive type inference.
    cycle_recovery: Option<Type<'db>>,

    /// The definitions that have some deferred parts.
    deferred: Box<[Definition<'db>]>,

    /// The diagnostics for this region.
    diagnostics: TypeCheckDiagnostics,

    /// For decorated function or class definitions, the type before applying decorators.
    undecorated_type: Option<Type<'db>>,

    /// Whether synthesized dictionary-key assignments derived from the right-hand side should be
    /// discarded.
    discards_dict_key_assignments: bool,

    /// Type qualifiers (`Required`, `NotRequired`, etc.) for annotation expressions.
    /// Only populated for expressions that have non-empty qualifiers.
    qualifiers: FrozenMap<ExpressionNodeKey, TypeQualifiers>,
}

impl<'db> DefinitionInferenceExtra<'db> {
    fn into_other(self) -> OtherDefinitionInferenceExtra<'db> {
        match self {
            Self::Qualifiers(qualifiers) => OtherDefinitionInferenceExtra {
                qualifiers,
                ..OtherDefinitionInferenceExtra::default()
            },
            Self::Deferred(deferred) => OtherDefinitionInferenceExtra {
                deferred,
                ..OtherDefinitionInferenceExtra::default()
            },
            Self::Diagnostics(diagnostics) => OtherDefinitionInferenceExtra {
                diagnostics: *diagnostics,
                ..OtherDefinitionInferenceExtra::default()
            },
            Self::Undecorated(undecorated_type) => OtherDefinitionInferenceExtra {
                undecorated_type: Some(*undecorated_type),
                ..OtherDefinitionInferenceExtra::default()
            },
            Self::DeferredAndUndecorated(extra) => {
                let DeferredAndUndecorated {
                    deferred,
                    undecorated_type,
                } = *extra;
                OtherDefinitionInferenceExtra {
                    deferred,
                    undecorated_type: Some(undecorated_type),
                    ..OtherDefinitionInferenceExtra::default()
                }
            }
            Self::CalledFunctions(called_functions) => OtherDefinitionInferenceExtra {
                called_functions,
                ..OtherDefinitionInferenceExtra::default()
            },
            Self::ExpectedTypes(expected_types) => OtherDefinitionInferenceExtra {
                expected_types,
                ..OtherDefinitionInferenceExtra::default()
            },
            Self::StringAnnotations(string_annotations) => OtherDefinitionInferenceExtra {
                string_annotations,
                ..OtherDefinitionInferenceExtra::default()
            },
            Self::DiscardsDictKeyAssignments => OtherDefinitionInferenceExtra {
                discards_dict_key_assignments: true,
                ..OtherDefinitionInferenceExtra::default()
            },
            Self::Other(extra) => *extra,
        }
    }

    fn collection_use_constraints(&self) -> Option<&CollectionUseConstraints<'db>> {
        match self {
            Self::Other(extra) => Some(&extra.collection_use_constraints),
            _ => None,
        }
    }
}

impl<'db> DefinitionInference<'db> {
    fn cycle_initial(
        db: &'db dyn Db,
        definition: Definition<'db>,
        cycle_recovery: Type<'db>,
    ) -> Self {
        let mut types = DefinitionTypes::Empty;

        // Eagerly store more precise types for collection literals to avoid an extra
        // cycle iteration, i.e., by inferring `list[Divergent]` instead of `Divergent`.
        if let DefinitionKind::Assignment(assignment) = definition.kind(db) {
            let module = parsed_module(db, definition.file(db)).load(db);
            let known_collection = match assignment.value(&module) {
                ast::Expr::Set(_) => Some(KnownClass::Set),
                ast::Expr::List(_) => Some(KnownClass::List),
                ast::Expr::Dict(_) => Some(KnownClass::Dict),
                _ => None,
            };

            if let Some(collection_class) = known_collection
                .and_then(|known_collection| known_collection.try_to_class_literal(db))
            {
                let divergent_collection = collection_class
                    .apply_specialization(db, |generic_context| {
                        generic_context.repeat_specialization(db, cycle_recovery)
                    });

                types = DefinitionTypes::Binding(Type::instance(db, divergent_collection));
            }
        }

        Self {
            expressions: FrozenMap::default(),
            types,
            #[cfg(debug_assertions)]
            scope: definition.scope(db),
            extra: Some(Box::new(DefinitionInferenceExtra::Other(Box::new(
                OtherDefinitionInferenceExtra {
                    cycle_recovery: Some(cycle_recovery),
                    ..OtherDefinitionInferenceExtra::default()
                },
            )))),
        }
    }

    fn cycle_normalized(
        mut self,
        db: &'db dyn Db,
        previous_inference: &DefinitionInference<'db>,
        cycle: &salsa::Cycle,
        definition: Definition<'db>,
    ) -> DefinitionInference<'db> {
        for (expr, ty) in &mut self.expressions {
            let previous_ty = previous_inference.expression_type(*expr);
            *ty = ty.cycle_normalized(db, previous_ty, cycle);
        }
        self.types = std::mem::take(&mut self.types).cycle_normalized(
            db,
            &previous_inference.types,
            cycle,
            definition,
        );

        if cycle.iteration() > crate::TAINTED_CYCLES
            && let Some(previous_constraints) = previous_inference
                .extra
                .as_deref()
                .and_then(DefinitionInferenceExtra::collection_use_constraints)
            && !previous_constraints.is_empty()
        {
            let mut extra = self
                .extra
                .take()
                .map_or_else(OtherDefinitionInferenceExtra::default, |extra| {
                    extra.into_other()
                });
            extend_collection_use_constraints(
                &mut extra.collection_use_constraints,
                previous_constraints,
            );
            self.extra = Some(Box::new(DefinitionInferenceExtra::Other(Box::new(extra))));
        }

        self
    }

    pub(crate) fn expression_type(&self, expression: impl Into<ExpressionNodeKey>) -> Type<'db> {
        self.try_expression_type(expression)
            .unwrap_or_else(Type::unknown)
    }

    pub(crate) fn try_expression_type(
        &self,
        expression: impl Into<ExpressionNodeKey>,
    ) -> Option<Type<'db>> {
        self.expressions
            .get(&expression.into())
            .copied()
            .or_else(|| self.fallback_type())
    }

    pub(crate) fn collection_use_constraints(
        &self,
        collection_def: Definition<'db>,
    ) -> Option<&FxIndexSet<Type<'db>>> {
        self.extra
            .as_deref()?
            .collection_use_constraints()?
            .get(&collection_def)
    }

    /// Get qualifiers for an annotation expression
    pub(crate) fn qualifiers(&self, expression: impl Into<ExpressionNodeKey>) -> TypeQualifiers {
        let expression = expression.into();
        match self.extra.as_deref() {
            Some(DefinitionInferenceExtra::Qualifiers(qualifiers)) => {
                qualifiers.get(&expression).copied().unwrap_or_default()
            }
            Some(DefinitionInferenceExtra::Other(extra)) => extra
                .qualifiers
                .get(&expression)
                .copied()
                .unwrap_or_default(),
            Some(_) | None => TypeQualifiers::default(),
        }
    }

    /// Get metadata for a type expression.
    pub(crate) fn type_expression_flags(
        &self,
        expression: impl Into<ExpressionNodeKey>,
    ) -> TypeExpressionFlags {
        match self.extra.as_deref() {
            Some(DefinitionInferenceExtra::Other(extra)) => extra
                .type_expression_flags
                .get(&expression.into())
                .copied()
                .unwrap_or_default(),
            Some(_) | None => TypeExpressionFlags::default(),
        }
    }

    #[track_caller]
    pub(crate) fn binding_type(&self, definition: Definition<'db>) -> Type<'db> {
        self.types
            .bindings(definition)
            .find_map(|(def, ty)| if def == definition { Some(ty) } else { None })
            .or_else(|| self.fallback_type())
            .expect(
                "definition should belong to this TypeInference region and \
                TypeInferenceBuilder should have inferred a type for it",
            )
    }

    fn bindings(
        &self,
        owner: Definition<'db>,
    ) -> impl ExactSizeIterator<Item = (Definition<'db>, Type<'db>)> {
        self.types.bindings(owner)
    }

    pub(crate) fn inferred_declaration(
        &self,
        definition: Definition<'db>,
    ) -> InferredDeclaration<'db> {
        self.types
            .declarations(definition)
            .find_map(|(def, declaration)| {
                if def == definition {
                    Some(InferredDeclaration::Declared(declaration))
                } else {
                    None
                }
            })
            .or_else(|| {
                self.fallback_type()
                    .map(TypeAndQualifiers::declared)
                    .map(InferredDeclaration::Declared)
            })
            .unwrap_or(InferredDeclaration::Rejected)
    }

    fn declarations(
        &self,
        owner: Definition<'db>,
    ) -> impl ExactSizeIterator<Item = (Definition<'db>, TypeAndQualifiers<'db>)> {
        self.types.declarations(owner)
    }

    fn declaration_types(&self) -> impl ExactSizeIterator<Item = TypeAndQualifiers<'db>> {
        self.types.declaration_types()
    }

    pub(crate) fn fallback_type(&self) -> Option<Type<'db>> {
        match self.extra.as_deref() {
            Some(DefinitionInferenceExtra::Other(extra)) => extra.cycle_recovery,
            Some(_) | None => None,
        }
    }

    pub(crate) fn discards_dict_key_assignments(&self) -> bool {
        match self.extra.as_deref() {
            Some(DefinitionInferenceExtra::DiscardsDictKeyAssignments) => true,
            Some(DefinitionInferenceExtra::Other(extra)) => extra.discards_dict_key_assignments,
            Some(_) | None => false,
        }
    }

    pub(crate) fn undecorated_type(&self) -> Option<Type<'db>> {
        match self.extra.as_deref() {
            Some(DefinitionInferenceExtra::Undecorated(undecorated_type)) => {
                Some(**undecorated_type)
            }
            Some(DefinitionInferenceExtra::DeferredAndUndecorated(extra)) => {
                Some(extra.undecorated_type)
            }
            Some(DefinitionInferenceExtra::Other(extra)) => extra.undecorated_type,
            Some(_) | None => None,
        }
    }

    pub(crate) fn function_type(&self, definition: Definition<'db>) -> Option<FunctionType<'db>> {
        let ty = if let Some(undecorated) = self.undecorated_type() {
            undecorated
        } else {
            self.inferred_declaration(definition)
                .declared()?
                .inner_type()
        };

        ty.as_function_literal()
    }
}

/// The inferred types for an expression region.
#[derive(Debug, Eq, PartialEq, get_size2::GetSize, salsa::SalsaValue)]
pub(crate) struct ExpressionInference<'db> {
    /// The types of every expression in this region.
    expressions: FrozenMap<ExpressionNodeKey, Type<'db>>,

    extra: Option<Box<ExpressionInferenceExtra<'db>>>,

    /// The scope this region is part of.
    #[cfg(debug_assertions)]
    scope: ScopeId<'db>,
}

/// Extra data that only exists for few inferred expression regions.
#[derive(Debug, Eq, PartialEq, get_size2::GetSize, Default, salsa::SalsaValue)]
struct ExpressionInferenceExtra<'db> {
    /// String annotations found in this region
    string_annotations: FrozenSet<ExpressionNodeKey>,

    /// Expected types for expression nodes tracked for IDE completion.
    expected_types: FrozenMap<ExpressionNodeKey, Type<'db>>,

    /// Metadata for type expressions in this region.
    type_expression_flags: FrozenMap<ExpressionNodeKey, TypeExpressionFlags>,

    /// The constraints on any collection initializers that are accessed in this region.
    collection_use_constraints: CollectionUseConstraints<'db>,

    /// The types of every binding in this expression region.
    ///
    /// Only very few expression regions have bindings (around 0.1%).
    bindings: Box<[(Definition<'db>, Type<'db>)]>,

    /// The diagnostics for this region.
    diagnostics: TypeCheckDiagnostics,

    /// The fallback type for missing expressions/bindings/declarations or recursive type inference.
    cycle_recovery: Option<Type<'db>>,
}

impl<'db> ExpressionInference<'db> {
    fn cycle_initial(scope: ScopeId<'db>, cycle_recovery: Type<'db>) -> Self {
        let _ = scope;
        Self {
            extra: Some(Box::new(ExpressionInferenceExtra {
                cycle_recovery: Some(cycle_recovery),
                ..ExpressionInferenceExtra::default()
            })),
            expressions: FrozenMap::default(),
            #[cfg(debug_assertions)]
            scope,
        }
    }

    fn cycle_normalized(
        mut self,
        db: &'db dyn Db,
        previous: &ExpressionInference<'db>,
        cycle: &salsa::Cycle,
    ) -> ExpressionInference<'db> {
        if let Some(extra) = self.extra.as_mut() {
            for (binding, binding_ty) in &mut extra.bindings {
                if let Some((_, previous_binding)) = previous.extra.as_deref().and_then(|extra| {
                    extra
                        .bindings
                        .iter()
                        .find(|(previous_binding, _)| previous_binding == binding)
                }) {
                    *binding_ty = binding_ty.cycle_normalized(db, *previous_binding, cycle);
                } else {
                    *binding_ty = binding_ty.recursive_type_normalized(db, cycle);
                }
            }
        }

        for (expr, ty) in &mut self.expressions {
            let previous_ty = previous.expression_type(*expr);
            *ty = ty.cycle_normalized(db, previous_ty, cycle);
        }

        if cycle.iteration() > crate::TAINTED_CYCLES
            && let Some(previous_extra) = previous.extra.as_deref()
            && !previous_extra.collection_use_constraints.is_empty()
        {
            let extra = self.extra.get_or_insert_default();
            extend_collection_use_constraints(
                &mut extra.collection_use_constraints,
                &previous_extra.collection_use_constraints,
            );
        }

        self
    }

    pub(crate) fn try_expression_type(
        &self,
        expression: impl Into<ExpressionNodeKey>,
    ) -> Option<Type<'db>> {
        self.expressions
            .get(&expression.into())
            .copied()
            .or_else(|| self.fallback_type())
    }

    pub(crate) fn expression_type(&self, expression: impl Into<ExpressionNodeKey>) -> Type<'db> {
        self.try_expression_type(expression)
            .unwrap_or_else(Type::unknown)
    }

    pub(crate) fn collection_use_constraints(
        &self,
        collection_def: Definition<'db>,
    ) -> Option<&FxIndexSet<Type<'db>>> {
        self.extra
            .as_ref()?
            .collection_use_constraints
            .get(&collection_def)
    }

    fn fallback_type(&self) -> Option<Type<'db>> {
        self.extra.as_ref().and_then(|extra| extra.cycle_recovery)
    }
}

/// The inferred types for a statement region.
///
/// Many statements can be treated directly as definitions or expressions,
/// and so simply wrapped the inference result of those regions.
#[derive(Debug, Eq, PartialEq, get_size2::GetSize)]
pub(crate) enum StatementInference<'db> {
    Expression(&'db ExpressionInference<'db>),
    Definition(Definition<'db>, &'db DefinitionInference<'db>),
    Other(&'db StatementInferenceInner<'db>),
}

impl<'db> StatementInference<'db> {
    pub(crate) fn expression_type(&self, expression: impl Into<ExpressionNodeKey>) -> Type<'db> {
        match self {
            StatementInference::Expression(inference) => inference.expression_type(expression),
            StatementInference::Definition(_, inference) => inference.expression_type(expression),
            StatementInference::Other(inference) => inference.expression_type(expression),
        }
    }

    pub(crate) fn collection_use_constraints(
        &self,
        collection_def: Definition<'db>,
    ) -> Option<&FxIndexSet<Type<'db>>> {
        match self {
            StatementInference::Expression(inference) => {
                inference.collection_use_constraints(collection_def)
            }
            StatementInference::Definition(_, inference) => {
                inference.collection_use_constraints(collection_def)
            }
            StatementInference::Other(inference) => {
                inference.collection_use_constraints(collection_def)
            }
        }
    }
}

/// The inferred types for a statement region.
#[derive(Debug, Eq, PartialEq, get_size2::GetSize, salsa::SalsaValue)]
pub(crate) struct StatementInferenceInner<'db> {
    /// The types of every expression in this region.
    expressions: FrozenMap<ExpressionNodeKey, Type<'db>>,

    /// The scope this region is part of.
    #[cfg(debug_assertions)]
    scope: ScopeId<'db>,

    /// The types of every binding in this region.
    bindings: Box<[(Definition<'db>, Type<'db>)]>,

    /// The types and type qualifiers of every declaration in this region.
    declarations: Box<[(Definition<'db>, TypeAndQualifiers<'db>)]>,

    /// The extra data that is only present for few inference regions.
    extra: Option<Box<StatementInferenceInnerExtra<'db>>>,
}

#[derive(Debug, Eq, PartialEq, get_size2::GetSize, Default, salsa::SalsaValue)]
struct StatementInferenceInnerExtra<'db> {
    /// String annotations found in this region
    string_annotations: FrozenSet<ExpressionNodeKey>,

    /// Expected types for expression nodes tracked for IDE completion.
    expected_types: FrozenMap<ExpressionNodeKey, Type<'db>>,

    /// Functions called while inferring this statement.
    called_functions: Box<[FunctionType<'db>]>,

    /// The returned types and their corresponding ranges, if this statement is in a function body.
    return_types_and_ranges: Box<[TypeAndRange<'db>]>,

    /// Metadata for type expressions in this region.
    type_expression_flags: FrozenMap<ExpressionNodeKey, TypeExpressionFlags>,

    /// The constraints on any collection initializers that are accessed in this region.
    collection_use_constraints: CollectionUseConstraints<'db>,

    /// The fallback type for missing expressions/bindings/declarations or recursive type inference.
    cycle_recovery: Option<Type<'db>>,

    /// The definitions that have some deferred parts.
    deferred: Box<[Definition<'db>]>,

    /// The diagnostics for this region.
    diagnostics: TypeCheckDiagnostics,

    /// Type qualifiers (`Required`, `NotRequired`, etc.) for annotation expressions.
    /// Only populated for expressions that have non-empty qualifiers.
    qualifiers: FrozenMap<ExpressionNodeKey, TypeQualifiers>,
}

impl<'db> StatementInferenceInner<'db> {
    fn cycle_initial(scope: ScopeId<'db>, cycle_recovery: Type<'db>) -> Self {
        let _ = scope;

        Self {
            expressions: FrozenMap::default(),
            bindings: Box::default(),
            declarations: Box::default(),
            #[cfg(debug_assertions)]
            scope,
            extra: Some(Box::new(StatementInferenceInnerExtra {
                cycle_recovery: Some(cycle_recovery),
                ..StatementInferenceInnerExtra::default()
            })),
        }
    }

    fn cycle_normalized(
        mut self,
        db: &'db dyn Db,
        previous_inference: &StatementInferenceInner<'db>,
        cycle: &salsa::Cycle,
    ) -> StatementInferenceInner<'db> {
        for (expr, ty) in &mut self.expressions {
            let previous_ty = previous_inference.expression_type(*expr);
            *ty = ty.cycle_normalized(db, previous_ty, cycle);
        }
        for (binding, binding_ty) in &mut self.bindings {
            if let Some((_, previous_binding)) = previous_inference
                .bindings
                .iter()
                .find(|(previous_binding, _)| previous_binding == binding)
            {
                *binding_ty = binding_ty.cycle_normalized(db, *previous_binding, cycle);
            } else {
                *binding_ty = binding_ty.recursive_type_normalized(db, cycle);
            }
        }
        for (declaration, declaration_ty) in &mut self.declarations {
            if let Some((_, previous_declaration)) = previous_inference
                .declarations
                .iter()
                .find(|(previous_declaration, _)| previous_declaration == declaration)
            {
                *declaration_ty = declaration_ty.map_type(|decl_ty| {
                    decl_ty.cycle_normalized(db, previous_declaration.inner_type(), cycle)
                });
            } else {
                *declaration_ty =
                    declaration_ty.map_type(|decl_ty| decl_ty.recursive_type_normalized(db, cycle));
            }
        }

        if cycle.iteration() > crate::TAINTED_CYCLES
            && let Some(previous_extra) = previous_inference.extra.as_deref()
            && !previous_extra.collection_use_constraints.is_empty()
        {
            let extra = self.extra.get_or_insert_default();
            extend_collection_use_constraints(
                &mut extra.collection_use_constraints,
                &previous_extra.collection_use_constraints,
            );
        }

        self
    }

    pub(crate) fn expression_type(&self, expression: impl Into<ExpressionNodeKey>) -> Type<'db> {
        self.try_expression_type(expression)
            .unwrap_or_else(Type::unknown)
    }

    pub(crate) fn try_expression_type(
        &self,
        expression: impl Into<ExpressionNodeKey>,
    ) -> Option<Type<'db>> {
        self.expressions
            .get(&expression.into())
            .copied()
            .or_else(|| self.fallback_type())
    }

    pub(crate) fn collection_use_constraints(
        &self,
        collection_def: Definition<'db>,
    ) -> Option<&FxIndexSet<Type<'db>>> {
        self.extra
            .as_ref()?
            .collection_use_constraints
            .get(&collection_def)
    }

    fn bindings(&self) -> impl ExactSizeIterator<Item = (Definition<'db>, Type<'db>)> {
        self.bindings.iter().copied()
    }

    fn declarations(
        &self,
    ) -> impl ExactSizeIterator<Item = (Definition<'db>, TypeAndQualifiers<'db>)> {
        self.declarations.iter().copied()
    }

    pub(crate) fn fallback_type(&self) -> Option<Type<'db>> {
        self.extra.as_ref().and_then(|extra| extra.cycle_recovery)
    }
}

bitflags::bitflags! {
    #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
    pub(crate) struct InferenceFlags: u16 {
        /// Whether to allow `ParamSpec` in type expressions.
        ///
        /// In most contexts inside type expressions, bare `ParamSpec`s are not allowed.
        /// They are specifically allowed as the first argument to `Callable`,
        /// the second argument to `Concatenate`, and certain other special cases.
        const ALLOW_PARAMSPEC_TYPE_EXPR = 1 << 0;

        /// Whether to check for unbound type variables in type expressions.
        /// This flag is set when processing annotation expressions, where unbound type variables
        /// are an error. It is unset in other contexts (e.g., `TypeVar` defaults, explicit class
        /// specialization) where unbound type variables are expected.
        const CHECK_UNBOUND_TYPEVARS = 1 << 1;

        /// Whether the visitor is currently visiting a vararg annotation
        /// (e.g., `*args: int` or `**kwargs: int` in a function definition).
        const IN_VARARG_ANNOTATION = 1 << 2;

        /// Whether the visitor is currently visiting a return-type annotation
        const IN_RETURN_TYPE = 1 << 3;

        /// Whether the visitor is currently visiting a type alias value expression
        const IN_TYPE_ALIAS = 1 << 4;

        /// Whether the visitor is currently visiting a parameter annotation
        const IN_PARAMETER_ANNOTATION = 1 << 5;

        /// Whether we are currently in a context where `Concatenate` can be legal
        const IN_VALID_CONCATENATE_CONTEXT = 1 << 6;

        /// Whether we're in the first pass of inferring a PEP-613 type alias.
        ///
        /// During this pass, `invalid-type-form` diagnostics are suppressed;
        /// these are emitted during the second, post-inference, pass.
        const IN_PEP_613_ALIAS_FIRST_PASS = 1 << 7;

        const IN_NO_TYPE_CHECK = 1 << 8;

        /// Whether the visitor is currently visiting a `**kwargs` annotation.
        const IN_KWARG_ANNOTATION = 1 << 9;

        /// Whether we're in a context where `Unpack` can be legal.
        const IN_VALID_UNPACK_CONTEXT = 1 << 10;

        /// Whether the visitor is currently visiting a type expression.
        const IN_TYPE_EXPRESSION = 1 << 12;

        /// Whether the visitor is currently visiting a nested position in a type expression.
        const IN_NESTED_TYPE_EXPRESSION = 1 << 13;

        /// Whether the visitor is currently visiting the argument to `Unpack[...]`.
        const IN_UNPACK_TYPE_ARGUMENT = 1 << 14;
    }
}

impl get_size2::GetSize for InferenceFlags {}

impl InferenceFlags {
    #[must_use = "Inference flags should always be restored to the original value after being temporarily modified"]
    fn replace(&mut self, other: Self, set_to: bool) -> bool {
        let previously_contained_flag = self.contains(other);
        self.set(other, set_to);
        previously_contained_flag
    }

    pub(super) const fn type_expression_context(self) -> &'static str {
        if self.contains(InferenceFlags::IN_RETURN_TYPE) {
            "return type annotation"
        } else if self.contains(InferenceFlags::IN_PARAMETER_ANNOTATION) {
            "parameter annotation"
        } else if self.contains(InferenceFlags::IN_TYPE_ALIAS) {
            "type alias value"
        } else {
            "type expression"
        }
    }
}
