Skip to content

Commit

Permalink
Auto merge of rust-lang#109915 - scottmcm:layout-indexvec, r=oli-obk
Browse files Browse the repository at this point in the history
Use `FieldIdx` in `FieldsShape`

Finally got to the main motivating example from rust-lang/compiler-team#606 :)
  • Loading branch information
bors committed Apr 6, 2023
2 parents 7f6edd3 + 21bb8ef commit 2824db3
Show file tree
Hide file tree
Showing 7 changed files with 132 additions and 100 deletions.
75 changes: 35 additions & 40 deletions compiler/rustc_abi/src/layout.rs
Original file line number Diff line number Diff line change
Expand Up @@ -8,19 +8,6 @@ use rand_xoshiro::Xoshiro128StarStar;

use tracing::debug;

// Invert a bijective mapping, i.e. `invert(map)[y] = x` if `map[x] = y`.
// 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.
fn invert_mapping(map: &[u32]) -> Vec<u32> {
let mut inverse = vec![0; map.len()];
for i in 0..map.len() {
inverse[map[i] as usize] = i as u32;
}
inverse
}

pub trait LayoutCalculator {
type TargetDataLayoutRef: Borrow<TargetDataLayout>;

Expand All @@ -45,8 +32,8 @@ pub trait LayoutCalculator {
LayoutS {
variants: Variants::Single { index: FIRST_VARIANT },
fields: FieldsShape::Arbitrary {
offsets: vec![Size::ZERO, b_offset],
memory_index: vec![0, 1],
offsets: [Size::ZERO, b_offset].into(),
memory_index: [0, 1].into(),
},
abi: Abi::ScalarPair(a, b),
largest_niche,
Expand All @@ -58,18 +45,18 @@ pub trait LayoutCalculator {
fn univariant(
&self,
dl: &TargetDataLayout,
fields: &[Layout<'_>],
fields: &IndexSlice<FieldIdx, Layout<'_>>,
repr: &ReprOptions,
kind: StructKind,
) -> Option<LayoutS> {
let pack = repr.pack;
let mut align = if pack.is_some() { dl.i8_align } else { dl.aggregate_align };
let mut inverse_memory_index: Vec<u32> = (0..fields.len() as u32).collect();
let mut inverse_memory_index: IndexVec<u32, FieldIdx> = fields.indices().collect();
let optimize = !repr.inhibit_struct_field_reordering_opt();
if optimize {
let end =
if let StructKind::MaybeUnsized = kind { fields.len() - 1 } else { fields.len() };
let optimizing = &mut inverse_memory_index[..end];
let optimizing = &mut inverse_memory_index.raw[..end];
let effective_field_align = |layout: Layout<'_>| {
if let Some(pack) = pack {
// return the packed alignment in bytes
Expand Down Expand Up @@ -105,7 +92,7 @@ pub trait LayoutCalculator {
// Place ZSTs first to avoid "interesting offsets",
// especially with only one or two non-ZST fields.
// Then place largest alignments first, largest niches within an alignment group last
let f = fields[x as usize];
let f = fields[x];
let niche_size = f.largest_niche().map_or(0, |n| n.available(dl));
(!f.0.is_zst(), cmp::Reverse(effective_field_align(f)), niche_size)
});
Expand All @@ -117,7 +104,7 @@ pub trait LayoutCalculator {
// And put the largest niche in an alignment group at the end
// so it can be used as discriminant in jagged enums
optimizing.sort_by_key(|&x| {
let f = fields[x as usize];
let f = fields[x];
let niche_size = f.largest_niche().map_or(0, |n| n.available(dl));
(effective_field_align(f), niche_size)
});
Expand All @@ -135,7 +122,7 @@ pub trait LayoutCalculator {
// At the bottom of this function, we invert `inverse_memory_index` to
// produce `memory_index` (see `invert_mapping`).
let mut sized = true;
let mut offsets = vec![Size::ZERO; fields.len()];
let mut offsets = IndexVec::from_elem(Size::ZERO, &fields);
let mut offset = Size::ZERO;
let mut largest_niche = None;
let mut largest_niche_available = 0;
Expand All @@ -146,7 +133,7 @@ pub trait LayoutCalculator {
offset = prefix_size.align_to(prefix_align);
}
for &i in &inverse_memory_index {
let field = &fields[i as usize];
let field = &fields[i];
if !sized {
self.delay_bug(&format!(
"univariant: field #{} comes after unsized field",
Expand All @@ -168,7 +155,7 @@ pub trait LayoutCalculator {
align = align.max(field_align);

debug!("univariant offset: {:?} field: {:#?}", offset, field);
offsets[i as usize] = offset;
offsets[i] = offset;

if let Some(mut niche) = field.largest_niche() {
let available = niche.available(dl);
Expand All @@ -192,14 +179,18 @@ pub trait LayoutCalculator {
// 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 { invert_mapping(&inverse_memory_index) } else { inverse_memory_index };
let memory_index = if optimize {
inverse_memory_index.invert_bijective_mapping()
} else {
debug_assert!(inverse_memory_index.iter().copied().eq(fields.indices()));
inverse_memory_index.into_iter().map(FieldIdx::as_u32).collect()
};
let size = min_size.align_to(align.abi);
let mut abi = Abi::Aggregate { sized };
// Unpack newtype ABIs and find scalar pairs.
if sized && size.bytes() > 0 {
// All other fields must be ZSTs.
let mut non_zst_fields = fields.iter().enumerate().filter(|&(_, f)| !f.0.is_zst());
let mut non_zst_fields = fields.iter_enumerated().filter(|&(_, f)| !f.0.is_zst());

match (non_zst_fields.next(), non_zst_fields.next(), non_zst_fields.next()) {
// We have exactly one non-ZST field.
Expand Down Expand Up @@ -238,13 +229,13 @@ pub trait LayoutCalculator {
let pair = self.scalar_pair(a, b);
let pair_offsets = match pair.fields {
FieldsShape::Arbitrary { ref offsets, ref memory_index } => {
assert_eq!(memory_index, &[0, 1]);
assert_eq!(memory_index.raw, [0, 1]);
offsets
}
_ => panic!(),
};
if offsets[i] == pair_offsets[0]
&& offsets[j] == pair_offsets[1]
if offsets[i] == pair_offsets[FieldIdx::from_usize(0)]
&& offsets[j] == pair_offsets[FieldIdx::from_usize(1)]
&& align == pair.align
&& size == pair.size
{
Expand Down Expand Up @@ -289,7 +280,7 @@ pub trait LayoutCalculator {
fn layout_of_struct_or_enum(
&self,
repr: &ReprOptions,
variants: &IndexSlice<VariantIdx, Vec<Layout<'_>>>,
variants: &IndexSlice<VariantIdx, IndexVec<FieldIdx, Layout<'_>>>,
is_enum: bool,
is_unsafe_cell: bool,
scalar_valid_range: (Bound<u128>, Bound<u128>),
Expand All @@ -312,7 +303,7 @@ pub trait LayoutCalculator {
// but *not* an encoding of the discriminant (e.g., a tag value).
// See issue #49298 for more details on the need to leave space
// for non-ZST uninhabited data (mostly partial initialization).
let absent = |fields: &[Layout<'_>]| {
let absent = |fields: &IndexSlice<FieldIdx, Layout<'_>>| {
let uninhabited = fields.iter().any(|f| f.abi().is_uninhabited());
let is_zst = fields.iter().all(|f| f.0.is_zst());
uninhabited && is_zst
Expand Down Expand Up @@ -510,7 +501,7 @@ pub trait LayoutCalculator {
// It'll fit, but we need to make some adjustments.
match layout.fields {
FieldsShape::Arbitrary { ref mut offsets, .. } => {
for (j, offset) in offsets.iter_mut().enumerate() {
for (j, offset) in offsets.iter_enumerated_mut() {
if !variants[i][j].0.is_zst() {
*offset += this_offset;
}
Expand Down Expand Up @@ -577,8 +568,8 @@ pub trait LayoutCalculator {
variants: IndexVec::new(),
},
fields: FieldsShape::Arbitrary {
offsets: vec![niche_offset],
memory_index: vec![0],
offsets: [niche_offset].into(),
memory_index: [0].into(),
},
abi,
largest_niche,
Expand Down Expand Up @@ -651,7 +642,8 @@ pub trait LayoutCalculator {
st.variants = Variants::Single { index: i };
// Find the first field we can't move later
// to make room for a larger discriminant.
for field in st.fields.index_by_increasing_offset().map(|j| &field_layouts[j]) {
for field_idx in st.fields.index_by_increasing_offset() {
let field = &field_layouts[FieldIdx::from_usize(field_idx)];
if !field.0.is_zst() || field.align().abi.bytes() != 1 {
start_align = start_align.min(field.align().abi);
break;
Expand Down Expand Up @@ -802,13 +794,13 @@ pub trait LayoutCalculator {
let pair = self.scalar_pair(tag, prim_scalar);
let pair_offsets = match pair.fields {
FieldsShape::Arbitrary { ref offsets, ref memory_index } => {
assert_eq!(memory_index, &[0, 1]);
assert_eq!(memory_index.raw, [0, 1]);
offsets
}
_ => panic!(),
};
if pair_offsets[0] == Size::ZERO
&& pair_offsets[1] == *offset
if pair_offsets[FieldIdx::from_u32(0)] == Size::ZERO
&& pair_offsets[FieldIdx::from_u32(1)] == *offset
&& align == pair.align
&& size == pair.size
{
Expand Down Expand Up @@ -844,7 +836,10 @@ pub trait LayoutCalculator {
tag_field: 0,
variants: IndexVec::new(),
},
fields: FieldsShape::Arbitrary { offsets: vec![Size::ZERO], memory_index: vec![0] },
fields: FieldsShape::Arbitrary {
offsets: [Size::ZERO].into(),
memory_index: [0].into(),
},
largest_niche,
abi,
align,
Expand Down Expand Up @@ -883,7 +878,7 @@ pub trait LayoutCalculator {
fn layout_of_union(
&self,
repr: &ReprOptions,
variants: &IndexSlice<VariantIdx, Vec<Layout<'_>>>,
variants: &IndexSlice<VariantIdx, IndexVec<FieldIdx, Layout<'_>>>,
) -> Option<LayoutS> {
let dl = self.current_data_layout();
let dl = dl.borrow();
Expand Down
23 changes: 11 additions & 12 deletions compiler/rustc_abi/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1108,7 +1108,7 @@ pub enum FieldsShape {
/// ordered to match the source definition order.
/// This vector does not go in increasing order.
// FIXME(eddyb) use small vector optimization for the common case.
offsets: Vec<Size>,
offsets: IndexVec<FieldIdx, Size>,

/// Maps source order field indices to memory order indices,
/// depending on how the fields were reordered (if at all).
Expand All @@ -1122,7 +1122,7 @@ pub enum FieldsShape {
///
// FIXME(eddyb) build a better abstraction for permutations, if possible.
// FIXME(camlorn) also consider small vector optimization here.
memory_index: Vec<u32>,
memory_index: IndexVec<FieldIdx, u32>,
},
}

Expand Down Expand Up @@ -1157,7 +1157,7 @@ impl FieldsShape {
assert!(i < count);
stride * i
}
FieldsShape::Arbitrary { ref offsets, .. } => offsets[i],
FieldsShape::Arbitrary { ref offsets, .. } => offsets[FieldIdx::from_usize(i)],
}
}

Expand All @@ -1168,28 +1168,27 @@ impl FieldsShape {
unreachable!("FieldsShape::memory_index: `Primitive`s have no fields")
}
FieldsShape::Union(_) | FieldsShape::Array { .. } => i,
FieldsShape::Arbitrary { ref memory_index, .. } => memory_index[i].try_into().unwrap(),
FieldsShape::Arbitrary { ref memory_index, .. } => {
memory_index[FieldIdx::from_usize(i)].try_into().unwrap()
}
}
}

/// Gets source indices of the fields by increasing offsets.
#[inline]
pub fn index_by_increasing_offset<'a>(&'a self) -> impl Iterator<Item = usize> + 'a {
let mut inverse_small = [0u8; 64];
let mut inverse_big = vec![];
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 i in 0..self.count() {
inverse_small[memory_index[i] as usize] = i as u8;
for (field_idx, &mem_idx) in memory_index.iter_enumerated() {
inverse_small[mem_idx as usize] = field_idx.as_u32() as u8;
}
} else {
inverse_big = vec![0; self.count()];
for i in 0..self.count() {
inverse_big[memory_index[i] as usize] = i as u32;
}
inverse_big = memory_index.invert_bijective_mapping();
}
}

Expand All @@ -1199,7 +1198,7 @@ impl FieldsShape {
if use_small {
inverse_small[i] as usize
} else {
inverse_big[i] as usize
inverse_big[i as u32].as_usize()
}
}
})
Expand Down
43 changes: 43 additions & 0 deletions compiler/rustc_index/src/vec.rs
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@ pub trait Idx: Copy + 'static + Eq + PartialEq + Debug + Hash {
}

#[inline]
#[must_use = "Use `increment_by` if you wanted to update the index in-place"]
fn plus(self, amount: usize) -> Self {
Self::new(self.index() + amount)
}
Expand Down Expand Up @@ -294,6 +295,11 @@ impl<I: Idx, T: Clone> ToOwned for IndexSlice<I, T> {
}

impl<I: Idx, T> IndexSlice<I, T> {
#[inline]
pub fn empty() -> &'static Self {
Default::default()
}

#[inline]
pub fn from_raw(raw: &[T]) -> &Self {
let ptr: *const [T] = raw;
Expand Down Expand Up @@ -409,6 +415,36 @@ impl<I: Idx, T> IndexSlice<I, T> {
}
}

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!(
self.iter().map(|x| x.index() as u128).sum::<u128>(),
(0..self.len() as u128).sum::<u128>(),
"The values aren't 0..N in input {self:?}",
);

let mut inverse = IndexVec::from_elem_n(Idx::new(0), self.len());
for (i1, &i2) in self.iter_enumerated() {
inverse[i2] = i1;
}

debug_assert_eq!(
inverse.iter().map(|x| x.index() as u128).sum::<u128>(),
(0..inverse.len() as u128).sum::<u128>(),
"The values aren't 0..N in result {self:?}",
);

inverse
}
}

/// `IndexVec` is often used as a map, so it provides some map-like APIs.
impl<I: Idx, T> IndexVec<I, Option<T>> {
#[inline]
Expand Down Expand Up @@ -513,6 +549,13 @@ impl<I: Idx, T> FromIterator<T> for IndexVec<I, T> {
}
}

impl<I: Idx, T, const N: usize> From<[T; N]> for IndexVec<I, T> {
#[inline]
fn from(array: [T; N]) -> Self {
IndexVec::from_raw(array.into())
}
}

impl<I: Idx, T> IntoIterator for IndexVec<I, T> {
type Item = T;
type IntoIter = vec::IntoIter<T>;
Expand Down
3 changes: 2 additions & 1 deletion compiler/rustc_middle/src/ty/layout.rs
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ use crate::ty::{self, ReprOptions, Ty, TyCtxt, TypeVisitableExt};
use rustc_errors::{DiagnosticBuilder, Handler, IntoDiagnostic};
use rustc_hir as hir;
use rustc_hir::def_id::DefId;
use rustc_index::vec::IndexVec;
use rustc_session::config::OptLevel;
use rustc_span::symbol::{sym, Symbol};
use rustc_span::{Span, DUMMY_SP};
Expand Down Expand Up @@ -635,7 +636,7 @@ where
variants: Variants::Single { index: variant_index },
fields: match NonZeroUsize::new(fields) {
Some(fields) => FieldsShape::Union(fields),
None => FieldsShape::Arbitrary { offsets: vec![], memory_index: vec![] },
None => FieldsShape::Arbitrary { offsets: IndexVec::new(), memory_index: IndexVec::new() },
},
abi: Abi::Uninhabited,
largest_niche: None,
Expand Down
2 changes: 1 addition & 1 deletion compiler/rustc_mir_transform/src/generator.rs
Original file line number Diff line number Diff line change
Expand Up @@ -287,7 +287,7 @@ impl<'tcx> TransformVisitor<'tcx> {
statements.push(Statement {
kind: StatementKind::Assign(Box::new((
Place::return_place(),
Rvalue::Aggregate(Box::new(kind), IndexVec::from_iter([val])),
Rvalue::Aggregate(Box::new(kind), [val].into()),
))),
source_info,
});
Expand Down
Loading

0 comments on commit 2824db3

Please sign in to comment.