reuse allocation used by always_storage_live_locals bitset

Author: klensy

compiler/rustc_const_eval/src/interpret/eval_context.rs

@@ -5,6 +5,7 @@ use std::mem;
 use either::{Either, Left, Right};
 
 use rustc_hir::{self as hir, def_id::DefId, definitions::DefPathData};
+use rustc_index::bit_set::GrowableBitSet;
 use rustc_index::vec::IndexVec;
 use rustc_middle::mir;
 use rustc_middle::mir::interpret::{ErrorHandled, InterpError, InvalidProgramInfo};
@@ -46,6 +47,9 @@ pub struct InterpCx<'mir, 'tcx, M: Machine<'mir, 'tcx>> {
 
     /// The recursion limit (cached from `tcx.recursion_limit(())`)
     pub recursion_limit: Limit,
+
+    // reuse allocation for bit set
+    always_live_locals_cache: GrowableBitSet<mir::Local>,
 }
 
 // The Phantomdata exists to prevent this type from being `Send`. If it were sent across a thread
@@ -408,6 +412,7 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
             param_env,
             memory: Memory::new(),
             recursion_limit: tcx.recursion_limit(),
+            always_live_locals_cache: Default::default(),
         }
     }
 
@@ -705,13 +710,19 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
         let mut locals = IndexVec::from_elem(dummy, &body.local_decls);
 
         // Now mark those locals as live that have no `Storage*` annotations.
-        let always_live = always_storage_live_locals(self.body());
+        // and take cached allocation for bitset ...
+        let mut always_live = mem::take(&mut self.always_live_locals_cache);
+        always_storage_live_locals(self.body(), &mut always_live);
+
         for local in locals.indices() {
             if always_live.contains(local) {
                 locals[local].value = LocalValue::Live(Operand::Immediate(Immediate::Uninit));
             }
         }
         // done
+        // ... and place it back
+        self.always_live_locals_cache = always_live;
+
         self.frame_mut().locals = locals;
         M::after_stack_push(self)?;
         self.frame_mut().loc = Left(mir::Location::START);

compiler/rustc_const_eval/src/transform/validate.rs

@@ -1,7 +1,7 @@
 //! Validates the MIR to ensure that invariants are upheld.
 
 use rustc_data_structures::fx::{FxHashMap, FxHashSet};
-use rustc_index::bit_set::BitSet;
+use rustc_index::bit_set::{BitSet, GrowableBitSet};
 use rustc_index::vec::IndexVec;
 use rustc_infer::traits::Reveal;
 use rustc_middle::mir::interpret::Scalar;
@@ -52,11 +52,13 @@ impl<'tcx> MirPass<'tcx> for Validator {
             Reveal::All => tcx.param_env_reveal_all_normalized(def_id),
         };
 
-        let always_live_locals = always_storage_live_locals(body);
-        let storage_liveness = MaybeStorageLive::new(std::borrow::Cow::Owned(always_live_locals))
-            .into_engine(tcx, body)
-            .iterate_to_fixpoint()
-            .into_results_cursor(body);
+        let mut always_live_locals = GrowableBitSet::new_empty();
+        always_storage_live_locals(body, &mut always_live_locals);
+        let storage_liveness =
+            MaybeStorageLive::new(std::borrow::Cow::Owned(always_live_locals.into()))
+                .into_engine(tcx, body)
+                .iterate_to_fixpoint()
+                .into_results_cursor(body);
 
         let mut checker = TypeChecker {
             when: &self.when,

compiler/rustc_index/src/bit_set.rs

@@ -1552,6 +1552,17 @@ impl<T: Idx> GrowableBitSet<T> {
     pub fn len(&self) -> usize {
         self.bit_set.count()
     }
+
+    // reuse allocation, set domain_size and fill
+    pub fn fill(&mut self, domain_size: usize) {
+        self.ensure(domain_size);
+        self.bit_set.domain_size = domain_size;
+        self.bit_set.insert_all();
+    }
+
+    pub fn as_bitset(&self) -> &BitSet<T> {
+        &self.bit_set
+    }
 }
 
 impl<T: Idx> From<BitSet<T>> for GrowableBitSet<T> {

compiler/rustc_mir_dataflow/src/storage.rs

@@ -1,11 +1,12 @@
-use rustc_index::bit_set::BitSet;
+use rustc_index::bit_set::GrowableBitSet;
 use rustc_middle::mir::{self, Local};
 
 /// The set of locals in a MIR body that do not have `StorageLive`/`StorageDead` annotations.
 ///
 /// These locals have fixed storage for the duration of the body.
-pub fn always_storage_live_locals(body: &mir::Body<'_>) -> BitSet<Local> {
-    let mut always_live_locals = BitSet::new_filled(body.local_decls.len());
+pub fn always_storage_live_locals(body: &mir::Body<'_>, locals: &mut GrowableBitSet<Local>) {
+    locals.fill(body.local_decls.len());
+    let always_live_locals = locals;
 
     for block in &*body.basic_blocks {
         for statement in &block.statements {
@@ -15,6 +16,4 @@ pub fn always_storage_live_locals(body: &mir::Body<'_>) -> BitSet<Local> {
             }
         }
     }
-
-    always_live_locals
 }

compiler/rustc_mir_transform/src/generator.rs

@@ -1401,9 +1401,11 @@ pub(crate) fn mir_generator_witnesses<'tcx>(
     };
 
     // When first entering the generator, move the resume argument into its new local.
-    let always_live_locals = always_storage_live_locals(&body);
+    let mut always_live_locals = GrowableBitSet::new_empty();
+    always_storage_live_locals(&body, &mut always_live_locals);
 
-    let liveness_info = locals_live_across_suspend_points(tcx, body, &always_live_locals, movable);
+    let liveness_info =
+        locals_live_across_suspend_points(tcx, body, always_live_locals.as_bitset(), movable);
 
     // Extract locals which are live across suspension point into `layout`
     // `remap` gives a mapping from local indices onto generator struct indices
@@ -1493,10 +1495,11 @@ impl<'tcx> MirPass<'tcx> for StateTransform {
             },
         );
 
-        let always_live_locals = always_storage_live_locals(&body);
+        let mut always_live_locals = GrowableBitSet::new_empty();
+        always_storage_live_locals(&body, &mut always_live_locals);
 
         let liveness_info =
-            locals_live_across_suspend_points(tcx, body, &always_live_locals, movable);
+            locals_live_across_suspend_points(tcx, body, always_live_locals.as_bitset(), movable);
 
         if tcx.sess.opts.unstable_opts.validate_mir {
             let mut vis = EnsureGeneratorFieldAssignmentsNeverAlias {
@@ -1533,7 +1536,7 @@ impl<'tcx> MirPass<'tcx> for StateTransform {
             state_substs,
             remap,
             storage_liveness,
-            always_live_locals,
+            always_live_locals: always_live_locals.into(),
             suspension_points: Vec::new(),
             new_ret_local,
             discr_ty,