Allow matching enums without needing allocas

Author: scottmcm

compiler/rustc_codegen_ssa/src/mir/analyze.rs

@@ -1,12 +1,13 @@
 //! An analysis to determine which locals require allocas and
 //! which do not.
 
+use rustc_abi as abi;
 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::ty::layout::{LayoutOf, TyAndLayout};
 use rustc_middle::{bug, span_bug};
 use tracing::debug;
 
@@ -99,63 +100,112 @@ impl<'a, 'b, 'tcx, Bx: BuilderMethods<'b, 'tcx>> LocalAnalyzer<'a, 'b, 'tcx, Bx>
         context: PlaceContext,
         location: Location,
     ) {
-        let cx = self.fx.cx;
+        const COPY_CONTEXT: PlaceContext =
+            PlaceContext::NonMutatingUse(NonMutatingUseContext::Copy);
+
+        // `PlaceElem::Index` is the only variant that can mention other `Local`s,
+        // so check for those up-front before any potential short-circuits.
+        for elem in place_ref.projection {
+            if let mir::PlaceElem::Index(index_local) = *elem {
+                self.visit_local(index_local, COPY_CONTEXT, location);
+            }
+        }
 
-        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)
-            };
+        // If our local is already memory, nothing can make it *more* memory
+        // so we don't need to bother checking the projections further.
+        if self.locals[place_ref.local] == LocalKind::Memory {
+            return;
+        }
 
-            // Allow uses of projections that are ZSTs or from scalar fields.
-            let is_consume = matches!(
-                context,
-                PlaceContext::NonMutatingUse(
-                    NonMutatingUseContext::Copy | NonMutatingUseContext::Move,
-                )
-            );
-            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 place_ref.is_indirect_first_projection() {
+            // If this starts with a `Deref`, we only need to record a read of the
+            // pointer being dereferenced, as all the subsequent projections are
+            // working on a place which is always supported. (And because we're
+            // looking at codegen MIR, it can only happen as the first projection.)
+            self.visit_local(place_ref.local, COPY_CONTEXT, location);
+            return;
+        }
+
+        if !place_ref.projection.is_empty() {
+            if context.is_mutating_use() {
+                // If it's a mutating use it doesn't matter what the projections are,
+                // if there are *any* then we need a place to write. (For example,
+                // `_1 = Foo()` works in SSA but `_2.0 = Foo()` does not.)
+                let mut_projection = PlaceContext::MutatingUse(MutatingUseContext::Projection);
+                self.visit_local(place_ref.local, mut_projection, location);
+                return;
+            }
+
+            // Scan through to ensure the only projections are those which
+            // `FunctionCx::maybe_codegen_consume_direct` can handle.
+            let base_ty = self.fx.monomorphized_place_ty(mir::PlaceRef::from(place_ref.local));
+            let mut layout = self.fx.cx.layout_of(base_ty);
+            for elem in place_ref.projection {
+                if layout.is_zst() {
+                    // Any further projections must still be ZSTs, so we're good.
+                    break;
                 }
 
-                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;
+                #[track_caller]
+                fn compatible_projection(src: TyAndLayout<'_>, tgt: TyAndLayout<'_>) -> bool {
+                    fn compatible_initness(a: abi::Scalar, b: abi::Scalar) -> bool {
+                        !a.is_uninit_valid() || b.is_uninit_valid()
+                    }
+
+                    use abi::BackendRepr::*;
+                    match (src.backend_repr, tgt.backend_repr) {
+                        _ if tgt.is_zst() => true,
+                        (Scalar(a), Scalar(b)) | (SimdVector { element: a, .. }, SimdVector { element: b, .. }) =>
+                            compatible_initness(a, b),
+                        (ScalarPair(a0, a1), Scalar(b)) =>
+                            compatible_initness(a0, b) &&
+                            compatible_initness(a1, b),
+                        (ScalarPair(a0, a1), ScalarPair(b0, b1)) =>
+                            compatible_initness(a0, b0) &&
+                            compatible_initness(a1, b1),
+                        // This arm is a hack; remove it as part of
+                        // <https://github.com/rust-lang/compiler-team/issues/838>
+                        (SimdVector { .. }, Memory { .. }) => true,
+                        _ => bug!("Mismatched layouts in analysis\nsrc: {src:?}\ntgt: {tgt:?}"),
                     }
                 }
-            }
 
-            if let mir::ProjectionElem::Deref = elem {
-                // Deref projections typically only read the pointer.
-                base_context = PlaceContext::NonMutatingUse(NonMutatingUseContext::Copy);
-            }
+                match *elem {
+                    mir::PlaceElem::Field(fidx, ..) => {
+                        let field_layout = layout.field(self.fx.cx, fidx.as_usize());
+                        if compatible_projection(layout, field_layout)
+                        {
+                            layout = field_layout;
+                            continue;
+                        }
+                    }
+                    mir::PlaceElem::Downcast(_, vidx) => {
+                        let variant_layout = layout.for_variant(self.fx.cx, vidx);
+                        if compatible_projection(layout, variant_layout) {
+                            layout = variant_layout;
+                            continue;
+                        }
+                    }
 
-            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 {
-                self.visit_local(
-                    local,
-                    PlaceContext::NonMutatingUse(NonMutatingUseContext::Copy),
-                    location,
-                );
+                    mir::PlaceElem::Index(..)
+                    | mir::PlaceElem::ConstantIndex { .. }
+                    | mir::PlaceElem::Subslice { .. }
+                    | mir::PlaceElem::OpaqueCast(..)
+                    | mir::PlaceElem::UnwrapUnsafeBinder(..)
+                    | mir::PlaceElem::Subtype(..) => {}
+
+                    mir::PlaceElem::Deref => bug!("Deref after first position"),
+                }
+
+                // If the above didn't `continue`, we can't handle the projection.
+                self.locals[place_ref.local] = LocalKind::Memory;
+                return;
             }
-        } else {
-            self.visit_local(place_ref.local, context, location);
         }
+
+        // Even with supported projections, we still need to have `visit_local`
+        // check for things that can't be done in SSA (like `SharedBorrow`).
+        self.visit_local(place_ref.local, context, location);
     }
 }
 

compiler/rustc_codegen_ssa/src/mir/locals.rs

@@ -12,6 +12,7 @@ use tracing::{debug, warn};
 use crate::mir::{FunctionCx, LocalRef};
 use crate::traits::BuilderMethods;
 
+#[derive(Debug)]
 pub(super) struct Locals<'tcx, V> {
     values: IndexVec<mir::Local, LocalRef<'tcx, V>>,
 }

compiler/rustc_codegen_ssa/src/mir/mod.rs

@@ -136,6 +136,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

@@ -126,7 +126,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)
     }
@@ -361,13 +361,17 @@ impl<'a, 'tcx, V: CodegenObject> OperandRef<'tcx, V> {
             let (in_scalar, imm) = 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)
+                        (a, a_llval)
                     } 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)
+                        (b, b_llval)
                     }
                 }
 
@@ -378,23 +382,12 @@ 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.
+                    transmute_scalar(bx, imm, in_scalar, out_scalar)
                 }
                 BackendRepr::ScalarPair(_, _) | BackendRepr::Memory { .. } => bug!(),
             })
@@ -875,7 +868,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(),
@@ -884,17 +877,21 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                             );
                             o = o.extract_field(self, bx, f.index());
                         }
-                        mir::ProjectionElem::Index(_)
-                        | mir::ProjectionElem::ConstantIndex { .. } => {
-                            // ZSTs don't require any actual memory access.
-                            // FIXME(eddyb) deduplicate this with the identical
-                            // checks in `codegen_consume` and `extract_field`.
-                            let elem = o.layout.field(bx.cx(), 0);
-                            if elem.is_zst() {
-                                o = OperandRef::zero_sized(elem);
-                            } else {
-                                return None;
+                        mir::ProjectionElem::Downcast(_sym, variant_idx) => {
+                            let layout = o.layout.for_variant(bx.cx(), variant_idx);
+                            let val = match layout.backend_repr {
+                                // 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.
+                                BackendRepr::Scalar(..) | BackendRepr::ScalarPair(..) => {
+                                    self.codegen_transmute_operand(bx, o, layout)
                             }
+                                BackendRepr::SimdVector { .. } | BackendRepr::Memory { .. } => {
+                                    o.val
+                                }
+                            };
+
+                            o = OperandRef { val, layout };
                         }
                         _ => return None,
                     }

tests/codegen/array-cmp.rs

@@ -39,6 +39,10 @@ pub fn array_of_tuple_le(a: &[(i16, u16); 2], b: &[(i16, u16); 2]) -> bool {
     // CHECK: %[[EQ00:.+]] = icmp eq i16 %[[A00]], %[[B00]]
     // CHECK-NEXT: br i1 %[[EQ00]], label %[[L01:.+]], label %[[EXIT_S:.+]]
 
+    // CHECK: [[EXIT_S]]:
+    // CHECK: %[[RET_S:.+]] = icmp sle i16
+    // CHECK: br label
+
     // CHECK: [[L01]]:
     // CHECK: %[[PA01:.+]] = getelementptr{{.+}}i8, ptr %a, {{i32|i64}} 2
     // CHECK: %[[PB01:.+]] = getelementptr{{.+}}i8, ptr %b, {{i32|i64}} 2
@@ -66,8 +70,12 @@ pub fn array_of_tuple_le(a: &[(i16, u16); 2], b: &[(i16, u16); 2]) -> bool {
     // CHECK: %[[EQ11:.+]] = icmp eq i16 %[[A11]], %[[B11]]
     // CHECK-NEXT: br i1 %[[EQ11]], label %[[DONE:.+]], label %[[EXIT_U]]
 
+    // CHECK: [[EXIT_U]]:
+    // CHECK: %[[RET_U:.+]] = icmp ule i16
+    // CHECK: br label
+
     // CHECK: [[DONE]]:
-    // CHECK: %[[RET:.+]] = phi i1 [ %{{.+}}, %[[EXIT_S]] ], [ %{{.+}}, %[[EXIT_U]] ], [ true, %[[L11]] ]
+    // CHECK: %[[RET:.+]] = phi i1 [ true, %[[L11]] ], [ %[[RET_S]], %[[EXIT_S]] ], [ %[[RET_U]], %[[EXIT_U]] ]
     // CHECK: ret i1 %[[RET]]
 
     a <= b

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(),