Optimize memory usage

Author: Peter554

src/import_discovery.rs

@@ -4,31 +4,33 @@ use rustpython_parser::{
     self,
     ast::{Mod, Stmt},
 };
+use std::collections::HashMap;
 use std::collections::HashSet;
 use std::fs;
-use std::{collections::HashMap, sync::Arc};
 
 use crate::ast_visit;
 use crate::indexing;
 use crate::package_discovery;
 
-pub type Imports = HashMap<String, HashSet<String>>;
+pub type Imports<'a> = HashMap<&'a str, HashSet<&'a str>>;
 
-pub fn discover_imports(modules_by_pypath: &indexing::ModulesByPypath) -> Result<Imports> {
+pub fn discover_imports<'a>(
+    modules_by_pypath: &'a indexing::ModulesByPypath,
+) -> Result<Imports<'a>> {
     modules_by_pypath
         .values()
         .par_bridge()
         .map(|module| {
-            let imports = get_imports_for_module(Arc::clone(module), modules_by_pypath)?;
-            Ok((module.pypath.clone(), imports))
+            let imports = get_imports_for_module(module, modules_by_pypath)?;
+            Ok((module.pypath.as_str(), imports))
         })
         .collect::<Result<Imports>>()
 }
 
-fn get_imports_for_module(
-    module: Arc<package_discovery::Module>,
-    modules_by_pypath: &indexing::ModulesByPypath,
-) -> Result<HashSet<String>> {
+fn get_imports_for_module<'a>(
+    module: &'a package_discovery::Module,
+    modules_by_pypath: &'a indexing::ModulesByPypath,
+) -> Result<HashSet<&'a str>> {
     let code = fs::read_to_string(&module.path)?;
     let ast = rustpython_parser::parse(
         &code,
@@ -44,7 +46,7 @@ fn get_imports_for_module(
     };
 
     let mut visitor = ImportVisitor {
-        module: Arc::clone(&module),
+        module,
         modules_by_pypath,
         imports: HashSet::new(),
     };
@@ -54,9 +56,9 @@ fn get_imports_for_module(
 }
 
 struct ImportVisitor<'a> {
-    module: Arc<package_discovery::Module>,
-    modules_by_pypath: &'a indexing::ModulesByPypath,
-    imports: HashSet<String>,
+    module: &'a package_discovery::Module,
+    modules_by_pypath: &'a indexing::ModulesByPypath<'a>,
+    imports: HashSet<&'a str>,
 }
 
 impl<'a> ast_visit::StatementVisitor for ImportVisitor<'a> {
@@ -65,9 +67,9 @@ impl<'a> ast_visit::StatementVisitor for ImportVisitor<'a> {
             Stmt::Import(stmt) => {
                 for name in stmt.names.iter() {
                     for pypath in [name.name.to_string(), format!("{}.__init__", name.name)] {
-                        match self.modules_by_pypath.get(&pypath) {
+                        match self.modules_by_pypath.get(pypath.as_str()) {
                             Some(imported_module) => {
-                                self.imports.insert(imported_module.pypath.clone());
+                                self.imports.insert(imported_module.pypath.as_str());
                             }
                             None => continue,
                         }
@@ -116,9 +118,9 @@ impl<'a> ast_visit::StatementVisitor for ImportVisitor<'a> {
                         format!("{}.{}.__init__", &pypath_prefix, name.name),
                         format!("{}.__init__", &pypath_prefix),
                     ] {
-                        match self.modules_by_pypath.get(&pypath) {
+                        match self.modules_by_pypath.get(pypath.as_str()) {
                             Some(imported_module) => {
-                                self.imports.insert(imported_module.pypath.clone());
+                                self.imports.insert(imported_module.pypath.as_str());
                                 break;
                             }
                             None => continue,
@@ -143,10 +145,10 @@ mod tests {
     fn test_get_imports_for_module() {
         let root_package_path = Path::new("./example");
         let root_package = package_discovery::discover_package(root_package_path).unwrap();
-        let modules_by_pypath = indexing::get_modules_by_pypath(Arc::clone(&root_package)).unwrap();
+        let modules_by_pypath = indexing::get_modules_by_pypath(&root_package).unwrap();
 
         let module = modules_by_pypath.get("example.__init__").unwrap();
-        let imports = get_imports_for_module(Arc::clone(module), &modules_by_pypath).unwrap();
+        let imports = get_imports_for_module(module, &modules_by_pypath).unwrap();
         assert_eq!(
             imports,
             [
@@ -167,12 +169,11 @@ mod tests {
                 "example.child5.__init__",
             ]
             .into_iter()
-            .map(|i| i.to_string())
             .collect::<HashSet<_>>()
         );
 
         let module = modules_by_pypath.get("example.child.__init__").unwrap();
-        let imports = get_imports_for_module(Arc::clone(module), &modules_by_pypath).unwrap();
+        let imports = get_imports_for_module(module, &modules_by_pypath).unwrap();
         assert_eq!(
             imports,
             [
@@ -193,12 +194,11 @@ mod tests {
                 "example.child5.__init__",
             ]
             .into_iter()
-            .map(|i| i.to_string())
             .collect::<HashSet<_>>()
         );
 
         let module = modules_by_pypath.get("example.z").unwrap();
-        let imports = get_imports_for_module(Arc::clone(module), &modules_by_pypath).unwrap();
+        let imports = get_imports_for_module(module, &modules_by_pypath).unwrap();
         assert_eq!(
             imports,
             [
@@ -219,12 +219,11 @@ mod tests {
                 "example.child5.__init__",
             ]
             .into_iter()
-            .map(|i| i.to_string())
             .collect::<HashSet<_>>()
         );
 
         let module = modules_by_pypath.get("example.child.c_z").unwrap();
-        let imports = get_imports_for_module(Arc::clone(module), &modules_by_pypath).unwrap();
+        let imports = get_imports_for_module(module, &modules_by_pypath).unwrap();
         assert_eq!(
             imports,
             [
@@ -245,7 +244,6 @@ mod tests {
                 "example.child5.__init__",
             ]
             .into_iter()
-            .map(|i| i.to_string())
             .collect::<HashSet<_>>()
         );
     }

src/indexing.rs

@@ -1,36 +1,34 @@
 use anyhow::Result;
-use std::{collections::HashMap, sync::Arc};
+use std::collections::HashMap;
 
 use crate::package_discovery;
 
-pub type PackagesByPypath = HashMap<String, Arc<package_discovery::Package>>;
-pub type ModulesByPypath = HashMap<String, Arc<package_discovery::Module>>;
+pub type PackagesByPypath<'a> = HashMap<&'a str, &'a package_discovery::Package>;
+pub type ModulesByPypath<'a> = HashMap<&'a str, &'a package_discovery::Module>;
 
 pub fn get_packages_by_pypath(
-    root_package: Arc<package_discovery::Package>,
+    root_package: &package_discovery::Package,
 ) -> Result<PackagesByPypath> {
     let mut m = HashMap::new();
     let mut q = vec![root_package];
     while let Some(package) = q.pop() {
-        m.insert(package.pypath.clone(), Arc::clone(&package));
+        m.insert(package.pypath.as_str(), package);
         for child in package.children.iter() {
-            q.push(Arc::clone(&child));
+            q.push(child);
         }
     }
     Ok(m)
 }
 
-pub fn get_modules_by_pypath(
-    root_package: Arc<package_discovery::Package>,
-) -> Result<ModulesByPypath> {
+pub fn get_modules_by_pypath(root_package: &package_discovery::Package) -> Result<ModulesByPypath> {
     let mut m = HashMap::new();
     let mut q = vec![root_package];
     while let Some(package) = q.pop() {
         for module in package.modules.iter() {
-            m.insert(module.pypath.clone(), Arc::clone(module));
+            m.insert(module.pypath.as_str(), module);
         }
         for child in package.children.iter() {
-            q.push(Arc::clone(&child));
+            q.push(child);
         }
     }
     Ok(m)
@@ -48,36 +46,29 @@ mod tests {
         let root_package_path = Path::new("./example");
         let root_package = package_discovery::discover_package(root_package_path).unwrap();
 
-        let mut packages_by_pypath = get_packages_by_pypath(Arc::clone(&root_package)).unwrap();
+        let packages_by_pypath = get_packages_by_pypath(&root_package).unwrap();
 
         assert_eq!(packages_by_pypath.len(), 6);
+        assert_eq!(packages_by_pypath.get("example").unwrap(), &&root_package);
         assert_eq!(
-            packages_by_pypath.remove("example").unwrap(),
-            Arc::clone(&root_package)
+            packages_by_pypath.get("example.child").unwrap(),
+            root_package
+                .children
+                .iter()
+                .filter(|child| child.pypath == "example.child")
+                .collect::<Vec<_>>()
+                .first()
+                .unwrap()
         );
         assert_eq!(
-            packages_by_pypath.remove("example.child").unwrap(),
-            Arc::clone(
-                &root_package
-                    .children
-                    .iter()
-                    .filter(|child| child.pypath == "example.child")
-                    .collect::<Vec<_>>()
-                    .first()
-                    .unwrap()
-            )
-        );
-        assert_eq!(
-            packages_by_pypath.remove("example.child2").unwrap(),
-            Arc::clone(
-                &root_package
-                    .children
-                    .iter()
-                    .filter(|child| child.pypath == "example.child2")
-                    .collect::<Vec<_>>()
-                    .first()
-                    .unwrap()
-            )
+            packages_by_pypath.get("example.child2").unwrap(),
+            root_package
+                .children
+                .iter()
+                .filter(|child| child.pypath == "example.child2")
+                .collect::<Vec<_>>()
+                .first()
+                .unwrap()
         );
     }
 
@@ -86,13 +77,13 @@ mod tests {
         let root_package_path = Path::new("./example");
         let root_package = package_discovery::discover_package(root_package_path).unwrap();
 
-        let mut modules_by_pypath = get_modules_by_pypath(Arc::clone(&root_package)).unwrap();
+        let modules_by_pypath = get_modules_by_pypath(&root_package).unwrap();
 
         assert_eq!(modules_by_pypath.len(), 18);
         for module in root_package.modules.iter() {
             assert_eq!(
-                modules_by_pypath.remove(&module.pypath).unwrap(),
-                Arc::clone(module)
+                modules_by_pypath.get(module.pypath.as_str()).unwrap(),
+                &module
             )
         }
         for child_module in root_package
@@ -106,8 +97,8 @@ mod tests {
             .iter()
         {
             assert_eq!(
-                modules_by_pypath.remove(&child_module.pypath).unwrap(),
-                Arc::clone(child_module)
+                modules_by_pypath.get(child_module.pypath.as_str()).unwrap(),
+                &child_module
             )
         }
     }

src/main.rs

@@ -1,5 +1,4 @@
 use anyhow::Result;
-use std::sync::Arc;
 use std::{env, path::Path};
 
 use pyimports::import_discovery;
@@ -11,8 +10,8 @@ pub fn main() -> Result<()> {
     let root_package_path = Path::new(&args[1]);
 
     let package = package_discovery::discover_package(root_package_path)?;
-    let packages_by_pypath = indexing::get_packages_by_pypath(Arc::clone(&package))?;
-    let modules_by_pypath = indexing::get_modules_by_pypath(Arc::clone(&package))?;
+    let packages_by_pypath = indexing::get_packages_by_pypath(&package)?;
+    let modules_by_pypath = indexing::get_modules_by_pypath(&package)?;
     let imports = import_discovery::discover_imports(&modules_by_pypath)?;
 
     // dbg!(&package);

src/package_discovery.rs

@@ -1,7 +1,6 @@
 use anyhow::Result;
 use rayon::prelude::*;
 use std::path::PathBuf;
-use std::sync::Arc;
 use std::{fs, io, path::Path};
 use thiserror::Error;
 
@@ -24,8 +23,8 @@ pub enum PackageDiscoveryError {
 pub struct Package {
     pub pypath: String,
     pub path: PathBuf,
-    pub children: Vec<Arc<Package>>,
-    pub modules: Vec<Arc<Module>>,
+    pub children: Vec<Package>,
+    pub modules: Vec<Module>,
 }
 
 impl Package {
@@ -38,11 +37,11 @@ impl Package {
         }
     }
 
-    pub fn add_child(&mut self, child: Arc<Package>) {
+    pub fn add_child(&mut self, child: Package) {
         self.children.push(child);
     }
 
-    pub fn add_module(&mut self, module: Arc<Module>) {
+    pub fn add_module(&mut self, module: Module) {
         self.modules.push(module);
     }
 }
@@ -59,11 +58,11 @@ impl Module {
     }
 }
 
-pub fn discover_package(package_path: &Path) -> Result<Arc<Package>> {
+pub fn discover_package(package_path: &Path) -> Result<Package> {
     _discover_package(package_path, package_path)
 }
 
-fn _discover_package(root_package_path: &Path, package_path: &Path) -> Result<Arc<Package>> {
+fn _discover_package(root_package_path: &Path, package_path: &Path) -> Result<Package> {
     let (is_package, files, dirs) = fs::read_dir(package_path)
         .map_err(PackageDiscoveryError::CannotReadDir)?
         .par_bridge()
@@ -101,18 +100,18 @@ fn _discover_package(root_package_path: &Path, package_path: &Path) -> Result<Ar
     }
 
     let pypath = get_pypath(root_package_path, package_path, false)?;
-    let mut package = Package::new(pypath, package_path.to_owned());
+    let mut package = Package::new(pypath, package_path.to_path_buf());
 
     for module in files
         .par_iter()
         .filter(|file| file.path().extension().unwrap_or_default() == "py")
         .map(|file| {
             let pypath = get_pypath(root_package_path, &file.path(), true)?;
-            Ok(Module::new(pypath, file.path().to_owned()))
+            Ok(Module::new(pypath, file.path().to_path_buf()))
         })
         .collect::<Result<Vec<_>>>()?
     {
-        package.add_module(Arc::new(module));
+        package.add_module(module);
     }
 
     for child in dirs
@@ -123,7 +122,7 @@ fn _discover_package(root_package_path: &Path, package_path: &Path) -> Result<Ar
     {
         match child {
             Ok(child) => {
-                package.add_child(Arc::clone(&child));
+                package.add_child(child);
             }
             Err(e) => match e.root_cause().downcast_ref::<PackageDiscoveryError>() {
                 Some(PackageDiscoveryError::NotAPythonPackage) => continue,
@@ -132,7 +131,7 @@ fn _discover_package(root_package_path: &Path, package_path: &Path) -> Result<Ar
         }
     }
 
-    Ok(Arc::new(package))
+    Ok(package)
 }
 
 fn get_pypath(root_package_path: &Path, path: &Path, is_file: bool) -> Result<String> {