layout: Store inverse memory index in FieldsShape::Arbitrary

All usages of `memory_index` start by calling `invert_bijective_mapping`, so
storing the inverted mapping directly saves some work and simplifies the code.

Author: moulins

compiler/rustc_abi/src/layout.rs

@@ -714,7 +714,7 @@ impl<Cx: HasDataLayout> LayoutCalculator<Cx> {
                 },
                 fields: FieldsShape::Arbitrary {
                     offsets: [niche_offset].into(),
-                    memory_index: [0].into(),
+                    inverse_memory_index: [FieldIdx::new(0)].into(),
                 },
                 backend_repr: abi,
                 largest_niche,
@@ -1008,8 +1008,8 @@ impl<Cx: HasDataLayout> LayoutCalculator<Cx> {
                 let pair =
                     LayoutData::<FieldIdx, VariantIdx>::scalar_pair(&self.cx, tag, prim_scalar);
                 let pair_offsets = match pair.fields {
-                    FieldsShape::Arbitrary { ref offsets, ref memory_index } => {
-                        assert_eq!(memory_index.raw, [0, 1]);
+                    FieldsShape::Arbitrary { ref offsets, ref inverse_memory_index } => {
+                        assert_eq!(inverse_memory_index.raw, [FieldIdx::new(0), FieldIdx::new(1)]);
                         offsets
                     }
                     _ => panic!("encountered a non-arbitrary layout during enum layout"),
@@ -1061,7 +1061,7 @@ impl<Cx: HasDataLayout> LayoutCalculator<Cx> {
             },
             fields: FieldsShape::Arbitrary {
                 offsets: [Size::ZERO].into(),
-                memory_index: [0].into(),
+                inverse_memory_index: [FieldIdx::new(0)].into(),
             },
             largest_niche,
             uninhabited,
@@ -1252,8 +1252,6 @@ impl<Cx: HasDataLayout> LayoutCalculator<Cx> {
         // That is, if field 5 has offset 0, the first element of inverse_memory_index is 5.
         // We now write field offsets to the corresponding offset slot;
         // field 5 with offset 0 puts 0 in offsets[5].
-        // At the bottom of this function, we invert `inverse_memory_index` to
-        // produce `memory_index` (see `invert_mapping`).
         let mut unsized_field = None::<&F>;
         let mut offsets = IndexVec::from_elem(Size::ZERO, fields);
         let mut offset = Size::ZERO;
@@ -1322,18 +1320,6 @@ impl<Cx: HasDataLayout> LayoutCalculator<Cx> {
 
         debug!("univariant min_size: {:?}", offset);
         let min_size = offset;
-        // As stated above, inverse_memory_index holds field indices by increasing offset.
-        // This makes it an already-sorted view of the offsets vec.
-        // To invert it, consider:
-        // If field 5 has offset 0, offsets[0] is 5, and memory_index[5] should be 0.
-        // Field 5 would be the first element, so memory_index is i:
-        // Note: if we didn't optimize, it's already right.
-        let memory_index = if optimize_field_order {
-            inverse_memory_index.invert_bijective_mapping()
-        } else {
-            debug_assert!(inverse_memory_index.iter().copied().eq(fields.indices()));
-            inverse_memory_index.into_iter().map(|it| it.index() as u32).collect()
-        };
         let size = min_size.align_to(align);
         // FIXME(oli-obk): deduplicate and harden these checks
         if size.bytes() >= dl.obj_size_bound() {
@@ -1389,8 +1375,14 @@ impl<Cx: HasDataLayout> LayoutCalculator<Cx> {
                             let pair =
                                 LayoutData::<FieldIdx, VariantIdx>::scalar_pair(&self.cx, a, b);
                             let pair_offsets = match pair.fields {
-                                FieldsShape::Arbitrary { ref offsets, ref memory_index } => {
-                                    assert_eq!(memory_index.raw, [0, 1]);
+                                FieldsShape::Arbitrary {
+                                    ref offsets,
+                                    ref inverse_memory_index,
+                                } => {
+                                    assert_eq!(
+                                        inverse_memory_index.raw,
+                                        [FieldIdx::new(0), FieldIdx::new(1)]
+                                    );
                                     offsets
                                 }
                                 FieldsShape::Primitive
@@ -1434,7 +1426,7 @@ impl<Cx: HasDataLayout> LayoutCalculator<Cx> {
 
         Ok(LayoutData {
             variants: Variants::Single { index: VariantIdx::new(0) },
-            fields: FieldsShape::Arbitrary { offsets, memory_index },
+            fields: FieldsShape::Arbitrary { offsets, inverse_memory_index },
             backend_repr: abi,
             largest_niche,
             uninhabited,
@@ -1530,7 +1522,10 @@ where
 
     Ok(LayoutData {
         variants: Variants::Single { index: VariantIdx::new(0) },
-        fields: FieldsShape::Arbitrary { offsets: [Size::ZERO].into(), memory_index: [0].into() },
+        fields: FieldsShape::Arbitrary {
+            offsets: [Size::ZERO].into(),
+            inverse_memory_index: [FieldIdx::new(0)].into(),
+        },
         backend_repr: repr,
         largest_niche: elt.largest_niche,
         uninhabited: false,

compiler/rustc_abi/src/layout/coroutine.rs

@@ -182,9 +182,7 @@ pub(super) fn layout<
     // CoroutineLayout.
     debug!("prefix = {:#?}", prefix);
     let (outer_fields, promoted_offsets, promoted_memory_index) = match prefix.fields {
-        FieldsShape::Arbitrary { mut offsets, memory_index } => {
-            let mut inverse_memory_index = memory_index.invert_bijective_mapping();
-
+        FieldsShape::Arbitrary { mut offsets, inverse_memory_index } => {
             // "a" (`0..b_start`) and "b" (`b_start..`) correspond to
             // "outer" and "promoted" fields respectively.
             let b_start = tag_index.plus(1);
@@ -194,21 +192,21 @@ pub(super) fn layout<
             // Disentangle the "a" and "b" components of `inverse_memory_index`
             // by preserving the order but keeping only one disjoint "half" each.
             // FIXME(eddyb) build a better abstraction for permutations, if possible.
-            let inverse_memory_index_b: IndexVec<u32, FieldIdx> = inverse_memory_index
-                .iter()
-                .filter_map(|&i| i.index().checked_sub(b_start.index()).map(FieldIdx::new))
-                .collect();
-            inverse_memory_index.raw.retain(|&i| i.index() < b_start.index());
-            let inverse_memory_index_a = inverse_memory_index;
-
-            // Since `inverse_memory_index_{a,b}` each only refer to their
-            // respective fields, they can be safely inverted
-            let memory_index_a = inverse_memory_index_a.invert_bijective_mapping();
-            let memory_index_b = inverse_memory_index_b.invert_bijective_mapping();
+            let mut inverse_memory_index_a = IndexVec::<u32, FieldIdx>::new();
+            let mut inverse_memory_index_b = IndexVec::<u32, FieldIdx>::new();
+            for i in inverse_memory_index {
+                if let Some(j) = i.index().checked_sub(b_start.index()) {
+                    inverse_memory_index_b.push(FieldIdx::new(j));
+                } else {
+                    inverse_memory_index_a.push(i);
+                }
+            }
 
-            let outer_fields =
-                FieldsShape::Arbitrary { offsets: offsets_a, memory_index: memory_index_a };
-            (outer_fields, offsets_b, memory_index_b)
+            let outer_fields = FieldsShape::Arbitrary {
+                offsets: offsets_a,
+                inverse_memory_index: inverse_memory_index_a,
+            };
+            (outer_fields, offsets_b, inverse_memory_index_b.invert_bijective_mapping())
         }
         _ => unreachable!(),
     };
@@ -236,7 +234,7 @@ pub(super) fn layout<
             )?;
             variant.variants = Variants::Single { index };
 
-            let FieldsShape::Arbitrary { offsets, memory_index } = variant.fields else {
+            let FieldsShape::Arbitrary { offsets, inverse_memory_index } = variant.fields else {
                 unreachable!();
             };
 
@@ -249,6 +247,7 @@ pub(super) fn layout<
             // promoted fields were being used, but leave the elements not in the
             // subset as `invalid_field_idx`, which we can filter out later to
             // obtain a valid (bijective) mapping.
+            let memory_index = inverse_memory_index.invert_bijective_mapping();
             let invalid_field_idx = promoted_memory_index.len() + memory_index.len();
             let mut combined_inverse_memory_index =
                 IndexVec::from_elem_n(FieldIdx::new(invalid_field_idx), invalid_field_idx);
@@ -273,14 +272,13 @@ pub(super) fn layout<
                 })
                 .collect();
 
-            // Remove the unused slots and invert the mapping to obtain the
-            // combined `memory_index` (also see previous comment).
+            // Remove the unused slots to obtain the combined `inverse_memory_index`
+            // (also see previous comment).
             combined_inverse_memory_index.raw.retain(|&i| i.index() != invalid_field_idx);
-            let combined_memory_index = combined_inverse_memory_index.invert_bijective_mapping();
 
             variant.fields = FieldsShape::Arbitrary {
                 offsets: combined_offsets,
-                memory_index: combined_memory_index,
+                inverse_memory_index: combined_inverse_memory_index,
             };
 
             size = size.max(variant.size);

compiler/rustc_abi/src/layout/simple.rs

@@ -16,7 +16,7 @@ impl<FieldIdx: Idx, VariantIdx: Idx> LayoutData<FieldIdx, VariantIdx> {
             variants: Variants::Single { index: VariantIdx::new(0) },
             fields: FieldsShape::Arbitrary {
                 offsets: IndexVec::new(),
-                memory_index: IndexVec::new(),
+                inverse_memory_index: IndexVec::new(),
             },
             backend_repr: BackendRepr::Memory { sized },
             largest_niche: None,
@@ -108,7 +108,7 @@ impl<FieldIdx: Idx, VariantIdx: Idx> LayoutData<FieldIdx, VariantIdx> {
             variants: Variants::Single { index: VariantIdx::new(0) },
             fields: FieldsShape::Arbitrary {
                 offsets: [Size::ZERO, b_offset].into(),
-                memory_index: [0, 1].into(),
+                inverse_memory_index: [FieldIdx::new(0), FieldIdx::new(1)].into(),
             },
             backend_repr: BackendRepr::ScalarPair(a, b),
             largest_niche,
@@ -133,7 +133,7 @@ impl<FieldIdx: Idx, VariantIdx: Idx> LayoutData<FieldIdx, VariantIdx> {
                 Some(fields) => FieldsShape::Union(fields),
                 None => FieldsShape::Arbitrary {
                     offsets: IndexVec::new(),
-                    memory_index: IndexVec::new(),
+                    inverse_memory_index: IndexVec::new(),
                 },
             },
             backend_repr: BackendRepr::Memory { sized: true },

compiler/rustc_abi/src/lib.rs

@@ -1636,19 +1636,14 @@ pub enum FieldsShape<FieldIdx: Idx> {
         // FIXME(eddyb) use small vector optimization for the common case.
         offsets: IndexVec<FieldIdx, Size>,
 
-        /// Maps source order field indices to memory order indices,
+        /// Maps memory order field indices to source order indices,
         /// depending on how the fields were reordered (if at all).
         /// This is a permutation, with both the source order and the
         /// memory order using the same (0..n) index ranges.
         ///
-        /// Note that during computation of `memory_index`, sometimes
-        /// it is easier to operate on the inverse mapping (that is,
-        /// from memory order to source order), and that is usually
-        /// named `inverse_memory_index`.
-        ///
         // FIXME(eddyb) build a better abstraction for permutations, if possible.
         // FIXME(camlorn) also consider small vector optimization here.
-        memory_index: IndexVec<FieldIdx, u32>,
+        inverse_memory_index: IndexVec<u32, FieldIdx>,
     },
 }
 
@@ -1682,50 +1677,18 @@ impl<FieldIdx: Idx> FieldsShape<FieldIdx> {
         }
     }
 
-    #[inline]
-    pub fn memory_index(&self, i: usize) -> usize {
-        match *self {
-            FieldsShape::Primitive => {
-                unreachable!("FieldsShape::memory_index: `Primitive`s have no fields")
-            }
-            FieldsShape::Union(_) | FieldsShape::Array { .. } => i,
-            FieldsShape::Arbitrary { ref memory_index, .. } => {
-                memory_index[FieldIdx::new(i)].try_into().unwrap()
-            }
-        }
-    }
-
     /// Gets source indices of the fields by increasing offsets.
     #[inline]
     pub fn index_by_increasing_offset(&self) -> impl ExactSizeIterator<Item = usize> {
-        let mut inverse_small = [0u8; 64];
-        let mut inverse_big = IndexVec::new();
-        let use_small = self.count() <= inverse_small.len();
-
-        // We have to write this logic twice in order to keep the array small.
-        if let FieldsShape::Arbitrary { ref memory_index, .. } = *self {
-            if use_small {
-                for (field_idx, &mem_idx) in memory_index.iter_enumerated() {
-                    inverse_small[mem_idx as usize] = field_idx.index() as u8;
-                }
-            } else {
-                inverse_big = memory_index.invert_bijective_mapping();
-            }
-        }
-
         // Primitives don't really have fields in the way that structs do,
         // but having this return an empty iterator for them is unhelpful
         // since that makes them look kinda like ZSTs, which they're not.
         let pseudofield_count = if let FieldsShape::Primitive = self { 1 } else { self.count() };
 
-        (0..pseudofield_count).map(move |i| match *self {
+        (0..pseudofield_count).map(move |i| match self {
             FieldsShape::Primitive | FieldsShape::Union(_) | FieldsShape::Array { .. } => i,
-            FieldsShape::Arbitrary { .. } => {
-                if use_small {
-                    inverse_small[i] as usize
-                } else {
-                    inverse_big[i as u32].index()
-                }
+            FieldsShape::Arbitrary { inverse_memory_index, .. } => {
+                inverse_memory_index[i as u32].index()
             }
         })
     }

compiler/rustc_const_eval/src/interpret/visitor.rs

@@ -4,7 +4,7 @@
 use std::num::NonZero;
 
 use rustc_abi::{FieldIdx, FieldsShape, VariantIdx, Variants};
-use rustc_index::IndexVec;
+use rustc_index::{Idx as _, IndexVec};
 use rustc_middle::mir::interpret::InterpResult;
 use rustc_middle::ty::{self, Ty};
 use tracing::trace;
@@ -27,13 +27,15 @@ pub trait ValueVisitor<'tcx, M: Machine<'tcx>>: Sized {
     /// This function provides the chance to reorder the order in which fields are visited for
     /// `FieldsShape::Aggregate`.
     ///
-    /// The default means we iterate in source declaration order; alternatively this can do some
-    /// work with `memory_index` to iterate in memory order.
+    /// The default means we iterate in source declaration order; alternatively this can use
+    /// `inverse_memory_index` to iterate in memory order.
     #[inline(always)]
     fn aggregate_field_iter(
-        memory_index: &IndexVec<FieldIdx, u32>,
-    ) -> impl Iterator<Item = FieldIdx> + 'static {
-        memory_index.indices()
+        inverse_memory_index: &IndexVec<u32, FieldIdx>,
+    ) -> impl Iterator<Item = FieldIdx> {
+        // Allow the optimizer to elide the bounds checking when creating each index.
+        let _ = FieldIdx::new(inverse_memory_index.len());
+        (0..inverse_memory_index.len()).map(FieldIdx::new)
     }
 
     // Recursive actions, ready to be overloaded.
@@ -168,8 +170,8 @@ pub trait ValueVisitor<'tcx, M: Machine<'tcx>>: Sized {
             &FieldsShape::Union(fields) => {
                 self.visit_union(v, fields)?;
             }
-            FieldsShape::Arbitrary { memory_index, .. } => {
-                for idx in Self::aggregate_field_iter(memory_index) {
+            FieldsShape::Arbitrary { inverse_memory_index, .. } => {
+                for idx in Self::aggregate_field_iter(inverse_memory_index) {
                     let field = self.ecx().project_field(v, idx)?;
                     self.visit_field(v, idx.as_usize(), &field)?;
                 }

compiler/rustc_index/src/slice.rs

@@ -181,9 +181,6 @@ impl<I: Idx, J: Idx> IndexSlice<I, J> {
     /// Invert a bijective mapping, i.e. `invert(map)[y] = x` if `map[x] = y`,
     /// assuming the values in `self` are a permutation of `0..self.len()`.
     ///
-    /// This is used to go between `memory_index` (source field order to memory order)
-    /// and `inverse_memory_index` (memory order to source field order).
-    /// See also `FieldsShape::Arbitrary::memory_index` for more details.
     // FIXME(eddyb) build a better abstraction for permutations, if possible.
     pub fn invert_bijective_mapping(&self) -> IndexVec<J, I> {
         debug_assert_eq!(

compiler/rustc_transmute/src/layout/tree.rs

@@ -491,7 +491,7 @@ pub(crate) mod rustc {
         ) -> Result<Self, Err> {
             // This constructor does not support non-`FieldsShape::Arbitrary`
             // layouts.
-            let FieldsShape::Arbitrary { offsets, memory_index } = layout.fields() else {
+            let FieldsShape::Arbitrary { offsets, inverse_memory_index } = layout.fields() else {
                 return Err(Err::NotYetSupported);
             };
 
@@ -519,7 +519,6 @@ pub(crate) mod rustc {
             }
 
             // Append the fields, in memory order, to the layout.
-            let inverse_memory_index = memory_index.invert_bijective_mapping();
             for &field_idx in inverse_memory_index.iter() {
                 // Add interfield padding.
                 let padding_needed = offsets[field_idx] - size;

compiler/rustc_ty_utils/src/layout.rs

@@ -9,7 +9,7 @@ use rustc_abi::{
 use rustc_hashes::Hash64;
 use rustc_hir::attrs::AttributeKind;
 use rustc_hir::find_attr;
-use rustc_index::IndexVec;
+use rustc_index::{Idx as _, IndexVec};
 use rustc_middle::bug;
 use rustc_middle::query::Providers;
 use rustc_middle::traits::ObligationCause;
@@ -374,7 +374,7 @@ fn layout_of_uncached<'tcx>(
             // specifically care about pattern types will have to handle it.
             layout.fields = FieldsShape::Arbitrary {
                 offsets: [Size::ZERO].into_iter().collect(),
-                memory_index: [0].into_iter().collect(),
+                inverse_memory_index: [FieldIdx::new(0)].into_iter().collect(),
             };
             tcx.mk_layout(layout)
         }

src/tools/miri/src/helpers.rs

@@ -584,10 +584,9 @@ pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
             }
 
             fn aggregate_field_iter(
-                memory_index: &IndexVec<FieldIdx, u32>,
-            ) -> impl Iterator<Item = FieldIdx> + 'static {
-                let inverse_memory_index = memory_index.invert_bijective_mapping();
-                inverse_memory_index.into_iter()
+                inverse_memory_index: &IndexVec<u32, FieldIdx>,
+            ) -> impl Iterator<Item = FieldIdx> {
+                inverse_memory_index.iter().copied()
             }
 
             // Hook to detect `UnsafeCell`.

tests/ui/abi/c-zst.aarch64-darwin.stderr

@@ -13,7 +13,7 @@ error: fn_abi_of(pass_zst) = FnAbi {
                            },
                            fields: Arbitrary {
                                offsets: [],
-                               memory_index: [],
+                               inverse_memory_index: [],
                            },
                            largest_niche: None,
                            uninhabited: false,
@@ -41,7 +41,7 @@ error: fn_abi_of(pass_zst) = FnAbi {
                        },
                        fields: Arbitrary {
                            offsets: [],
-                           memory_index: [],
+                           inverse_memory_index: [],
                        },
                        largest_niche: None,
                        uninhabited: false,