compiler/rustc_mir_dataflow/src/impls/precise_liveness.rs

//! Computes the points where each local must have a distinct allocation.
//!
//! The result is a [`SparseIntervalMatrix`] with one row per local. Two locals
//! may share the same address only if their rows are disjoint. To model MIR
//! statements where a source operand and destination place may share an
//! address, each statement and terminator is split into an early point, where
//! operands are read, and a late point, where destinations are written.
//!
//! A local live range starts at the late point of any statement or terminator
//! that writes to it without a `Deref` projection. It ends at `StorageLive` or
//! `StorageDead`, or at the early point of a whole-local move operand or the
//! last use of that local on a control-flow path (only for locals whose address
//! is never observed). This happens at `StorageLive` because it resets a local
//! to an uninitialized state.
//!
//! `Call` terminators are handled specially: move operands are kept live
//! through the late point of the terminator so they conflict with each other
//! and with the destination place. This matches the runtime behavior where the
//! place is donated to the callee for the duration of the call.
//!
//! The analysis consists of three passes:
//!
//! 1. A backward value-liveness analysis which identifies the last point at
//!    which each local is directly used in each control flow path. This is used
//!    to shorten the lifetime of unborrowed locals to their last use.
//! 2. A forward allocation-liveness analysis which uses that information
//!    together with writes, moves, and storage markers to compute which
//!    allocations may be live on entry to each block.
//! 3. A final walk using those entry states which constructs the matrix and
//!    distinguishes whether each effect happens at the early or late point of a
//!    statement or terminator.

use rustc_index::IndexVec;
use rustc_index::bit_set::DenseBitSet;
use rustc_index::interval::SparseIntervalMatrix;
use rustc_middle::mir::visit::{
    MutatingUseContext, NonMutatingUseContext, NonUseContext, PlaceContext, VisitPlacesWith,
    Visitor,
};
use rustc_middle::mir::{self, BasicBlock, Local, Location, MirDumper, PassWhere, Place};
use rustc_middle::ty::TyCtxt;
use tracing::trace;

use crate::impls::{DefUse, MaybeLiveLocals, borrowed_locals};
use crate::points::{DenseLocationMap, PointIndex};
use crate::{Analysis, GenKill, ResultsVisitor, visit_results};

/// Visits locals that are implicitly allocated by a statement or terminator as
/// part of destination place evaluation.
pub struct VisitAllocatedLocalsWith<F>(pub F);

impl<'tcx, F> Visitor<'tcx> for VisitAllocatedLocalsWith<F>
where
    F: FnMut(Local, PlaceContext),
{
    fn visit_place(&mut self, place: &Place<'tcx>, context: PlaceContext, _: Location) {
        match DefUse::for_place(*place, context) {
            DefUse::Def | DefUse::PartialWrite => (self.0)(place.local, context),
            DefUse::Use | DefUse::NonUse => {}
        }
    }
}

/// Visits locals that are implicitly freed by being evaluated as a bare-local
/// move operand or through a storage statement.
pub struct VisitDeallocatedLocalsWith<F>(pub F);

impl<'tcx, F> Visitor<'tcx> for VisitDeallocatedLocalsWith<F>
where
    F: FnMut(Local, PlaceContext),
{
    fn visit_statement(&mut self, statement: &mir::Statement<'tcx>, location: Location) {
        match statement.kind {
            // StorageLive resets a local to the live-but-unallocated state,
            // which also has the effect of freeing its allocation.
            mir::StatementKind::StorageLive(local) => {
                (self.0)(local, PlaceContext::NonUse(NonUseContext::StorageLive));
            }
            mir::StatementKind::StorageDead(local) => {
                (self.0)(local, PlaceContext::NonUse(NonUseContext::StorageDead));
            }
            _ => self.super_statement(statement, location),
        }
    }

    fn visit_operand(&mut self, operand: &mir::Operand<'tcx>, _: Location) {
        if let mir::Operand::Move(place) = operand
            && let Some(local) = place.as_local()
        {
            (self.0)(local, PlaceContext::NonMutatingUse(NonMutatingUseContext::Move));
        }
    }
}

////////////////////////////////////////////////////////////////////////////////
// Backward dataflow pass
//
// This pass computes "kill points" for each local, indicating the location of
// their last use in a particular control flow branch. These are later used in
// the forward pass to end the live range of locals that are never borrowed at
// their last direct use.
//
// This pass ignores borrowed locals: we can't determine their last use from
// syntactic uses of the base local, so we conservatively treat them as always
// live in the backwards pass and don't emit kill points for them. The main
// forward pass still kills at `StorageDead` and whole-local moves, which is
// when their allocation is freed.
//
// This pass has 2 outputs: a set of kill points that mark the last use
// locations of locals and a per-block bitset indicating which locals are live
// on entry to that block. The latter acts as a set of kill points at block
// entry for all locals not in the bit set.

struct KillPoints<'a> {
    live_on_entry: IndexVec<BasicBlock, DenseBitSet<Local>>,
    kill_points_map: IndexVec<PointIndex, &'a [(Local, Location)]>,
}

impl<'a> KillPoints<'a> {
    fn compute<'tcx>(
        tcx: TyCtxt<'tcx>,
        body: &mir::Body<'tcx>,
        pass_name: Option<&'static str>,
        points: &DenseLocationMap,
        kill_points: &'a mut Vec<(Local, Location)>,
    ) -> Self {
        let maybe_live_locals = MaybeLiveLocals.iterate_to_fixpoint(tcx, body, pass_name);
        let borrowed_locals = borrowed_locals(body);

        // Initialize all borrowed locals as live on entry. We never try to kill
        // those.
        let mut live_on_entry =
            IndexVec::from_elem_n(borrowed_locals.clone(), body.basic_blocks.len());

        // Collect kill points and live-on-entry states from the results of
        // MaybeLiveLocals.
        kill_points.clear();
        let mut visitor = KillPointsVisitor {
            kill_points,
            live_on_entry: &mut live_on_entry,
            borrowed_locals: &borrowed_locals,
        };
        visit_results(
            body,
            mir::traversal::reachable(body).map(|(block, _)| block),
            &maybe_live_locals,
            &mut visitor,
        );
        trace!(?kill_points);
        trace!(?live_on_entry);

        // Create a mapping of `PointIndex` to the set of killed locals at that
        // location.
        let mut kill_points_map = IndexVec::from_elem_n(&[][..], points.num_points());
        for chunk in kill_points.chunk_by(|a, b| a.1 == b.1) {
            let point = points.point_from_location(chunk[0].1);
            trace!("Kill points at {:?}: {:?}", chunk[0].1, chunk);
            debug_assert!(kill_points_map[point].is_empty());
            kill_points_map[point] = chunk;
        }

        Self { live_on_entry, kill_points_map }
    }
}

struct KillPointsVisitor<'a> {
    kill_points: &'a mut Vec<(Local, Location)>,
    live_on_entry: &'a mut IndexVec<BasicBlock, DenseBitSet<Local>>,
    borrowed_locals: &'a DenseBitSet<Local>,
}

impl<'tcx> ResultsVisitor<'tcx, MaybeLiveLocals> for KillPointsVisitor<'_> {
    fn visit_block_exit(&mut self, state: &DenseBitSet<Local>, block: BasicBlock) {
        // Borrowed locals are already marked as live when live_on_entry was
        // initialized. This adds the non-borrowed locals that we have
        // determined are live on entry to this block.
        self.live_on_entry[block].union(state);
    }

    fn visit_after_early_statement_effect(
        &mut self,
        state: &DenseBitSet<Local>,
        statement: &mir::Statement<'tcx>,
        location: Location,
    ) {
        VisitPlacesWith(|place: Place<'tcx>, ctxt| {
            // Ignore non-uses.
            match ctxt {
                PlaceContext::NonMutatingUse(_) | PlaceContext::MutatingUse(_) => {}
                PlaceContext::NonUse(_) => return,
            }

            // If a local is used in a statement but is dead after it then this
            // location is a kill point. Don't emit a kill point for borrowed
            // locals.
            if !state.contains(place.local) && !self.borrowed_locals.contains(place.local) {
                self.kill_points.push((place.local, location));
            }
        })
        .visit_statement(statement, location);
    }

    fn visit_after_early_terminator_effect(
        &mut self,
        state: &DenseBitSet<Local>,
        terminator: &mir::Terminator<'tcx>,
        location: Location,
    ) {
        VisitPlacesWith(|place: Place<'tcx>, ctxt| {
            // Ignore non-uses (they don't do anything) and edge uses
            // (implicitly killed though live_on_entry at the start of the
            // corresponding successor).
            match ctxt {
                PlaceContext::MutatingUse(
                    MutatingUseContext::AsmOutput
                    | MutatingUseContext::Call
                    | MutatingUseContext::Yield,
                )
                | PlaceContext::NonUse(_) => return,
                PlaceContext::NonMutatingUse(_) | PlaceContext::MutatingUse(_) => {}
            }

            // If a local is used in a terminator but is dead after it then this
            // location is a kill point. Don't emit a kill point for borrowed
            // locals.
            if !state.contains(place.local) && !self.borrowed_locals.contains(place.local) {
                self.kill_points.push((place.local, location));
            }
        })
        .visit_terminator(terminator, location);
    }
}

////////////////////////////////////////////////////////////////////////////////
// Forward dataflow pass
//
// This is a simpler version of the main pass in `liveness_matrix`. It only
// needs to determine local liveness at block edges and doesn't distinguish
// between "early" and "late" effects. However it must still preserve the
// ordering of effects: moved operands are killed before destinations are
// generated.

struct PreciseLiveness<'a> {
    kill_points: &'a KillPoints<'a>,
    points: &'a DenseLocationMap,
}

impl PreciseLiveness<'_> {
    fn apply_block_start_effect(&self, state: &mut DenseBitSet<Local>, block: BasicBlock) {
        // Kill any unborrowed locals that are dead at the start of this block.
        // Notably this kills any dead results produced by a predecessor's
        // terminator. Borrowed locals are always included in live_on_entry.
        state.intersect(&self.kill_points.live_on_entry[block]);
    }
}

impl<'tcx> Analysis<'tcx> for PreciseLiveness<'_> {
    type Domain = DenseBitSet<Local>;

    const NAME: &'static str = "precise_liveness";

    fn bottom_value(&self, body: &mir::Body<'tcx>) -> DenseBitSet<Local> {
        DenseBitSet::new_empty(body.local_decls.len())
    }

    fn initialize_start_block(&self, body: &mir::Body<'tcx>, state: &mut DenseBitSet<Local>) {
        // Function arguments start out as live.
        for arg in body.args_iter() {
            state.gen_(arg);
        }
    }

    fn apply_primary_statement_effect(
        &self,
        state: &mut DenseBitSet<Local>,
        statement: &mir::Statement<'tcx>,
        location: Location,
    ) {
        if location.statement_index == 0 {
            self.apply_block_start_effect(state, location.block);
        }

        VisitDeallocatedLocalsWith(|local, _| state.kill(local))
            .visit_statement(statement, location);
        VisitAllocatedLocalsWith(|local, _| state.gen_(local)).visit_statement(statement, location);

        // Apply kill points at this statement: if a variable is dead then it
        // doesn't need storage.
        let point = self.points.point_from_location(location);
        for &(local, _) in self.kill_points.kill_points_map[point] {
            state.kill(local);
        }
    }

    fn apply_primary_terminator_effect(
        &self,
        state: &mut DenseBitSet<Local>,
        terminator: &mir::Terminator<'tcx>,
        location: Location,
    ) {
        if location.statement_index == 0 {
            self.apply_block_start_effect(state, location.block);
        }

        VisitDeallocatedLocalsWith(|local, _| state.kill(local))
            .visit_terminator(terminator, location);

        // Gen destination places.
        VisitAllocatedLocalsWith(|local, ctxt| {
            // The resume place of a Yield is only evaluated after the coroutine
            // is resumed. This is in contrast to call destinations and inline
            // asm outputs which are evaluated before the call and therefore
            // remain allocated on both normal and unwind edges.
            if ctxt != PlaceContext::MutatingUse(MutatingUseContext::Yield) {
                state.gen_(local);
            }
        })
        .visit_terminator(terminator, location);
    }

    fn apply_call_return_effect(
        &self,
        state: &mut DenseBitSet<Local>,
        _block: BasicBlock,
        return_places: mir::CallReturnPlaces<'_, 'tcx>,
    ) {
        if let mir::CallReturnPlaces::Yield(place) = return_places
            && !place.is_indirect()
        {
            state.gen_(place.local);
        }
    }
}

////////////////////////////////////////////////////////////////////////////////
// Matrix construction

/// Different "phases" of a single MIR statement, used to describe how
/// overlapping operands are handled.
///
/// As a general rule, source operands are read in the `Early` phase and
/// destination places are written in the `Late` phase.
#[derive(Copy, Clone, Debug, PartialEq, Eq)]
pub enum SplitPointEffect {
    Early = 0,
    Late = 1,
}

rustc_index::newtype_index! {
    /// A `PointIndex` with the lower bit encoding early/late inside a
    /// statement.
    ///
    /// This is used to model overlap constraints within a MIR statement: if a
    /// source/destination are allowed to overlap then the source is read in
    /// `SplitPointEffect::Early` and the write is done in
    /// `SplitPointEffect::Late`.
    #[orderable]
    #[debug_format = "SplitPointIndex({})"]
    pub struct SplitPointIndex {}
}

impl SplitPointIndex {
    pub fn new(point: PointIndex, effect: SplitPointEffect) -> SplitPointIndex {
        let index = (point.as_u32() << 1) | (effect as u32);
        SplitPointIndex::from_u32(index)
    }

    pub fn point(self) -> PointIndex {
        PointIndex::from_u32(self.as_u32() >> 1)
    }

    pub fn effect(self) -> SplitPointEffect {
        match self.as_u32() & 1 {
            0 => SplitPointEffect::Early,
            1 => SplitPointEffect::Late,
            _ => unreachable!(),
        }
    }
}

/// Helper type to construct a `SparseIntervalMatrix`.
struct MatrixBuilder {
    matrix: SparseIntervalMatrix<Local, SplitPointIndex>,
    range_start: IndexVec<Local, Option<SplitPointIndex>>,

    // Track locals that have been live at any point in a block so that at the
    // end of a block we don't need to iterate over all locals. This
    // significantly speeds up matrix building.
    maybe_live_locals: Vec<Local>,
}

impl MatrixBuilder {
    fn gen_(&mut self, local: Local, point: PointIndex, effect: SplitPointEffect) {
        let split_point = SplitPointIndex::new(point, effect);

        // No-op if the local is already live.
        if self.range_start[local].is_none() {
            self.range_start[local] = Some(split_point);
            self.maybe_live_locals.push(local);
        }
    }

    fn kill(&mut self, local: Local, point: PointIndex, effect: SplitPointEffect) {
        let end = SplitPointIndex::new(point, effect);

        // No-op if the local is already dead.
        if let Some(start) = self.range_start[local].take() {
            debug_assert!(end >= start);
            self.matrix.append_range(local, start..=end);
        }
    }

    fn kill_all(&mut self, point: PointIndex, effect: SplitPointEffect) {
        while let Some(local) = self.maybe_live_locals.pop() {
            self.kill(local, point, effect);
        }
    }

    fn kill_all_except(&mut self, except: Local, point: PointIndex, effect: SplitPointEffect) {
        while let Some(local) = self.maybe_live_locals.pop() {
            if local != except {
                self.kill(local, point, effect);
            }
        }
        self.maybe_live_locals.push(except);
    }
}

pub fn liveness_matrix<'tcx>(
    tcx: TyCtxt<'tcx>,
    body: &mir::Body<'tcx>,
    points: &DenseLocationMap,
    pass_name: Option<&'static str>,
) -> SparseIntervalMatrix<Local, SplitPointIndex> {
    let mut kill_points_vec = vec![];
    let kill_points = KillPoints::compute(tcx, body, pass_name, points, &mut kill_points_vec);
    let mut results = PreciseLiveness { kill_points: &kill_points, points }
        .iterate_to_fixpoint(tcx, body, pass_name);

    let mut builder = MatrixBuilder {
        matrix: SparseIntervalMatrix::new(points.num_points() * 2),
        range_start: IndexVec::from_elem_n(None, body.local_decls.len()),
        maybe_live_locals: Vec::new(),
    };
    for (block, block_data) in body.basic_blocks.iter_enumerated() {
        // We can mutate the state in-place since we're not using it any more
        // after this point.
        let state = &mut results.entry_states[block];

        // Notably this kills any dead results produced by a predecessor's
        // terminator.
        state.intersect(&kill_points.live_on_entry[block]);

        // Gen any locals that are live at the start of the block. If this block
        // only consists of a return terminator then only gen the return place
        // instead. This ensures that StorageDead for all other locals are
        // inserted before the return terminator.
        let terminator = block_data.terminator();
        if let mir::TerminatorKind::Return = terminator.kind
            && block_data.statements.is_empty()
        {
            if state.contains(mir::RETURN_PLACE) {
                builder.gen_(mir::RETURN_PLACE, points.entry_point(block), SplitPointEffect::Early);
            }
        } else {
            for local in state.iter() {
                builder.gen_(local, points.entry_point(block), SplitPointEffect::Early);
            }
        }

        for (statement_index, statement) in block_data.statements.iter().enumerate() {
            let location = Location { block, statement_index };
            let point = points.point_from_location(location);

            // Kill moved locals at the early point.
            VisitDeallocatedLocalsWith(|local, _| {
                builder.kill(local, point, SplitPointEffect::Early);
            })
            .visit_statement(statement, location);

            // Kill any locals which are no longer used after this statement.
            for &(local, _) in kill_points.kill_points_map[point] {
                builder.kill(local, point, SplitPointEffect::Early);
            }

            // Gen destination places at the late point.
            VisitAllocatedLocalsWith(|local, _| {
                builder.gen_(local, point, SplitPointEffect::Late);
            })
            .visit_statement(statement, location);

            // Kill any dead destination places: they will only appear at the
            // late point of the statement they are generated in, which is
            // sufficient for determining overlap.
            for &(local, _) in kill_points.kill_points_map[point] {
                builder.kill(local, point, SplitPointEffect::Late);
            }
        }

        // If this block ends in a return terminator, end all live ranges before
        // the terminator so that StorageDead statements are inserted before it.
        //
        // This is useful after inlining so that the lifetime of locals in the
        // inlined callee don't extend past the call in the caller.
        if let mir::TerminatorKind::Return = terminator.kind
            && !block_data.statements.is_empty()
        {
            // Blocks with only a return terminator are handled above.
            let location = Location { block, statement_index: block_data.statements.len() - 1 };
            let point = points.point_from_location(location);
            builder.kill_all_except(mir::RETURN_PLACE, point, SplitPointEffect::Late);
        }

        let location = Location { block, statement_index: block_data.statements.len() };
        let point = points.point_from_location(location);

        // Kill moved locals at the early point.
        VisitDeallocatedLocalsWith(|local, _| {
            builder.kill(local, point, SplitPointEffect::Early);
        })
        .visit_terminator(terminator, location);

        // Kill any locals which are no longer used after this terminator.
        for &(local, _) in kill_points.kill_points_map[point] {
            builder.kill(local, point, SplitPointEffect::Early);
        }

        // Gen destination places at the late point.
        VisitAllocatedLocalsWith(|local, _| {
            builder.gen_(local, point, SplitPointEffect::Late);
        })
        .visit_terminator(terminator, location);

        // Move arguments to a call are treated specially: the place that they
        // represent is passed directly to the callee, which means that they are
        // not allowed to alias any other move operand or the destination place.
        // This is represented here by extending their live range to the late
        // part, making it overlap with that of the destination place.
        //
        // Notably, this *doesn't* apply to TailCall.
        if let mir::TerminatorKind::Call {
            func: _,
            args,
            destination: _,
            target: _,
            unwind: _,
            call_source: _,
            fn_span: _,
        } = &terminator.kind
        {
            for arg in args {
                if let mir::Operand::Move(place) = arg.node
                    && !place.is_indirect()
                {
                    builder.gen_(place.local, point, SplitPointEffect::Late);
                }
            }
        }

        // End the lifetimes of all locals at the end of the block. Successor
        // blocks (which may not be contiguous in the index space!) will
        // initialize the lifetimes again from their entry state.
        builder.kill_all(point, SplitPointEffect::Late);
    }

    builder.matrix
}

pub fn dump_liveness_matrix<'tcx>(
    tcx: TyCtxt<'tcx>,
    body: &mir::Body<'tcx>,
    pass_name: &'static str,
    points: &DenseLocationMap,
    matrix: &SparseIntervalMatrix<Local, SplitPointIndex>,
) {
    let locals_live_at = |split_point| {
        matrix.rows().filter(|&r| matrix.contains(r, split_point)).collect::<Vec<_>>()
    };

    if let Some(dumper) = MirDumper::new(tcx, pass_name, body) {
        let extra_data = &|pass_where, w: &mut dyn std::io::Write| {
            if let PassWhere::BeforeLocation(loc) = pass_where {
                let point = points.point_from_location(loc);
                let split_point = SplitPointIndex::new(point, SplitPointEffect::Early);
                let live = locals_live_at(split_point);
                writeln!(w, "        // {loc:?}-early => {live:?}")?;
                let split_point = SplitPointIndex::new(point, SplitPointEffect::Late);
                let live = locals_live_at(split_point);
                writeln!(w, "        // {loc:?}-late => {live:?}")?;
            }
            Ok(())
        };

        dumper.set_extra_data(extra_data).dump_mir(body)
    }
}