DestinationPropagation: compute liveness as ranges instead of traveling bitsets

compiler/rustc_index/src/interval.rs

@@ -164,6 +164,38 @@ impl<I: Idx> IntervalSet<I> {
         );
     }
 
+    /// Specialized version of `insert_range` when we know that the inserted point is *after* any
+    /// contained.
+    pub fn append_range(&mut self, range: impl RangeBounds<I> + Clone) {
+        let start = inclusive_start(range.clone());
+        let Some(end) = inclusive_end(self.domain, range) else {
+            // empty range
+            return;
+        };
+        if start > end {
+            return;
+        }
+
+        let start = start.index() as u32;
+        let end = end.index() as u32;
+        if let Some((_, last_end)) = self.map.last_mut() {
+            assert!(*last_end <= start);
+            // The start is already adjacent to the set.
+            if start <= *last_end + 1 {
+                *last_end = end;
+            } else {
+                self.map.push((start, end));
+            }
+        } else {
+            self.map.push((start, end));
+        }
+
+        debug_assert!(
+            self.check_invariants(),
+            "wrong intervals after append {start:?}..={end:?} to {self:?}"
+        );
+    }
+
     pub fn contains(&self, needle: I) -> bool {
         let needle = needle.index() as u32;
         let Some(last) = self.map.partition_point(|r| r.0 <= needle).checked_sub(1) else {
@@ -379,6 +411,10 @@ impl<R: Idx, C: Step + Idx> SparseIntervalMatrix<R, C> {
         self.ensure_row(row).append(point)
     }
 
+    pub fn append_range(&mut self, row: R, point: impl RangeBounds<C> + Clone) {
+        self.ensure_row(row).append_range(point)
+    }
+
     pub fn contains(&self, row: R, point: C) -> bool {
         self.row(row).is_some_and(|r| r.contains(point))
     }

compiler/rustc_mir_dataflow/src/impls/liveness.rs

@@ -25,6 +25,12 @@ use crate::{Analysis, Backward, GenKill};
 /// [liveness]: https://en.wikipedia.org/wiki/Live_variable_analysis
 pub struct MaybeLiveLocals;
 
+impl MaybeLiveLocals {
+    pub fn transfer_function<I>(state: &mut I) -> TransferFunction<'_, I> {
+        TransferFunction(state)
+    }
+}
+
 impl<'tcx> Analysis<'tcx> for MaybeLiveLocals {
     type Domain = DenseBitSet<Local>;
     type Direction = Backward;
@@ -81,9 +87,12 @@ impl<'tcx> Analysis<'tcx> for MaybeLiveLocals {
     }
 }
 
-pub struct TransferFunction<'a>(pub &'a mut DenseBitSet<Local>);
+pub struct TransferFunction<'a, I>(pub &'a mut I);
 
-impl<'tcx> Visitor<'tcx> for TransferFunction<'_> {
+impl<'tcx, I> Visitor<'tcx> for TransferFunction<'_, I>
+where
+    I: GenKill<Local>,
+{
     fn visit_place(&mut self, place: &mir::Place<'tcx>, context: PlaceContext, location: Location) {
         if let PlaceContext::MutatingUse(MutatingUseContext::Yield) = context {
             // The resume place is evaluated and assigned to only after coroutine resumes, so its
@@ -143,7 +152,7 @@ pub enum DefUse {
 }
 
 impl DefUse {
-    fn apply(state: &mut DenseBitSet<Local>, place: Place<'_>, context: PlaceContext) {
+    fn apply(state: &mut impl GenKill<Local>, place: Place<'_>, context: PlaceContext) {
         match DefUse::for_place(place, context) {
             DefUse::Def => state.kill(place.local),
             DefUse::Use => state.gen_(place.local),

compiler/rustc_mir_transform/src/dest_prop.rs

@@ -146,7 +146,7 @@ use rustc_middle::mir::*;
 use rustc_middle::ty::TyCtxt;
 use rustc_mir_dataflow::impls::{DefUse, MaybeLiveLocals};
 use rustc_mir_dataflow::points::DenseLocationMap;
-use rustc_mir_dataflow::{Analysis, EntryStates};
+use rustc_mir_dataflow::{Analysis, EntryStates, GenKill};
 use tracing::{debug, trace};
 
 pub(super) struct DestinationPropagation;
@@ -502,25 +502,88 @@ impl TwoStepIndex {
 }
 
 /// Add points depending on the result of the given dataflow analysis.
+#[tracing::instrument(level = "trace", skip(elements, body))]
 fn save_as_intervals<'tcx>(
     elements: &DenseLocationMap,
     body: &Body<'tcx>,
     relevant: &RelevantLocals,
     entry_states: EntryStates<DenseBitSet<Local>>,
 ) -> SparseIntervalMatrix<RelevantLocal, TwoStepIndex> {
-    let mut values = SparseIntervalMatrix::new(2 * elements.num_points());
-    let mut state = MaybeLiveLocals.bottom_value(body);
-    let reachable_blocks = traversal::reachable_as_bitset(body);
+    /// Generalized dataflow state for use inside a given block.
+    struct GenKillIntervalMatrix<'a> {
+        values: SparseIntervalMatrix<RelevantLocal, TwoStepIndex>,
+        relevant: &'a RelevantLocals,
+        /// If a local is live, this stores the start of the live range.
+        /// If a local is dead, this stores `None`.
+        pending: IndexVec<RelevantLocal, Option<TwoStepIndex>>,
+        /// The current position of the cursor inside the MIR body.
+        current: TwoStepIndex,
+    }
 
-    let two_step_loc = |location, effect| TwoStepIndex::new(elements, location, effect);
-    let append_at =
-        |values: &mut SparseIntervalMatrix<_, _>, state: &DenseBitSet<Local>, twostep| {
-            for (relevant, &original) in relevant.original.iter_enumerated() {
-                if state.contains(original) {
-                    values.append(relevant, twostep);
+    impl GenKill<Local> for GenKillIntervalMatrix<'_> {
+        fn gen_(&mut self, elem: Local) {
+            let Some(elem) = self.relevant.shrink[elem] else { return };
+            // If the local was already live, do not overwrite the start position.
+            let _ = self.pending[elem].get_or_insert(self.current);
+        }
+
+        fn kill(&mut self, elem: Local) {
+            // Ensure we only kill for odd positions, so `insert_single` is well-behaved.
+            debug_assert!(self.current.as_u32() & 1 == 1);
+            let Some(elem) = self.relevant.shrink[elem] else { return };
+            if let Some(start) = self.pending[elem].take() {
+                debug_assert!(start <= self.current);
+                // The local is live since `start`.
+                // We are killing it, so it won't be after `current`, hence an exclusive range.
+                self.values.append_range(elem, start..self.current);
+            }
+        }
+    }
+
+    impl GenKillIntervalMatrix<'_> {
+        /// Insert a singleton range. This can be used for dead locals to mark conflicts, for
+        /// instance `move` operands in function calls or partial writes.
+        fn insert_single(&mut self, elem: RelevantLocal) {
+            // If we have a set pending, we will insert it when killing it, so nothing more to do.
+            // Kills only happen for odd positions, so we don't risk `kill` to insert
+            // a range excluding `self.current`.
+            debug_assert!(self.current.as_u32() & 1 == 0);
+            if self.pending[elem].is_none() {
+                self.values.append(elem, self.current);
+            }
+        }
+
+        fn start_block(&mut self, entry_state: &DenseBitSet<Local>) {
+            debug_assert!(self.pending.iter().all(Option::is_none));
+            for local in entry_state.iter() {
+                if let Some(elem) = self.relevant.shrink[local] {
+                    self.pending[elem] = Some(self.current);
                 }
             }
-        };
+        }
+
+        fn end_block(&mut self) {
+            for (elem, start) in self.pending.iter_enumerated_mut() {
+                if let Some(start) = start.take() {
+                    debug_assert!(start <= self.current);
+                    // We are ending a block, mark all live locals as live up to `current`,
+                    // including that position (which is still inside the block).
+                    self.values.append_range(elem, start..=self.current);
+                }
+            }
+        }
+    }
+
+    let reachable_blocks = traversal::reachable_as_bitset(body);
+    let two_step_loc = |location, effect| TwoStepIndex::new(elements, location, effect);
+
+    let mut state = GenKillIntervalMatrix {
+        values: SparseIntervalMatrix::new(2 * elements.num_points()),
+        relevant,
+        pending: IndexVec::from_elem(None, &relevant.original),
+        // Dummy value.
+        current: TwoStepIndex::from_u32(0),
+    };
 
     // Iterate blocks in decreasing order, to visit locations in decreasing order. This
     // allows to use the more efficient `append` method to interval sets.
@@ -529,14 +592,15 @@ fn save_as_intervals<'tcx>(
             continue;
         }
 
-        state.clone_from(&entry_states[block]);
-
         let block_data = &body.basic_blocks[block];
         let loc = Location { block, statement_index: block_data.statements.len() };
+        state.current = two_step_loc(loc, Effect::After);
+
+        // Setup the new block.
+        state.start_block(&entry_states[block]);
 
         let term = block_data.terminator();
-        let mut twostep = two_step_loc(loc, Effect::After);
-        append_at(&mut values, &state, twostep);
+
         // Ensure we have a non-zero live range even for dead stores. This is done by marking all
         // the written-to locals as live in the second half of the statement.
         // We also ensure that operands read by terminators conflict with writes by that terminator.
@@ -545,25 +609,23 @@ fn save_as_intervals<'tcx>(
             if let Some(relevant) = relevant.shrink[place.local] {
                 match DefUse::for_place(place, ctxt) {
                     DefUse::Def | DefUse::Use | DefUse::PartialWrite => {
-                        values.insert(relevant, twostep);
+                        state.insert_single(relevant);
                     }
                     DefUse::NonUse => {}
                 }
             }
         })
         .visit_terminator(term, loc);
 
-        twostep = TwoStepIndex::from_u32(twostep.as_u32() + 1);
-        debug_assert_eq!(twostep, two_step_loc(loc, Effect::Before));
-        MaybeLiveLocals.apply_early_terminator_effect(&mut state, term, loc);
-        MaybeLiveLocals.apply_primary_terminator_effect(&mut state, term, loc);
-        append_at(&mut values, &state, twostep);
+        state.current = state.current + 1;
+        debug_assert_eq!(state.current, two_step_loc(loc, Effect::Before));
+        MaybeLiveLocals::transfer_function(&mut state).visit_terminator(term, loc);
 
         for (statement_index, stmt) in block_data.statements.iter().enumerate().rev() {
             let loc = Location { block, statement_index };
-            twostep = TwoStepIndex::from_u32(twostep.as_u32() + 1);
-            debug_assert_eq!(twostep, two_step_loc(loc, Effect::After));
-            append_at(&mut values, &state, twostep);
+            state.current = state.current + 1;
+            debug_assert_eq!(state.current, two_step_loc(loc, Effect::After));
+
             // Like terminators, ensure we have a non-zero live range even for dead stores.
             // Some rvalues interleave reads and writes, for instance `Rvalue::Aggregate`, see
             // https://github.com/rust-lang/rust/issues/146383. By precaution, treat statements
@@ -582,26 +644,27 @@ fn save_as_intervals<'tcx>(
                 if let Some(relevant) = relevant.shrink[place.local] {
                     match DefUse::for_place(place, ctxt) {
                         DefUse::Def | DefUse::PartialWrite => {
-                            values.insert(relevant, twostep);
+                            state.insert_single(relevant);
                         }
                         DefUse::Use if !is_simple_assignment => {
-                            values.insert(relevant, twostep);
+                            state.insert_single(relevant);
                         }
                         DefUse::Use | DefUse::NonUse => {}
                     }
                 }
             })
             .visit_statement(stmt, loc);
 
-            twostep = TwoStepIndex::from_u32(twostep.as_u32() + 1);
-            debug_assert_eq!(twostep, two_step_loc(loc, Effect::Before));
-            MaybeLiveLocals.apply_early_statement_effect(&mut state, stmt, loc);
-            MaybeLiveLocals.apply_primary_statement_effect(&mut state, stmt, loc);
             // ... but reads from operands are marked as live here so they do not conflict with
             // the all the writes we manually marked as live in the second half of the statement.
-            append_at(&mut values, &state, twostep);
+            state.current = TwoStepIndex::from_u32(state.current.as_u32() + 1);
+            debug_assert_eq!(state.current, two_step_loc(loc, Effect::Before));
+            MaybeLiveLocals::transfer_function(&mut state).visit_statement(stmt, loc);
         }
+
+        // Cleanup the current block for the next one.
+        state.end_block();
     }
 
-    values
+    state.values
 }