Don't alloca for unused locals

Author: saethlin

compiler/rustc_codegen_ssa/src/mir/analyze.rs

@@ -15,13 +15,18 @@ use crate::traits::*;
 
 pub(crate) fn non_ssa_locals<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
     fx: &FunctionCx<'a, 'tcx, Bx>,
+    traversal_order: &[mir::BasicBlock],
+    reachable_locals: &BitSet<mir::Local>,
 ) -> BitSet<mir::Local> {
     let mir = fx.mir;
     let dominators = mir.basic_blocks.dominators();
     let locals = mir
         .local_decls
-        .iter()
-        .map(|decl| {
+        .iter_enumerated()
+        .map(|(local, decl)| {
+            if !reachable_locals.contains(local) {
+                return LocalKind::Unused;
+            }
             let ty = fx.monomorphize(decl.ty);
             let layout = fx.cx.spanned_layout_of(ty, decl.source_info.span);
             if layout.is_zst() {
@@ -44,7 +49,8 @@ pub(crate) fn non_ssa_locals<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
     // If there exists a local definition that dominates all uses of that local,
     // the definition should be visited first. Traverse blocks in an order that
     // is a topological sort of dominance partial order.
-    for (bb, data) in traversal::reverse_postorder(mir) {
+    for bb in traversal_order.iter().copied() {
+        let data = &mir.basic_blocks[bb];
         analyzer.visit_basic_block_data(bb, data);
     }
 

compiler/rustc_codegen_ssa/src/mir/mod.rs

@@ -192,6 +192,9 @@ pub fn codegen_mir<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
             })
             .collect();
 
+    let (traversal_order, reachable_locals) =
+        traversal::mono_reachable_reverse_postorder(mir, cx.tcx(), instance);
+
     let mut fx = FunctionCx {
         instance,
         mir,
@@ -218,7 +221,7 @@ pub fn codegen_mir<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
 
     fx.per_local_var_debug_info = fx.compute_per_local_var_debug_info(&mut start_bx);
 
-    let memory_locals = analyze::non_ssa_locals(&fx);
+    let memory_locals = analyze::non_ssa_locals(&fx, &traversal_order, &reachable_locals);
 
     // Allocate variable and temp allocas
     let local_values = {
@@ -277,17 +280,14 @@ pub fn codegen_mir<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
     // So drop the builder of `start_llbb` to avoid having two at the same time.
     drop(start_bx);
 
-    let reachable_blocks = traversal::mono_reachable_as_bitset(mir, cx.tcx(), instance);
-
-    // Codegen the body of each block using reverse postorder
-    for (bb, _) in traversal::reverse_postorder(mir) {
-        if reachable_blocks.contains(bb) {
-            fx.codegen_block(bb);
-        } else {
-            // We want to skip this block, because it's not reachable. But we still create
-            // the block so terminators in other blocks can reference it.
-            fx.codegen_block_as_unreachable(bb);
-        }
+    let mut unreached_blocks = BitSet::new_filled(mir.basic_blocks.len());
+    // Codegen the body of each reachable block using our reverse postorder list.
+    for bb in traversal_order {
+        fx.codegen_block(bb);
+        unreached_blocks.remove(bb);
+    }
+    for bb in unreached_blocks.iter() {
+        fx.codegen_block_as_unreachable(bb);
     }
 }
 

compiler/rustc_middle/src/mir/mod.rs

@@ -1424,6 +1424,19 @@ impl<'tcx> BasicBlockData<'tcx> {
     pub fn is_empty_unreachable(&self) -> bool {
         self.statements.is_empty() && matches!(self.terminator().kind, TerminatorKind::Unreachable)
     }
+
+    /// Like [`Terminator::successors`] but tries to use information available from the [`Instance`]
+    /// to skip successors like the `false` side of an `if const {`.
+    ///
+    /// This is used to implement [`traversal::mono_reachable`] and
+    /// [`traversal::mono_reachable_reverse_postorder`].
+    pub fn mono_successors(&self, tcx: TyCtxt<'tcx>, instance: Instance<'tcx>) -> Successors<'_> {
+        if let Some((bits, targets)) = Body::try_const_mono_switchint(tcx, instance, self) {
+            targets.successors_for_value(bits)
+        } else {
+            self.terminator().successors()
+        }
+    }
 }
 
 ///////////////////////////////////////////////////////////////////////////

compiler/rustc_middle/src/mir/terminator.rs

@@ -413,6 +413,17 @@ mod helper {
     use super::*;
     pub type Successors<'a> = impl DoubleEndedIterator<Item = BasicBlock> + 'a;
     pub type SuccessorsMut<'a> = impl DoubleEndedIterator<Item = &'a mut BasicBlock> + 'a;
+
+    impl SwitchTargets {
+        /// Like [`SwitchTargets::target_for_value`], but returning the same type as
+        /// [`Terminator::successors`].
+        #[inline]
+        pub fn successors_for_value(&self, value: u128) -> Successors<'_> {
+            let target = self.target_for_value(value);
+            (&[]).into_iter().copied().chain(Some(target))
+        }
+    }
+
     impl<'tcx> TerminatorKind<'tcx> {
         #[inline]
         pub fn successors(&self) -> Successors<'_> {

compiler/rustc_middle/src/mir/traversal.rs

@@ -1,4 +1,5 @@
 use super::*;
+use crate::mir::visit::Visitor;
 
 /// Preorder traversal of a graph.
 ///
@@ -232,6 +233,90 @@ pub fn postorder<'a, 'tcx>(
     reverse_postorder(body).rev()
 }
 
+struct UsedLocals(BitSet<Local>);
+
+impl<'tcx> Visitor<'tcx> for UsedLocals {
+    fn visit_local(
+        &mut self,
+        local: Local,
+        _ctx: crate::mir::visit::PlaceContext,
+        _location: Location,
+    ) {
+        self.0.insert(local);
+    }
+}
+
+struct MonoReachablePostorder<'a, 'tcx> {
+    basic_blocks: &'a IndexSlice<BasicBlock, BasicBlockData<'tcx>>,
+    visited: BitSet<BasicBlock>,
+    visit_stack: Vec<(BasicBlock, Successors<'a>)>,
+    locals: UsedLocals,
+    tcx: TyCtxt<'tcx>,
+    instance: Instance<'tcx>,
+}
+
+impl<'a, 'tcx> MonoReachablePostorder<'a, 'tcx> {
+    fn new(
+        body: &'a Body<'tcx>,
+        tcx: TyCtxt<'tcx>,
+        instance: Instance<'tcx>,
+    ) -> MonoReachablePostorder<'a, 'tcx> {
+        let basic_blocks = &body.basic_blocks;
+        let mut po = MonoReachablePostorder {
+            basic_blocks,
+            visited: BitSet::new_empty(basic_blocks.len()),
+            visit_stack: Vec::new(),
+            locals: UsedLocals(BitSet::new_empty(body.local_decls.len())),
+            tcx,
+            instance,
+        };
+
+        po.visit(START_BLOCK);
+        po.traverse_successor();
+        po
+    }
+
+    fn visit(&mut self, bb: BasicBlock) {
+        if !self.visited.insert(bb) {
+            return;
+        }
+        let data = &self.basic_blocks[bb];
+        self.locals.visit_basic_block_data(bb, data);
+        let successors = data.mono_successors(self.tcx, self.instance);
+        self.visit_stack.push((bb, successors));
+    }
+
+    fn traverse_successor(&mut self) {
+        while let Some(bb) = self.visit_stack.last_mut().and_then(|(_, iter)| iter.next_back()) {
+            self.visit(bb);
+        }
+    }
+}
+
+impl<'tcx> Iterator for MonoReachablePostorder<'_, 'tcx> {
+    type Item = BasicBlock;
+
+    fn next(&mut self) -> Option<BasicBlock> {
+        let (bb, _) = self.visit_stack.pop()?;
+        self.traverse_successor();
+        Some(bb)
+    }
+}
+
+pub fn mono_reachable_reverse_postorder<'a, 'tcx>(
+    body: &'a Body<'tcx>,
+    tcx: TyCtxt<'tcx>,
+    instance: Instance<'tcx>,
+) -> (Vec<BasicBlock>, BitSet<Local>) {
+    let mut iter = MonoReachablePostorder::new(body, tcx, instance);
+    let mut items = Vec::with_capacity(body.basic_blocks.len());
+    while let Some(block) = iter.next() {
+        items.push(block);
+    }
+    items.reverse();
+    (items, iter.locals.0)
+}
+
 /// Returns an iterator over all basic blocks reachable from the `START_BLOCK` in no particular
 /// order.
 ///
@@ -358,14 +443,8 @@ impl<'a, 'tcx> Iterator for MonoReachable<'a, 'tcx> {
 
             let data = &self.body[idx];
 
-            if let Some((bits, targets)) =
-                Body::try_const_mono_switchint(self.tcx, self.instance, data)
-            {
-                let target = targets.target_for_value(bits);
-                self.add_work([target]);
-            } else {
-                self.add_work(data.terminator().successors());
-            }
+            let targets = data.mono_successors(self.tcx, self.instance);
+            self.add_work(targets);
 
             return Some((idx, data));
         }

tests/codegen/constant-branch.rs

@@ -41,10 +41,8 @@ pub fn if_constant_match() {
         _ => 4,
     };
 
-    // CHECK: br label %[[MINUS1:.+]]
+    // CHECK: br label %{{.+}}
     _ = match -1 {
-        // CHECK: [[MINUS1]]:
-        // CHECK: store i32 1
         -1 => 1,
         _ => 0,
     }

tests/codegen/no-alloca-inside-if-false.rs

@@ -0,0 +1,23 @@
+//@ compile-flags: -Cno-prepopulate-passes -Copt-level=0
+
+#![crate_type = "lib"]
+
+#[inline(never)]
+fn test<const SIZE: usize>() {
+    // CHECK-LABEL: no_alloca_inside_if_false::test
+    // CHECK: start:
+    // CHECK-NEXT: %0 = alloca
+    // CHECK-NEXT: %vec = alloca
+    // CHECK-NOT: %arr = alloca
+    if const { SIZE < 4096 } {
+        let arr = [0u8; SIZE];
+        std::hint::black_box(&arr);
+    } else {
+        let vec = vec![0u8; SIZE];
+        std::hint::black_box(&vec);
+    }
+}
+
+pub fn main() {
+    test::<8192>();
+}