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(),
+                    in_memory_order: [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 in_memory_order } => {
+                        assert_eq!(in_memory_order.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(),
+                in_memory_order: [FieldIdx::new(0)].into(),
             },
             largest_niche,
             uninhabited,
@@ -1110,10 +1110,10 @@ impl<Cx: HasDataLayout> LayoutCalculator<Cx> {
         let pack = repr.pack;
         let mut align = if pack.is_some() { dl.i8_align } else { dl.aggregate_align };
         let mut max_repr_align = repr.align;
-        let mut inverse_memory_index: IndexVec<u32, FieldIdx> = fields.indices().collect();
+        let mut in_memory_order: IndexVec<u32, FieldIdx> = fields.indices().collect();
         let optimize_field_order = !repr.inhibit_struct_field_reordering();
         let end = if let StructKind::MaybeUnsized = kind { fields.len() - 1 } else { fields.len() };
-        let optimizing = &mut inverse_memory_index.raw[..end];
+        let optimizing = &mut in_memory_order.raw[..end];
         let fields_excluding_tail = &fields.raw[..end];
         // unsizable tail fields are excluded so that we use the same seed for the sized and unsized layouts.
         let field_seed = fields_excluding_tail
@@ -1248,12 +1248,10 @@ impl<Cx: HasDataLayout> LayoutCalculator<Cx> {
                 //                 regardless of the status of `-Z randomize-layout`
             }
         }
-        // inverse_memory_index holds field indices by increasing memory offset.
-        // That is, if field 5 has offset 0, the first element of inverse_memory_index is 5.
+        // in_memory_order holds field indices by increasing memory offset.
+        // That is, if field 5 has offset 0, the first element of in_memory_order 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;
@@ -1265,7 +1263,7 @@ impl<Cx: HasDataLayout> LayoutCalculator<Cx> {
             align = align.max(prefix_align);
             offset = prefix_size.align_to(prefix_align);
         }
-        for &i in &inverse_memory_index {
+        for &i in &in_memory_order {
             let field = &fields[i];
             if let Some(unsized_field) = unsized_field {
                 return Err(LayoutCalculatorError::UnexpectedUnsized(*unsized_field));
@@ -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,11 @@ 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 in_memory_order } => {
+                                    assert_eq!(
+                                        in_memory_order.raw,
+                                        [FieldIdx::new(0), FieldIdx::new(1)]
+                                    );
                                     offsets
                                 }
                                 FieldsShape::Primitive
@@ -1434,7 +1423,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, in_memory_order },
             backend_repr: abi,
             largest_niche,
             uninhabited,
@@ -1530,7 +1519,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(),
+            in_memory_order: [FieldIdx::new(0)].into(),
+        },
         backend_repr: repr,
         largest_niche: elt.largest_niche,
         uninhabited: false,

compiler/rustc_abi/src/layout/coroutine.rs

@@ -182,33 +182,29 @@ 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, in_memory_order } => {
             // "a" (`0..b_start`) and "b" (`b_start..`) correspond to
             // "outer" and "promoted" fields respectively.
             let b_start = tag_index.plus(1);
             let offsets_b = IndexVec::from_raw(offsets.raw.split_off(b_start.index()));
             let offsets_a = offsets;
 
-            // Disentangle the "a" and "b" components of `inverse_memory_index`
+            // Disentangle the "a" and "b" components of `in_memory_order`
             // 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 in_memory_order_a = IndexVec::<u32, FieldIdx>::new();
+            let mut in_memory_order_b = IndexVec::<u32, FieldIdx>::new();
+            for i in in_memory_order {
+                if let Some(j) = i.index().checked_sub(b_start.index()) {
+                    in_memory_order_b.push(FieldIdx::new(j));
+                } else {
+                    in_memory_order_a.push(i);
+                }
+            }
 
             let outer_fields =
-                FieldsShape::Arbitrary { offsets: offsets_a, memory_index: memory_index_a };
-            (outer_fields, offsets_b, memory_index_b)
+                FieldsShape::Arbitrary { offsets: offsets_a, in_memory_order: in_memory_order_a };
+            (outer_fields, offsets_b, in_memory_order_b.invert_bijective_mapping())
         }
         _ => unreachable!(),
     };
@@ -236,7 +232,7 @@ pub(super) fn layout<
             )?;
             variant.variants = Variants::Single { index };
 
-            let FieldsShape::Arbitrary { offsets, memory_index } = variant.fields else {
+            let FieldsShape::Arbitrary { offsets, in_memory_order } = variant.fields else {
                 unreachable!();
             };
 
@@ -249,8 +245,9 @@ 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 = in_memory_order.invert_bijective_mapping();
             let invalid_field_idx = promoted_memory_index.len() + memory_index.len();
-            let mut combined_inverse_memory_index =
+            let mut combined_in_memory_order =
                 IndexVec::from_elem_n(FieldIdx::new(invalid_field_idx), invalid_field_idx);
 
             let mut offsets_and_memory_index = iter::zip(offsets, memory_index);
@@ -268,19 +265,18 @@ pub(super) fn layout<
                             (promoted_offsets[field_idx], promoted_memory_index[field_idx])
                         }
                     };
-                    combined_inverse_memory_index[memory_index] = i;
+                    combined_in_memory_order[memory_index] = i;
                     offset
                 })
                 .collect();
 
-            // Remove the unused slots and invert the mapping to obtain the
-            // combined `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();
+            // Remove the unused slots to obtain the combined `in_memory_order`
+            // (also see previous comment).
+            combined_in_memory_order.raw.retain(|&i| i.index() != invalid_field_idx);
 
             variant.fields = FieldsShape::Arbitrary {
                 offsets: combined_offsets,
-                memory_index: combined_memory_index,
+                in_memory_order: combined_in_memory_order,
             };
 
             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(),
+                in_memory_order: 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(),
+                in_memory_order: [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(),
+                    in_memory_order: 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>,
+        in_memory_order: IndexVec<u32, FieldIdx>,
     },
 }
 
@@ -1682,51 +1677,17 @@ 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 { in_memory_order, .. } => in_memory_order[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
+    /// `in_memory_order` to iterate in memory order.
     #[inline(always)]
     fn aggregate_field_iter(
-        memory_index: &IndexVec<FieldIdx, u32>,
-    ) -> impl Iterator<Item = FieldIdx> + 'static {
-        memory_index.indices()
+        in_memory_order: &IndexVec<u32, FieldIdx>,
+    ) -> impl Iterator<Item = FieldIdx> {
+        // Allow the optimizer to elide the bounds checking when creating each index.
+        let _ = FieldIdx::new(in_memory_order.len());
+        (0..in_memory_order.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 { in_memory_order, .. } => {
+                for idx in Self::aggregate_field_iter(in_memory_order) {
                     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, in_memory_order } = layout.fields() else {
                 return Err(Err::NotYetSupported);
             };
 
@@ -519,8 +519,7 @@ 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() {
+            for &field_idx in in_memory_order.iter() {
                 // Add interfield padding.
                 let padding_needed = offsets[field_idx] - size;
                 let padding = Self::padding(padding_needed.bytes_usize());