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

crates/oxc_ast/src/lib.rs

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

tasks/ast_codegen/src/defs.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,24 @@ 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 size_32: usize,
+    pub align_32: usize,
 }
 
 #[derive(Debug, Serialize)]
@@ -23,6 +36,10 @@ 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 size_32: usize,
+    pub align_32: usize,
 }
 
 #[derive(Debug, Serialize)]
@@ -57,21 +74,47 @@ impl From<&RType> for Option<TypeDef> {
 
 impl From<&REnum> for EnumDef {
     fn from(it @ REnum { item, meta }: &REnum) -> Self {
+        let (size_64, align_64) = meta
+            .layout_64
+            .layout()
+            .map_or_else(|| panic!("Uncalculated layout on {}!", item.ident), KnownLayout::unpack);
+        let (size_32, align_32) = meta
+            .layout_32
+            .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,
+            size_32,
+            align_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) = meta
+            .layout_64
+            .layout()
+            .map_or_else(|| panic!("Uncalculated layout on {}!", item.ident), KnownLayout::unpack);
+        let (size_32, align_32) = meta
+            .layout_32
+            .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,
+            size_32,
+            align_32,
         }
     }
 }

tasks/ast_codegen/src/generators/assert_layouts.rs

@@ -0,0 +1,75 @@
+use proc_macro2::TokenStream;
+use quote::quote;
+use syn::{parse_quote, PathArguments, Type};
+
+use crate::{
+    defs::{EnumDef, StructDef, 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() {
+                                if let PathArguments::AngleBracketed(args) = &mut seg.arguments {
+                                    *args = parse_quote!(<'static>);
+                                }
+                            }
+                        }
+                        ty
+                    });
+                match def {
+                    TypeDef::Struct(StructDef { size_64, align_64, size_32, align_32, .. })
+                    | TypeDef::Enum(EnumDef { size_64, align_64, size_32, align_32, .. }) => (
+                        quote! {
+                            assert!(size_of::<#typ>() == #size_64);
+                            assert!(align_of::<#typ>() == #align_64);
+                        },
+                        quote! {
+                            assert!(size_of::<#typ>() == #size_32);
+                            assert!(align_of::<#typ>() == #align_32);
+                        },
+                    ),
+                }
+            })
+            .collect::<(Vec<TokenStream>, Vec<TokenStream>)>();
+
+        let header = generated_header!();
+
+        GeneratorOutput::Stream((
+            output(crate::AST_CRATE, "assert_layouts.rs"),
+            quote! {
+                #header
+
+                use crate::ast::*;
+
+                endl!();
+
+                #[cfg(target_pointer_width = "64")]
+                const _: () = { #(#assertions_64)* };
+                #[cfg(target_pointer_width = "32")]
+                const _: () = { #(#assertions_32)* };
+                #[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");
+            },
+        ))
+    }
+}

tasks/ast_codegen/src/generators/mod.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;

tasks/ast_codegen/src/layout.rs

@@ -0,0 +1,166 @@
+#[derive(Debug, Default, Clone, Copy)]
+pub enum Layout {
+    #[default]
+    Unknown,
+    Layout(KnownLayout),
+}
+
+impl Layout {
+    pub const fn known(size: usize, align: usize, niches: u128) -> Self {
+        Self::Layout(KnownLayout { size, align, niches })
+    }
+
+    pub fn layout(self) -> Option<KnownLayout> {
+        if let Self::Layout(layout) = self {
+            Some(layout)
+        } else {
+            None
+        }
+    }
+}
+
+impl From<KnownLayout> for Layout {
+    fn from(layout: KnownLayout) -> Self {
+        Self::Layout(layout)
+    }
+}
+
+#[derive(Debug, Default, Clone, Copy)]
+pub struct KnownLayout {
+    size: usize,
+    align: usize,
+    /// number of available niches
+    niches: u128,
+}
+
+impl KnownLayout {
+    pub const fn new(size: usize, align: usize, niches: u128) -> Self {
+        Self { size, align, niches }
+    }
+    #[inline]
+    pub fn size(&self) -> usize {
+        self.size
+    }
+
+    #[inline]
+    pub fn align(&self) -> usize {
+        self.align
+    }
+
+    /// number of available niches
+    #[inline]
+    pub fn niches(&self) -> u128 {
+        self.niches
+    }
+
+    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;
+    }
+
+    /// Panics
+    /// if doesn't have enough viable space and `can_resize` is false
+    pub fn consume_niches(&mut self, n: u128, can_resize: bool) {
+        if self.niches() >= n {
+            self.niches -= 1;
+        } else if can_resize {
+            let align = self.align();
+            self.size += align;
+            self.niches += match align {
+                1 => u8::MAX as u128,
+                2 => u16::MAX as u128,
+                4 => u32::MAX as u128,
+                8 => u64::MAX as u128,
+                16 => u128::MAX,
+                _ => unreachable!("We do not support paddings bigger than 16 bytes."),
+            };
+            self.consume_niches(n, can_resize);
+        } else {
+            panic!("`{}` called on a layout without enough space.", stringify!(consume_niches));
+        }
+    }
+
+    pub fn unpack(self) -> (/* size */ usize, /* align */ usize) {
+        let Self { size, align, .. } = self;
+        (size, align)
+    }
+}
+
+impl Layout {
+    /// # Panics
+    /// If alignment of `T` is higher than one byte.
+    pub const fn of<T>() -> Self {
+        // TODO: find a better way of calculating this.
+        struct N1<T>(Option<T>);
+        struct N2<T>(N1<N1<T>>);
+        struct N3<T>(N1<N2<T>>);
+        struct N4<T>(N1<N3<T>>);
+        struct N5<T>(N1<N4<T>>);
+        struct N6<T>(N1<N5<T>>);
+        struct N7<T>(N1<N6<T>>);
+        struct N8<T>(N1<N7<T>>);
+
+        let size = size_of::<T>();
+        let align = align_of::<T>();
+        let niches = if size_of::<N1<T>>() > size {
+            0
+        } else if size_of::<N2<T>>() > size {
+            1
+        } else if size_of::<N3<T>>() > size {
+            2
+        } else if size_of::<N4<T>>() > size {
+            3
+        } else if size_of::<N5<T>>() > size {
+            4
+        } else if size_of::<N6<T>>() > size {
+            5
+        } else if size_of::<N7<T>>() > size {
+            6
+        } else if size_of::<N8<T>>() > size {
+            7
+        } else if size_of::<N8<T>>() == size {
+            8
+        } else {
+            panic!(
+                "Alignment of `T` is bigger than what this method can calculate the headroom for."
+            );
+        };
+        // NOTE: some or all of `niches` might be `trailing_pad` but we don't need to
+        // distinguish between them. This method is only used to get layout info of simple types.
+        // most of them are builtin primitives.
+        Self::known(size, align, niches)
+    }
+
+    pub const fn zero() -> Self {
+        #[repr(C)]
+        struct Empty;
+        Self::of::<Empty>()
+    }
+
+    pub const fn ptr_32() -> Self {
+        Layout::known(4, 4, 0)
+    }
+
+    pub const fn ptr_64() -> Self {
+        Layout::known(8, 8, 0)
+    }
+
+    pub const fn wide_ptr_32() -> Self {
+        Layout::known(8, 4, 1)
+    }
+
+    pub const fn wide_ptr_64() -> Self {
+        Layout::of::<&str>()
+    }
+
+    pub fn is_unknown(&self) -> bool {
+        matches!(self, Self::Unknown)
+    }
+}

tasks/ast_codegen/src/main.rs

@@ -1,8 +1,11 @@
 const AST_CRATE: &str = "crates/oxc_ast";
+#[allow(dead_code)]
+const AST_MACROS_CRATE: &str = "crates/oxc_ast_macros";
 
 mod defs;
 mod fmt;
 mod generators;
+mod layout;
 mod markers;
 mod passes;
 mod schema;
@@ -13,13 +16,14 @@ use std::{cell::RefCell, collections::HashMap, io::Read, path::PathBuf, rc::Rc};
 use bpaf::{Bpaf, Parser};
 use fmt::{cargo_fmt, pprint};
 use itertools::Itertools;
-use passes::{BuildSchema, Linker, Pass};
+use passes::{BuildSchema, CalcLayout, Linker, Pass};
 use proc_macro2::TokenStream;
 use syn::parse_file;
 
 use defs::TypeDef;
 use generators::{
-    AstBuilderGenerator, AstKindGenerator, Generator, VisitGenerator, VisitMutGenerator,
+    AssertLayouts, AstBuilderGenerator, AstKindGenerator, Generator, VisitGenerator,
+    VisitMutGenerator,
 };
 use schema::{Module, REnum, RStruct, RType, Schema};
 use util::{write_all_to, NormalizeError};
@@ -258,7 +262,9 @@ fn main() -> std::result::Result<(), Box<dyn std::error::Error>> {
     let CodegenResult { outputs, schema } = files()
         .fold(AstCodegen::default(), AstCodegen::add_file)
         .pass(Linker)
+        .pass(CalcLayout)
         .pass(BuildSchema)
+        .gen(AssertLayouts)
         .gen(AstKindGenerator)
         .gen(AstBuilderGenerator)
         .gen(ImplGetSpanGenerator)