//! Deeply normalize types using the old trait solver.

use rustc_data_structures::stack::ensure_sufficient_stack;
use rustc_errors::msg;
use rustc_infer::infer::at::At;
use rustc_infer::infer::{InferCtxt, InferOk};
use rustc_infer::traits::{
    FromSolverError, Normalized, Obligation, PredicateObligations, TraitEngine,
};
use rustc_macros::extension;
use rustc_middle::span_bug;
use rustc_middle::traits::{ObligationCause, ObligationCauseCode};
use rustc_middle::ty::{
    self, AliasTerm, Term, Ty, TyCtxt, TypeFoldable, TypeFolder, TypeSuperFoldable, TypeVisitable,
    TypeVisitableExt, TypingMode, Unnormalized,
};
use tracing::{debug, instrument};

use super::{BoundVarReplacer, PlaceholderReplacer, SelectionContext, project};
use crate::error_reporting::InferCtxtErrorExt;
use crate::error_reporting::traits::OverflowCause;
use crate::solve::NextSolverError;

#[extension(pub trait NormalizeExt<'tcx>)]
impl<'tcx> At<'_, 'tcx> {
    /// Normalize a value using the `AssocTypeNormalizer`.
    ///
    /// This normalization should be used when the type contains inference variables or the
    /// projection may be fallible.
    fn normalize<T: TypeFoldable<TyCtxt<'tcx>>>(
        &self,
        value: Unnormalized<'tcx, T>,
    ) -> InferOk<'tcx, T> {
        if self.infcx.next_trait_solver() {
            let Normalized { value, obligations } = crate::solve::normalize(*self, value);
            InferOk { value, obligations }
        } else {
            let value = value.skip_normalization();
            let mut selcx = SelectionContext::new(self.infcx);
            let Normalized { value, obligations } =
                normalize_with_depth(&mut selcx, self.param_env, self.cause.clone(), 0, value);
            InferOk { value, obligations }
        }
    }

    /// Deeply normalizes `value`, replacing all aliases which can by normalized in
    /// the current environment. In the new solver this errors in case normalization
    /// fails or is ambiguous.
    ///
    /// In the old solver this simply uses `normalizes` and adds the nested obligations
    /// to the `fulfill_cx`. This is necessary as we otherwise end up recomputing the
    /// same goals in both a temporary and the shared context which negatively impacts
    /// performance as these don't share caching.
    ///
    /// FIXME(-Znext-solver=no): For performance reasons, we currently reuse an existing
    /// fulfillment context in the old solver. Once we have removed the old solver, we
    /// can remove the `fulfill_cx` parameter on this function.
    fn deeply_normalize<T, E>(
        self,
        value: Unnormalized<'tcx, T>,
        fulfill_cx: &mut dyn TraitEngine<'tcx, E>,
    ) -> Result<T, Vec<E>>
    where
        T: TypeFoldable<TyCtxt<'tcx>>,
        E: FromSolverError<'tcx, NextSolverError<'tcx>>,
    {
        if self.infcx.next_trait_solver() {
            crate::solve::deeply_normalize(self, value)
        } else {
            if fulfill_cx.has_pending_obligations() {
                let pending_obligations = fulfill_cx.pending_obligations();
                span_bug!(
                    pending_obligations[0].cause.span,
                    "deeply_normalize should not be called with pending obligations: \
                    {pending_obligations:#?}"
                );
            }
            let value = self
                .normalize(value)
                .into_value_registering_obligations(self.infcx, &mut *fulfill_cx);
            let errors = fulfill_cx.evaluate_obligations_error_on_ambiguity(self.infcx);
            let value = self.infcx.resolve_vars_if_possible(value);
            if errors.is_empty() {
                Ok(value)
            } else {
                // Drop pending obligations, since deep normalization may happen
                // in a loop and we don't want to trigger the assertion on the next
                // iteration due to pending ambiguous obligations we've left over.
                let _ = fulfill_cx.collect_remaining_errors(self.infcx);
                Err(errors)
            }
        }
    }
}

/// As `normalize`, but with a custom depth.
pub(crate) fn normalize_with_depth<'a, 'b, 'tcx, T>(
    selcx: &'a mut SelectionContext<'b, 'tcx>,
    param_env: ty::ParamEnv<'tcx>,
    cause: ObligationCause<'tcx>,
    depth: usize,
    value: T,
) -> Normalized<'tcx, T>
where
    T: TypeFoldable<TyCtxt<'tcx>>,
{
    let mut obligations = PredicateObligations::new();
    let value = normalize_with_depth_to(selcx, param_env, cause, depth, value, &mut obligations);
    Normalized { value, obligations }
}

#[instrument(level = "info", skip(selcx, param_env, cause, obligations))]
pub(crate) fn normalize_with_depth_to<'a, 'b, 'tcx, T>(
    selcx: &'a mut SelectionContext<'b, 'tcx>,
    param_env: ty::ParamEnv<'tcx>,
    cause: ObligationCause<'tcx>,
    depth: usize,
    value: T,
    obligations: &mut PredicateObligations<'tcx>,
) -> T
where
    T: TypeFoldable<TyCtxt<'tcx>>,
{
    debug!(obligations.len = obligations.len());
    let mut normalizer = AssocTypeNormalizer::new(selcx, param_env, cause, depth, obligations);
    let result = ensure_sufficient_stack(|| AssocTypeNormalizer::fold(&mut normalizer, value));
    debug!(?result, obligations.len = normalizer.obligations.len());
    debug!(?normalizer.obligations,);
    result
}

pub(super) fn needs_normalization<'tcx, T: TypeVisitable<TyCtxt<'tcx>>>(
    infcx: &InferCtxt<'tcx>,
    value: &T,
) -> bool {
    let mut flags = ty::TypeFlags::HAS_ALIAS;

    // Opaques are treated as rigid outside of `TypingMode::PostAnalysis`,
    // so we can ignore those.
    match infcx.typing_mode_raw().assert_not_erased() {
        // FIXME(#132279): We likely want to reveal opaques during post borrowck analysis
        TypingMode::Coherence
        | TypingMode::Typeck { .. }
        | TypingMode::PostTypeckUntilBorrowck { .. }
        | TypingMode::PostBorrowck { .. } => flags.remove(ty::TypeFlags::HAS_TY_OPAQUE),
        TypingMode::PostAnalysis | TypingMode::Codegen => {}
    }

    value.has_type_flags(flags)
}

struct AssocTypeNormalizer<'a, 'b, 'tcx> {
    selcx: &'a mut SelectionContext<'b, 'tcx>,
    param_env: ty::ParamEnv<'tcx>,
    cause: ObligationCause<'tcx>,
    obligations: &'a mut PredicateObligations<'tcx>,
    depth: usize,
    universes: Vec<Option<ty::UniverseIndex>>,
}

impl<'a, 'b, 'tcx> AssocTypeNormalizer<'a, 'b, 'tcx> {
    fn new(
        selcx: &'a mut SelectionContext<'b, 'tcx>,
        param_env: ty::ParamEnv<'tcx>,
        cause: ObligationCause<'tcx>,
        depth: usize,
        obligations: &'a mut PredicateObligations<'tcx>,
    ) -> AssocTypeNormalizer<'a, 'b, 'tcx> {
        debug_assert!(!selcx.infcx.next_trait_solver());
        AssocTypeNormalizer { selcx, param_env, cause, obligations, depth, universes: vec![] }
    }

    fn fold<T: TypeFoldable<TyCtxt<'tcx>>>(&mut self, value: T) -> T {
        let value = self.selcx.infcx.resolve_vars_if_possible(value);
        debug!(?value);

        assert!(
            !value.has_escaping_bound_vars(),
            "Normalizing {value:?} without wrapping in a `Binder`"
        );

        if !needs_normalization(self.selcx.infcx, &value) { value } else { value.fold_with(self) }
    }

    // FIXME(mgca): While this supports constants, it is only used for types by default right now
    #[instrument(level = "debug", skip(self), ret)]
    fn normalize_trait_projection(&mut self, proj: AliasTerm<'tcx>) -> Term<'tcx> {
        if !proj.has_escaping_bound_vars() {
            // When we don't have escaping bound vars we can normalize ambig aliases
            // to inference variables (done in `normalize_projection_ty`). This would
            // be wrong if there were escaping bound vars as even if we instantiated
            // the bound vars with placeholders, we wouldn't be able to map them back
            // after normalization succeeded.
            //
            // Also, as an optimization: when we don't have escaping bound vars, we don't
            // need to replace them with placeholders (see branch below).
            let proj = proj.fold_with(self);
            project::normalize_projection_term(
                self.selcx,
                self.param_env,
                proj,
                self.cause.clone(),
                self.depth,
                self.obligations,
            )
        } else {
            // If there are escaping bound vars, we temporarily replace the
            // bound vars with placeholders. Note though, that in the case
            // that we still can't project for whatever reason (e.g. self
            // type isn't known enough), we *can't* register an obligation
            // and return an inference variable (since then that obligation
            // would have bound vars and that's a can of worms). Instead,
            // we just give up and fall back to pretending like we never tried!
            //
            // Note: this isn't necessarily the final approach here; we may
            // want to figure out how to register obligations with escaping vars
            // or handle this some other way.
            let infcx = self.selcx.infcx;
            let (proj, mapped_regions, mapped_types, mapped_consts) =
                BoundVarReplacer::replace_bound_vars(infcx, &mut self.universes, proj);
            let proj = proj.fold_with(self);
            let normalized_term = project::opt_normalize_projection_term(
                self.selcx,
                self.param_env,
                proj,
                self.cause.clone(),
                self.depth,
                self.obligations,
            )
            .ok()
            .flatten()
            .unwrap_or_else(|| proj.to_term(infcx.tcx, ty::IsRigid::No));

            PlaceholderReplacer::replace_placeholders(
                infcx,
                mapped_regions,
                mapped_types,
                mapped_consts,
                &self.universes,
                normalized_term,
            )
        }
    }

    // FIXME(mgca): While this supports constants, it is only used for types by default right now
    #[instrument(level = "debug", skip(self), ret)]
    fn normalize_inherent_projection(&mut self, inherent: AliasTerm<'tcx>) -> Term<'tcx> {
        if !inherent.has_escaping_bound_vars() {
            // When we don't have escaping bound vars we can normalize ambig aliases
            // to inference variables (done in `normalize_projection_ty`). This would
            // be wrong if there were escaping bound vars as even if we instantiated
            // the bound vars with placeholders, we wouldn't be able to map them back
            // after normalization succeeded.
            //
            // Also, as an optimization: when we don't have escaping bound vars, we don't
            // need to replace them with placeholders (see branch below).

            let inherent = inherent.fold_with(self);
            project::normalize_inherent_projection(
                self.selcx,
                self.param_env,
                inherent,
                self.cause.clone(),
                self.depth,
                self.obligations,
            )
        } else {
            let infcx = self.selcx.infcx;
            let (inherent, mapped_regions, mapped_types, mapped_consts) =
                BoundVarReplacer::replace_bound_vars(infcx, &mut self.universes, inherent);
            let inherent = inherent.fold_with(self);
            let inherent = project::normalize_inherent_projection(
                self.selcx,
                self.param_env,
                inherent,
                self.cause.clone(),
                self.depth,
                self.obligations,
            );

            PlaceholderReplacer::replace_placeholders(
                infcx,
                mapped_regions,
                mapped_types,
                mapped_consts,
                &self.universes,
                inherent,
            )
        }
    }

    // FIXME(mgca): While this supports constants, it is only used for types by default right now
    #[instrument(level = "debug", skip(self), ret)]
    fn normalize_free_alias(&mut self, free: AliasTerm<'tcx>) -> Term<'tcx> {
        let recursion_limit = self.cx().recursion_limit();
        if !recursion_limit.value_within_limit(self.depth) {
            self.selcx.infcx.err_ctxt().report_overflow_error(
                OverflowCause::DeeplyNormalize(free),
                self.cause.span,
                false,
                |diag| {
                    diag.note(msg!("in case this is a recursive type alias, consider using a struct, enum, or union instead"));
                },
            );
        }

        let def_id = free.expect_free_def_id();

        // We don't replace bound vars in the generic arguments of the free alias with
        // placeholders. This doesn't cause any issues as instantiating parameters with
        // bound variables is special-cased to rewrite the debruijn index to be higher
        // whenever we fold through a binder.
        //
        // However, we do replace any escaping bound vars in the resulting goals with
        // placeholders as the trait solver does not expect to encounter escaping bound
        // vars in obligations.
        //
        // FIXME(lazy_type_alias): Check how much this actually matters for perf before
        // stabilization. This is a bit weird and generally not how we handle binders in
        // the compiler so ideally we'd do the same boundvar->placeholder->boundvar dance
        // that other kinds of normalization do.
        let infcx = self.selcx.infcx;
        self.obligations.extend(
            infcx
                .tcx
                .predicates_of(def_id)
                .instantiate_own(infcx.tcx, free.args)
                .map(|(pred, span)| (pred.skip_norm_wip(), span))
                .map(|(mut predicate, span)| {
                    if free.has_escaping_bound_vars() {
                        (predicate, ..) = BoundVarReplacer::replace_bound_vars(
                            infcx,
                            &mut self.universes,
                            predicate,
                        );
                    }
                    let mut cause = self.cause.clone();
                    cause.map_code(|code| ObligationCauseCode::TypeAlias(code, span, def_id));
                    Obligation::new(infcx.tcx, cause, self.param_env, predicate)
                }),
        );
        self.depth += 1;
        let res: ty::Term<'tcx> = if free.kind.is_type() {
            infcx
                .tcx
                .type_of(def_id)
                .instantiate(infcx.tcx, free.args)
                .skip_norm_wip()
                .fold_with(self)
                .into()
        } else {
            infcx
                .tcx
                .const_of_item(def_id)
                .instantiate(infcx.tcx, free.args)
                .skip_norm_wip()
                .fold_with(self)
                .into()
        };
        // When normalizing a free const alias, register a `ConstArgHasType`
        // obligation to ensure the const value's type matches the declared type.
        if let Some(ct) = res.as_const() {
            let expected_ty =
                infcx.tcx.type_of(def_id).instantiate(infcx.tcx, free.args).skip_norm_wip();
            self.obligations.push(Obligation::with_depth(
                infcx.tcx,
                self.cause.clone(),
                self.depth,
                self.param_env,
                ty::ClauseKind::ConstArgHasType(ct, expected_ty),
            ));
        }
        self.depth -= 1;
        res
    }
}

impl<'a, 'b, 'tcx> TypeFolder<TyCtxt<'tcx>> for AssocTypeNormalizer<'a, 'b, 'tcx> {
    fn cx(&self) -> TyCtxt<'tcx> {
        self.selcx.tcx()
    }

    fn fold_binder<T: TypeFoldable<TyCtxt<'tcx>>>(
        &mut self,
        t: ty::Binder<'tcx, T>,
    ) -> ty::Binder<'tcx, T> {
        self.universes.push(None);
        let t = t.super_fold_with(self);
        self.universes.pop();
        t
    }

    fn fold_ty(&mut self, ty: Ty<'tcx>) -> Ty<'tcx> {
        // We don't use the rigid marker in old solver.
        debug_assert!(!ty.has_rigid_aliases());

        if !needs_normalization(self.selcx.infcx, &ty) {
            return ty;
        }

        let ty::Alias(_, data) = *ty.kind() else { return ty.super_fold_with(self) };

        // We try to be a little clever here as a performance optimization in
        // cases where there are nested projections under binders.
        // For example:
        // ```
        // for<'a> fn(<T as Foo>::One<'a, Box<dyn Bar<'a, Item=<T as Foo>::Two<'a>>>>)
        // ```
        // We normalize the args on the projection before the projecting, but
        // if we're naive, we'll
        //   replace bound vars on inner, project inner, replace placeholders on inner,
        //   replace bound vars on outer, project outer, replace placeholders on outer
        //
        // However, if we're a bit more clever, we can replace the bound vars
        // on the entire type before normalizing nested projections, meaning we
        //   replace bound vars on outer, project inner,
        //   project outer, replace placeholders on outer
        //
        // This is possible because the inner `'a` will already be a placeholder
        // when we need to normalize the inner projection
        //
        // On the other hand, this does add a bit of complexity, since we only
        // replace bound vars if the current type is a `Projection` and we need
        // to make sure we don't forget to fold the args regardless.

        match data.kind {
            ty::Opaque { def_id } => {
                // Only normalize `impl Trait` outside of type inference, usually in codegen.
                match self.selcx.typing_mode() {
                    // FIXME(#132279): We likely want to reveal opaques during post borrowck analysis
                    TypingMode::Coherence
                    | TypingMode::Typeck { .. }
                    | TypingMode::PostTypeckUntilBorrowck { .. }
                    | TypingMode::PostBorrowck { .. } => ty.super_fold_with(self),
                    TypingMode::PostAnalysis | TypingMode::Codegen => {
                        let recursion_limit = self.cx().recursion_limit();
                        if !recursion_limit.value_within_limit(self.depth) {
                            self.selcx.infcx.err_ctxt().report_overflow_error(
                                OverflowCause::DeeplyNormalize(data.into()),
                                self.cause.span,
                                true,
                                |_| {},
                            );
                        }

                        let args = data.args.fold_with(self);
                        let generic_ty = self.cx().type_of(def_id);
                        let concrete_ty = generic_ty.instantiate(self.cx(), args).skip_norm_wip();
                        self.depth += 1;
                        let folded_ty = self.fold_ty(concrete_ty);
                        self.depth -= 1;
                        folded_ty
                    }
                }
            }

            ty::Projection { .. } => self.normalize_trait_projection(data.into()).expect_type(),
            ty::Inherent { .. } => self.normalize_inherent_projection(data.into()).expect_type(),
            ty::Free { .. } => self.normalize_free_alias(data.into()).expect_type(),
        }
    }

    #[instrument(skip(self), level = "debug")]
    fn fold_const(&mut self, ct: ty::Const<'tcx>) -> ty::Const<'tcx> {
        // We don't use the rigid marker in old solver.
        debug_assert!(!ct.has_rigid_aliases());

        let tcx = self.selcx.tcx();

        if tcx.features().generic_const_exprs()
            // Normalize type_const items even with feature `generic_const_exprs`.
            && !matches!(ct.kind(), ty::ConstKind::Alias(_, alias_const) if alias_const.kind.is_type_const(tcx))
            || !needs_normalization(self.selcx.infcx, &ct)
        {
            return ct;
        }

        let alias_const = match ct.kind() {
            ty::ConstKind::Alias(_, alias_const) => alias_const,
            _ => return ct.super_fold_with(self),
        };

        // Note that the Projection/Inherent/Free cases are unreachable on stable,
        // unless a `min_generic_const_args` feature gate error has already
        // been emitted earlier in compilation.
        //
        // That's because we can only end up with an Alias ty::Const for a const item
        // if it was marked with `type const`. Using this attribute without the mgca
        // feature gate causes a parse error.
        let ct = match alias_const.kind {
            ty::AliasConstKind::Projection { .. } => {
                self.normalize_trait_projection(alias_const.into()).expect_const()
            }
            ty::AliasConstKind::Inherent { .. } => {
                self.normalize_inherent_projection(alias_const.into()).expect_const()
            }
            ty::AliasConstKind::Free { .. } => {
                self.normalize_free_alias(alias_const.into()).expect_const()
            }
            ty::AliasConstKind::Anon { .. } => {
                let ct = ct.super_fold_with(self);
                super::with_replaced_escaping_bound_vars(
                    self.selcx.infcx,
                    &mut self.universes,
                    ct,
                    |ct| super::evaluate_const(self.selcx.infcx, ct, self.param_env),
                )
            }
        };

        // We re-fold the normalized const as the `ty` field on `ConstKind::Value` may be
        // unnormalized after const evaluation returns.
        ct.super_fold_with(self)
    }

    #[inline]
    fn fold_predicate(&mut self, p: ty::Predicate<'tcx>) -> ty::Predicate<'tcx> {
        if p.allow_normalization() && needs_normalization(self.selcx.infcx, &p) {
            p.super_fold_with(self)
        } else {
            p
        }
    }
}