Auto merge of rust-lang#138582 - scottmcm:option-ssa-2, r=<try>

Don't require `alloca`s for consuming simple enums

For example, if you pass an `Option<u32>` to a function, don't immediately write it to an `alloca` then read it again.

Author: bors

compiler/rustc_codegen_ssa/src/mir/analyze.rs

@@ -5,8 +5,8 @@ use rustc_data_structures::graph::dominators::Dominators;
 use rustc_index::bit_set::DenseBitSet;
 use rustc_index::{IndexSlice, IndexVec};
 use rustc_middle::mir::visit::{MutatingUseContext, NonMutatingUseContext, PlaceContext, Visitor};
-use rustc_middle::mir::{self, DefLocation, Location, TerminatorKind, traversal};
-use rustc_middle::ty::layout::{HasTyCtxt, LayoutOf};
+use rustc_middle::mir::{self, DefLocation, Location, PlaceElem, TerminatorKind, traversal};
+use rustc_middle::ty::layout::LayoutOf;
 use rustc_middle::{bug, span_bug};
 use tracing::debug;
 
@@ -96,65 +96,92 @@ impl<'a, 'b, 'tcx, Bx: BuilderMethods<'b, 'tcx>> LocalAnalyzer<'a, 'b, 'tcx, Bx>
     fn process_place(
         &mut self,
         place_ref: &mir::PlaceRef<'tcx>,
-        context: PlaceContext,
+        mut context: PlaceContext,
         location: Location,
     ) {
-        let cx = self.fx.cx;
-
-        if let Some((place_base, elem)) = place_ref.last_projection() {
-            let mut base_context = if context.is_mutating_use() {
-                PlaceContext::MutatingUse(MutatingUseContext::Projection)
-            } else {
-                PlaceContext::NonMutatingUse(NonMutatingUseContext::Projection)
-            };
-
-            // Allow uses of projections that are ZSTs or from scalar fields.
-            let is_consume = matches!(
-                context,
-                PlaceContext::NonMutatingUse(
-                    NonMutatingUseContext::Copy | NonMutatingUseContext::Move,
-                )
+        let maybe_local = if place_ref.is_indirect_first_projection() {
+            // After we deref a pointer, the local *of that pointer* is no
+            // longer interesting for the rest of the projection chain.
+            self.visit_local(
+                place_ref.local,
+                PlaceContext::NonMutatingUse(NonMutatingUseContext::Copy),
+                location,
             );
-            if is_consume {
-                let base_ty = place_base.ty(self.fx.mir, cx.tcx());
-                let base_ty = self.fx.monomorphize(base_ty);
-
-                // ZSTs don't require any actual memory access.
-                let elem_ty = base_ty.projection_ty(cx.tcx(), self.fx.monomorphize(elem)).ty;
-                let span = self.fx.mir.local_decls[place_ref.local].source_info.span;
-                if cx.spanned_layout_of(elem_ty, span).is_zst() {
-                    return;
-                }
-
-                if let mir::ProjectionElem::Field(..) = elem {
-                    let layout = cx.spanned_layout_of(base_ty.ty, span);
-                    if cx.is_backend_immediate(layout) || cx.is_backend_scalar_pair(layout) {
-                        // Recurse with the same context, instead of `Projection`,
-                        // potentially stopping at non-operand projections,
-                        // which would trigger `not_ssa` on locals.
-                        base_context = context;
-                    }
-                }
-            }
-
-            if let mir::ProjectionElem::Deref = elem {
-                // Deref projections typically only read the pointer.
-                base_context = PlaceContext::NonMutatingUse(NonMutatingUseContext::Copy);
-            }
+            None
+        } else {
+            Some(place_ref.local)
+        };
 
-            self.process_place(&place_base, base_context, location);
-            // HACK(eddyb) this emulates the old `visit_projection_elem`, this
-            // entire `visit_place`-like `process_place` method should be rewritten,
-            // now that we have moved to the "slice of projections" representation.
-            if let mir::ProjectionElem::Index(local) = elem {
+        let mut projection: &[PlaceElem<'tcx>] = place_ref.projection;
+        loop {
+            // Index projections are the only ones with another local, so handle
+            // that special case before the normal projection match.
+            if let [PlaceElem::Index(index_local), ..] = *projection {
                 self.visit_local(
-                    local,
+                    index_local,
                     PlaceContext::NonMutatingUse(NonMutatingUseContext::Copy),
                     location,
                 );
             }
-        } else {
-            self.visit_local(place_ref.local, context, location);
+
+            projection = match projection {
+                // No more projections means we're done looping.
+                [] => break,
+                // The only deref allowed in a Runtime-phase place is at the
+                // beginning, which we checked before the loop.
+                [PlaceElem::Deref, rest @ ..] => {
+                    assert_eq!(maybe_local, None);
+                    rest
+                }
+                // Making SSA locals useful for non-primitives heavily depends on
+                // not forcing stack allocation for basic newtypes and simple
+                // enums like `Option<u32>` or `Result<bool, Box<MyError>>`.
+                [PlaceElem::Downcast { .. }, PlaceElem::Field { .. }, rest @ ..]
+                | [PlaceElem::Field { .. }, rest @ ..] => {
+                    if let PlaceContext::NonMutatingUse(
+                        NonMutatingUseContext::Copy
+                        | NonMutatingUseContext::Move
+                        | NonMutatingUseContext::Inspect,
+                    ) = context
+                    {
+                        // Reading fields (or pseudo-fields) in operands can stay SSA
+                    } else {
+                        // But storing into a projection needs memory, especially for function returns
+                        context = PlaceContext::MutatingUse(MutatingUseContext::Projection);
+                    }
+                    rest
+                }
+                [PlaceElem::Downcast { .. }, ..] => {
+                    span_bug!(self.fx.mir.span, "Non-field downcast in {place_ref:?}");
+                }
+                // FIXME: These are type-changing, but not layout-affecting, so
+                // they probably needn't force memory, but for now they do since
+                // `maybe_codegen_consume_direct` doesn't handle them.
+                [
+                    PlaceElem::OpaqueCast { .. }
+                    | PlaceElem::UnwrapUnsafeBinder { .. }
+                    | PlaceElem::Subtype { .. },
+                    rest @ ..,
+                ] => {
+                    context = PlaceContext::MutatingUse(MutatingUseContext::Projection);
+                    rest
+                }
+                // The various types of indexing use address arithmetic, so we
+                // need to force the local to Memory like a borrow would.
+                [
+                    PlaceElem::Index { .. }
+                    | PlaceElem::ConstantIndex { .. }
+                    | PlaceElem::Subslice { .. },
+                    rest @ ..,
+                ] => {
+                    context = PlaceContext::MutatingUse(MutatingUseContext::Projection);
+                    rest
+                }
+            };
+        }
+
+        if let Some(local) = maybe_local {
+            self.visit_local(local, context, location);
         }
     }
 }

compiler/rustc_codegen_ssa/src/mir/mod.rs

@@ -135,6 +135,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
     }
 }
 
+#[derive(Debug)]
 enum LocalRef<'tcx, V> {
     Place(PlaceRef<'tcx, V>),
     /// `UnsizedPlace(p)`: `p` itself is a thin pointer (indirect place).

compiler/rustc_codegen_ssa/src/mir/operand.rs

@@ -142,7 +142,7 @@ pub struct OperandRef<'tcx, V> {
     pub layout: TyAndLayout<'tcx>,
 }
 
-impl<V: CodegenObject> fmt::Debug for OperandRef<'_, V> {
+impl<V: fmt::Debug> fmt::Debug for OperandRef<'_, V> {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         write!(f, "OperandRef({:?} @ {:?})", self.val, self.layout)
     }
@@ -372,16 +372,22 @@ impl<'a, 'tcx, V: CodegenObject> OperandRef<'tcx, V> {
                 )
             })
         } else {
-            let (in_scalar, imm) = match (self.val, self.layout.backend_repr) {
+            let (imm, in_scalar, in_bty) = match (self.val, self.layout.backend_repr) {
                 // Extract a scalar component from a pair.
                 (OperandValue::Pair(a_llval, b_llval), BackendRepr::ScalarPair(a, b)) => {
-                    if offset.bytes() == 0 {
+                    // This needs to look at `offset`, rather than `i`, because
+                    // for a type like `Option<u32>`, the first thing in the pair
+                    // is the tag, so `(_2 as Some).0` needs to read the *second*
+                    // thing in the pair despite it being "field zero".
+                    if offset == Size::ZERO {
                         assert_eq!(field.size, a.size(bx.cx()));
-                        (Some(a), a_llval)
+                        let bty = bx.scalar_pair_element_backend_type(self.layout, 0, false);
+                        (a_llval, a, bty)
                     } else {
                         assert_eq!(offset, a.size(bx.cx()).align_to(b.align(bx.cx()).abi));
                         assert_eq!(field.size, b.size(bx.cx()));
-                        (Some(b), b_llval)
+                        let bty = bx.scalar_pair_element_backend_type(self.layout, 1, false);
+                        (b_llval, b, bty)
                     }
                 }
 
@@ -392,23 +398,13 @@ impl<'a, 'tcx, V: CodegenObject> OperandRef<'tcx, V> {
             OperandValue::Immediate(match field.backend_repr {
                 BackendRepr::SimdVector { .. } => imm,
                 BackendRepr::Scalar(out_scalar) => {
-                    let Some(in_scalar) = in_scalar else {
-                        span_bug!(
-                            fx.mir.span,
-                            "OperandRef::extract_field({:?}): missing input scalar for output scalar",
-                            self
-                        )
-                    };
-                    if in_scalar != out_scalar {
-                        // If the backend and backend_immediate types might differ,
-                        // flip back to the backend type then to the new immediate.
-                        // This avoids nop truncations, but still handles things like
-                        // Bools in union fields needs to be truncated.
-                        let backend = bx.from_immediate(imm);
-                        bx.to_immediate_scalar(backend, out_scalar)
-                    } else {
-                        imm
-                    }
+                    // For a type like `Result<usize, &u32>` the layout is `Pair(i64, ptr)`.
+                    // But if we're reading the `Ok` payload, we need to turn that `ptr`
+                    // back into an integer. To avoid repeating logic we do that by
+                    // calling the transmute code, which is legal thanks to the size
+                    // assert we did when pulling it out of the pair.
+                    let out_bty = bx.backend_type(field);
+                    fx.transmute_immediate(bx, imm, in_scalar, in_bty, out_scalar, out_bty)
                 }
                 BackendRepr::ScalarPair(_, _) | BackendRepr::Memory { .. } => bug!(),
             })
@@ -712,7 +708,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
             LocalRef::Operand(mut o) => {
                 // Moves out of scalar and scalar pair fields are trivial.
                 for elem in place_ref.projection.iter() {
-                    match elem {
+                    match *elem {
                         mir::ProjectionElem::Field(f, _) => {
                             assert!(
                                 !o.layout.ty.is_any_ptr(),
@@ -721,6 +717,16 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                             );
                             o = o.extract_field(self, bx, f.index());
                         }
+                        mir::ProjectionElem::Downcast(_sym, variant_idx) => {
+                            let layout = o.layout.for_variant(bx.cx(), variant_idx);
+                            // The transmute here handles cases like `Result<bool, u8>`
+                            // where the immediate values need to change for
+                            // the specific types in the cast-to variant.
+                            let Some(val) = self.codegen_transmute_operand(bx, o, layout) else {
+                                bug!("Couldn't transmute in downcast to {variant_idx:?} of {o:?}");
+                            };
+                            o = OperandRef { val, layout };
+                        }
                         mir::ProjectionElem::Index(_)
                         | mir::ProjectionElem::ConstantIndex { .. } => {
                             // ZSTs don't require any actual memory access.

compiler/rustc_codegen_ssa/src/mir/rvalue.rs

@@ -281,13 +281,13 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                     && in_a.size(self.cx) == out_a.size(self.cx)
                     && in_b.size(self.cx) == out_b.size(self.cx)
                 {
-                    let in_a_ibty = bx.scalar_pair_element_backend_type(operand.layout, 0, false);
-                    let in_b_ibty = bx.scalar_pair_element_backend_type(operand.layout, 1, false);
-                    let out_a_ibty = bx.scalar_pair_element_backend_type(cast, 0, false);
-                    let out_b_ibty = bx.scalar_pair_element_backend_type(cast, 1, false);
+                    let in_a_bty = bx.scalar_pair_element_backend_type(operand.layout, 0, false);
+                    let in_b_bty = bx.scalar_pair_element_backend_type(operand.layout, 1, false);
+                    let out_a_bty = bx.scalar_pair_element_backend_type(cast, 0, false);
+                    let out_b_bty = bx.scalar_pair_element_backend_type(cast, 1, false);
                     Some(OperandValue::Pair(
-                        self.transmute_immediate(bx, imm_a, in_a, in_a_ibty, out_a, out_a_ibty),
-                        self.transmute_immediate(bx, imm_b, in_b, in_b_ibty, out_b, out_b_ibty),
+                        self.transmute_immediate(bx, imm_a, in_a, in_a_bty, out_a, out_a_bty),
+                        self.transmute_immediate(bx, imm_b, in_b, in_b_bty, out_b, out_b_bty),
                     ))
                 } else {
                     None
@@ -353,7 +353,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
     ///
     /// `to_backend_ty` must be the *non*-immediate backend type (so it will be
     /// `i8`, not `i1`, for `bool`-like types.)
-    fn transmute_immediate(
+    pub(crate) fn transmute_immediate(
         &self,
         bx: &mut Bx,
         mut imm: Bx::Value,
@@ -371,8 +371,8 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
             return imm;
         }
 
-        use abi::Primitive::*;
         imm = bx.from_immediate(imm);
+        debug_assert_eq!(bx.cx().val_ty(imm), from_backend_ty);
 
         // If we have a scalar, we must already know its range. Either
         //
@@ -385,8 +385,10 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
         // <https://github.com/rust-lang/rust/pull/135610#issuecomment-2599275182>
         self.assume_scalar_range(bx, imm, from_scalar, from_backend_ty);
 
+        use abi::Primitive::*;
         imm = match (from_scalar.primitive(), to_scalar.primitive()) {
-            (Int(..) | Float(_), Int(..) | Float(_)) => bx.bitcast(imm, to_backend_ty),
+            (Int(..), Int(..)) | (Float(_), Float(_)) => imm,
+            (Int(..), Float(_)) | (Float(_), Int(..)) => bx.bitcast(imm, to_backend_ty),
             (Pointer(..), Pointer(..)) => bx.pointercast(imm, to_backend_ty),
             (Int(..), Pointer(..)) => bx.ptradd(bx.const_null(bx.type_ptr()), imm),
             (Pointer(..), Int(..)) => {
@@ -411,6 +413,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
         // constraint that the `transmute` introduced is to `assume` it.
         self.assume_scalar_range(bx, imm, to_scalar, to_backend_ty);
 
+        debug_assert_eq!(bx.cx().val_ty(imm), to_backend_ty);
         imm = bx.to_immediate_scalar(imm, to_scalar);
         imm
     }
@@ -1157,8 +1160,8 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                     mir::AggregateKind::Coroutine(..) | mir::AggregateKind::CoroutineClosure(..) => false,
                 };
                 allowed_kind && {
-                let ty = rvalue.ty(self.mir, self.cx.tcx());
-                let ty = self.monomorphize(ty);
+                    let ty = rvalue.ty(self.mir, self.cx.tcx());
+                    let ty = self.monomorphize(ty);
                     let layout = self.cx.spanned_layout_of(ty, span);
                     !self.cx.is_backend_ref(layout)
                 }

tests/codegen/common_prim_int_ptr.rs

@@ -40,9 +40,13 @@ pub unsafe fn extract_int(x: Result<usize, Box<()>>) -> usize {
 }
 
 // CHECK-LABEL: @extract_box
-// CHECK-SAME: (i{{[0-9]+}} {{[^%]+}} [[DISCRIMINANT:%[0-9]+]], ptr {{[^%]+}} [[PAYLOAD:%[0-9]+]])
+// CHECK-SAME: (i{{[0-9]+}} {{[^%]+}} [[DISCRIMINANT:%x.0]], ptr {{[^%]+}} [[PAYLOAD:%x.1]])
 #[no_mangle]
 pub unsafe fn extract_box(x: Result<usize, Box<i32>>) -> Box<i32> {
+    // CHECK: [[NOT_OK:%.+]] = icmp ne i{{[0-9]+}} [[DISCRIMINANT]], 0
+    // CHECK: call void @llvm.assume(i1 [[NOT_OK]])
+    // CHECK: [[NOT_NULL:%.+]] = icmp ne ptr [[PAYLOAD]], null
+    // CHECK: call void @llvm.assume(i1 [[NOT_NULL]])
     // CHECK: ret ptr [[PAYLOAD]]
     match x {
         Ok(_) => std::intrinsics::unreachable(),