//! Various utilities to load MIR.
//! Allow to easily load the MIR code generated by a specific pass.
use std::panic;

use crate::hax;
use crate::hax::UnderOwnerState;
use rustc_middle::mir;
use rustc_middle::ty;

use charon_lib::ast::*;
use charon_lib::options::MirLevel;

use super::translate_ctx::ItemTransCtx;

impl<'tcx> ItemTransCtx<'tcx, '_> {
    pub fn get_mir(
        &mut self,
        item_ref: &hax::ItemRef,
        span: Span,
    ) -> Result<Option<mir::Body<'tcx>>, Error> {
        // Stopgap measure because there are still many panics in charon and hax.
        let mut this = panic::AssertUnwindSafe(&mut *self);
        let res = panic::catch_unwind(move || this.get_mir_inner(item_ref));
        match res {
            Ok(Ok(body)) => Ok(body),
            // Translation error
            Ok(Err(e)) => Err(e),
            Err(_) => {
                raise_error!(self, span, "Thread panicked when extracting body.");
            }
        }
    }

    fn get_mir_inner(&mut self, item_ref: &hax::ItemRef) -> Result<Option<mir::Body<'tcx>>, Error> {
        let tcx = self.t_ctx.tcx;
        let mir_level = self.t_ctx.options.mir_level;
        let def_id = &item_ref.def_id;
        Ok(match get_mir_for_def_id_and_level(tcx, def_id, mir_level) {
            Some(body) => {
                let s = self.hax_state_with_id();
                Some(if self.monomorphize() {
                    let typing_env = s.typing_env();
                    let args = item_ref.rustc_args(s);
                    hax::substitute(tcx, typing_env, Some(args), body)
                } else {
                    body
                })
            }
            None => None,
        })
    }
}

/// Query the MIR for a function at a specific level. Return `None` in the case of a foreign body
/// with no MIR available.
#[tracing::instrument(skip(tcx))]
fn get_mir_for_def_id_and_level<'tcx>(
    tcx: ty::TyCtxt<'tcx>,
    def_id: &hax::DefId,
    level: MirLevel,
) -> Option<mir::Body<'tcx>> {
    match def_id.base {
        hax::DefIdBase::Real(rust_def_id) => {
            if let Some(local_def_id) = rust_def_id.as_local() {
                match level {
                    MirLevel::Built => {
                        let body = tcx.mir_built(local_def_id);
                        if !body.is_stolen() {
                            return Some(body.borrow().clone());
                        }
                    }
                    MirLevel::Promoted => {
                        let (body, _) = tcx.mir_promoted(local_def_id);
                        if !body.is_stolen() {
                            return Some(body.borrow().clone());
                        }
                    }
                    MirLevel::Elaborated => {
                        let body = tcx.mir_drops_elaborated_and_const_checked(local_def_id);
                        if !body.is_stolen() {
                            return Some(body.borrow().clone());
                        }
                    }
                    MirLevel::Optimized => {}
                }
                // We fall back to optimized MIR if the requested body was stolen.
            }

            // There are only two MIRs we can fetch for non-local bodies: CTFE mir for globals and const
            // fns, and optimized MIR for functions.
            //
            // We pass `-Zalways-encode-mir` so that we get MIR for all the dependencies we compiled
            // ourselves. This doesn't apply to the stdlib; there we only get MIR for const items and
            // generic or inlineable functions.
            let is_global = matches!(
                def_id.kind,
                hax::DefKind::Const { .. }
                    | hax::DefKind::AnonConst
                    | hax::DefKind::AssocConst { .. }
                    | hax::DefKind::InlineConst
            );
            let is_static = matches!(def_id.kind, hax::DefKind::Static { .. });
            let mir_available = tcx.is_mir_available(rust_def_id);

            if mir_available && !is_global && !is_static {
                Some(tcx.optimized_mir(rust_def_id).clone())
            } else if (is_global && !tcx.is_trivial_const(rust_def_id))
                || (is_static && rust_def_id.is_local())
                || tcx.is_const_fn(rust_def_id)
            {
                Some(tcx.mir_for_ctfe(rust_def_id).clone())
            } else {
                trace!("mir not available for {:?}", rust_def_id);
                None
            }
        }
        hax::DefIdBase::Promoted(rust_def_id, promoted_id) => {
            Some(hax::get_promoted_mir(tcx, rust_def_id, promoted_id))
        }
        _ => unreachable!(),
    }
}
