frontend 0.4.1

//! A pass that eliminates branches on uninhabited or unreachable enum variants.

// `#![no_std]`: these arrive with the standard prelude and name no path, so a `std::`
// search cannot see them - and a `#[derive]` can use them without the name appearing
// in this file at all, which is why they are not trimmed by inspection.
use alloc::borrow::ToOwned;
use alloc::boxed::Box;
use alloc::format;
use alloc::string::{String, ToString};
use alloc::vec;
use alloc::vec::Vec;

use crate::rustc_abi::{VariantIdx, Variants};
use crate::rustc_data_structures::fx::FxHashSet;
use crate::rustc_index::Idx;
use crate::bug;
use crate::rustc_middle::mir::{
    BasicBlockData, Body, Local, Operand, Rvalue, StatementKind, TerminatorKind,
};
use crate::rustc_middle::ty::layout::TyAndLayout;
use crate::rustc_middle::ty::{Ty, TyCtxt};
use tracing::trace;

use crate::rustc_mir_transform::PassPolicy;
use crate::rustc_mir_transform::patch::MirPatch;

pub(super) struct UnreachableEnumBranching;

fn get_discriminant_local(terminator: &TerminatorKind<'_>) -> Option<Local> {
    if let TerminatorKind::SwitchInt { discr: Operand::Move(p), .. } = terminator {
        p.as_local()
    } else {
        None
    }
}

/// If the basic block terminates by switching on a discriminant, this returns the `Ty` the
/// discriminant is read from. Otherwise, returns None.
fn get_switched_on_type<'tcx>(
    block_data: &BasicBlockData<'tcx>,
    tcx: TyCtxt<'tcx>,
    body: &Body<'tcx>,
) -> Option<Ty<'tcx>> {
    let terminator = block_data.terminator();

    // Only bother checking blocks which terminate by switching on a local.
    let local = get_discriminant_local(&terminator.kind)?;

    let stmt_before_term = block_data.statements.last()?;

    if let StatementKind::Assign(ref assign) = stmt_before_term.kind
        && let (l, Rvalue::Discriminant(place)) = **assign
        && l.as_local() == Some(local)
    {
        let ty = place.ty(body, tcx).ty;
        if ty.is_enum() {
            return Some(ty);
        }
    }

    None
}

fn variant_discriminants<'tcx>(
    layout: &TyAndLayout<'tcx>,
    ty: Ty<'tcx>,
    tcx: TyCtxt<'tcx>,
) -> FxHashSet<u128> {
    match &layout.variants {
        Variants::Empty => {
            // Uninhabited, no valid discriminant.
            FxHashSet::default()
        }
        Variants::Single { index } => {
            let mut res = FxHashSet::default();
            res.insert(
                ty.discriminant_for_variant(tcx, *index)
                    .map_or(index.as_u32() as u128, |discr| discr.val),
            );
            res
        }
        Variants::Multiple { variants, .. } => variants
            .iter_enumerated()
            .filter_map(|(idx, layout)| {
                (!layout.is_uninhabited())
                    .then(|| ty.discriminant_for_variant(tcx, idx).unwrap().val)
            })
            .collect(),
    }
}

impl<'tcx> crate::rustc_mir_transform::MirPass<'tcx> for UnreachableEnumBranching {
    fn policy(&self, sess: &crate::rustc_session::Session) -> PassPolicy {
        PassPolicy::optimization(sess.mir_opt_level() > 0)
    }

    fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
        trace!("UnreachableEnumBranching starting for {:?}", body.source);

        let mut unreachable_targets = Vec::new();
        let mut patch = MirPatch::new(body);

        for (bb, bb_data) in body.basic_blocks.iter_enumerated() {
            trace!("processing block {:?}", bb);

            if bb_data.is_cleanup {
                continue;
            }

            let Some(discriminant_ty) = get_switched_on_type(bb_data, tcx, body) else { continue };

            let layout = tcx.layout_of(body.typing_env(tcx).as_query_input(discriminant_ty));

            let mut allowed_variants = if let Ok(layout) = layout {
                // Find allowed variants based on uninhabited.
                variant_discriminants(&layout, discriminant_ty, tcx)
            } else if let Some(variant_range) = discriminant_ty.variant_range(tcx) {
                // If there are some generics, we can still get the allowed variants.
                (variant_range.start.index()..variant_range.end.index())
                    .map(VariantIdx::new)
                    .map(|variant| {
                        discriminant_ty.discriminant_for_variant(tcx, variant).unwrap().val
                    })
                    .collect()
            } else {
                continue;
            };

            trace!("allowed_variants = {:?}", allowed_variants);

            unreachable_targets.clear();
            let TerminatorKind::SwitchInt { targets, discr } = &bb_data.terminator().kind else {
                bug!()
            };

            for (index, (val, _)) in targets.iter().enumerate() {
                if !allowed_variants.remove(&val) {
                    unreachable_targets.push(index);
                }
            }

            // If and only if there is a variant that does not have a branch set, change the
            // current of otherwise as the variant branch and set otherwise to unreachable. It
            // transforms the following code
            // ```rust
            // match c {
            //     Ordering::Less => 1,
            //     Ordering::Equal => 2,
            //     _ => 3,
            // }
            // ```
            // to
            // ```rust
            // match c {
            //     Ordering::Less => 1,
            //     Ordering::Equal => 2,
            //     Ordering::Greater => 3,
            // }
            // ```
            let replace_otherwise_to_unreachable = allowed_variants.len() <= 1
                && !body.basic_blocks[targets.otherwise()].is_empty_unreachable();
            if unreachable_targets.is_empty() && !replace_otherwise_to_unreachable {
                continue;
            }

            let unreachable_block = patch.unreachable_no_cleanup_block();
            let mut targets = targets.clone();
            if replace_otherwise_to_unreachable {
                let otherwise_is_last_variant = allowed_variants.len() == 1;
                if otherwise_is_last_variant {
                    // We have checked that `allowed_variants` has only one element.
                    let last_variant = *allowed_variants.iter().next().unwrap();
                    targets.add_target(last_variant, targets.otherwise());
                }
                unreachable_targets.push(targets.iter().count());
            }
            for index in unreachable_targets.iter() {
                targets.all_targets_mut()[*index] = unreachable_block;
            }
            patch.patch_terminator(bb, TerminatorKind::SwitchInt { targets, discr: discr.clone() });
        }

        patch.apply(body);
    }
}