//! Monomorphization of mir, which is used in mir interpreter and const eval.
//!
//! The job of monomorphization is:
//! * Monomorphization. That is, replacing `Option<T>` with `Option<i32>` where `T:=i32` substitution
//!   is provided
//! * Normalizing types, for example replacing RPIT of other functions called in this body.
//!
//! So the monomorphization should be called even if the substitution is empty.

use hir_def::DefWithBodyId;
use rustc_type_ir::inherent::IntoKind;
use rustc_type_ir::{
    FallibleTypeFolder, TypeFlags, TypeFoldable, TypeSuperFoldable, TypeVisitableExt,
};
use triomphe::Arc;

use crate::{
    ParamEnvAndCrate,
    next_solver::{
        Allocation, AllocationData, Const, ConstKind, Region, RegionKind, StoredConst,
        StoredGenericArgs, StoredTy,
    },
    traits::StoredParamEnvAndCrate,
};
use crate::{
    db::{HirDatabase, InternedClosureId},
    next_solver::{
        DbInterner, GenericArgs, Ty, TyKind, TypingMode,
        infer::{DbInternerInferExt, InferCtxt, traits::ObligationCause},
        obligation_ctxt::ObligationCtxt,
        references_non_lt_error,
    },
};

use super::{MirBody, MirLowerError, Operand, OperandKind, Rvalue, StatementKind, TerminatorKind};

struct Filler<'db> {
    infcx: InferCtxt<'db>,
    trait_env: ParamEnvAndCrate<'db>,
    subst: GenericArgs<'db>,
}

impl<'db> FallibleTypeFolder<DbInterner<'db>> for Filler<'db> {
    type Error = MirLowerError;

    fn cx(&self) -> DbInterner<'db> {
        self.infcx.interner
    }

    fn try_fold_ty(&mut self, ty: Ty<'db>) -> Result<Ty<'db>, Self::Error> {
        if !ty.has_type_flags(TypeFlags::HAS_ALIAS | TypeFlags::HAS_PARAM) {
            return Ok(ty);
        }

        match ty.kind() {
            TyKind::Alias(..) => {
                // First instantiate params.
                let ty = ty.try_super_fold_with(self)?;

                let mut ocx = ObligationCtxt::new(&self.infcx);
                let ty = ocx
                    .structurally_normalize_ty(
                        &ObligationCause::dummy(),
                        self.trait_env.param_env,
                        ty,
                    )
                    .map_err(|_| MirLowerError::NotSupported("can't normalize alias".to_owned()))?;
                ty.try_super_fold_with(self)
            }
            TyKind::Param(param) => Ok(self
                .subst
                .as_slice()
                .get(param.index as usize)
                .and_then(|arg| arg.ty())
                .ok_or_else(|| {
                    MirLowerError::GenericArgNotProvided(param.id.into(), self.subst.store())
                })?),
            _ => ty.try_super_fold_with(self),
        }
    }

    fn try_fold_const(&mut self, ct: Const<'db>) -> Result<Const<'db>, Self::Error> {
        let ConstKind::Param(param) = ct.kind() else {
            return ct.try_super_fold_with(self);
        };
        self.subst.as_slice().get(param.index as usize).and_then(|arg| arg.konst()).ok_or_else(
            || MirLowerError::GenericArgNotProvided(param.id.into(), self.subst.store()),
        )
    }

    fn try_fold_region(&mut self, region: Region<'db>) -> Result<Region<'db>, Self::Error> {
        let RegionKind::ReEarlyParam(param) = region.kind() else {
            return Ok(region);
        };
        self.subst.as_slice().get(param.index as usize).and_then(|arg| arg.region()).ok_or_else(
            || MirLowerError::GenericArgNotProvided(param.id.into(), self.subst.store()),
        )
    }
}

impl<'db> Filler<'db> {
    fn new(db: &'db dyn HirDatabase, env: ParamEnvAndCrate<'db>, subst: GenericArgs<'db>) -> Self {
        let interner = DbInterner::new_with(db, env.krate);
        let infcx = interner.infer_ctxt().build(TypingMode::PostAnalysis);
        Self { infcx, trait_env: env, subst }
    }

    fn fill_ty(&mut self, t: &mut StoredTy) -> Result<(), MirLowerError> {
        // Can't deep normalized as that'll try to normalize consts and fail.
        *t = t.as_ref().try_fold_with(self)?.store();
        if references_non_lt_error(&t.as_ref()) {
            Err(MirLowerError::NotSupported("monomorphization resulted in errors".to_owned()))
        } else {
            Ok(())
        }
    }

    fn fill_const(&mut self, t: &mut StoredConst) -> Result<(), MirLowerError> {
        // Can't deep normalized as that'll try to normalize consts and fail.
        *t = t.as_ref().try_fold_with(self)?.store();
        if references_non_lt_error(&t.as_ref()) {
            Err(MirLowerError::NotSupported("monomorphization resulted in errors".to_owned()))
        } else {
            Ok(())
        }
    }

    fn fill_args(&mut self, t: &mut StoredGenericArgs) -> Result<(), MirLowerError> {
        // Can't deep normalized as that'll try to normalize consts and fail.
        *t = t.as_ref().try_fold_with(self)?.store();
        if references_non_lt_error(&t.as_ref()) {
            Err(MirLowerError::NotSupported("monomorphization resulted in errors".to_owned()))
        } else {
            Ok(())
        }
    }

    fn fill_operand(&mut self, op: &mut Operand) -> Result<(), MirLowerError> {
        match &mut op.kind {
            OperandKind::Constant { konst, ty } => {
                self.fill_const(konst)?;
                self.fill_ty(ty)?;
            }
            OperandKind::Allocation { allocation } => {
                let alloc = allocation.as_ref();
                let mut ty = alloc.ty.store();
                self.fill_ty(&mut ty)?;
                *allocation = Allocation::new(AllocationData {
                    ty: ty.as_ref(),
                    memory: alloc.memory.clone(),
                    // FIXME: Do we need to fill the memory map too?
                    memory_map: alloc.memory_map.clone(),
                })
                .store();
            }
            OperandKind::Copy(_) | OperandKind::Move(_) | OperandKind::Static(_) => (),
        }
        Ok(())
    }

    fn fill_body(&mut self, body: &mut MirBody) -> Result<(), MirLowerError> {
        for (_, l) in body.locals.iter_mut() {
            self.fill_ty(&mut l.ty)?;
        }
        for (_, bb) in body.basic_blocks.iter_mut() {
            for statement in &mut bb.statements {
                match &mut statement.kind {
                    StatementKind::Assign(_, r) => match r {
                        Rvalue::Aggregate(ak, ops) => {
                            for op in &mut **ops {
                                self.fill_operand(op)?;
                            }
                            match ak {
                                super::AggregateKind::Array(ty)
                                | super::AggregateKind::Tuple(ty)
                                | super::AggregateKind::Closure(ty) => self.fill_ty(ty)?,
                                super::AggregateKind::Adt(_, subst) => self.fill_args(subst)?,
                                super::AggregateKind::Union(_, _) => (),
                            }
                        }
                        Rvalue::ShallowInitBox(_, ty) | Rvalue::ShallowInitBoxWithAlloc(ty) => {
                            self.fill_ty(ty)?;
                        }
                        Rvalue::Use(op) => {
                            self.fill_operand(op)?;
                        }
                        Rvalue::Repeat(op, len) => {
                            self.fill_operand(op)?;
                            self.fill_const(len)?;
                        }
                        Rvalue::Ref(_, _)
                        | Rvalue::Len(_)
                        | Rvalue::Cast(_, _, _)
                        | Rvalue::CheckedBinaryOp(_, _, _)
                        | Rvalue::UnaryOp(_, _)
                        | Rvalue::Discriminant(_)
                        | Rvalue::CopyForDeref(_) => (),
                        Rvalue::ThreadLocalRef(n)
                        | Rvalue::AddressOf(n)
                        | Rvalue::BinaryOp(n)
                        | Rvalue::NullaryOp(n) => match *n {},
                    },
                    StatementKind::Deinit(_)
                    | StatementKind::FakeRead(_)
                    | StatementKind::StorageLive(_)
                    | StatementKind::StorageDead(_)
                    | StatementKind::Nop => (),
                }
            }
            if let Some(terminator) = &mut bb.terminator {
                match &mut terminator.kind {
                    TerminatorKind::Call { func, args, .. } => {
                        self.fill_operand(func)?;
                        for op in &mut **args {
                            self.fill_operand(op)?;
                        }
                    }
                    TerminatorKind::SwitchInt { discr, .. } => {
                        self.fill_operand(discr)?;
                    }
                    TerminatorKind::Goto { .. }
                    | TerminatorKind::UnwindResume
                    | TerminatorKind::Abort
                    | TerminatorKind::Return
                    | TerminatorKind::Unreachable
                    | TerminatorKind::Drop { .. }
                    | TerminatorKind::DropAndReplace { .. }
                    | TerminatorKind::Assert { .. }
                    | TerminatorKind::Yield { .. }
                    | TerminatorKind::CoroutineDrop
                    | TerminatorKind::FalseEdge { .. }
                    | TerminatorKind::FalseUnwind { .. } => (),
                }
            }
        }
        Ok(())
    }
}

pub fn monomorphized_mir_body_query(
    db: &dyn HirDatabase,
    owner: DefWithBodyId,
    subst: StoredGenericArgs,
    trait_env: StoredParamEnvAndCrate,
) -> Result<Arc<MirBody>, MirLowerError> {
    let mut filler = Filler::new(db, trait_env.as_ref(), subst.as_ref());
    let body = db.mir_body(owner)?;
    let mut body = (*body).clone();
    filler.fill_body(&mut body)?;
    Ok(Arc::new(body))
}

pub(crate) fn monomorphized_mir_body_cycle_result(
    _db: &dyn HirDatabase,
    _: salsa::Id,
    _: DefWithBodyId,
    _: StoredGenericArgs,
    _: StoredParamEnvAndCrate,
) -> Result<Arc<MirBody>, MirLowerError> {
    Err(MirLowerError::Loop)
}

pub fn monomorphized_mir_body_for_closure_query(
    db: &dyn HirDatabase,
    closure: InternedClosureId,
    subst: StoredGenericArgs,
    trait_env: StoredParamEnvAndCrate,
) -> Result<Arc<MirBody>, MirLowerError> {
    let mut filler = Filler::new(db, trait_env.as_ref(), subst.as_ref());
    let body = db.mir_body_for_closure(closure)?;
    let mut body = (*body).clone();
    filler.fill_body(&mut body)?;
    Ok(Arc::new(body))
}