feat(ast_codegen): add alignment and size data to the schema. (#4615)

This PR generates the layouts in the schema but doesn't use it to reorder.

Author: rzvxa

crates/oxc_ast/src/generated/assert_layouts.rs

@@ -37,6 +37,8 @@ pub mod syntax_directed_operations;
 mod trivia;
 
 mod generated {
+    #[cfg(test)]
+    pub mod assert_layouts;
     pub mod ast_builder;
     pub mod ast_kind;
     pub mod span;

crates/oxc_ast/src/lib.rs

@@ -1,5 +1,5 @@
 use super::{REnum, RStruct, RType};
-use crate::{schema::Inherit, util::TypeExt, TypeName};
+use crate::{layout::KnownLayout, schema::Inherit, util::TypeExt, TypeName};
 use quote::ToTokens;
 use serde::Serialize;
 
@@ -9,11 +9,26 @@ pub enum TypeDef {
     Enum(EnumDef),
 }
 
+impl TypeDef {
+    pub fn name(&self) -> &String {
+        match self {
+            Self::Struct(it) => &it.name,
+            Self::Enum(it) => &it.name,
+        }
+    }
+}
+
 #[derive(Debug, Serialize)]
 pub struct StructDef {
     pub name: TypeName,
     pub fields: Vec<FieldDef>,
     pub has_lifetime: bool,
+    pub size_64: usize,
+    pub align_64: usize,
+    pub offsets_64: Option<Vec<usize>>,
+    pub size_32: usize,
+    pub align_32: usize,
+    pub offsets_32: Option<Vec<usize>>,
 }
 
 #[derive(Debug, Serialize)]
@@ -23,6 +38,12 @@ pub struct EnumDef {
     /// For `@inherits` inherited enum variants
     pub inherits: Vec<EnumInheritDef>,
     pub has_lifetime: bool,
+    pub size_64: usize,
+    pub align_64: usize,
+    pub offsets_64: Option<Vec<usize>>,
+    pub size_32: usize,
+    pub align_32: usize,
+    pub offsets_32: Option<Vec<usize>>,
 }
 
 #[derive(Debug, Serialize)]
@@ -57,21 +78,55 @@ impl From<&RType> for Option<TypeDef> {
 
 impl From<&REnum> for EnumDef {
     fn from(it @ REnum { item, meta }: &REnum) -> Self {
+        let (size_64, align_64, offsets_64) = meta
+            .layout_64
+            .clone()
+            .layout()
+            .map_or_else(|| panic!("Uncalculated layout on {}!", item.ident), KnownLayout::unpack);
+        let (size_32, align_32, offsets_32) = meta
+            .layout_32
+            .clone()
+            .layout()
+            .map_or_else(|| panic!("Uncalculated layout on {}!", item.ident), KnownLayout::unpack);
         Self {
             name: it.ident().to_string(),
             variants: item.variants.iter().map(Into::into).collect(),
-            has_lifetime: item.generics.lifetimes().count() > 0,
             inherits: meta.inherits.iter().map(Into::into).collect(),
+            has_lifetime: item.generics.lifetimes().count() > 0,
+
+            size_64,
+            align_64,
+            offsets_64,
+            size_32,
+            align_32,
+            offsets_32,
         }
     }
 }
 
 impl From<&RStruct> for StructDef {
-    fn from(it @ RStruct { item, .. }: &RStruct) -> Self {
+    fn from(it @ RStruct { item, meta }: &RStruct) -> Self {
+        let (size_64, align_64, offsets_64) = meta
+            .layout_64
+            .clone()
+            .layout()
+            .map_or_else(|| panic!("Uncalculated layout on {}!", item.ident), KnownLayout::unpack);
+        let (size_32, align_32, offsets_32) = meta
+            .layout_32
+            .clone()
+            .layout()
+            .map_or_else(|| panic!("Uncalculated layout on {}!", item.ident), KnownLayout::unpack);
         Self {
             name: it.ident().to_string(),
             fields: item.fields.iter().map(Into::into).collect(),
             has_lifetime: item.generics.lifetimes().count() > 0,
+
+            size_64,
+            align_64,
+            offsets_64,
+            size_32,
+            align_32,
+            offsets_32,
         }
     }
 }

tasks/ast_codegen/src/defs.rs

@@ -0,0 +1,118 @@
+use proc_macro2::TokenStream;
+use quote::{format_ident, quote};
+use syn::{PathArguments, Type};
+
+use crate::{
+    defs::{FieldDef, TypeDef},
+    output, CodegenCtx, Generator, GeneratorOutput,
+};
+
+use super::{define_generator, generated_header};
+
+define_generator! {
+    pub struct AssertLayouts;
+}
+
+impl Generator for AssertLayouts {
+    fn name(&self) -> &'static str {
+        stringify!(AssertLayouts)
+    }
+
+    fn generate(&mut self, ctx: &CodegenCtx) -> GeneratorOutput {
+        let (assertions_64, assertions_32) = ctx
+            .schema
+            .borrow()
+            .definitions
+            .iter()
+            .map(|def| {
+                let typ =
+                    ctx.find(def.name()).and_then(|ty| ty.borrow().as_type()).map(|mut ty| {
+                        if let Type::Path(ty) = &mut ty {
+                            if let Some(seg) = ty.path.segments.first_mut() {
+                                seg.arguments = PathArguments::None;
+                            }
+                        }
+                        ty
+                    });
+                let typ = typ.unwrap();
+                assert_type(&typ, def)
+            })
+            .collect::<(Vec<TokenStream>, Vec<TokenStream>)>();
+
+        let header = generated_header!();
+
+        GeneratorOutput::Stream((
+            output(crate::AST_CRATE, "assert_layouts.rs"),
+            quote! {
+                #header
+
+                use std::mem::{align_of, offset_of, size_of};
+
+                endl!();
+
+                use crate::ast::*;
+
+                endl!();
+
+                #[cfg(target_pointer_width = "64")]
+                const _: () = { #(#assertions_64)* };
+                endl!();
+
+                #[cfg(target_pointer_width = "32")]
+                const _: () = { #(#assertions_32)* };
+                endl!();
+
+                #[cfg(not(any(target_pointer_width = "64", target_pointer_width = "32")))]
+                const _: () = panic!("Platforms with pointer width other than 64 or 32 bit are not supported");
+            },
+        ))
+    }
+}
+
+fn assert_type(ty: &Type, def: &TypeDef) -> (TokenStream, TokenStream) {
+    match def {
+        TypeDef::Struct(def) => (
+            with_offsets_assertion(
+                assert_size_align(ty, def.size_64, def.align_64),
+                ty,
+                &def.fields,
+                def.offsets_64.as_deref(),
+            ),
+            with_offsets_assertion(
+                assert_size_align(ty, def.size_32, def.align_32),
+                ty,
+                &def.fields,
+                def.offsets_64.as_deref(),
+            ),
+        ),
+        TypeDef::Enum(def) => (
+            assert_size_align(ty, def.size_64, def.align_64),
+            assert_size_align(ty, def.size_32, def.align_32),
+        ),
+    }
+}
+
+fn assert_size_align(ty: &Type, size: usize, align: usize) -> TokenStream {
+    quote! {
+        assert!(size_of::<#ty>() == #size);
+        assert!(align_of::<#ty>() == #align);
+    }
+}
+
+fn with_offsets_assertion(
+    mut tk: TokenStream,
+    ty: &Type,
+    fields: &[FieldDef],
+    offsets: Option<&[usize]>,
+) -> TokenStream {
+    let Some(offsets) = offsets else { return tk };
+
+    let assertions = fields.iter().zip(offsets).map(|(field, offset)| {
+        let field = field.name.as_ref().map(|it| format_ident!("{it}"));
+        quote! {
+            assert!(offset_of!(#ty, #field) == #offset);
+        }
+    });
+    tk.extend(assertions);
+    tk
+}

tasks/ast_codegen/src/generators/assert_layouts.rs

@@ -1,3 +1,4 @@
+mod assert_layouts;
 mod ast_builder;
 mod ast_kind;
 mod impl_get_span;
@@ -39,6 +40,7 @@ macro_rules! generated_header {
 pub(crate) use generated_header;
 pub(crate) use insert;
 
+pub use assert_layouts::AssertLayouts;
 pub use ast_builder::AstBuilderGenerator;
 pub use ast_kind::AstKindGenerator;
 pub use impl_get_span::ImplGetSpanGenerator;