Lazily calculate MaybeTransitiveLiveLocals

Author: dianqk

compiler/rustc_mir_transform/src/merge_branches.rs

@@ -24,11 +24,7 @@ impl<'tcx> crate::MirPass<'tcx> for MergeBranchSimplification {
         let param_env = tcx.param_env_reveal_all_normalized(def_id);
 
         let borrowed_locals = borrowed_locals(body);
-        let mut maybe_live: ResultsCursor<'_, '_, MaybeTransitiveLiveLocals<'_>> =
-            MaybeTransitiveLiveLocals::new(&borrowed_locals)
-                .into_engine(tcx, body)
-                .iterate_to_fixpoint()
-                .into_results_cursor(body);
+        let mut stmt_live_result = StatementLiveResult::new(tcx, body, &borrowed_locals);
         for i in 0..body.basic_blocks.len() {
             let bbs = &*body.basic_blocks;
             let switch_bb_idx = BasicBlock::from_usize(i);
@@ -75,8 +71,7 @@ impl<'tcx> crate::MirPass<'tcx> for MergeBranchSimplification {
                 targets,
                 src_place,
                 src_ty,
-                &borrowed_locals,
-                &mut maybe_live,
+                &mut stmt_live_result,
             ) {
                 let statement_index = bbs[switch_bb_idx].statements.len();
                 let parent_end = Location { block: switch_bb_idx, statement_index };
@@ -93,6 +88,33 @@ impl<'tcx> crate::MirPass<'tcx> for MergeBranchSimplification {
     }
 }
 
+struct StatementLiveResult<'tcx, 'mir, 'a> {
+    tcx: TyCtxt<'tcx>,
+    body: &'mir Body<'tcx>,
+    result: Option<ResultsCursor<'mir, 'tcx, MaybeTransitiveLiveLocals<'a>>>,
+    borrowed_locals: &'a BitSet<Local>,
+}
+
+impl<'tcx, 'mir, 'a> StatementLiveResult<'tcx, 'mir, 'a> {
+    fn new(tcx: TyCtxt<'tcx>, body: &'mir Body<'tcx>, borrowed_locals: &'a BitSet<Local>) -> Self {
+        Self { tcx, body, result: None, borrowed_locals }
+    }
+
+    fn is_live(&mut self, loc: Location, local: Local) -> bool {
+        if self.borrowed_locals.contains(local) {
+            return true;
+        }
+        let maybe_live = self.result.get_or_insert_with(|| {
+            MaybeTransitiveLiveLocals::new(&self.borrowed_locals)
+                .into_engine(self.tcx, self.body)
+                .iterate_to_fixpoint()
+                .into_results_cursor(self.body)
+        });
+        maybe_live.seek_before_primary_effect(loc);
+        maybe_live.get().contains(local)
+    }
+}
+
 /// The GVN simplified
 /// ```ignore (syntax-highlighting-only)
 /// match a {
@@ -116,22 +138,11 @@ fn can_simplify_to_copy<'tcx>(
     targets: &SwitchTargets,
     src_place: Place<'tcx>,
     src_ty: tcx::PlaceTy<'tcx>,
-    borrowed_locals: &BitSet<Local>,
-    maybe_live: &mut ResultsCursor<'_, 'tcx, MaybeTransitiveLiveLocals<'_>>,
+    stmt_live_result: &mut StatementLiveResult<'tcx, '_, '_>,
 ) -> Option<Place<'tcx>> {
     let mut targets_iter = targets.iter();
     let dest_place = targets_iter.next().and_then(|(index, target)| {
-        find_copy_assign(
-            tcx,
-            param_env,
-            body,
-            index,
-            target,
-            src_place,
-            src_ty,
-            borrowed_locals,
-            maybe_live,
-        )
+        find_copy_assign(tcx, param_env, body, index, target, src_place, src_ty, stmt_live_result)
     })?;
     let dest_ty = dest_place.ty(body.local_decls(), tcx);
     if dest_ty.ty != src_ty.ty || dest_ty.variant_index.is_some() {
@@ -147,8 +158,7 @@ fn can_simplify_to_copy<'tcx>(
                 other_target,
                 src_place,
                 src_ty,
-                borrowed_locals,
-                maybe_live,
+                stmt_live_result,
             )
     }) {
         return None;
@@ -164,8 +174,7 @@ fn find_copy_assign<'tcx>(
     target_block: BasicBlock,
     src_place: Place<'tcx>,
     src_ty: tcx::PlaceTy<'tcx>,
-    borrowed_locals: &BitSet<Local>,
-    maybe_live: &mut ResultsCursor<'_, 'tcx, MaybeTransitiveLiveLocals<'_>>,
+    stmt_live_result: &mut StatementLiveResult<'tcx, '_, '_>,
 ) -> Option<Place<'tcx>> {
     let statements = &body.basic_blocks[target_block].statements;
     if statements.is_empty() {
@@ -190,11 +199,10 @@ fn find_copy_assign<'tcx>(
                 StatementKind::Assign(box (dest_place, _))
                 | StatementKind::SetDiscriminant { place: box dest_place, .. }
                 | StatementKind::Deinit(box dest_place) => {
-                    if dest_place.is_indirect() || borrowed_locals.contains(dest_place.local) {
+                    if dest_place.is_indirect() {
                         return None;
                     }
-                    maybe_live.seek_before_primary_effect(loc);
-                    if !maybe_live.get().contains(dest_place.local) {
+                    if !stmt_live_result.is_live(loc, dest_place.local) {
                         lived_stmts.remove(statement_index);
                     } else if matches!(statement.kind, StatementKind::Assign(_))
                         && expected_assign_stmt.is_none()