fix(ast_tools): fix miscalculation of enum niches (#6743)

Follow-on after #5774. Correct the logic for calculating niches in enums.

It's still not quite correct - number of niches depends on how many spare discriminant "slots" there are at *start or end* of the range, not in total. But this is closer to correct than it was - we now don't take into account whether enum variant payloads have niches or not, which is not relevant.
This commit is contained in:
overlookmotel 2024-10-21 14:27:45 +00:00
parent 202c7f66c7
commit 5b41eeaa9d
2 changed files with 45 additions and 68 deletions

View file

@ -61,18 +61,6 @@ impl KnownLayout {
self.offsets.as_ref()
}
pub unsafe fn set_size_unchecked(&mut self, size: usize) {
self.size = size;
}
pub unsafe fn set_align_unchecked(&mut self, align: usize) {
self.align = align;
}
pub unsafe fn set_niches_unchecked(&mut self, niches: u128) {
self.niches = niches;
}
pub fn with_offsets(mut self, offsets: Vec<usize>) -> Self {
self.offsets = Some(offsets);
self

View file

@ -1,3 +1,5 @@
use std::cmp::max;
use itertools::Itertools;
use lazy_static::lazy_static;
use quote::ToTokens;
@ -81,69 +83,56 @@ pub fn calc_layout(ty: &mut AstType, ctx: &EarlyCtx) -> Result<bool> {
}
fn calc_enum_layout(ty: &mut Enum, ctx: &EarlyCtx) -> Result<PlatformLayout> {
fn collect_variant_layouts(ty: &Enum, ctx: &EarlyCtx) -> Result<Vec<PlatformLayout>> {
// all unit variants?
if ty.item.variants.iter().all(|var| var.fields.is_empty()) {
// all AST enums are `repr(u8)` so it would have a 1 byte layout/alignment,
// if it holds no data
let layout = KnownLayout::new(0, 1, 0);
let layout = Layout::Layout(layout);
Ok(vec![PlatformLayout(layout.clone(), layout)])
} else {
ty.item
.variants
.iter()
.map(|var| {
let typ =
var.fields.iter().exactly_one().map(|f| f.ty.analyze(ctx)).normalize()?;
calc_type_layout(&typ, ctx)
})
.collect()
}
struct SizeAlign {
size: usize,
align: usize,
}
#[expect(clippy::needless_pass_by_value)]
fn fold_layout(mut acc: KnownLayout, layout: KnownLayout) -> KnownLayout {
// SAFETY: we are folding valid layouts so it is safe.
unsafe {
// max alignment
if layout.align() > acc.align() {
acc.set_align_unchecked(layout.align());
}
// max size
if layout.size() > acc.size() {
acc.set_size_unchecked(layout.size());
}
// min niches
if layout.niches() < acc.niches() {
acc.set_niches_unchecked(layout.niches());
}
// Get max size and align of variants
let mut size_align_64 = SizeAlign { size: 0, align: 1 };
let mut size_align_32 = SizeAlign { size: 0, align: 1 };
for variant in &ty.item.variants {
if variant.fields.is_empty() {
continue;
}
acc
let field = variant.fields.iter().exactly_one().normalize().unwrap();
let typ = field.ty.analyze(ctx);
let PlatformLayout(variant_layout_64, variant_layout_32) = calc_type_layout(&typ, ctx)?;
let variant_layout_64 = variant_layout_64.layout().unwrap();
size_align_64.size = max(size_align_64.size, variant_layout_64.size());
size_align_64.align = max(size_align_64.align, variant_layout_64.align());
let variant_layout_32 = variant_layout_32.layout().unwrap();
size_align_32.size = max(size_align_32.size, variant_layout_32.size());
size_align_32.align = max(size_align_32.align, variant_layout_32.align());
}
let with_tag = |layout: KnownLayout| -> Result<KnownLayout> {
let niches = layout.niches();
let layout = std::alloc::Layout::from_size_align(layout.size(), layout.align())
.normalize()?
.pad_to_align();
let mut layout = KnownLayout::new(layout.size(), layout.align(), niches);
layout.consume_niches(ty.item.variants.len() as u128, true);
Ok(layout)
};
// Round up size to largest variant alignment.
// Largest variant is not necessarily the most highly aligned e.g. `enum { A([u8; 50]), B(u64) }`
size_align_64.size = size_align_64.size.next_multiple_of(size_align_64.align);
size_align_32.size = size_align_32.size.next_multiple_of(size_align_32.align);
let layouts = collect_variant_layouts(ty, ctx)?;
let (layouts_x64, layouts_x32): (Vec<KnownLayout>, Vec<KnownLayout>) = layouts
.into_iter()
.map(|PlatformLayout(x64, x32)| {
x64.layout().and_then(|x64| x32.layout().map(|x32| (x64, x32)))
})
.collect::<Option<_>>()
.expect("already checked.");
// Add discriminant.
// All enums are `#[repr(C, u8)]` (fieldful) or `#[repr(u8)]` (fieldless), so disriminant is 1 byte.
// But padding is inserted to align all payloads to largest alignment of any variant.
size_align_64.size += size_align_64.align;
size_align_32.size += size_align_32.align;
let x32 = with_tag(layouts_x32.into_iter().fold(KnownLayout::default(), fold_layout))?;
let x64 = with_tag(layouts_x64.into_iter().fold(KnownLayout::default(), fold_layout))?;
Ok(PlatformLayout(Layout::from(x64), Layout::from(x32)))
// Variant payloads are not relevant in niche calculation for `#[repr(u8)]` / `#[repr(C, u8)]` enums.
// The niche optimization for Option-like enums is disabled by `#[repr(u8)]`.
// https://doc.rust-lang.org/nightly/nomicon/other-reprs.html#repru-repri
// So number of niches only depends on the number of discriminants.
// TODO: This isn't quite correct. Number of niches depends only on how many unused discriminant
// values at *start* or *end* of range.
// https://github.com/oxc-project/oxc/pull/5774#pullrequestreview-2306334340
let niches = (256 - ty.item.variants.len()) as u128;
let layout_64 = KnownLayout::new(size_align_64.size, size_align_64.align, niches);
let layout_32 = KnownLayout::new(size_align_32.size, size_align_32.align, niches);
Ok(PlatformLayout(Layout::from(layout_64), Layout::from(layout_32)))
}
fn calc_struct_layout(ty: &mut Struct, ctx: &EarlyCtx) -> Result<PlatformLayout> {