diff --git a/tasks/ast_codegen/src/codegen.rs b/tasks/ast_codegen/src/codegen.rs
deleted file mode 100644
index b3c0f8066bccb..0000000000000
--- a/tasks/ast_codegen/src/codegen.rs
+++ /dev/null
@@ -1,156 +0,0 @@
-use std::{cell::RefCell, collections::HashMap, path::PathBuf};
-
-use itertools::Itertools;
-
-use crate::{
-    generators::{Generator, GeneratorOutput},
-    passes::Pass,
-    rust_ast::{self, AstRef},
-    schema::{lower_ast_types, Schema, TypeDef},
-    Result, TypeId,
-};
-
-#[derive(Default)]
-pub struct AstCodegen {
-    files: Vec<PathBuf>,
-    passes: Vec<Box<dyn Runner<Output = (), Context = EarlyCtx>>>,
-    generators: Vec<Box<dyn Runner<Output = GeneratorOutput, Context = LateCtx>>>,
-}
-
-pub struct AstCodegenResult {
-    pub schema: Schema,
-    pub outputs: Vec<(/* generator name */ &'static str, /* output */ GeneratorOutput)>,
-}
-
-pub trait Runner {
-    type Context;
-    type Output;
-    fn name(&self) -> &'static str;
-    fn run(&mut self, ctx: &Self::Context) -> Result<Self::Output>;
-}
-
-pub struct EarlyCtx {
-    ty_table: Vec<AstRef>,
-    ident_table: HashMap<String, TypeId>,
-    mods: RefCell<Vec<rust_ast::Module>>,
-}
-
-impl EarlyCtx {
-    fn new(mods: Vec<rust_ast::Module>) -> Self {
-        // worst case len
-        let len = mods.iter().fold(0, |acc, it| acc + it.items.len());
-        let adts = mods.iter().flat_map(|it| it.items.iter());
-
-        let mut ty_table = Vec::with_capacity(len);
-        let mut ident_table = HashMap::with_capacity(len);
-        for adt in adts {
-            if let Some(ident) = adt.borrow().ident() {
-                let ident = ident.to_string();
-                let type_id = ty_table.len();
-                ty_table.push(AstRef::clone(adt));
-                ident_table.insert(ident, type_id);
-            }
-        }
-
-        Self { ty_table, ident_table, mods: RefCell::new(mods) }
-    }
-
-    pub fn chronological_idents(&self) -> impl Iterator<Item = &String> {
-        self.ident_table.iter().sorted_by_key(|it| it.1).map(|it| it.0)
-    }
-
-    pub fn mods(&self) -> &RefCell<Vec<rust_ast::Module>> {
-        &self.mods
-    }
-
-    pub fn find(&self, key: &String) -> Option<AstRef> {
-        self.type_id(key).map(|id| AstRef::clone(&self.ty_table[id]))
-    }
-
-    pub fn type_id(&self, key: &String) -> Option<TypeId> {
-        self.ident_table.get(key).copied()
-    }
-
-    pub fn ast_ref(&self, id: TypeId) -> AstRef {
-        AstRef::clone(&self.ty_table[id])
-    }
-
-    fn into_late_ctx(self) -> LateCtx {
-        let schema = lower_ast_types(&self);
-
-        LateCtx { schema }
-    }
-}
-
-pub struct LateCtx {
-    schema: Schema,
-}
-
-impl LateCtx {
-    pub fn schema(&self) -> &Schema {
-        &self.schema
-    }
-
-    pub fn type_def(&self, id: TypeId) -> Option<&TypeDef> {
-        self.schema.get(id)
-    }
-}
-
-impl AstCodegen {
-    #[must_use]
-    pub fn add_file<P>(mut self, path: P) -> Self
-    where
-        P: AsRef<str>,
-    {
-        self.files.push(path.as_ref().into());
-        self
-    }
-
-    #[must_use]
-    pub fn pass<P>(mut self, pass: P) -> Self
-    where
-        P: Pass + Runner<Output = (), Context = EarlyCtx> + 'static,
-    {
-        self.passes.push(Box::new(pass));
-        self
-    }
-
-    #[must_use]
-    pub fn gen<G>(mut self, generator: G) -> Self
-    where
-        G: Generator + Runner<Output = GeneratorOutput, Context = LateCtx> + 'static,
-    {
-        self.generators.push(Box::new(generator));
-        self
-    }
-
-    pub fn generate(self) -> Result<AstCodegenResult> {
-        let modules = self
-            .files
-            .into_iter()
-            .map(rust_ast::Module::from)
-            .map(rust_ast::Module::load)
-            .map_ok(rust_ast::Module::expand)
-            .map_ok(|it| it.map(rust_ast::Module::analyze))
-            .collect::<Result<Result<Result<Vec<_>>>>>()???;
-
-        // early passes
-        let ctx = {
-            let ctx = EarlyCtx::new(modules);
-            _ = self
-                .passes
-                .into_iter()
-                .map(|mut runner| runner.run(&ctx).map(|res| (runner.name(), res)))
-                .collect::<Result<Vec<_>>>()?;
-            ctx.into_late_ctx()
-        };
-
-        let outputs = self
-            .generators
-            .into_iter()
-            .map(|mut runner| runner.run(&ctx).map(|res| (runner.name(), res)))
-            .collect::<Result<Vec<_>>>()?;
-
-        Ok(AstCodegenResult { outputs, schema: ctx.schema })
-    }
-}
diff --git a/tasks/ast_codegen/src/generators/assert_layouts.rs b/tasks/ast_codegen/src/generators/assert_layouts.rs
index 28d5df5527bcc..aa8bdfecc6575 100644
--- a/tasks/ast_codegen/src/generators/assert_layouts.rs
+++ b/tasks/ast_codegen/src/generators/assert_layouts.rs
@@ -3,10 +3,9 @@ use quote::{format_ident, quote};
 use syn::Type;
 
 use crate::{
-    codegen::LateCtx,
     output,
     schema::{FieldDef, ToType, TypeDef},
-    Generator, GeneratorOutput,
+    Generator, GeneratorOutput, LateCtx,
 };
 
 use super::{define_generator, generated_header};
@@ -22,8 +21,9 @@ impl Generator for AssertLayouts {
 
     fn generate(&mut self, ctx: &LateCtx) -> GeneratorOutput {
         let (assertions_64, assertions_32) = ctx
-            .schema()
-            .into_iter()
+            .schema
+            .definitions
+            .iter()
             .map(|def| {
                 let typ = def.to_type_elide();
                 assert_type(&typ, def)
diff --git a/tasks/ast_codegen/src/generators/ast_builder.rs b/tasks/ast_codegen/src/generators/ast_builder.rs
index 6f8a66a0c73d0..47055e4295d3e 100644
--- a/tasks/ast_codegen/src/generators/ast_builder.rs
+++ b/tasks/ast_codegen/src/generators/ast_builder.rs
@@ -9,14 +9,13 @@ use quote::{format_ident, quote, ToTokens};
 use syn::{parse_quote, Ident, Type};
 
 use crate::{
-    codegen::LateCtx,
     generators::generated_header,
     output,
     schema::{
         EnumDef, FieldDef, GetIdent, InheritDef, StructDef, ToType, TypeDef, TypeName, VariantDef,
     },
     util::{TypeAnalysis, TypeWrapper},
-    Generator, GeneratorOutput,
+    Generator, GeneratorOutput, LateCtx,
 };
 
 use super::define_generator;
@@ -32,8 +31,9 @@ impl Generator for AstBuilderGenerator {
 
     fn generate(&mut self, ctx: &LateCtx) -> GeneratorOutput {
         let fns = ctx
-            .schema()
-            .into_iter()
+            .schema
+            .definitions
+            .iter()
             .filter(|it| it.visitable())
             .map(|it| generate_builder_fn(it, ctx))
             .collect_vec();
diff --git a/tasks/ast_codegen/src/generators/ast_kind.rs b/tasks/ast_codegen/src/generators/ast_kind.rs
index 5b1368fa2b346..f75ce47592d81 100644
--- a/tasks/ast_codegen/src/generators/ast_kind.rs
+++ b/tasks/ast_codegen/src/generators/ast_kind.rs
@@ -3,11 +3,10 @@ use quote::quote;
 use syn::{parse_quote, Arm, Ident, Type, Variant};
 
 use crate::{
-    codegen::LateCtx,
     output,
     schema::{GetIdent, ToType, TypeDef},
     util::ToIdent,
-    Generator, GeneratorOutput,
+    Generator, GeneratorOutput, LateCtx,
 };
 
 use super::{define_generator, generated_header};
@@ -142,8 +141,8 @@ impl Generator for AstKindGenerator {
 
     fn generate(&mut self, ctx: &LateCtx) -> GeneratorOutput {
         let have_kinds: Vec<(Ident, Type)> = ctx
-            .schema()
-            .into_iter()
+            .schema.definitions
+            .iter()
             .filter(|it| it.visitable())
             .filter(
                 |maybe_kind| matches!(maybe_kind, kind @ (TypeDef::Enum(_) | TypeDef::Struct(_)) if kind.visitable())
diff --git a/tasks/ast_codegen/src/generators/derive_clone_in.rs b/tasks/ast_codegen/src/generators/derive_clone_in.rs
index 1fae59a78d71b..f3316c14cf902 100644
--- a/tasks/ast_codegen/src/generators/derive_clone_in.rs
+++ b/tasks/ast_codegen/src/generators/derive_clone_in.rs
@@ -4,10 +4,9 @@ use quote::{format_ident, quote};
 use syn::Ident;
 
 use crate::{
-    codegen::LateCtx,
     output,
     schema::{EnumDef, GetIdent, StructDef, TypeDef},
-    GeneratorOutput,
+    GeneratorOutput, LateCtx,
 };
 
 use super::{define_generator, generated_header, Generator};
@@ -23,8 +22,9 @@ impl Generator for DeriveCloneIn {
 
     fn generate(&mut self, ctx: &LateCtx) -> GeneratorOutput {
         let impls: Vec<TokenStream> = ctx
-            .schema()
-            .into_iter()
+            .schema
+            .definitions
+            .iter()
             .filter(|def| def.generates_derive("CloneIn"))
             .map(|def| match &def {
                 TypeDef::Enum(it) => derive_enum(it),
diff --git a/tasks/ast_codegen/src/generators/derive_get_span.rs b/tasks/ast_codegen/src/generators/derive_get_span.rs
index a1fee556a549a..ab30bd4b0b544 100644
--- a/tasks/ast_codegen/src/generators/derive_get_span.rs
+++ b/tasks/ast_codegen/src/generators/derive_get_span.rs
@@ -3,11 +3,10 @@ use quote::{format_ident, quote};
 use syn::{Generics, Ident, Type};
 
 use crate::{
-    codegen::LateCtx,
     output,
     schema::{EnumDef, GetGenerics, StructDef, ToType, TypeDef},
     util::ToIdent,
-    Generator, GeneratorOutput,
+    Generator, GeneratorOutput, LateCtx,
 };
 
 use super::{define_generator, generated_header};
@@ -89,8 +88,9 @@ fn derive<const MUT: bool>(ctx: &LateCtx) -> TokenStream {
         )
     };
     let impls: Vec<TokenStream> = ctx
-        .schema()
-        .into_iter()
+        .schema
+        .definitions
+        .iter()
         .filter(|def| def.visitable())
         .map(|def| match &def {
             TypeDef::Enum(it) => derive_enum(it),
diff --git a/tasks/ast_codegen/src/generators/mod.rs b/tasks/ast_codegen/src/generators/mod.rs
index 3cac362d08072..bc7a039ee9374 100644
--- a/tasks/ast_codegen/src/generators/mod.rs
+++ b/tasks/ast_codegen/src/generators/mod.rs
@@ -38,8 +38,6 @@ macro_rules! generated_header {
     }};
 }
 
-use std::path::PathBuf;
-
 pub(crate) use generated_header;
 pub(crate) use insert;
 
@@ -48,101 +46,27 @@ pub use ast_builder::AstBuilderGenerator;
 pub use ast_kind::AstKindGenerator;
 pub use derive_clone_in::DeriveCloneIn;
 pub use derive_get_span::{DeriveGetSpan, DeriveGetSpanMut};
-use proc_macro2::TokenStream;
 pub use visit::{VisitGenerator, VisitMutGenerator};
 
-use crate::codegen::LateCtx;
+use crate::{GeneratorOutput, LateCtx};
 
 pub trait Generator {
     fn name(&self) -> &'static str;
     fn generate(&mut self, ctx: &LateCtx) -> GeneratorOutput;
 }
 
-pub type GeneratedTokenStream = (/* output path */ PathBuf, TokenStream);
-pub type GeneratedDataStream = (/* output path */ PathBuf, Vec<u8>);
-
-// TODO: remove me
-#[allow(dead_code)]
-#[derive(Debug, Clone)]
-pub enum GeneratorOutput {
-    None,
-    Info(Vec<u8>),
-    Data(GeneratedDataStream),
-    Stream(GeneratedTokenStream),
-}
-
-// TODO: remove me
-#[allow(dead_code)]
-impl GeneratorOutput {
-    pub fn is_none(&self) -> bool {
-        matches!(self, Self::None)
-    }
-
-    pub fn expect_none(&self) {
-        assert!(self.is_none());
-    }
-
-    pub fn to_info(&self) -> &[u8] {
-        if let Self::Info(it) = self {
-            it
-        } else {
-            panic!();
-        }
-    }
-
-    pub fn to_data(&self) -> &GeneratedDataStream {
-        if let Self::Data(it) = self {
-            it
-        } else {
-            panic!();
-        }
-    }
-
-    pub fn to_stream(&self) -> &GeneratedTokenStream {
-        if let Self::Stream(it) = self {
-            it
-        } else {
-            panic!();
-        }
-    }
-
-    pub fn into_info(self) -> Vec<u8> {
-        if let Self::Info(it) = self {
-            it
-        } else {
-            panic!();
-        }
-    }
-
-    pub fn into_data(self) -> GeneratedDataStream {
-        if let Self::Data(it) = self {
-            it
-        } else {
-            panic!();
-        }
-    }
-
-    pub fn into_stream(self) -> GeneratedTokenStream {
-        if let Self::Stream(it) = self {
-            it
-        } else {
-            panic!();
-        }
-    }
-}
-
 macro_rules! define_generator {
     ($vis:vis struct $ident:ident $($lifetime:lifetime)? $($rest:tt)*) => {
         $vis struct $ident $($lifetime)? $($rest)*
-        impl $($lifetime)? $crate::codegen::Runner for $ident $($lifetime)? {
-            type Context = $crate::codegen::LateCtx;
+        impl $($lifetime)? $crate::Runner for $ident $($lifetime)? {
+            type Context = $crate::LateCtx;
             type Output = $crate::GeneratorOutput;
 
             fn name(&self) -> &'static str {
                 $crate::Generator::name(self)
             }
 
-            fn run(&mut self, ctx: &$crate::codegen::LateCtx) -> $crate::Result<Self::Output> {
+            fn run(&mut self, ctx: &$crate::LateCtx) -> $crate::Result<Self::Output> {
                 Ok(self.generate(ctx))
             }
         }
diff --git a/tasks/ast_codegen/src/generators/visit.rs b/tasks/ast_codegen/src/generators/visit.rs
index 2527421f8edc4..de400836e0d0e 100644
--- a/tasks/ast_codegen/src/generators/visit.rs
+++ b/tasks/ast_codegen/src/generators/visit.rs
@@ -7,13 +7,12 @@ use quote::{format_ident, quote, ToTokens};
 use syn::{parse_quote, Ident};
 
 use crate::{
-    codegen::LateCtx,
     generators::{ast_kind::BLACK_LIST as KIND_BLACK_LIST, insert},
     markers::{ScopeMarkers, VisitArg, VisitMarkers},
     output,
     schema::{EnumDef, GetIdent, StructDef, ToType, TypeDef},
     util::{StrExt, ToIdent, TokenStreamExt, TypeWrapper},
-    Generator, GeneratorOutput,
+    Generator, GeneratorOutput, LateCtx,
 };
 
 use super::{define_generator, generated_header};
@@ -171,8 +170,9 @@ impl<'a> VisitBuilder<'a> {
     fn build(mut self) -> (/* visits */ Vec<TokenStream>, /* walks */ Vec<TokenStream>) {
         let program = self
             .ctx
-            .schema()
-            .into_iter()
+            .schema
+            .definitions
+            .iter()
             .filter(|it| it.visitable())
             .find(|it| it.name() == "Program")
             .expect("Couldn't find the `Program` type!");
diff --git a/tasks/ast_codegen/src/main.rs b/tasks/ast_codegen/src/main.rs
index ad8e281216b58..bbf25b9cdcf48 100644
--- a/tasks/ast_codegen/src/main.rs
+++ b/tasks/ast_codegen/src/main.rs
@@ -1,11 +1,10 @@
-use std::{cell::RefCell, io::Read, path::PathBuf, rc::Rc};
+use std::{cell::RefCell, collections::HashMap, io::Read, path::PathBuf, rc::Rc};
 
 use bpaf::{Bpaf, Parser};
-use codegen::{AstCodegen, AstCodegenResult};
 use itertools::Itertools;
+use proc_macro2::TokenStream;
 use syn::parse_file;
 
-mod codegen;
 mod fmt;
 mod generators;
 mod layout;
@@ -18,10 +17,11 @@ mod util;
 use fmt::{cargo_fmt, pprint};
 use generators::{
     AssertLayouts, AstBuilderGenerator, AstKindGenerator, DeriveCloneIn, DeriveGetSpan,
-    DeriveGetSpanMut, GeneratedDataStream, GeneratedTokenStream, Generator, GeneratorOutput,
-    VisitGenerator, VisitMutGenerator,
+    DeriveGetSpanMut, Generator, VisitGenerator, VisitMutGenerator,
 };
 use passes::{CalcLayout, Linker, Pass};
+use rust_ast::AstRef;
+use schema::{lower_ast_types, Schema, TypeDef};
 use util::{write_all_to, NormalizeError};
 
 static SOURCE_PATHS: &[&str] = &[
@@ -36,9 +36,14 @@ static SOURCE_PATHS: &[&str] = &[
 ];
 
 const AST_CRATE: &str = "crates/oxc_ast";
+#[allow(dead_code)]
+const AST_MACROS_CRATE: &str = "crates/oxc_ast_macros";
 
 type Result<R> = std::result::Result<R, String>;
 type TypeId = usize;
+type TypeTable = Vec<AstRef>;
+type DefTable = Vec<TypeDef>;
+type IdentTable = HashMap<String, TypeId>;
 
 #[derive(Debug, Bpaf)]
 pub struct CliOptions {
@@ -52,10 +57,217 @@ pub struct CliOptions {
     schema: Option<std::path::PathBuf>,
 }
 
+#[derive(Default)]
+struct AstCodegen {
+    files: Vec<PathBuf>,
+    passes: Vec<Box<dyn Runner<Output = (), Context = EarlyCtx>>>,
+    builders: Vec<Box<dyn Runner<Output = GeneratorOutput, Context = LateCtx>>>,
+}
+
+type GeneratedStream = (/* output path */ PathBuf, TokenStream);
+type DataStream = (/* output path */ PathBuf, Vec<u8>);
+
+// TODO: remove me
+#[allow(dead_code)]
+#[derive(Debug, Clone)]
+enum GeneratorOutput {
+    None,
+    Info(Vec<u8>),
+    Data(DataStream),
+    Stream(GeneratedStream),
+}
+
+// TODO: remove me
+#[allow(dead_code)]
+impl GeneratorOutput {
+    pub fn is_none(&self) -> bool {
+        matches!(self, Self::None)
+    }
+
+    pub fn expect_none(&self) {
+        assert!(self.is_none());
+    }
+
+    pub fn to_info(&self) -> &[u8] {
+        if let Self::Info(it) = self {
+            it
+        } else {
+            panic!();
+        }
+    }
+
+    pub fn to_data(&self) -> &DataStream {
+        if let Self::Data(it) = self {
+            it
+        } else {
+            panic!();
+        }
+    }
+
+    pub fn to_stream(&self) -> &GeneratedStream {
+        if let Self::Stream(it) = self {
+            it
+        } else {
+            panic!();
+        }
+    }
+
+    pub fn into_info(self) -> Vec<u8> {
+        if let Self::Info(it) = self {
+            it
+        } else {
+            panic!();
+        }
+    }
+
+    pub fn into_data(self) -> DataStream {
+        if let Self::Data(it) = self {
+            it
+        } else {
+            panic!();
+        }
+    }
+
+    pub fn into_stream(self) -> GeneratedStream {
+        if let Self::Stream(it) = self {
+            it
+        } else {
+            panic!();
+        }
+    }
+}
+
+struct EarlyCtx {
+    ty_table: TypeTable,
+    ident_table: IdentTable,
+    mods: RefCell<Vec<rust_ast::Module>>,
+}
+
+impl EarlyCtx {
+    fn new(mods: Vec<rust_ast::Module>) -> Self {
+        // worst case len
+        let len = mods.iter().fold(0, |acc, it| acc + it.items.len());
+        let adts = mods.iter().flat_map(|it| it.items.iter());
+
+        let mut ty_table = TypeTable::with_capacity(len);
+        let mut ident_table = IdentTable::with_capacity(len);
+        for adt in adts {
+            if let Some(ident) = adt.borrow().ident() {
+                let ident = ident.to_string();
+                let type_id = ty_table.len();
+                ty_table.push(AstRef::clone(adt));
+                ident_table.insert(ident, type_id);
+            }
+        }
+
+        Self { ty_table, ident_table, mods: RefCell::new(mods) }
+    }
+
+    fn into_late_ctx(self) -> LateCtx {
+        let schema = lower_ast_types(&self);
+
+        LateCtx { schema }
+    }
+
+    fn find(&self, key: &String) -> Option<AstRef> {
+        self.type_id(key).map(|id| AstRef::clone(&self.ty_table[id]))
+    }
+
+    fn type_id(&self, key: &String) -> Option<TypeId> {
+        self.ident_table.get(key).copied()
+    }
+
+    fn ast_ref(&self, id: TypeId) -> AstRef {
+        AstRef::clone(&self.ty_table[id])
+    }
+}
+
+struct LateCtx {
+    schema: Schema,
+}
+
+struct CodegenResult {
+    schema: Schema,
+    outputs: Vec<(/* generator name */ &'static str, /* output */ GeneratorOutput)>,
+}
+
+impl LateCtx {
+    fn type_def(&self, id: TypeId) -> Option<&TypeDef> {
+        self.schema.definitions.get(id)
+    }
+}
+
+trait Runner {
+    type Context;
+    type Output;
+    fn name(&self) -> &'static str;
+    fn run(&mut self, ctx: &Self::Context) -> Result<Self::Output>;
+}
+
+impl AstCodegen {
+    #[must_use]
+    fn add_file<P>(mut self, path: P) -> Self
+    where
+        P: AsRef<str>,
+    {
+        self.files.push(path.as_ref().into());
+        self
+    }
+
+    #[must_use]
+    fn pass<P>(mut self, pass: P) -> Self
+    where
+        P: Pass + Runner<Output = (), Context = EarlyCtx> + 'static,
+    {
+        self.passes.push(Box::new(pass));
+        self
+    }
+
+    #[must_use]
+    fn gen<G>(mut self, generator: G) -> Self
+    where
+        G: Generator + Runner<Output = GeneratorOutput, Context = LateCtx> + 'static,
+    {
+        self.builders.push(Box::new(generator));
+        self
+    }
+
+    fn generate(self) -> Result<CodegenResult> {
+        let modules = self
+            .files
+            .into_iter()
+            .map(rust_ast::Module::from)
+            .map(rust_ast::Module::load)
+            .map_ok(rust_ast::Module::expand)
+            .map_ok(|it| it.map(rust_ast::Module::analyze))
+            .collect::<Result<Result<Result<Vec<_>>>>>()???;
+
+        // early passes
+        let ctx = {
+            let ctx = EarlyCtx::new(modules);
+            _ = self
+                .passes
+                .into_iter()
+                .map(|mut runner| runner.run(&ctx).map(|res| (runner.name(), res)))
+                .collect::<Result<Vec<_>>>()?;
+            ctx.into_late_ctx()
+        };
+
+        let outputs = self
+            .builders
+            .into_iter()
+            .map(|mut runner| runner.run(&ctx).map(|res| (runner.name(), res)))
+            .collect::<Result<Vec<_>>>()?;
+
+        Ok(CodegenResult { outputs, schema: ctx.schema })
+    }
+}
+
+#[allow(clippy::print_stdout)]
 fn main() -> std::result::Result<(), Box<dyn std::error::Error>> {
     let cli_options = cli_options().run();
 
-    let AstCodegenResult { outputs, schema } = SOURCE_PATHS
+    let CodegenResult { outputs, schema } = SOURCE_PATHS
         .iter()
         .fold(AstCodegen::default(), AstCodegen::add_file)
         .pass(Linker)
@@ -104,7 +316,7 @@ fn output(krate: &str, path: &str) -> PathBuf {
 
 /// Writes all streams and returns a vector pointing to side-effects written on the disk
 fn write_generated_streams(
-    streams: impl IntoIterator<Item = GeneratedTokenStream>,
+    streams: impl IntoIterator<Item = GeneratedStream>,
 ) -> std::io::Result<Vec<String>> {
     streams
         .into_iter()
@@ -120,7 +332,7 @@ fn write_generated_streams(
 
 /// Writes all streams and returns a vector pointing to side-effects written on the disk
 fn write_data_streams(
-    streams: impl IntoIterator<Item = GeneratedDataStream>,
+    streams: impl IntoIterator<Item = DataStream>,
 ) -> std::io::Result<Vec<String>> {
     streams
         .into_iter()
diff --git a/tasks/ast_codegen/src/passes/calc_layout.rs b/tasks/ast_codegen/src/passes/calc_layout.rs
index cdee71abbbb3a..f0ba11dc9ed72 100644
--- a/tasks/ast_codegen/src/passes/calc_layout.rs
+++ b/tasks/ast_codegen/src/passes/calc_layout.rs
@@ -6,11 +6,10 @@ use quote::ToTokens;
 use syn::Type;
 
 use crate::{
-    codegen::EarlyCtx,
     layout::{KnownLayout, Layout},
     rust_ast::{AstRef, AstType, Enum, Struct},
     util::{NormalizeError, TypeAnalysis, TypeExt, TypeWrapper},
-    Result,
+    EarlyCtx, Result,
 };
 
 use super::{define_pass, Pass};
diff --git a/tasks/ast_codegen/src/passes/linker.rs b/tasks/ast_codegen/src/passes/linker.rs
index 3060dc53deca9..860e1d8b60903 100644
--- a/tasks/ast_codegen/src/passes/linker.rs
+++ b/tasks/ast_codegen/src/passes/linker.rs
@@ -2,7 +2,7 @@ use std::borrow::Cow;
 
 use syn::parse_quote;
 
-use crate::{codegen::EarlyCtx, rust_ast::Inherit, util::NormalizeError};
+use crate::{rust_ast::Inherit, util::NormalizeError, EarlyCtx};
 
 use super::{define_pass, AstType, Pass, Result};
 
diff --git a/tasks/ast_codegen/src/passes/mod.rs b/tasks/ast_codegen/src/passes/mod.rs
index 0151940abc8c1..ce1d3694ba653 100644
--- a/tasks/ast_codegen/src/passes/mod.rs
+++ b/tasks/ast_codegen/src/passes/mod.rs
@@ -1,6 +1,8 @@
 use std::collections::VecDeque;
 
-use crate::{codegen::EarlyCtx, rust_ast::AstType, Result};
+use itertools::Itertools;
+
+use crate::{rust_ast::AstType, EarlyCtx, Result};
 
 mod calc_layout;
 mod linker;
@@ -28,13 +30,13 @@ pub trait Pass {
 
         // call each
         // we sort by `TypeId` so we always have the same ordering as how it is written in the rust.
-        let mut unresolved = ctx.chronological_idents().collect::<VecDeque<_>>();
+        let mut unresolved =
+            ctx.ident_table.iter().sorted_by_key(|it| it.1).map(|it| it.0).collect::<VecDeque<_>>();
 
         while let Some(next) = unresolved.pop_back() {
             let next_id = ctx.type_id(next).unwrap();
 
-            let ast_ref = ctx.ast_ref(next_id);
-            let val = &mut ast_ref.borrow_mut();
+            let val = &mut ctx.ty_table[next_id].borrow_mut();
 
             if !self.each(val, ctx)? {
                 unresolved.push_front(next);
@@ -47,8 +49,8 @@ pub trait Pass {
 macro_rules! define_pass {
     ($vis:vis struct $ident:ident $($lifetime:lifetime)? $($rest:tt)*) => {
         $vis struct $ident $($lifetime)? $($rest)*
-        impl $($lifetime)? $crate::codegen::Runner for $ident $($lifetime)? {
-            type Context = $crate::codegen::EarlyCtx;
+        impl $($lifetime)? $crate::Runner for $ident $($lifetime)? {
+            type Context = $crate::EarlyCtx;
             type Output = ();
             fn name(&self) -> &'static str {
                 $crate::Pass::name(self)
diff --git a/tasks/ast_codegen/src/schema/mod.rs b/tasks/ast_codegen/src/schema/mod.rs
index e0828f80eaeac..d588e9a8fb62e 100644
--- a/tasks/ast_codegen/src/schema/mod.rs
+++ b/tasks/ast_codegen/src/schema/mod.rs
@@ -2,12 +2,11 @@ use quote::ToTokens;
 use serde::Serialize;
 
 use crate::{
-    codegen,
     layout::KnownLayout,
     markers::{get_scope_attr, get_scope_markers, get_visit_markers},
     rust_ast as rust,
     util::{unexpanded_macro_err, TypeExt},
-    Result, TypeId,
+    DefTable, Result,
 };
 
 mod defs;
@@ -84,22 +83,7 @@ impl<'a> From<crate::util::TypeIdentResult<'a>> for TypeName {
 
 #[derive(Debug, Default, serde::Serialize)]
 pub struct Schema {
-    defs: Vec<TypeDef>,
-}
-
-impl Schema {
-    pub fn get(&self, id: TypeId) -> Option<&TypeDef> {
-        self.defs.get(id)
-    }
-}
-
-impl<'a> IntoIterator for &'a Schema {
-    type Item = &'a TypeDef;
-    type IntoIter = std::slice::Iter<'a, TypeDef>;
-
-    fn into_iter(self) -> Self::IntoIter {
-        self.defs.iter()
-    }
+    pub definitions: DefTable,
 }
 
 fn parse_outer_markers(attrs: &Vec<syn::Attribute>) -> Result<OuterMarkers> {
@@ -115,18 +99,18 @@ fn parse_inner_markers(attrs: &Vec<syn::Attribute>) -> Result<InnerMarkers> {
 }
 
 // lower `AstType` to `TypeDef`.
-pub fn lower_ast_types(ctx: &codegen::EarlyCtx) -> Schema {
-    let defs = ctx
-        .mods()
+pub fn lower_ast_types(ctx: &crate::EarlyCtx) -> Schema {
+    let definitions = ctx
+        .mods
         .borrow()
         .iter()
         .flat_map(|it| &it.items)
         .map(|it| lower_ast_type(&it.borrow(), ctx))
         .collect();
-    Schema { defs }
+    Schema { definitions }
 }
 
-fn lower_ast_type(ty: &rust::AstType, ctx: &codegen::EarlyCtx) -> TypeDef {
+fn lower_ast_type(ty: &rust::AstType, ctx: &crate::EarlyCtx) -> TypeDef {
     match ty {
         rust::AstType::Enum(it) => TypeDef::Enum(lower_ast_enum(it, ctx)),
         rust::AstType::Struct(it) => TypeDef::Struct(lower_ast_struct(it, ctx)),
@@ -134,7 +118,7 @@ fn lower_ast_type(ty: &rust::AstType, ctx: &codegen::EarlyCtx) -> TypeDef {
     }
 }
 
-fn lower_ast_enum(it @ rust::Enum { item, meta }: &rust::Enum, ctx: &codegen::EarlyCtx) -> EnumDef {
+fn lower_ast_enum(it @ rust::Enum { item, meta }: &rust::Enum, ctx: &crate::EarlyCtx) -> EnumDef {
     let (size_64, align_64, offsets_64) = meta
         .layout_64
         .clone()
@@ -171,7 +155,7 @@ fn lower_ast_enum(it @ rust::Enum { item, meta }: &rust::Enum, ctx: &codegen::Ea
 
 fn lower_ast_struct(
     it @ rust::Struct { item, meta }: &rust::Struct,
-    ctx: &codegen::EarlyCtx,
+    ctx: &crate::EarlyCtx,
 ) -> StructDef {
     let (size_64, align_64, offsets_64) = meta
         .layout_64
@@ -201,7 +185,7 @@ fn lower_ast_struct(
     }
 }
 
-fn lower_variant<F>(variant: &syn::Variant, enum_dbg_name: F, ctx: &codegen::EarlyCtx) -> VariantDef
+fn lower_variant<F>(variant: &syn::Variant, enum_dbg_name: F, ctx: &crate::EarlyCtx) -> VariantDef
 where
     F: Fn() -> String,
 {
@@ -221,7 +205,7 @@ where
     }
 }
 
-fn lower_inherit(inherit: &rust::Inherit, ctx: &codegen::EarlyCtx) -> InheritDef {
+fn lower_inherit(inherit: &rust::Inherit, ctx: &crate::EarlyCtx) -> InheritDef {
     match inherit {
         rust::Inherit::Linked { super_, variants } => InheritDef {
             super_: create_type_ref(super_, ctx),
@@ -236,7 +220,7 @@ fn lower_inherit(inherit: &rust::Inherit, ctx: &codegen::EarlyCtx) -> InheritDef
     }
 }
 
-fn lower_field(field: &syn::Field, ctx: &codegen::EarlyCtx) -> FieldDef {
+fn lower_field(field: &syn::Field, ctx: &crate::EarlyCtx) -> FieldDef {
     FieldDef {
         name: field.ident.as_ref().map(ToString::to_string),
         vis: Visibility::from(&field.vis),
@@ -246,7 +230,7 @@ fn lower_field(field: &syn::Field, ctx: &codegen::EarlyCtx) -> FieldDef {
     }
 }
 
-fn create_type_ref(ty: &syn::Type, ctx: &codegen::EarlyCtx) -> TypeRef {
+fn create_type_ref(ty: &syn::Type, ctx: &crate::EarlyCtx) -> TypeRef {
     let ident = ty.get_ident();
     let id = ident.as_ident().and_then(|id| ctx.type_id(&id.to_string()));
     let transparent_id = ctx.type_id(&ident.inner_ident().to_string());
diff --git a/tasks/ast_codegen/src/util.rs b/tasks/ast_codegen/src/util.rs
index 0610ea72e02c8..8ebbde3efc229 100644
--- a/tasks/ast_codegen/src/util.rs
+++ b/tasks/ast_codegen/src/util.rs
@@ -4,7 +4,7 @@ use quote::{format_ident, ToTokens};
 use serde::Serialize;
 use syn::{spanned::Spanned, GenericArgument, Ident, ItemMacro, PathArguments, Type, TypePath};
 
-use crate::{codegen::EarlyCtx, TypeId};
+use crate::{EarlyCtx, TypeId};
 
 pub trait NormalizeError<T> {
     fn normalize(self) -> crate::Result<T>;