Fast references impl

Author: bagel897

Cargo.lock

@@ -20,7 +20,7 @@ codegen-sdk-resolution = { workspace = true}
 sysinfo = "0.33.1"
 rkyv.workspace = true
 [features]
-python = [ "codegen-sdk-analyzer/python", "codegen-sdk-python"] # TODO: Add python support
+python = [ "codegen-sdk-analyzer/python", "codegen-sdk-python"]
 typescript = [ "codegen-sdk-analyzer/typescript", "codegen-sdk-typescript"]
 tsx = [ "codegen-sdk-analyzer/tsx"]
 jsx = [ "codegen-sdk-analyzer/jsx"]

Cargo.toml

@@ -88,12 +88,14 @@ impl Codebase {
     pub fn execute_op_with_progress<T: Send + Sync>(
         &self,
         name: &str,
-        op: fn(&dyn Db, File) -> T,
+        parallel: bool,
+        op: fn(&dyn Db, codegen_sdk_common::FileNodeId<'_>) -> T,
     ) -> Vec<T> {
         execute_op_with_progress(
             self._db(),
             codegen_sdk_resolution::files(self._db()),
             name,
+            parallel,
             op,
         )
     }
@@ -118,8 +120,9 @@ impl CodebaseContext for Codebase {
     fn get_file<'a>(&'a self, path: PathBuf) -> Option<&'a Self::File<'a>> {
         if let Ok(path) = path.canonicalize() {
             let file = self.db.files.get(&path);
-            if let Some(file) = file {
-                return parse_file(&self.db, file.clone()).file(&self.db).as_ref();
+            if let Some(_) = file {
+                let file_id = codegen_sdk_common::FileNodeId::new(&self.db, path);
+                return parse_file(&self.db, file_id).file(&self.db).as_ref();
             }
         }
         None

codegen-sdk-analyzer/src/codebase.rs

@@ -1,17 +1,23 @@
+use std::path::PathBuf;
+
 use codegen_sdk_ast::{Definitions, References};
 #[cfg(feature = "serialization")]
 use codegen_sdk_common::serialize::Cache;
-use codegen_sdk_cst::File;
 use codegen_sdk_resolution::{Db, Scope};
 use indicatif::{ProgressBar, ProgressStyle};
 
 use super::discovery::{FilesToParse, log_languages};
 use crate::{ParsedFile, database::CodegenDatabase, parser::parse_file};
-pub fn execute_op_with_progress<Database: Db + ?Sized + 'static, T: Send + Sync>(
+pub fn execute_op_with_progress<
+    Database: Db + ?Sized + 'static,
+    Input: Send + Sync,
+    T: Send + Sync,
+>(
     db: &Database,
-    files: codegen_sdk_common::hash::FxHashSet<File>,
+    files: codegen_sdk_common::hash::FxHashSet<Input>,
     name: &str,
-    op: fn(&Database, File) -> T,
+    parallel: bool,
+    op: fn(&Database, Input) -> T,
 ) -> Vec<T> {
     let multi = db.multi_progress();
     let style = ProgressStyle::with_template(
@@ -27,19 +33,35 @@ pub fn execute_op_with_progress<Database: Db + ?Sized + 'static, T: Send + Sync>
         .into_iter()
         .map(|file| (&pg, file, op))
         .collect::<Vec<_>>();
-    let results: Vec<T> = salsa::par_map(db, inputs, move |db, input| {
-        let (pg, file, op) = input;
-        let res = op(
-            db,
-            #[cfg(feature = "serialization")]
-            &cache,
-            file,
-        );
-        pg.inc(1);
-        res
-    });
+    let results: Vec<_> = if parallel {
+        salsa::par_map(db, inputs, move |db, input| {
+            let (pg, file, op) = input;
+            let res = op(
+                db,
+                #[cfg(feature = "serialization")]
+                &cache,
+                file,
+            );
+            pg.inc(1);
+            res
+        })
+    } else {
+        inputs
+            .into_iter()
+            .map(|input| {
+                let (pg, file, op) = input;
+                let res = op(
+                    db,
+                    #[cfg(feature = "serialization")]
+                    &cache,
+                    file,
+                );
+                pg.inc(1);
+                res
+            })
+            .collect()
+    };
     pg.finish();
-    multi.remove(&pg);
     results
 }
 // #[salsa::tracked]
@@ -50,7 +72,12 @@ pub fn execute_op_with_progress<Database: Db + ?Sized + 'static, T: Send + Sync>
 // }
 #[salsa::tracked]
 fn parse_files_definitions_par(db: &dyn Db, files: FilesToParse) {
-    let _: Vec<_> = execute_op_with_progress(db, files.files(db), "Parsing Files", |db, input| {
+    let ids = files
+        .files(db)
+        .iter()
+        .map(|input| codegen_sdk_common::FileNodeId::new(db, input.path(db)))
+        .collect::<codegen_sdk_common::hash::FxHashSet<_>>();
+    let _: Vec<_> = execute_op_with_progress(db, ids, "Parsing Files", true, |db, input| {
         let file = parse_file(db, input.clone());
         if let Some(parsed) = file.file(db) {
             #[cfg(feature = "typescript")]
@@ -62,12 +89,47 @@ fn parse_files_definitions_par(db: &dyn Db, files: FilesToParse) {
             if let ParsedFile::Python(parsed) = parsed {
                 parsed.definitions(db);
                 parsed.references(db);
-                codegen_sdk_python::ast::dependencies(db, input);
+                // let deps = codegen_sdk_python::ast::dependencies(db, input);
+                // for dep in deps.dependencies(db).keys() {
+                //     codegen_sdk_resolution::ast::references_impl(db, dep);
+                // }
             }
         }
         ()
     });
 }
+#[salsa::tracked]
+fn compute_dependencies_par(db: &dyn Db, files: FilesToParse) {
+    let ids = files
+        .files(db)
+        .iter()
+        .map(|input| codegen_sdk_common::FileNodeId::new(db, input.path(db)))
+        .collect::<codegen_sdk_common::hash::FxHashSet<_>>();
+    let targets: codegen_sdk_common::hash::FxHashSet<(PathBuf, String)> =
+        execute_op_with_progress(db, ids, "Computing Dependencies", true, |db, input| {
+            let file = parse_file(db, input.clone());
+            if let Some(parsed) = file.file(db) {
+                #[cfg(feature = "python")]
+                if let ParsedFile::Python(parsed) = parsed {
+                    let deps = codegen_sdk_python::ast::dependency_keys(db, input);
+                    return deps
+                        .iter()
+                        .map(|dep| (dep.path(db).path(db).clone(), dep.name(db).clone()))
+                        .collect::<Vec<_>>();
+                }
+            }
+            Vec::new()
+        })
+        .into_iter()
+        .flatten()
+        .collect();
+    // let _: Vec<_> = execute_op_with_progress(db, targets, "Finding Usages", true, |db, input: (PathBuf, String)| {
+    //     let file_node_id = codegen_sdk_common::FileNodeId::new(db, input.0);
+    //     let fully_qualified_name = codegen_sdk_resolution::FullyQualifiedName::new(db, file_node_id, input.1);
+    //     codegen_sdk_python::ast::references_impl(db, fully_qualified_name);
+    // });
+}
+
 pub fn parse_files<'db>(
     db: &'db CodegenDatabase,
     #[cfg(feature = "serialization")] cache: &'db Cache,
@@ -89,6 +151,12 @@ pub fn parse_files<'db>(
         &cache,
         files_to_parse,
     );
+    compute_dependencies_par(
+        db,
+        #[cfg(feature = "serialization")]
+        &cache,
+        files_to_parse,
+    );
     #[cfg(feature = "serialization")]
     report_cached_count(cached, &files_to_parse.files(db));
 }

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

@@ -20,7 +20,7 @@ use crate::progress::get_multi_progress;
 // Basic Database implementation for Query generation. This is not used for anything else.
 pub struct CodegenDatabase {
     storage: salsa::Storage<Self>,
-    pub files: DashMap<PathBuf, File>,
+    pub files: Arc<DashMap<PathBuf, File>>,
     dirs: Vec<PathBuf>,
     multi_progress: MultiProgress,
     file_watcher: Arc<Mutex<Debouncer<RecommendedWatcher>>>,
@@ -41,7 +41,7 @@ impl CodegenDatabase {
             file_watcher: get_watcher(tx),
             storage: salsa::Storage::default(),
             multi_progress,
-            files: DashMap::new(),
+            files: Arc::new(DashMap::new()),
             dirs: Vec::new(),
             root,
         }
@@ -72,10 +72,10 @@ impl salsa::Database for CodegenDatabase {
 }
 #[salsa::db]
 impl Db for CodegenDatabase {
-    fn files(&self) -> codegen_sdk_common::hash::FxHashSet<codegen_sdk_cst::File> {
+    fn files(&self) -> codegen_sdk_common::hash::FxHashSet<codegen_sdk_common::FileNodeId<'_>> {
         self.files
             .iter()
-            .map(|entry| entry.value().clone())
+            .map(|entry| codegen_sdk_common::FileNodeId::new(self, entry.key().clone()))
             .collect()
     }
     fn watch_dir(&mut self, path: PathBuf) -> anyhow::Result<()> {

codegen-sdk-analyzer/src/database.rs

@@ -14,7 +14,10 @@ pub struct Parsed<'db> {
     pub file: Option<ParsedFile<'db>>,
 }
 #[salsa::tracked(return_ref)]
-pub fn parse_file(db: &dyn salsa::Database, file: codegen_sdk_cst::File) -> Parsed<'_> {
+pub fn parse_file<'db>(
+    db: &'db dyn codegen_sdk_resolution::Db,
+    file: codegen_sdk_common::FileNodeId<'db>,
+) -> Parsed<'db> {
     parse_language!();
-    Parsed::new(db, FileNodeId::new(db, file.path(db)), None)
+    Parsed::new(db, file, None)
 }

codegen-sdk-analyzer/src/parser.rs

@@ -72,7 +72,7 @@ pub fn generate_ast(language: &Language) -> anyhow::Result<TokenStream> {
         pub id: codegen_sdk_common::FileNodeId<'db>,
     }
     impl<'db> codegen_sdk_resolution::Parse<'db> for #language_struct_name<'db> {
-        fn parse(db: &'db dyn salsa::Database, input: codegen_sdk_cst::File) -> &'db Self {
+        fn parse(db: &'db dyn codegen_sdk_resolution::Db, input: codegen_sdk_common::FileNodeId<'db>) -> &'db Self {
             parse(db, input)
         }
     }
@@ -82,7 +82,8 @@ pub fn generate_ast(language: &Language) -> anyhow::Result<TokenStream> {
     //     }}
     // }}
     #[salsa::tracked(return_ref)]
-    pub fn parse(db: &dyn salsa::Database, input: codegen_sdk_cst::File) -> #language_struct_name<'_> {
+    pub fn parse<'db>(db: &'db dyn codegen_sdk_resolution::Db, input: codegen_sdk_common::FileNodeId<'db>) -> #language_struct_name<'db> {
+        let input = db.input(input.path(db)).unwrap();
         log::debug!("Parsing {} file: {}", input.path(db).display(), #language_name_str);
         let ast = crate::cst::parse_program_raw(db, input);
         let file_id = codegen_sdk_common::FileNodeId::new(db, input.path(db).clone());

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

@@ -662,7 +662,7 @@ impl<'a> Query<'a> {
         let symbol_name = self.symbol_name();
         return quote! {
             let fully_qualified_name = codegen_sdk_resolution::FullyQualifiedName::new(db, node.file_id(),#name.source());
-            let symbol = #symbol_name::new(db, fully_qualified_name, id, node.clone(), #(#args.clone().into()),*);
+            let symbol = #symbol_name::new(db, fully_qualified_name, id, #(#args.clone().into()),*);
             #to_append.entry(#name.source()).or_default().push(symbol);
         };
     }
@@ -913,6 +913,8 @@ impl<'a> Query<'a> {
             }
             let name_ident = format_ident!("{}", name);
             fields.push(parse_quote!(
+                #[tracked]
+                #[return_ref]
                 pub #name_ident: crate::cst::#type_name<'db>
             ));
         }

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

@@ -84,24 +84,30 @@ pub fn generate_visitor<'db>(
                 _fully_qualified_name: codegen_sdk_resolution::FullyQualifiedName<'db>,
                 #[id]
                 node_id: indextree::NodeId,
-                #[tracked]
-                #[return_ref]
-                pub node: crate::cst::#type_name<'db>,
+                // #[tracked]
+                // #[return_ref]
+                // pub node: crate::cst::#type_name<'db>,
                 #(#fields),*
             }
+            impl<'db> #variant<'db> {
+                pub fn node(&self, db: &'db dyn codegen_sdk_resolution::Db) -> &'db crate::cst::#type_name<'db> {
+                    let file = self.file(db);
+                    let tree = file.tree(db);
+                    tree.get(&self.node_id(db)).unwrap().as_ref().try_into().unwrap()
+                }
+            }
             impl<'db> codegen_sdk_resolution::HasFile<'db> for #variant<'db> {
                 type File<'db1> = #language_struct<'db1>;
                 fn file(&self, db: &'db dyn codegen_sdk_resolution::Db) -> &'db Self::File<'db> {
-                    let path = self.node(db).id().file(db).path(db);
-                    let input = db.get_file(path).unwrap();
-                    parse(db, input)
+                    let path = self._fully_qualified_name(db).path(db);
+                    parse(db, path)
                 }
-                fn root_path(&self, db: &'db dyn salsa::Database) -> PathBuf {
+                fn root_path(&self, db: &'db dyn codegen_sdk_resolution::Db) -> PathBuf {
                     self.node(db).id().root(db).path(db)
                 }
             }
             impl<'db> codegen_sdk_resolution::HasId<'db> for #variant<'db> {
-                fn fully_qualified_name(&self, db: &'db dyn codegen_sdk_resolution::Db) -> codegen_sdk_resolution::FullyQualifiedName<'db> {
+                fn fully_qualified_name(&self, db: &'db dyn salsa::Database) -> codegen_sdk_resolution::FullyQualifiedName<'db> {
                     self._fully_qualified_name(db)
                 }
             }
@@ -125,14 +131,14 @@ pub fn generate_visitor<'db>(
                         #(Self::#symbol_names(symbol) => symbol.file(db),)*
                     }
                 }
-                fn root_path(&self, db: &'db dyn salsa::Database) -> PathBuf {
+                fn root_path(&self, db: &'db dyn codegen_sdk_resolution::Db) -> PathBuf {
                     match self {
                         #(Self::#symbol_names(symbol) => symbol.root_path(db),)*
                     }
                 }
             }
             impl<'db> codegen_sdk_resolution::HasId<'db> for #symbol_name<'db> {
-                fn fully_qualified_name(&self, db: &'db dyn codegen_sdk_resolution::Db) -> codegen_sdk_resolution::FullyQualifiedName<'db> {
+                fn fully_qualified_name(&self, db: &'db dyn salsa::Database) -> codegen_sdk_resolution::FullyQualifiedName<'db> {
                     match self {
                         #(Self::#symbol_names(symbol) => symbol.fully_qualified_name(db),)*
                     }

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

@@ -41,6 +41,7 @@ salsa = { workspace = true }
 indextree = { workspace = true }
 indexmap = { workspace = true }
 rustc-hash = "2.1.1"
+hashbrown = { version = "0.15.2", features = ["rayon"] }
 [dev-dependencies]
 test-log = { workspace = true }
 [features]