//! The general point of the optimizations provided here is to simplify something like:
//!
//! ```rust
//! match x {
//!     Ok(x) => Ok(x),
//!     Err(x) => Err(x)
//! }
//! ```
//!
//! into just `x`.

use crate::transform::{simplify, MirPass, MirSource};
use itertools::Itertools as _;
use rustc_index::{bit_set::BitSet, vec::IndexVec};
use rustc_middle::mir::visit::{NonUseContext, PlaceContext, Visitor};
use rustc_middle::mir::*;
use rustc_middle::ty::{List, Ty, TyCtxt};
use rustc_target::abi::VariantIdx;
use std::iter::{Enumerate, Peekable};
use std::slice::Iter;

/// Simplifies arms of form `Variant(x) => Variant(x)` to just a move.
///
/// This is done by transforming basic blocks where the statements match:
///
/// ```rust
/// _LOCAL_TMP = ((_LOCAL_1 as Variant ).FIELD: TY );
/// _TMP_2 = _LOCAL_TMP;
/// ((_LOCAL_0 as Variant).FIELD: TY) = move _TMP_2;
/// discriminant(_LOCAL_0) = VAR_IDX;
/// ```
///
/// into:
///
/// ```rust
/// _LOCAL_0 = move _LOCAL_1
/// ```
pub struct SimplifyArmIdentity;

#[derive(Debug)]
struct ArmIdentityInfo<'tcx> {
    /// Storage location for the variant's field
    local_temp_0: Local,
    /// Storage location holding the variant being read from
    local_1: Local,
    /// The variant field being read from
    vf_s0: VarField<'tcx>,
    /// Index of the statement which loads the variant being read
    get_variant_field_stmt: usize,

    /// Tracks each assignment to a temporary of the variant's field
    field_tmp_assignments: Vec<(Local, Local)>,

    /// Storage location holding the variant's field that was read from
    local_tmp_s1: Local,
    /// Storage location holding the enum that we are writing to
    local_0: Local,
    /// The variant field being written to
    vf_s1: VarField<'tcx>,

    /// Storage location that the discriminant is being written to
    set_discr_local: Local,
    /// The variant being written
    set_discr_var_idx: VariantIdx,

    /// Index of the statement that should be overwritten as a move
    stmt_to_overwrite: usize,
    /// SourceInfo for the new move
    source_info: SourceInfo,

    /// Indices of matching Storage{Live,Dead} statements encountered.
    /// (StorageLive index,, StorageDead index, Local)
    storage_stmts: Vec<(usize, usize, Local)>,

    /// The statements that should be removed (turned into nops)
    stmts_to_remove: Vec<usize>,

    /// Indices of debug variables that need to be adjusted to point to
    // `{local_0}.{dbg_projection}`.
    dbg_info_to_adjust: Vec<usize>,

    /// The projection used to rewrite debug info.
    dbg_projection: &'tcx List<PlaceElem<'tcx>>,
}

fn get_arm_identity_info<'a, 'tcx>(
    stmts: &'a [Statement<'tcx>],
    locals_count: usize,
    debug_info: &'a [VarDebugInfo<'tcx>],
) -> Option<ArmIdentityInfo<'tcx>> {
    // This can't possibly match unless there are at least 3 statements in the block
    // so fail fast on tiny blocks.
    if stmts.len() < 3 {
        return None;
    }

    let mut tmp_assigns = Vec::new();
    let mut nop_stmts = Vec::new();
    let mut storage_stmts = Vec::new();
    let mut storage_live_stmts = Vec::new();
    let mut storage_dead_stmts = Vec::new();

    type StmtIter<'a, 'tcx> = Peekable<Enumerate<Iter<'a, Statement<'tcx>>>>;

    fn is_storage_stmt<'tcx>(stmt: &Statement<'tcx>) -> bool {
        matches!(stmt.kind, StatementKind::StorageLive(_) | StatementKind::StorageDead(_))
    }

    /// Eats consecutive Statements which match `test`, performing the specified `action` for each.
    /// The iterator `stmt_iter` is not advanced if none were matched.
    fn try_eat<'a, 'tcx>(
        stmt_iter: &mut StmtIter<'a, 'tcx>,
        test: impl Fn(&'a Statement<'tcx>) -> bool,
        mut action: impl FnMut(usize, &'a Statement<'tcx>),
    ) {
        while stmt_iter.peek().map(|(_, stmt)| test(stmt)).unwrap_or(false) {
            let (idx, stmt) = stmt_iter.next().unwrap();

            action(idx, stmt);
        }
    }

    /// Eats consecutive `StorageLive` and `StorageDead` Statements.
    /// The iterator `stmt_iter` is not advanced if none were found.
    fn try_eat_storage_stmts<'a, 'tcx>(
        stmt_iter: &mut StmtIter<'a, 'tcx>,
        storage_live_stmts: &mut Vec<(usize, Local)>,
        storage_dead_stmts: &mut Vec<(usize, Local)>,
    ) {
        try_eat(stmt_iter, is_storage_stmt, |idx, stmt| {
            if let StatementKind::StorageLive(l) = stmt.kind {
                storage_live_stmts.push((idx, l));
            } else if let StatementKind::StorageDead(l) = stmt.kind {
                storage_dead_stmts.push((idx, l));
            }
        })
    }

    fn is_tmp_storage_stmt<'tcx>(stmt: &Statement<'tcx>) -> bool {
        use rustc_middle::mir::StatementKind::Assign;
        if let Assign(box (place, Rvalue::Use(Operand::Copy(p) | Operand::Move(p)))) = &stmt.kind {
            place.as_local().is_some() && p.as_local().is_some()
        } else {
            false
        }
    }

    /// Eats consecutive `Assign` Statements.
    // The iterator `stmt_iter` is not advanced if none were found.
    fn try_eat_assign_tmp_stmts<'a, 'tcx>(
        stmt_iter: &mut StmtIter<'a, 'tcx>,
        tmp_assigns: &mut Vec<(Local, Local)>,
        nop_stmts: &mut Vec<usize>,
    ) {
        try_eat(stmt_iter, is_tmp_storage_stmt, |idx, stmt| {
            use rustc_middle::mir::StatementKind::Assign;
            if let Assign(box (place, Rvalue::Use(Operand::Copy(p) | Operand::Move(p)))) =
                &stmt.kind
            {
                tmp_assigns.push((place.as_local().unwrap(), p.as_local().unwrap()));
                nop_stmts.push(idx);
            }
        })
    }

    fn find_storage_live_dead_stmts_for_local<'tcx>(
        local: Local,
        stmts: &[Statement<'tcx>],
    ) -> Option<(usize, usize)> {
        trace!("looking for {:?}", local);
        let mut storage_live_stmt = None;
        let mut storage_dead_stmt = None;
        for (idx, stmt) in stmts.iter().enumerate() {
            if stmt.kind == StatementKind::StorageLive(local) {
                storage_live_stmt = Some(idx);
            } else if stmt.kind == StatementKind::StorageDead(local) {
                storage_dead_stmt = Some(idx);
            }
        }

        Some((storage_live_stmt?, storage_dead_stmt.unwrap_or(usize::MAX)))
    }

    // Try to match the expected MIR structure with the basic block we're processing.
    // We want to see something that looks like:
    // ```
    // (StorageLive(_) | StorageDead(_));*
    // _LOCAL_INTO = ((_LOCAL_FROM as Variant).FIELD: TY);
    // (StorageLive(_) | StorageDead(_));*
    // (tmp_n+1 = tmp_n);*
    // (StorageLive(_) | StorageDead(_));*
    // (tmp_n+1 = tmp_n);*
    // ((LOCAL_FROM as Variant).FIELD: TY) = move tmp;
    // discriminant(LOCAL_FROM) = VariantIdx;
    // (StorageLive(_) | StorageDead(_));*
    // ```
    let mut stmt_iter = stmts.iter().enumerate().peekable();

    try_eat_storage_stmts(&mut stmt_iter, &mut storage_live_stmts, &mut storage_dead_stmts);

    let (get_variant_field_stmt, stmt) = stmt_iter.next()?;
    let (local_tmp_s0, local_1, vf_s0, dbg_projection) = match_get_variant_field(stmt)?;

    try_eat_storage_stmts(&mut stmt_iter, &mut storage_live_stmts, &mut storage_dead_stmts);

    try_eat_assign_tmp_stmts(&mut stmt_iter, &mut tmp_assigns, &mut nop_stmts);

    try_eat_storage_stmts(&mut stmt_iter, &mut storage_live_stmts, &mut storage_dead_stmts);

    try_eat_assign_tmp_stmts(&mut stmt_iter, &mut tmp_assigns, &mut nop_stmts);

    let (idx, stmt) = stmt_iter.next()?;
    let (local_tmp_s1, local_0, vf_s1) = match_set_variant_field(stmt)?;
    nop_stmts.push(idx);

    let (idx, stmt) = stmt_iter.next()?;
    let (set_discr_local, set_discr_var_idx) = match_set_discr(stmt)?;
    let discr_stmt_source_info = stmt.source_info;
    nop_stmts.push(idx);

    try_eat_storage_stmts(&mut stmt_iter, &mut storage_live_stmts, &mut storage_dead_stmts);

    for (live_idx, live_local) in storage_live_stmts {
        if let Some(i) = storage_dead_stmts.iter().rposition(|(_, l)| *l == live_local) {
            let (dead_idx, _) = storage_dead_stmts.swap_remove(i);
            storage_stmts.push((live_idx, dead_idx, live_local));

            if live_local == local_tmp_s0 {
                nop_stmts.push(get_variant_field_stmt);
            }
        }
    }

    nop_stmts.sort();

    // Use one of the statements we're going to discard between the point
    // where the storage location for the variant field becomes live and
    // is killed.
    let (live_idx, dead_idx) = find_storage_live_dead_stmts_for_local(local_tmp_s0, stmts)?;
    let stmt_to_overwrite =
        nop_stmts.iter().find(|stmt_idx| live_idx < **stmt_idx && **stmt_idx < dead_idx);

    let mut tmp_assigned_vars = BitSet::new_empty(locals_count);
    for (l, r) in &tmp_assigns {
        tmp_assigned_vars.insert(*l);
        tmp_assigned_vars.insert(*r);
    }

    let mut dbg_info_to_adjust = Vec::new();
    for (i, var_info) in debug_info.iter().enumerate() {
        if tmp_assigned_vars.contains(var_info.place.local) {
            dbg_info_to_adjust.push(i);
        }
    }

    Some(ArmIdentityInfo {
        local_temp_0: local_tmp_s0,
        local_1,
        vf_s0,
        get_variant_field_stmt,
        field_tmp_assignments: tmp_assigns,
        local_tmp_s1,
        local_0,
        vf_s1,
        set_discr_local,
        set_discr_var_idx,
        stmt_to_overwrite: *stmt_to_overwrite?,
        source_info: discr_stmt_source_info,
        storage_stmts,
        stmts_to_remove: nop_stmts,
        dbg_info_to_adjust,
        dbg_projection,
    })
}

fn optimization_applies<'tcx>(
    opt_info: &ArmIdentityInfo<'tcx>,
    local_decls: &IndexVec<Local, LocalDecl<'tcx>>,
    local_uses: &IndexVec<Local, usize>,
    var_debug_info: &[VarDebugInfo<'tcx>],
) -> bool {
    trace!("testing if optimization applies...");

    // FIXME(wesleywiser): possibly relax this restriction?
    if opt_info.local_0 == opt_info.local_1 {
        trace!("NO: moving into ourselves");
        return false;
    } else if opt_info.vf_s0 != opt_info.vf_s1 {
        trace!("NO: the field-and-variant information do not match");
        return false;
    } else if local_decls[opt_info.local_0].ty != local_decls[opt_info.local_1].ty {
        // FIXME(Centril,oli-obk): possibly relax to same layout?
        trace!("NO: source and target locals have different types");
        return false;
    } else if (opt_info.local_0, opt_info.vf_s0.var_idx)
        != (opt_info.set_discr_local, opt_info.set_discr_var_idx)
    {
        trace!("NO: the discriminants do not match");
        return false;
    }

    // Verify the assigment chain consists of the form b = a; c = b; d = c; etc...
    if opt_info.field_tmp_assignments.is_empty() {
        trace!("NO: no assignments found");
        return false;
    }
    let mut last_assigned_to = opt_info.field_tmp_assignments[0].1;
    let source_local = last_assigned_to;
    for (l, r) in &opt_info.field_tmp_assignments {
        if *r != last_assigned_to {
            trace!("NO: found unexpected assignment {:?} = {:?}", l, r);
            return false;
        }

        last_assigned_to = *l;
    }

    // Check that the first and last used locals are only used twice
    // since they are of the form:
    //
    // ```
    // _first = ((_x as Variant).n: ty);
    // _n = _first;
    // ...
    // ((_y as Variant).n: ty) = _n;
    // discriminant(_y) = z;
    // ```
    for (l, r) in &opt_info.field_tmp_assignments {
        if local_uses[*l] != 2 {
            warn!("NO: FAILED assignment chain local {:?} was used more than twice", l);
            return false;
        } else if local_uses[*r] != 2 {
            warn!("NO: FAILED assignment chain local {:?} was used more than twice", r);
            return false;
        }
    }

    // Check that debug info only points to full Locals and not projections.
    for dbg_idx in &opt_info.dbg_info_to_adjust {
        let dbg_info = &var_debug_info[*dbg_idx];
        if !dbg_info.place.projection.is_empty() {
            trace!("NO: debug info for {:?} had a projection {:?}", dbg_info.name, dbg_info.place);
            return false;
        }
    }

    if source_local != opt_info.local_temp_0 {
        trace!(
            "NO: start of assignment chain does not match enum variant temp: {:?} != {:?}",
            source_local,
            opt_info.local_temp_0
        );
        return false;
    } else if last_assigned_to != opt_info.local_tmp_s1 {
        trace!(
            "NO: end of assignemnt chain does not match written enum temp: {:?} != {:?}",
            last_assigned_to,
            opt_info.local_tmp_s1
        );
        return false;
    }

    trace!("SUCCESS: optimization applies!");
    return true;
}

impl<'tcx> MirPass<'tcx> for SimplifyArmIdentity {
    fn run_pass(&self, tcx: TyCtxt<'tcx>, source: MirSource<'tcx>, body: &mut Body<'tcx>) {
        if tcx.sess.opts.debugging_opts.mir_opt_level < 2 {
            return;
        }

        trace!("running SimplifyArmIdentity on {:?}", source);
        let local_uses = LocalUseCounter::get_local_uses(body);
        let (basic_blocks, local_decls, debug_info) =
            body.basic_blocks_local_decls_mut_and_var_debug_info();
        for bb in basic_blocks {
            if let Some(opt_info) =
                get_arm_identity_info(&bb.statements, local_decls.len(), debug_info)
            {
                trace!("got opt_info = {:#?}", opt_info);
                if !optimization_applies(&opt_info, local_decls, &local_uses, &debug_info) {
                    debug!("optimization skipped for {:?}", source);
                    continue;
                }

                // Also remove unused Storage{Live,Dead} statements which correspond
                // to temps used previously.
                for (live_idx, dead_idx, local) in &opt_info.storage_stmts {
                    // The temporary that we've read the variant field into is scoped to this block,
                    // so we can remove the assignment.
                    if *local == opt_info.local_temp_0 {
                        bb.statements[opt_info.get_variant_field_stmt].make_nop();
                    }

                    for (left, right) in &opt_info.field_tmp_assignments {
                        if local == left || local == right {
                            bb.statements[*live_idx].make_nop();
                            bb.statements[*dead_idx].make_nop();
                        }
                    }
                }

                // Right shape; transform
                for stmt_idx in opt_info.stmts_to_remove {
                    bb.statements[stmt_idx].make_nop();
                }

                let stmt = &mut bb.statements[opt_info.stmt_to_overwrite];
                stmt.source_info = opt_info.source_info;
                stmt.kind = StatementKind::Assign(box (
                    opt_info.local_0.into(),
                    Rvalue::Use(Operand::Move(opt_info.local_1.into())),
                ));

                bb.statements.retain(|stmt| stmt.kind != StatementKind::Nop);

                // Fix the debug info to point to the right local
                for dbg_index in opt_info.dbg_info_to_adjust {
                    let dbg_info = &mut debug_info[dbg_index];
                    assert!(dbg_info.place.projection.is_empty());
                    dbg_info.place.local = opt_info.local_0;
                    dbg_info.place.projection = opt_info.dbg_projection;
                }

                trace!("block is now {:?}", bb.statements);
            }
        }
    }
}

struct LocalUseCounter {
    local_uses: IndexVec<Local, usize>,
}

impl LocalUseCounter {
    fn get_local_uses<'tcx>(body: &Body<'tcx>) -> IndexVec<Local, usize> {
        let mut counter = LocalUseCounter { local_uses: IndexVec::from_elem(0, &body.local_decls) };
        counter.visit_body(body);
        counter.local_uses
    }
}

impl<'tcx> Visitor<'tcx> for LocalUseCounter {
    fn visit_local(&mut self, local: &Local, context: PlaceContext, _location: Location) {
        if context.is_storage_marker()
            || context == PlaceContext::NonUse(NonUseContext::VarDebugInfo)
        {
            return;
        }

        self.local_uses[*local] += 1;
    }
}

/// Match on:
/// ```rust
/// _LOCAL_INTO = ((_LOCAL_FROM as Variant).FIELD: TY);
/// ```
fn match_get_variant_field<'tcx>(
    stmt: &Statement<'tcx>,
) -> Option<(Local, Local, VarField<'tcx>, &'tcx List<PlaceElem<'tcx>>)> {
    match &stmt.kind {
        StatementKind::Assign(box (place_into, rvalue_from)) => match rvalue_from {
            Rvalue::Use(Operand::Copy(pf) | Operand::Move(pf)) => {
                let local_into = place_into.as_local()?;
                let (local_from, vf) = match_variant_field_place(*pf)?;
                Some((local_into, local_from, vf, pf.projection))
            }
            _ => None,
        },
        _ => None,
    }
}

/// Match on:
/// ```rust
/// ((_LOCAL_FROM as Variant).FIELD: TY) = move _LOCAL_INTO;
/// ```
fn match_set_variant_field<'tcx>(stmt: &Statement<'tcx>) -> Option<(Local, Local, VarField<'tcx>)> {
    match &stmt.kind {
        StatementKind::Assign(box (place_from, rvalue_into)) => match rvalue_into {
            Rvalue::Use(Operand::Move(place_into)) => {
                let local_into = place_into.as_local()?;
                let (local_from, vf) = match_variant_field_place(*place_from)?;
                Some((local_into, local_from, vf))
            }
            _ => None,
        },
        _ => None,
    }
}

/// Match on:
/// ```rust
/// discriminant(_LOCAL_TO_SET) = VAR_IDX;
/// ```
fn match_set_discr<'tcx>(stmt: &Statement<'tcx>) -> Option<(Local, VariantIdx)> {
    match &stmt.kind {
        StatementKind::SetDiscriminant { place, variant_index } => {
            Some((place.as_local()?, *variant_index))
        }
        _ => None,
    }
}

#[derive(PartialEq, Debug)]
struct VarField<'tcx> {
    field: Field,
    field_ty: Ty<'tcx>,
    var_idx: VariantIdx,
}

/// Match on `((_LOCAL as Variant).FIELD: TY)`.
fn match_variant_field_place<'tcx>(place: Place<'tcx>) -> Option<(Local, VarField<'tcx>)> {
    match place.as_ref() {
        PlaceRef {
            local,
            projection: &[ProjectionElem::Downcast(_, var_idx), ProjectionElem::Field(field, ty)],
        } => Some((local, VarField { field, field_ty: ty, var_idx })),
        _ => None,
    }
}

/// Simplifies `SwitchInt(_) -> [targets]`,
/// where all the `targets` have the same form,
/// into `goto -> target_first`.
pub struct SimplifyBranchSame;

impl<'tcx> MirPass<'tcx> for SimplifyBranchSame {
    fn run_pass(&self, _: TyCtxt<'tcx>, _: MirSource<'tcx>, body: &mut Body<'tcx>) {
        let mut did_remove_blocks = false;
        let bbs = body.basic_blocks_mut();
        for bb_idx in bbs.indices() {
            let targets = match &bbs[bb_idx].terminator().kind {
                TerminatorKind::SwitchInt { targets, .. } => targets,
                _ => continue,
            };

            let mut iter_bbs_reachable = targets
                .iter()
                .map(|idx| (*idx, &bbs[*idx]))
                .filter(|(_, bb)| {
                    // Reaching `unreachable` is UB so assume it doesn't happen.
                    bb.terminator().kind != TerminatorKind::Unreachable
                    // But `asm!(...)` could abort the program,
                    // so we cannot assume that the `unreachable` terminator itself is reachable.
                    // FIXME(Centril): use a normalization pass instead of a check.
                    || bb.statements.iter().any(|stmt| match stmt.kind {
                        StatementKind::LlvmInlineAsm(..) => true,
                        _ => false,
                    })
                })
                .peekable();

            // We want to `goto -> bb_first`.
            let bb_first = iter_bbs_reachable.peek().map(|(idx, _)| *idx).unwrap_or(targets[0]);

            // All successor basic blocks should have the exact same form.
            let all_successors_equivalent =
                iter_bbs_reachable.map(|(_, bb)| bb).tuple_windows().all(|(bb_l, bb_r)| {
                    bb_l.is_cleanup == bb_r.is_cleanup
                        && bb_l.terminator().kind == bb_r.terminator().kind
                        && bb_l.statements.iter().eq_by(&bb_r.statements, |x, y| x.kind == y.kind)
                });

            if all_successors_equivalent {
                // Replace `SwitchInt(..) -> [bb_first, ..];` with a `goto -> bb_first;`.
                bbs[bb_idx].terminator_mut().kind = TerminatorKind::Goto { target: bb_first };
                did_remove_blocks = true;
            }
        }

        if did_remove_blocks {
            // We have dead blocks now, so remove those.
            simplify::remove_dead_blocks(body);
        }
    }
}