Hopefully this is less work...

1. Don't downcast both parts of a pair to read just one
2. Avoid the transmute calls if it's not downcasting

Author: scottmcm

compiler/rustc_abi/src/lib.rs

@@ -1821,7 +1821,7 @@ pub enum Variants<FieldIdx: Idx, VariantIdx: Idx> {
 
     /// Single enum variants, structs/tuples, unions, and all non-ADTs.
     Single {
-        /// Always `0` for types that cannot have multiple variants.
+        /// Always [`FIRST_VARIANT`] for types that cannot have multiple variants.
         index: VariantIdx,
     },
 

compiler/rustc_codegen_ssa/src/mir/analyze.rs

@@ -100,33 +100,33 @@ impl<'a, 'b, 'tcx, Bx: BuilderMethods<'b, 'tcx>> LocalAnalyzer<'a, 'b, 'tcx, Bx>
         context: PlaceContext,
         location: Location,
     ) {
-        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 !place_ref.projection.is_empty() {
+            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 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;
-        }
+            // 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;
+            }
 
-        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.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,
@@ -199,6 +199,10 @@ impl<'a, 'b, 'tcx, Bx: BuilderMethods<'b, 'tcx>> LocalAnalyzer<'a, 'b, 'tcx, Bx>
                 self.locals[place_ref.local] = LocalKind::Memory;
                 return;
             }
+            debug_assert!(
+                !self.fx.cx.is_backend_ref(layout),
+                "Post-projection {place_ref:?} layout should be non-Ref, but it's {layout:?}",
+            );
         }
 
         // Even with supported projections, we still need to have `visit_local`

compiler/rustc_codegen_ssa/src/mir/block.rs

@@ -1,6 +1,6 @@
 use std::cmp;
 
-use rustc_abi::{Align, BackendRepr, ExternAbi, HasDataLayout, Reg, Size, WrappingRange};
+use rustc_abi::{Align, BackendRepr, ExternAbi, FieldIdx, HasDataLayout, Reg, Size, WrappingRange};
 use rustc_ast as ast;
 use rustc_ast::{InlineAsmOptions, InlineAsmTemplatePiece};
 use rustc_data_structures::packed::Pu128;
@@ -1038,7 +1038,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                             let (idx, _) = op.layout.non_1zst_field(bx).expect(
                                 "not exactly one non-1-ZST field in a `DispatchFromDyn` type",
                             );
-                            op = op.extract_field(self, bx, idx.as_usize());
+                            op = op.extract_field(self, bx, None, idx);
                         }
 
                         // Now that we have `*dyn Trait` or `&dyn Trait`, split it up into its
@@ -1598,7 +1598,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
         } else {
             // If the tuple is immediate, the elements are as well.
             for i in 0..tuple.layout.fields.count() {
-                let op = tuple.extract_field(self, bx, i);
+                let op = tuple.extract_field(self, bx, None, FieldIdx::from_usize(i));
                 self.codegen_argument(bx, op, llargs, &args[i], lifetime_ends_after_call);
             }
         }

compiler/rustc_codegen_ssa/src/mir/operand.rs

@@ -319,16 +319,40 @@ impl<'a, 'tcx, V: CodegenObject> OperandRef<'tcx, V> {
         OperandRef { val, layout }
     }
 
+    /// Extracts field `field_idx` from `self`, after downcasting to
+    /// `downcast_variant` if it's specified.
+    ///
+    /// Reading from things like tuples or structs, which are always single-variant,
+    /// don't need to pass a downcast variant since downcasting them to
+    /// `FIRST_VARIANT` doesn't actually change anything.
+    /// Things like enums and coroutines, though, must pass the variant from which
+    /// they want to read unless they're specifically reading the tag field
+    /// (which is the index at the type level, not inside one of the variants).
     pub(crate) fn extract_field<Bx: BuilderMethods<'a, 'tcx, Value = V>>(
         &self,
         fx: &mut FunctionCx<'a, 'tcx, Bx>,
         bx: &mut Bx,
-        i: usize,
+        downcast_variant: Option<VariantIdx>,
+        field_idx: FieldIdx,
     ) -> Self {
-        let field = self.layout.field(bx.cx(), i);
-        let offset = self.layout.fields.offset(i);
+        let layout = if let Some(vidx) = downcast_variant {
+            self.layout.for_variant(bx.cx(), vidx)
+        } else {
+            debug_assert!(
+                match self.layout.variants {
+                    Variants::Empty => true,
+                    Variants::Single { index } => index == FIRST_VARIANT,
+                    Variants::Multiple { tag_field, .. } => tag_field == field_idx,
+                },
+                "Should have specified a variant to read field {field_idx:?} of {self:?}",
+            );
+            self.layout
+        };
 
-        if !bx.is_backend_ref(self.layout) && bx.is_backend_ref(field) {
+        let field = layout.field(bx.cx(), field_idx.as_usize());
+        let offset = layout.fields.offset(field_idx.as_usize());
+
+        if !bx.is_backend_ref(layout) && bx.is_backend_ref(field) {
             // Part of https://github.com/rust-lang/compiler-team/issues/838
             span_bug!(
                 fx.mir.span,
@@ -338,18 +362,22 @@ impl<'a, 'tcx, V: CodegenObject> OperandRef<'tcx, V> {
 
         let val = if field.is_zst() {
             OperandValue::ZeroSized
-        } else if let BackendRepr::SimdVector { .. } = self.layout.backend_repr {
+        } else if let BackendRepr::SimdVector { .. } = layout.backend_repr {
             // codegen_transmute_operand doesn't support SIMD, but since the previous
             // check handled ZSTs, the only possible field access into something SIMD
             // is to the `non_1zst_field` that's the same SIMD. (Other things, even
             // just padding, would change the wrapper's representation type.)
-            assert_eq!(field.size, self.layout.size);
+            assert_eq!(field.size, layout.size);
             self.val
-        } else if field.size == self.layout.size {
-            assert_eq!(offset.bytes(), 0);
-            fx.codegen_transmute_operand(bx, *self, field)
+        } else if field.size == layout.size {
+            debug_assert_eq!(offset.bytes(), 0);
+            if downcast_variant.is_some() {
+                fx.codegen_transmute_operand(bx, *self, field)
+            } else {
+                self.val
+            }
         } else {
-            let (in_scalar, imm) = match (self.val, self.layout.backend_repr) {
+            let (in_scalar, imm) = match (self.val, layout.backend_repr) {
                 // Extract a scalar component from a pair.
                 (OperandValue::Pair(a_llval, b_llval), BackendRepr::ScalarPair(a, b)) => {
                     // This needs to look at `offset`, rather than `i`, because
@@ -371,19 +399,18 @@ impl<'a, 'tcx, V: CodegenObject> OperandRef<'tcx, V> {
                 }
             };
             OperandValue::Immediate(match field.backend_repr {
-                BackendRepr::SimdVector { .. } => imm,
-                BackendRepr::Scalar(out_scalar) => {
+                BackendRepr::Scalar(out_scalar) if downcast_variant.is_some() => {
                     // 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::Scalar(_) | BackendRepr::SimdVector { .. } => imm,
                 BackendRepr::ScalarPair(_, _) | BackendRepr::Memory { .. } => bug!(),
             })
         };
-
         OperandRef { val, layout: field }
     }
 
@@ -431,7 +458,7 @@ impl<'a, 'tcx, V: CodegenObject> OperandRef<'tcx, V> {
         let tag_op = match self.val {
             OperandValue::ZeroSized => bug!(),
             OperandValue::Immediate(_) | OperandValue::Pair(_, _) => {
-                self.extract_field(fx, bx, tag_field.as_usize())
+                self.extract_field(fx, bx, None, tag_field)
             }
             OperandValue::Ref(place) => {
                 let tag = place.with_type(self.layout).project_field(bx, tag_field.as_usize());
@@ -922,6 +949,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
     ) -> Option<OperandRef<'tcx, Bx::Value>> {
         debug!("maybe_codegen_consume_direct(place_ref={:?})", place_ref);
 
+        let mut downcast_variant = None;
         match self.locals[place_ref.local] {
             LocalRef::Operand(mut o) => {
                 // Moves out of scalar and scalar pair fields are trivial.
@@ -933,27 +961,16 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                                 "Bad PlaceRef: destructing pointers should use cast/PtrMetadata, \
                                  but tried to access field {f:?} of pointer {o:?}",
                             );
-                            o = o.extract_field(self, bx, f.index());
+                            o = o.extract_field(self, bx, downcast_variant, f);
+                            downcast_variant = 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 };
+                            downcast_variant = Some(variant_idx);
                         }
                         _ => return None,
                     }
                 }
+                debug_assert_eq!(downcast_variant, None);
 
                 Some(o)
             }