Auto merge of #129283 - saethlin:unreachable-allocas, r=<try>

Don't alloca for unused locals

We already have a concept of mono-unreachable basic blocks; this is primarily useful for ensuring that we do not compile code under an `if false`. But since we never gave locals the same analysis, a large local only used under an `if false` will still have stack space allocated for it.

There are 3 places we traverse MIR during monomorphization: Inside the collector, `non_ssa_locals`, and the walk to generate code. Unfortunately, #129283 (comment) indicates that we cannot afford the expense of tracking reachable locals during the collector's traversal, so we do need at least two mono-reachable traversals. And of course caching is of no help here because the benchmarks that regress are incr-unchanged; they don't do any codegen.

This fixes the second problem in #129282, and brings us anther step toward `const if` at home.

compiler/rustc_codegen_ssa/src/mir/analyze.rs

@@ -15,13 +15,18 @@ use crate::traits::*;
 
 pub 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 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/terminator.rs

@@ -413,6 +413,15 @@ 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 {
+        #[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

@@ -232,6 +232,155 @@ pub fn postorder<'a, 'tcx>(
     reverse_postorder(body).rev()
 }
 
+pub struct MonoReachablePostorder<'a, 'tcx> {
+    basic_blocks: &'a IndexSlice<BasicBlock, BasicBlockData<'tcx>>,
+    visited: BitSet<BasicBlock>,
+    visit_stack: Vec<(BasicBlock, Successors<'a>)>,
+    locals: BitSet<Local>,
+    tcx: TyCtxt<'tcx>,
+    instance: Instance<'tcx>,
+}
+
+impl<'a, 'tcx> MonoReachablePostorder<'a, 'tcx> {
+    pub 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: BitSet::new_empty(body.local_decls.len()),
+            tcx,
+            instance,
+        };
+
+        let root = START_BLOCK;
+        let data = &po.basic_blocks[root];
+
+        UsedLocals { locals: &mut po.locals }.visit_basic_block_data(root, data);
+        if let Some(ref term) = data.terminator {
+            po.visited.insert(root);
+
+            let successors = if let Some((bits, targets)) =
+                Body::try_const_mono_switchint(tcx, instance, data)
+            {
+                targets.successors_for_value(bits)
+            } else {
+                term.successors()
+            };
+
+            po.visit_stack.push((root, successors));
+            po.traverse_successor();
+        }
+
+        po
+    }
+
+    fn traverse_successor(&mut self) {
+        // This is quite a complex loop due to 1. the borrow checker not liking it much
+        // and 2. what exactly is going on is not clear
+        //
+        // It does the actual traversal of the graph, while the `next` method on the iterator
+        // just pops off of the stack. `visit_stack` is a stack containing pairs of nodes and
+        // iterators over the successors of those nodes. Each iteration attempts to get the next
+        // node from the top of the stack, then pushes that node and an iterator over the
+        // successors to the top of the stack. This loop only grows `visit_stack`, stopping when
+        // we reach a child that has no children that we haven't already visited.
+        //
+        // For a graph that looks like this:
+        //
+        //         A
+        //        / \
+        //       /   \
+        //      B     C
+        //      |     |
+        //      |     |
+        //      |     D
+        //       \   /
+        //        \ /
+        //         E
+        //
+        // The state of the stack starts out with just the root node (`A` in this case);
+        //     [(A, [B, C])]
+        //
+        // When the first call to `traverse_successor` happens, the following happens:
+        //
+        //     [(C, [D]),  // `C` taken from the successors of `A`, pushed to the
+        //                 // top of the stack along with the successors of `C`
+        //      (A, [B])]
+        //
+        //     [(D, [E]),  // `D` taken from successors of `C`, pushed to stack
+        //      (C, []),
+        //      (A, [B])]
+        //
+        //     [(E, []),   // `E` taken from successors of `D`, pushed to stack
+        //      (D, []),
+        //      (C, []),
+        //      (A, [B])]
+        //
+        // Now that the top of the stack has no successors we can traverse, each item will
+        // be popped off during iteration until we get back to `A`. This yields [E, D, C].
+        //
+        // When we yield `C` and call `traverse_successor`, we push `B` to the stack, but
+        // since we've already visited `E`, that child isn't added to the stack. The last
+        // two iterations yield `B` and finally `A` for a final traversal of [E, D, C, B, A]
+        while let Some(bb) = self.visit_stack.last_mut().and_then(|(_, iter)| iter.next_back()) {
+            if self.visited.insert(bb) {
+                let data = &self.basic_blocks[bb];
+                UsedLocals { locals: &mut self.locals }.visit_basic_block_data(bb, data);
+
+                let Some(term) = &data.terminator else {
+                    continue;
+                };
+
+                let successors = if let Some((bits, targets)) =
+                    Body::try_const_mono_switchint(self.tcx, self.instance, data)
+                {
+                    targets.successors_for_value(bits)
+                } else {
+                    term.successors()
+                };
+
+                self.visit_stack.push((bb, successors));
+            }
+        }
+    }
+}
+
+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)
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        // All the blocks, minus the number of blocks we've visited.
+        let remaining = self.basic_blocks.len() - self.visited.count();
+        (remaining, Some(remaining))
+    }
+}
+
+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)
+}
+
 /// Returns an iterator over all basic blocks reachable from the `START_BLOCK` in no particular
 /// order.
 ///
@@ -320,6 +469,22 @@ pub struct MonoReachable<'a, 'tcx> {
     worklist: BitSet<BasicBlock>,
 }
 
+struct UsedLocals<'a> {
+    locals: &'a mut BitSet<Local>,
+}
+
+use crate::mir::visit::Visitor;
+impl<'a, 'tcx> Visitor<'tcx> for UsedLocals<'a> {
+    fn visit_local(
+        &mut self,
+        local: Local,
+        _ctx: crate::mir::visit::PlaceContext,
+        _location: Location,
+    ) {
+        self.locals.insert(local);
+    }
+}
+
 impl<'a, 'tcx> MonoReachable<'a, 'tcx> {
     pub fn new(
         body: &'a Body<'tcx>,

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>();
+}