Skip to content

Commit

Permalink
feat(ast_codegen): support for generating VisitMut. (#4006)
Browse files Browse the repository at this point in the history
  • Loading branch information
@rzvxa
rzvxa committed Jul 2, 2024
1 parent 7a3cb83 commit 272df6e
Show file tree
Hide file tree
Showing 3 changed files with 146 additions and 31 deletions.
9 changes: 4 additions & 5 deletions tasks/ast_codegen/src/generators/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,8 +15,8 @@ macro_rules! endl {
/// used outside and accepts expressions.
/// Wraps the result of the given expression in `insert!({value here});` and outputs it as `TokenStream`.
macro_rules! insert {
($txt:expr) => {{
let txt: &str = &*$txt;
($fmt:literal $(, $args:expr)*) => {{
let txt = format!($fmt, $($args)*);
format!(r#"insert!("{}");"#, txt).parse::<proc_macro2::TokenStream>().unwrap()
}};
}
Expand All @@ -25,9 +25,8 @@ macro_rules! insert {
macro_rules! generated_header {
() => {{
let file = file!().replace("\\", "/");
let edit_comment = $crate::generators::insert!(format!(
"// To edit this generated file you have to edit `{file}`"
));
let edit_comment =
$crate::generators::insert!("// To edit this generated file you have to edit `{file}`");
// TODO add generation date, AST source hash, etc here.
quote::quote! {
insert!("// Auto-generated code, DO NOT EDIT DIRECTLY!");
Expand Down
156 changes: 135 additions & 21 deletions tasks/ast_codegen/src/generators/visit.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -37,11 +37,20 @@ impl Generator for VisitGenerator {

fn generate(&mut self, ctx: &CodegenCtx) -> GeneratorOutput {
let visit = (String::from("visit"), generate_visit(ctx));
let visit_mut = (String::from("visit_mut"), generate_visit_mut(ctx));

GeneratorOutput::Many(HashMap::from_iter(vec![visit]))
GeneratorOutput::Many(HashMap::from_iter(vec![visit, visit_mut]))
}
}

static CLIPPY_ALLOW: &str = "\
unused_variables,\
clippy::extra_unused_type_parameters,\
clippy::explicit_iter_loop,\
clippy::self_named_module_files,\
clippy::semicolon_if_nothing_returned,\
clippy::match_wildcard_for_single_variants";

fn generate_visit(ctx: &CodegenCtx) -> TokenStream {
let header = generated_header!();
// we evaluate it outside of quote to take advantage of expression evaluation
Expand All @@ -54,12 +63,13 @@ fn generate_visit(ctx: &CodegenCtx) -> TokenStream {
//! * [rustc visitor](https://github.com/rust-lang/rust/blob/master/compiler/rustc_ast/src/visit.rs)\n\
"};

let (visits, walks) = VisitBuilder::new(ctx).build();
let (visits, walks) = VisitBuilder::new(ctx, false).build();
let clippy_attr = insert!("#![allow({})]", CLIPPY_ALLOW);

quote! {
#header
#file_docs
insert!("#![allow(clippy::self_named_module_files, clippy::semicolon_if_nothing_returned, clippy::match_wildcard_for_single_variants)]");
#clippy_attr

endl!();

Expand All @@ -78,14 +88,11 @@ fn generate_visit(ctx: &CodegenCtx) -> TokenStream {

/// Syntax tree traversal
pub trait Visit<'a>: Sized {
#[allow(unused_variables)]
fn enter_node(&mut self, kind: AstKind<'a>) {}
#[allow(unused_variables)]
fn leave_node(&mut self, kind: AstKind<'a>) {}

endl!();

#[allow(unused_variables)]
fn enter_scope(&mut self, flags: ScopeFlags) {}
fn leave_scope(&mut self) {}

Expand Down Expand Up @@ -115,16 +122,80 @@ fn generate_visit(ctx: &CodegenCtx) -> TokenStream {
}
}

fn generate_visit_mut(ctx: &CodegenCtx) -> TokenStream {
let header = generated_header!();
// we evaluate it outside of quote to take advantage of expression evaluation
// otherwise the `\n\` wouldn't work!
let file_docs = insert! {"\
//! Visitor Pattern\n\
//!\n\
//! See:\n\
//! * [visitor pattern](https://rust-unofficial.github.io/patterns/patterns/behavioural/visitor.html)\n\
//! * [rustc visitor](https://github.com/rust-lang/rust/blob/master/compiler/rustc_ast/src/visit.rs)\n\
"};

let (visits, walks) = VisitBuilder::new(ctx, true).build();
let clippy_attr = insert!("#![allow({})]", CLIPPY_ALLOW);

quote! {
#header
#file_docs
#clippy_attr

endl!();

use oxc_allocator::Vec;
use oxc_syntax::scope::ScopeFlags;

endl!();

use crate::{ast::*, ast_kind::AstType};

endl!();

use walk_mut::*;

endl!();

/// Syntax tree traversal to mutate an exclusive borrow of a syntax tree in place.
pub trait VisitMut<'a>: Sized {
fn enter_node(&mut self, ty: AstType) {}
fn leave_node(&mut self, ty: AstType) {}

endl!();

fn enter_scope(&mut self, flags: ScopeFlags) {}
fn leave_scope(&mut self) {}

endl!();

#(#visits)*
}

endl!();

pub mod walk_mut {
use super::*;

#(#walks)*

}
}
}

struct VisitBuilder<'a> {
ctx: &'a CodegenCtx,

is_mut: bool,

visits: Vec<TokenStream>,
walks: Vec<TokenStream>,
cache: HashMap<Ident, [Option<Cow<'a, Ident>>; 2]>,
}

impl<'a> VisitBuilder<'a> {
fn new(ctx: &'a CodegenCtx) -> Self {
Self { ctx, visits: Vec::new(), walks: Vec::new(), cache: HashMap::new() }
fn new(ctx: &'a CodegenCtx, is_mut: bool) -> Self {
Self { ctx, is_mut, visits: Vec::new(), walks: Vec::new(), cache: HashMap::new() }
}

fn build(mut self) -> (/* visits */ Vec<TokenStream>, /* walks */ Vec<TokenStream>) {
Expand All @@ -143,6 +214,33 @@ impl<'a> VisitBuilder<'a> {
(self.visits, self.walks)
}

fn with_ref_pat<T>(&self, tk: T) -> TokenStream
where
T: ToTokens,
{
if self.is_mut {
quote!(&mut #tk)
} else {
quote!(&#tk)
}
}

fn kind_type(&self, ident: &Ident) -> TokenStream {
if self.is_mut {
quote!(AstType::#ident)
} else {
quote!(AstKind::#ident(visitor.alloc(it)))
}
}

fn get_iter(&self) -> TokenStream {
if self.is_mut {
quote!(iter_mut)
} else {
quote!(iter)
}
}

fn get_visitor(
&mut self,
ty: &TypeRef,
Expand Down Expand Up @@ -189,7 +287,7 @@ impl<'a> VisitBuilder<'a> {
format_ident!("{it}")
};

let as_param_type = quote!(&#as_type);
let as_param_type = self.with_ref_pat(&as_type);
let (extra_params, extra_args) = if ident == "Function" {
(quote!(, flags: Option<ScopeFlags>,), quote!(, flags))
} else {
Expand Down Expand Up @@ -223,8 +321,9 @@ impl<'a> VisitBuilder<'a> {

let walk_body = if collection {
let singular_visit = self.get_visitor(ty, false, None);
let iter = self.get_iter();
quote! {
for el in it {
for el in it.#iter() {
visitor.#singular_visit(el);
}
}
Expand All @@ -237,9 +336,10 @@ impl<'a> VisitBuilder<'a> {
};

// replace the placeholder walker with the actual one!
let visit_trait = if self.is_mut { quote!(VisitMut) } else { quote!(Visit) };
self.walks[this_walker] = quote! {
endl!();
pub fn #walk_name <'a, V: Visit<'a>>(visitor: &mut V, it: #as_param_type #extra_params) {
pub fn #walk_name <'a, V: #visit_trait<'a>>(visitor: &mut V, it: #as_param_type #extra_params) {
#walk_body
}
};
Expand Down Expand Up @@ -330,7 +430,11 @@ impl<'a> VisitBuilder<'a> {
} else {
None
};
let to_child = format_ident!("to_{snake_name}");
let to_child = if self.is_mut {
format_ident!("to_{snake_name}_mut")
} else {
format_ident!("to_{snake_name}")
};
let visit = self.get_visitor(&typ, false, visit_as.as_ref());
Some(quote!(#match_macro => visitor.#visit(it.#to_child())))
} else {
Expand All @@ -345,8 +449,9 @@ impl<'a> VisitBuilder<'a> {
if KIND_BLACK_LIST.contains(&ident.to_string().as_str()) {
tk
} else {
let kind = self.kind_type(ident);
quote! {
let kind = AstKind::#ident(visitor.alloc(it));
let kind = #kind;
visitor.enter_node(kind);
#tk
visitor.leave_node(kind);
Expand Down Expand Up @@ -437,19 +542,23 @@ impl<'a> VisitBuilder<'a> {
visit_as.as_ref(),
);
let name = it.ident.as_ref().expect("expected named fields!");
let borrowed_field = self.with_ref_pat(quote!(it.#name));
let mut result = match typ_wrapper {
TypeWrapper::Opt | TypeWrapper::OptBox | TypeWrapper::OptVec => quote! {
if let Some(ref #name) = it.#name {
if let Some(#name) = #borrowed_field {
visitor.#visit(#name #(#args)*);
}
},
TypeWrapper::VecOpt => quote! {
for #name in (&it.#name).into_iter().flatten() {
visitor.#visit(#name #(#args)*);
TypeWrapper::VecOpt => {
let iter = self.get_iter();
quote! {
for #name in it.#name.#iter().flatten() {
visitor.#visit(#name #(#args)*);
}
}
},
}
_ => quote! {
visitor.#visit(&it.#name #(#args)*);
visitor.#visit(#borrowed_field #(#args)*);
},
};
if have_enter_scope {
Expand All @@ -476,14 +585,19 @@ impl<'a> VisitBuilder<'a> {
let body = if KIND_BLACK_LIST.contains(&ident.to_string().as_str()) {
let unused =
if fields_visits.is_empty() { Some(quote!(let _ = (visitor, it);)) } else { None };
let note = insert!(
"// NOTE: {} doesn't exists!",
if self.is_mut { "AstType" } else { "AstKind" }
);
quote! {
insert!("// NOTE: AstKind doesn't exists!");
#note
#(#fields_visits)*
#unused
}
} else {
let kind = self.kind_type(ident);
quote! {
let kind = AstKind::#ident(visitor.alloc(it));
let kind = #kind;
visitor.enter_node(kind);
#(#fields_visits)*
visitor.leave_node(kind);
Expand Down
12 changes: 7 additions & 5 deletions tasks/ast_codegen/src/main.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -221,14 +221,16 @@ fn main() -> std::result::Result<(), Box<dyn std::error::Error>> {
}

{
// write `visit.rs` file
// write `visit.rs` and `visit_mut.rs` files
let output = outputs[VisitGenerator.name()].as_many();
let span_content = pprint(&output["visit"]);
let content = pprint(&output["visit"]);
let content_mut = pprint(&output["visit_mut"]);

let path = format!("{output_dir}/visit.rs");
let mut file = fs::File::create(path)?;
let mut visit = fs::File::create(format!("{output_dir}/visit.rs"))?;
let mut visit_mut = fs::File::create(format!("{output_dir}/visit_mut.rs"))?;

file.write_all(span_content.as_bytes())?;
visit.write_all(content.as_bytes())?;
visit_mut.write_all(content_mut.as_bytes())?;
}

cargo_fmt(".")?;
Expand Down

0 comments on commit 272df6e

Please sign in to comment.