use std::cmp;
use rustc_abi::{Align, Size};
use rustc_data_structures::fx::{FxHashMap, FxHashSet};
use rustc_data_structures::sync::Lock;
use rustc_span::Symbol;
use rustc_span::def_id::DefId;
#[derive(Clone, PartialEq, Eq, Hash, Debug)]
pub struct VariantInfo {
pub name: Option<Symbol>,
pub kind: SizeKind,
pub size: u64,
pub align: u64,
pub fields: Vec<FieldInfo>,
}
#[derive(Copy, Clone, PartialEq, Eq, Hash, Debug)]
pub enum SizeKind {
Exact,
Min,
}
#[derive(Copy, Clone, PartialEq, Eq, Hash, Debug)]
pub enum FieldKind {
AdtField,
Upvar,
CoroutineLocal,
}
impl std::fmt::Display for FieldKind {
fn fmt(&self, w: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
FieldKind::AdtField => write!(w, "field"),
FieldKind::Upvar => write!(w, "upvar"),
FieldKind::CoroutineLocal => write!(w, "local"),
}
}
}
#[derive(Copy, Clone, PartialEq, Eq, Hash, Debug)]
pub struct FieldInfo {
pub kind: FieldKind,
pub name: Symbol,
pub offset: u64,
pub size: u64,
pub align: u64,
pub type_name: Option<Symbol>,
}
#[derive(Copy, Clone, PartialEq, Eq, Hash, Debug)]
pub enum DataTypeKind {
Struct,
Union,
Enum,
Closure,
Coroutine,
}
#[derive(PartialEq, Eq, Hash, Debug)]
pub struct TypeSizeInfo {
pub kind: DataTypeKind,
pub type_description: String,
pub align: u64,
pub overall_size: u64,
pub packed: bool,
pub opt_discr_size: Option<u64>,
pub variants: Vec<VariantInfo>,
}
pub struct VTableSizeInfo {
pub trait_name: String,
pub entries: usize,
pub entries_ignoring_upcasting: usize,
pub entries_for_upcasting: usize,
pub upcasting_cost_percent: f64,
}
#[derive(Default)]
pub struct CodeStats {
type_sizes: Lock<FxHashSet<TypeSizeInfo>>,
vtable_sizes: Lock<FxHashMap<DefId, VTableSizeInfo>>,
}
impl CodeStats {
pub fn record_type_size<S: ToString>(
&self,
kind: DataTypeKind,
type_desc: S,
align: Align,
overall_size: Size,
packed: bool,
opt_discr_size: Option<Size>,
mut variants: Vec<VariantInfo>,
) {
if kind != DataTypeKind::Coroutine {
variants.sort_by_key(|info| cmp::Reverse(info.size));
}
let info = TypeSizeInfo {
kind,
type_description: type_desc.to_string(),
align: align.bytes(),
overall_size: overall_size.bytes(),
packed,
opt_discr_size: opt_discr_size.map(|s| s.bytes()),
variants,
};
self.type_sizes.borrow_mut().insert(info);
}
pub fn record_vtable_size(&self, trait_did: DefId, trait_name: &str, info: VTableSizeInfo) {
let prev = self.vtable_sizes.lock().insert(trait_did, info);
assert!(
prev.is_none(),
"size of vtable for `{trait_name}` ({trait_did:?}) is already recorded"
);
}
pub fn print_type_sizes(&self) {
let type_sizes = self.type_sizes.borrow();
#[allow(rustc::potential_query_instability)]
let mut sorted: Vec<_> = type_sizes.iter().collect();
sorted.sort_by_key(|info| (cmp::Reverse(info.overall_size), &info.type_description));
for info in sorted {
let TypeSizeInfo { type_description, overall_size, align, kind, variants, .. } = info;
println!(
"print-type-size type: `{type_description}`: {overall_size} bytes, alignment: {align} bytes"
);
let indent = " ";
let discr_size = if let Some(discr_size) = info.opt_discr_size {
println!("print-type-size {indent}discriminant: {discr_size} bytes");
discr_size
} else {
0
};
let mut max_variant_size = discr_size;
let struct_like = match kind {
DataTypeKind::Struct | DataTypeKind::Closure => true,
DataTypeKind::Enum | DataTypeKind::Union | DataTypeKind::Coroutine => false,
};
for (i, variant_info) in variants.into_iter().enumerate() {
let VariantInfo { ref name, kind: _, align: _, size, ref fields } = *variant_info;
let indent = if !struct_like {
let name = match name.as_ref() {
Some(name) => name.to_string(),
None => i.to_string(),
};
println!(
"print-type-size {indent}variant `{name}`: {diff} bytes",
diff = size - discr_size
);
" "
} else {
assert!(i < 1);
" "
};
max_variant_size = cmp::max(max_variant_size, size);
let mut min_offset = discr_size;
let mut fields = fields.clone();
fields.sort_by_key(|f| (f.offset, f.size));
for field in fields {
let FieldInfo { kind, ref name, offset, size, align, type_name } = field;
if offset > min_offset {
let pad = offset - min_offset;
println!("print-type-size {indent}padding: {pad} bytes");
}
if offset < min_offset {
print!(
"print-type-size {indent}{kind} `.{name}`: {size} bytes, \
offset: {offset} bytes, \
alignment: {align} bytes"
);
} else if info.packed || offset == min_offset {
print!("print-type-size {indent}{kind} `.{name}`: {size} bytes");
} else {
print!(
"print-type-size {indent}{kind} `.{name}`: {size} bytes, \
alignment: {align} bytes"
);
}
if let Some(type_name) = type_name {
println!(", type: {type_name}");
} else {
println!();
}
min_offset = offset + size;
}
}
match overall_size.checked_sub(max_variant_size) {
None => panic!("max_variant_size {max_variant_size} > {overall_size} overall_size"),
Some(diff @ 1..) => println!("print-type-size {indent}end padding: {diff} bytes"),
Some(0) => {}
}
}
}
pub fn print_vtable_sizes(&self, crate_name: Symbol) {
#[allow(rustc::potential_query_instability)]
let mut infos =
std::mem::take(&mut *self.vtable_sizes.lock()).into_values().collect::<Vec<_>>();
infos.sort_by(|a, b| {
a.upcasting_cost_percent
.total_cmp(&b.upcasting_cost_percent)
.reverse()
.then_with(|| a.trait_name.cmp(&b.trait_name))
});
for VTableSizeInfo {
trait_name,
entries,
entries_ignoring_upcasting,
entries_for_upcasting,
upcasting_cost_percent,
} in infos
{
println!(
r#"print-vtable-sizes {{ "crate_name": "{crate_name}", "trait_name": "{trait_name}", "entries": "{entries}", "entries_ignoring_upcasting": "{entries_ignoring_upcasting}", "entries_for_upcasting": "{entries_for_upcasting}", "upcasting_cost_percent": "{upcasting_cost_percent}" }}"#
);
}
}
}