refactor(ast_codegen): declutter the main file.

same as #4741

tasks/ast_codegen/src/codegen.rs

@@ -0,0 +1,156 @@
+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 })
+ }
+}

tasks/ast_codegen/src/generators/assert_layouts.rs

@@ -3,9 +3,10 @@ use quote::{format_ident, quote};
 use syn::Type;
 
 use crate::{
+ codegen::LateCtx,
  output,
  schema::{FieldDef, ToType, TypeDef},
- Generator, GeneratorOutput, LateCtx,
+ Generator, GeneratorOutput,
 };
 
 use super::{define_generator, generated_header};
@@ -17,9 +18,8 @@ define_generator! {
 impl Generator for AssertLayouts {
  fn generate(&mut self, ctx: &LateCtx) -> GeneratorOutput {
  let (assertions_64, assertions_32) = ctx
- .schema
- .definitions
- .iter()
+ .schema()
+ .into_iter()
  .map(|def| {
  let typ = def.to_type_elide();
  assert_type(&typ, def)

tasks/ast_codegen/src/generators/ast_builder.rs

@@ -9,13 +9,14 @@ 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, LateCtx,
+ Generator, GeneratorOutput,
 };
 
 use super::define_generator;
@@ -27,9 +28,8 @@ define_generator! {
 impl Generator for AstBuilderGenerator {
  fn generate(&mut self, ctx: &LateCtx) -> GeneratorOutput {
  let fns = ctx
- .schema
- .definitions
- .iter()
+ .schema()
+ .into_iter()
  .filter(|it| it.visitable())
  .map(|it| generate_builder_fn(it, ctx))
  .collect_vec();

tasks/ast_codegen/src/generators/ast_kind.rs

@@ -3,10 +3,11 @@ use quote::quote;
 use syn::{parse_quote, Arm, Ident, Type, Variant};
 
 use crate::{
+ codegen::LateCtx,
  output,
  schema::{GetIdent, ToType, TypeDef},
  util::ToIdent,
- Generator, GeneratorOutput, LateCtx,
+ Generator, GeneratorOutput,
 };
 
 use super::{define_generator, generated_header};
@@ -137,8 +138,8 @@ pub fn process_types(def: &TypeDef, _: &LateCtx) -> Vec<(Ident, Type)> {
 impl Generator for AstKindGenerator {
  fn generate(&mut self, ctx: &LateCtx) -> GeneratorOutput {
  let have_kinds: Vec<(Ident, Type)> = ctx
- .schema.definitions
- .iter()
+ .schema()
+ .into_iter()
  .filter(|it| it.visitable())
  .filter(
  |maybe_kind| matches!(maybe_kind, kind @ (TypeDef::Enum(_) | TypeDef::Struct(_)) if kind.visitable())

tasks/ast_codegen/src/generators/derive_clone_in.rs

@@ -4,9 +4,10 @@ use quote::{format_ident, quote};
 use syn::Ident;
 
 use crate::{
+ codegen::LateCtx,
  output,
  schema::{EnumDef, GetIdent, StructDef, TypeDef},
- GeneratorOutput, LateCtx,
+ GeneratorOutput,
 };
 
 use super::{define_generator, generated_header, Generator};
@@ -18,9 +19,8 @@ define_generator! {
 impl Generator for DeriveCloneIn {
  fn generate(&mut self, ctx: &LateCtx) -> GeneratorOutput {
  let impls: Vec<TokenStream> = ctx
- .schema
- .definitions
- .iter()
+ .schema()
+ .into_iter()
  .filter(|def| def.generates_derive("CloneIn"))
  .map(|def| match &def {
  TypeDef::Enum(it) => derive_enum(it),

tasks/ast_codegen/src/generators/derive_get_span.rs

@@ -3,10 +3,11 @@ use quote::{format_ident, quote};
 use syn::Ident;
 
 use crate::{
+ codegen::LateCtx,
  output,
  schema::{EnumDef, GetGenerics, StructDef, ToType, TypeDef},
  util::ToIdent,
- Generator, GeneratorOutput, LateCtx,
+ Generator, GeneratorOutput,
 };
 
 use super::{define_generator, generated_header};
@@ -54,9 +55,8 @@ fn derive(
  ctx: &LateCtx,
 ) -> TokenStream {
  let impls: Vec<TokenStream> = ctx
- .schema
- .definitions
- .iter()
+ .schema()
+ .into_iter()
  .filter(|def| def.visitable())
  .map(|def| match &def {
  TypeDef::Enum(def) => derive_enum(def, trait_name, method_name, self_type, result_type),