perf: Avoid cloning entire tree

Cargo.toml

@@ -37,6 +37,9 @@ stable = ["json", "toml", "typescript", "tsx", "jsx", "go", "python", "yaml", "j
 default = []
 [dev-dependencies]
 test-log = { workspace = true }
+criterion = { version = "0.5", features = ["html_reports"] }
+git2 = { version = "0.20.0" }
+tempfile = {workspace = true}
 [workspace]
 members = [
   "codegen-sdk-analyzer",
@@ -109,40 +112,48 @@ codegen-sdk-json = { path = "languages/codegen-sdk-json" }
 tempfile = "3.16.0"
 quote = "1.0.38"
 proc-macro2 = "1.0.93"
-derive_generic_visitor = "0.1.1"
 insta = "1.42.1"
 prettyplease = "0.2.29"
-syn = { version = "2.0.98", features = ["proc-macro"] }
+syn = { version = "2.0.98", features = ["proc-macro", "full"] }
 derive_more = { version = "2.0.1", features = ["debug", "display"] }
-salsa = "0.18.0"
+salsa = {git = "https://github.com/salsa-rs/salsa", branch = "master"}
 subenum = {git = "https://github.com/mrenow/subenum", branch = "main"}
 indicatif-log-bridge = "0.2.3"
 indicatif = { version = "0.17.11", features = ["rayon"] }
 crossbeam-channel = "0.5.11"
 rstest = "0.25.0"
+indextree = "4.7.3"
+thiserror = "2.0.11"
+
 [profile.dev]
 # codegen-backend = "cranelift"
-split-debuginfo = "unpacked"
+# split-debuginfo = "unpacked"
 
 [profile.dev.package]
 insta.opt-level = 3
 similar.opt-level = 3
 syn.opt-level = 3
 convert_case.opt-level = 3
-codegen-sdk-cst-generator.opt-level = 3
-codegen-sdk-common.opt-level = 3
+# codegen-sdk-cst-generator.opt-level = 3
+# codegen-sdk-common.opt-level = 3
 
 [profile.dev.build-override]
 opt-level = 3
+debug = true
 
-[profile.test.package."codegen-sdk-cst"]
-inherits = "dev"
-opt-level = 0
-debug = 0
-strip = "none"
-codegen-units = 256
-incremental = true
+# [profile.test.package."codegen-sdk-cst"]
+# inherits = "dev"
+# opt-level = 0
+# debug = 0
+# strip = "none"
+# codegen-units = 256
+# incremental = true
 # codegen-backend = "cranelift"
 
 [profile.test]
 lto = false
+
+[[bench]]
+name = "parse"
+harness = false
+required-features = ["stable"]

benches/parse.rs

@@ -0,0 +1,36 @@
+use std::{hint::black_box, path::PathBuf};
+
+use codegen_sdk_analyzer::Codebase;
+use criterion::{Criterion, criterion_group, criterion_main};
+fn clone_repo(url: String, name: String, tmp_dir: &tempfile::TempDir) -> PathBuf {
+    let repo_path = tmp_dir.path().join(name);
+    if !repo_path.exists() {
+        log::info!("Cloning repo: {} to {}", url, repo_path.display());
+        let mut fetch_opts = git2::FetchOptions::new();
+        fetch_opts.depth(1);
+        let _ = git2::build::RepoBuilder::new()
+            .fetch_options(fetch_opts)
+            .clone(&url, &repo_path)
+            .unwrap();
+    }
+    repo_path
+}
+fn parse_nest(path: &PathBuf) {
+    let _ = Codebase::new(path.clone());
+}
+
+fn criterion_benchmark(c: &mut Criterion) {
+    env_logger::init();
+    let temp_dir = tempfile::tempdir().unwrap();
+    let repo_path = clone_repo(
+        "https://github.com/nestjs/nest".to_string(),
+        "nest".to_string(),
+        &temp_dir,
+    );
+    c.bench_function("parse_nest", |b| {
+        b.iter(|| parse_nest(black_box(&repo_path)))
+    });
+}
+
+criterion_group!(benches, criterion_benchmark);
+criterion_main!(benches);

codegen-sdk-analyzer/src/codebase/discovery.rs

@@ -22,14 +22,15 @@ pub fn log_languages() {
 
 pub fn collect_files(db: &CodegenDatabase, dir: &PathBuf) -> FilesToParse {
     let mut files = Vec::new();
+    let dir = dir.canonicalize().unwrap();
     for language in LANGUAGES.iter() {
         for extension in language.file_extensions.iter() {
             files.extend(
-                glob(&format!(
-                    "{dir}**/*.{extension}",
-                    extension = extension,
-                    dir = dir.display()
-                ))
+                glob(
+                    &dir.join(format!("**/*.{extension}", extension = extension))
+                        .to_str()
+                        .unwrap(),
+                )
                 .unwrap(),
             );
         }

codegen-sdk-analyzer/src/parser.rs

@@ -1,18 +1,18 @@
-use std::path::PathBuf;
-
+use codegen_sdk_common::FileNodeId;
 use codegen_sdk_cst::CSTLanguage;
 use codegen_sdk_macros::{languages_ast, parse_language};
 languages_ast!();
 
 #[salsa::tracked]
 pub struct Parsed<'db> {
     #[id]
-    path: PathBuf,
+    path: FileNodeId<'db>,
     #[return_ref]
+    #[tracked]
     pub file: Option<ParsedFile<'db>>,
 }
-#[salsa::tracked]
+#[salsa::tracked(return_ref)]
 pub fn parse_file(db: &dyn salsa::Database, file: codegen_sdk_ast::input::File) -> Parsed<'_> {
     parse_language!();
-    Parsed::new(db, file.path(db), None)
+    Parsed::new(db, FileNodeId::new(db, file.path(db)), None)
 }

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

@@ -8,8 +8,8 @@ fn get_definitions_impl(language: &Language) -> TokenStream {
 
         impl<'db> codegen_sdk_ast::Definitions<'db> for #language_struct_name<'db> {
             type Definitions = ();
-            fn definitions(self, _db: &'db dyn salsa::Database) -> Self::Definitions{
-                ()
+            fn definitions(self, _db: &'db dyn salsa::Database) -> &'db Self::Definitions{
+                &()
             }
         }
         };
@@ -18,13 +18,13 @@ fn get_definitions_impl(language: &Language) -> TokenStream {
         #[salsa::tracked]
         impl<'db> codegen_sdk_ast::Definitions<'db> for #language_struct_name<'db> {
             type Definitions = Definitions<'db>;
-            #[salsa::tracked]
+            #[salsa::tracked(return_ref)]
             fn definitions(self, db: &'db dyn salsa::Database) -> Self::Definitions {
-                let mut definitions = Definitions::default();
                 if let Some(program) = self.node(db) {
-                    definitions = definitions.visit_by_val_infallible(&program);
+                    return Definitions::visit(db, program);
+                } else {
+                    return Definitions::default(db);
                 }
-                definitions
             }
         }
     }
@@ -35,8 +35,8 @@ fn get_references_impl(language: &Language) -> TokenStream {
         return quote! {
             impl<'db> codegen_sdk_ast::References<'db> for #language_struct_name<'db> {
                 type References = ();
-                fn references(self, _db: &'db dyn salsa::Database) -> Self::References {
-                    ()
+                fn references(self, _db: &'db dyn salsa::Database) -> &'db Self::References {
+                    &()
                 }
             }
         };
@@ -45,13 +45,13 @@ fn get_references_impl(language: &Language) -> TokenStream {
         #[salsa::tracked]
         impl<'db> codegen_sdk_ast::References<'db> for #language_struct_name<'db> {
             type References = References<'db>;
-            #[salsa::tracked]
+            #[salsa::tracked(return_ref)]
             fn references(self, db: &'db dyn salsa::Database) -> Self::References {
-                let mut references = References::default();
-            if let Some(program) = self.node(db) {
-                references = references.visit_by_val_infallible(&program);
-            }
-                references
+                if let Some(program) = self.node(db) {
+                    return References::visit(db, program);
+                } else {
+                    return References::default(db);
+                }
             }
         }
     }
@@ -61,14 +61,15 @@ pub fn generate_ast(language: &Language) -> anyhow::Result<TokenStream> {
     let language_name_str = language.name();
     let definitions_impl = get_definitions_impl(language);
     let references_impl = get_references_impl(language);
-    let program_id = format_ident!("{}", language.root_node());
+    let root_node_name = format_ident!("{}", language.root_node());
     let content = quote! {
     #[salsa::tracked]
     pub struct #language_struct_name<'db> {
+        #[tracked]
         #[return_ref]
-        node: Option<crate::cst::#program_id<'db>>,
+        pub node: Option<crate::cst::Parsed<'db>>,
         #[id]
-        pub path: PathBuf,
+        pub id: codegen_sdk_common::FileNodeId<'db>,
     }
     // impl<'db> File for {language_struct_name}File<'db> {{
     //     fn path(&self) -> &PathBuf {{
@@ -77,15 +78,24 @@ pub fn generate_ast(language: &Language) -> anyhow::Result<TokenStream> {
     // }}
     pub fn parse(db: &dyn salsa::Database, input: codegen_sdk_ast::input::File) -> #language_struct_name<'_> {
         log::debug!("Parsing {} file: {}", input.path(db).display(), #language_name_str);
-        let ast = crate::cst::parse_program_raw(db, input.contents(db));
-        #language_struct_name::new(db, ast, input.path(db).clone())
+        let ast = crate::cst::parse_program_raw(db, input.contents(db), input.path(db).clone());
+        let file_id = codegen_sdk_common::FileNodeId::new(db, input.path(db).clone());
+        #language_struct_name::new(db, ast, file_id)
     }
-    #[salsa::tracked]
+    #[salsa::tracked(return_ref)]
     pub fn parse_query(db: &dyn salsa::Database, input: codegen_sdk_ast::input::File) -> #language_struct_name<'_> {
         parse(db, input)
     }
 
-
+    impl<'db> #language_struct_name<'db> {
+        pub fn tree(&self, db: &'db dyn salsa::Database) -> &'db codegen_sdk_common::Tree<crate::cst::NodeTypes<'db>> {
+            self.node(db).unwrap().tree(db)
+        }
+        pub fn root(&self, db: &'db dyn salsa::Database) -> &'db crate::cst::#root_node_name<'db> {
+            let tree = self.tree(db);
+            tree.get(&self.node(db).unwrap().program(db)).unwrap().as_ref().try_into().unwrap()
+        }
+    }
     #definitions_impl
     #references_impl
     // impl<'db> HasNode for {language_struct_name}File<'db> {

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

@@ -10,11 +10,11 @@ use syn::parse_quote;
 pub fn generate_ast(language: &Language) -> anyhow::Result<()> {
     let db = CSTDatabase::default();
     let imports = quote! {
-        use derive_generic_visitor::{Visitor, Drive, Visit};
         use codegen_sdk_common::*;
         use std::path::PathBuf;
         use codegen_sdk_cst::CSTLanguage;
         use std::collections::BTreeMap;
+        use std::sync::mpsc::Sender;
     };
     let ast = generator::generate_ast(language)?;
     let definition_visitor = visitor::generate_visitor(&db, language, "definition");