ssa_range_prop.rs

//! A pass that propagates the known ranges of SSA locals.
//! We can know the ranges of SSA locals in certain locations for the following code:
//! ```
//! fn foo(a: u32) {
//!   let b = a < 9; // the integer representation of b is within the full range [0, 2).
//!   if b {
//!     let c = b; // c is true since b is within the range [1, 2).
//!     let d = a < 8; // d is true since a is within the range [0, 9).
//!   }
//! }
//! ```
use rustc_abi::WrappingRange;
use rustc_const_eval::interpret::Scalar;
use rustc_data_structures::fx::FxHashMap;
use rustc_data_structures::graph::dominators::Dominators;
use rustc_index::bit_set::DenseBitSet;
use rustc_middle::mir::visit::MutVisitor;
use rustc_middle::mir::{BasicBlock, Body, Location, Operand, Place, TerminatorKind, *};
use rustc_middle::ty::{TyCtxt, TypingEnv};
use rustc_span::DUMMY_SP;

use crate::ssa::SsaLocals;

pub(super) struct SsaRangePropagation;

impl<'tcx> crate::MirPass<'tcx> for SsaRangePropagation {
    fn is_enabled(&self, sess: &rustc_session::Session) -> bool {
        sess.mir_opt_level() > 1
    }

    fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
        let typing_env = body.typing_env(tcx);
        let ssa = SsaLocals::new(tcx, body, typing_env);
        // Clone dominators because we need them while mutating the body.
        let dominators = body.basic_blocks.dominators().clone();
        let mut range_set =
            RangeSet::new(tcx, typing_env, body, &ssa, &body.local_decls, dominators);

        let reverse_postorder = body.basic_blocks.reverse_postorder().to_vec();
        for bb in reverse_postorder {
            let data = &mut body.basic_blocks.as_mut_preserves_cfg()[bb];
            range_set.visit_basic_block_data(bb, data);
        }
    }

    fn is_required(&self) -> bool {
        false
    }
}

struct RangeSet<'tcx, 'body, 'a> {
    tcx: TyCtxt<'tcx>,
    typing_env: TypingEnv<'tcx>,
    ssa: &'a SsaLocals,
    local_decls: &'body LocalDecls<'tcx>,
    dominators: Dominators<BasicBlock>,
    /// Known ranges at each locations.
    ranges: FxHashMap<Place<'tcx>, Vec<(Location, WrappingRange)>>,
    /// Determines if the basic block has a single unique predecessor.
    unique_predecessors: DenseBitSet<BasicBlock>,
}

impl<'tcx, 'body, 'a> RangeSet<'tcx, 'body, 'a> {
    fn new(
        tcx: TyCtxt<'tcx>,
        typing_env: TypingEnv<'tcx>,
        body: &Body<'tcx>,
        ssa: &'a SsaLocals,
        local_decls: &'body LocalDecls<'tcx>,
        dominators: Dominators<BasicBlock>,
    ) -> Self {
        let predecessors = body.basic_blocks.predecessors();
        let mut unique_predecessors = DenseBitSet::new_empty(body.basic_blocks.len());
        for bb in body.basic_blocks.indices() {
            if predecessors[bb].len() == 1 {
                unique_predecessors.insert(bb);
            }
        }
        RangeSet {
            tcx,
            typing_env,
            ssa,
            local_decls,
            dominators,
            ranges: FxHashMap::default(),
            unique_predecessors,
        }
    }

    /// Create a new known range at the location.
    fn insert_range(&mut self, place: Place<'tcx>, location: Location, range: WrappingRange) {
        assert!(self.is_ssa(place));
        self.ranges.entry(place).or_default().push((location, range));
    }

    /// Get the known range at the location.
    fn get_range(&self, place: &Place<'tcx>, location: Location) -> Option<WrappingRange> {
        let Some(ranges) = self.ranges.get(place) else {
            return None;
        };
        // FIXME: This should use the intersection of all valid ranges.
        let (_, range) =
            ranges.iter().find(|(range_loc, _)| range_loc.dominates(location, &self.dominators))?;
        Some(*range)
    }

    fn try_as_constant(
        &mut self,
        place: Place<'tcx>,
        location: Location,
    ) -> Option<ConstOperand<'tcx>> {
        if let Some(range) = self.get_range(&place, location)
            && range.start == range.end
        {
            let ty = place.ty(self.local_decls, self.tcx).ty;
            let layout = self.tcx.layout_of(self.typing_env.as_query_input(ty)).ok()?;
            let value = ConstValue::Scalar(Scalar::from_uint(range.start, layout.size));
            let const_ = Const::Val(value, ty);
            return Some(ConstOperand { span: DUMMY_SP, user_ty: None, const_ });
        }
        None
    }

    fn is_ssa(&self, place: Place<'tcx>) -> bool {
        self.ssa.is_ssa(place.local) && place.is_stable_offset()
    }
}

impl<'tcx> MutVisitor<'tcx> for RangeSet<'tcx, '_, '_> {
    fn tcx(&self) -> TyCtxt<'tcx> {
        self.tcx
    }

    fn visit_operand(&mut self, operand: &mut Operand<'tcx>, location: Location) {
        // Attempts to simplify an operand to a constant value.
        if let Some(place) = operand.place()
            && let Some(const_) = self.try_as_constant(place, location)
        {
            *operand = Operand::Constant(Box::new(const_));
        };
    }

    fn visit_statement(&mut self, statement: &mut Statement<'tcx>, location: Location) {
        self.super_statement(statement, location);
        match &statement.kind {
            StatementKind::Intrinsic(box NonDivergingIntrinsic::Assume(operand)) => {
                if let Some(place) = operand.place()
                    && self.is_ssa(place)
                {
                    let successor = location.successor_within_block();
                    let range = WrappingRange { start: 1, end: 1 };
                    self.insert_range(place, successor, range);
                }
            }
            _ => {}
        }
    }

    fn visit_terminator(&mut self, terminator: &mut Terminator<'tcx>, location: Location) {
        self.super_terminator(terminator, location);
        match &terminator.kind {
            TerminatorKind::Assert { cond, expected, target, .. } => {
                if let Some(place) = cond.place()
                    && self.is_ssa(place)
                {
                    let successor = Location { block: *target, statement_index: 0 };
                    if location.dominates(successor, &self.dominators) {
                        assert_ne!(location.block, successor.block);
                        let val = *expected as u128;
                        let range = WrappingRange { start: val, end: val };
                        self.insert_range(place, successor, range);
                    }
                }
            }
            TerminatorKind::SwitchInt { discr, targets } => {
                if let Some(place) = discr.place()
                    && self.is_ssa(place)
                    // Reduce the potential compile-time overhead.
                    && targets.all_targets().len() < 16
                {
                    let mut distinct_targets: FxHashMap<BasicBlock, u64> = FxHashMap::default();
                    for (_, target) in targets.iter() {
                        let targets = distinct_targets.entry(target).or_default();
                        *targets += 1;
                    }
                    for (val, target) in targets.iter() {
                        if distinct_targets[&target] != 1 {
                            // FIXME: For multiple targets, the range can be the union of their values.
                            continue;
                        }
                        let successor = Location { block: target, statement_index: 0 };
                        if self.unique_predecessors.contains(successor.block) {
                            assert_ne!(location.block, successor.block);
                            let range = WrappingRange { start: val, end: val };
                            self.insert_range(place, successor, range);
                        }
                    }

                    // FIXME: The range for the otherwise target be extend to more types.
                    // For instance, `val` is within the range [4, 1) at the otherwise target of `matches!(val, 1 | 2 | 3)`.
                    let otherwise = Location { block: targets.otherwise(), statement_index: 0 };
                    if place.ty(self.local_decls, self.tcx).ty.is_bool()
                        && let [val] = targets.all_values()
                        && self.unique_predecessors.contains(otherwise.block)
                    {
                        assert_ne!(location.block, otherwise.block);
                        let range = if val.get() == 0 {
                            WrappingRange { start: 1, end: 1 }
                        } else {
                            WrappingRange { start: 0, end: 0 }
                        };
                        self.insert_range(place, otherwise, range);
                    }
                }
            }
            _ => {}
        }
    }
}