Generate references queries

codegen-sdk-ast-generator/src/generator.rs

@@ -6,7 +6,8 @@ pub fn generate_ast(language: &Language) -> anyhow::Result<String> {
     pub struct {language_struct_name}File {{
         node: {language_name}::{root_node_name},
         path: PathBuf,
-        pub visitor: QueryExecutor
+        pub references: References,
+        pub definitions: Definitions
     }}
     impl File for {language_struct_name}File {{
         fn path(&self) -> &PathBuf {{
@@ -15,9 +16,11 @@ pub fn generate_ast(language: &Language) -> anyhow::Result<String> {
         fn parse(path: &PathBuf) -> Result<Self, ParseError> {{
             log::debug!(\"Parsing {language_name} file: {{}}\", path.display());
             let ast = {language_name}::{language_struct_name}::parse_file(path)?;
-            let mut visitor = QueryExecutor::default();
-            ast.drive(&mut visitor);
-            Ok({language_struct_name}File {{ node: ast, path: path.clone(), visitor }})
+            let mut references = References::default();
+            let mut definitions = Definitions::default();
+            ast.drive(&mut definitions);
+            ast.drive(&mut references);
+            Ok({language_struct_name}File {{ node: ast, path: path.clone(), references, definitions }})
         }}
     }}
     impl HasNode for {language_struct_name}File {{

codegen-sdk-ast-generator/src/lib.rs

@@ -3,7 +3,7 @@
 use codegen_sdk_common::{generator::format_code, language::Language};
 use quote::quote;
 
 use crate::query::HasQuery;
 mod generator;
 mod query;
 mod visitor;
@@ -15,16 +15,9 @@
         use codegen_sdk_cst::CSTLanguage;
     };
     let mut ast = generator::generate_ast(language)?;
-    let visitor = visitor::generate_visitor(
-        &language
-            .definitions()
-            .values()
-            .into_iter()
-            .flatten()
-            .collect(),
-        language,
-    );
-    ast = imports.to_string() + &ast + &visitor.to_string();
+    let definitions = visitor::generate_visitor(language, "definitions");
+    let references = visitor::generate_visitor(language, "references");
+    ast = imports.to_string() + &ast + &definitions.to_string() + &references.to_string();
     ast = format_code(&ast).unwrap();
     let out_dir = std::env::var("OUT_DIR")?;
     let out_file = format!("{}/{}.rs", out_dir, language.name());

codegen-sdk-ast-generator/src/query.rs

@@ -412,12 +412,6 @@ pub trait HasQuery {
         }
         queries
     }
-    fn definitions(&self) -> BTreeMap<String, Vec<Query<'_>>> {
-        self.queries_with_prefix("definition")
-    }
-    // fn references(&self) -> BTreeMap<String, Vec<Query<'_>>> {
-    //     self.queries_with_prefix("reference")
-    // }
 }
 impl HasQuery for Language {
     fn queries(&self) -> BTreeMap<String, Query<'_>> {

codegen-sdk-ast-generator/src/visitor.rs

@@ -8,8 +8,15 @@ use proc_macro2::TokenStream;
 use quote::{format_ident, quote};
 
 use super::query::Query;
-pub fn generate_visitor(queries: &Vec<&Query>, language: &Language) -> TokenStream {
-    log::info!("Generating visitor for language: {}", language.name());
+use crate::query::HasQuery;
+pub fn generate_visitor(language: &Language, name: &str) -> TokenStream {
+    log::info!(
+        "Generating visitor for language: {} for {}",
+        language.name(),
+        name
+    );
+    let raw_queries = language.queries_with_prefix(&format!("@{}", name));
+    let queries: Vec<&Query> = raw_queries.values().flatten().collect();
     let language_name = format_ident!("{}", language.name());
     let mut names = Vec::new();
     let mut types = Vec::new();
@@ -52,7 +59,7 @@ pub fn generate_visitor(queries: &Vec<&Query>, language: &Language) -> TokenStre
             }
         });
     }
-    let name = format_ident!("QueryExecutor");
+    let name = format_ident!("{}", name.to_case(Case::Pascal));
     quote! {
         #[derive(Visitor, Default, Debug, Clone)]
         #[visitor(
@@ -77,9 +84,7 @@ mod tests {
     #[test_log::test]
     fn test_generate_visitor() {
         let language = &Typescript;
-        let queries = language.definitions();
-        log::info!("Gathered {} queries", queries.len());
-        let visitor = generate_visitor(&queries.values().into_iter().flatten().collect(), language);
+        let visitor = generate_visitor(language, "definitions");
         insta::assert_snapshot!(
             codegen_sdk_common::generator::format_code(&visitor.to_string()).unwrap()
         );