Ensure fair benchmarks

rust/Cargo.toml

@@ -8,14 +8,25 @@ name = "_rustgrimp"
 crate-type = ["cdylib", "rlib"]
 
 [dependencies]
-log = "0.4.19"
-pyo3-log = "0.11.0"
-serde_json = "1.0.103"
 rayon = "1.10"
+bimap = "0.6.3"
+slotmap = "1.0.7"
+getset = "0.1.3"
+derive-new = "0.7.0"
+lazy_static = "1.5.0"
+string-interner = "0.18.0"
+thiserror = "2.0.11"
+itertools = "0.14.0"
+tap = "1.0.1"
+rustc-hash = "2.1.0"
+indexmap = "2.7.1"
 
 [dependencies.pyo3]
-version = "0.22.4"
+version = "0.23.4"
 
 [features]
 extension-module = ["pyo3/extension-module"]
 default = ["extension-module"]
+
+[dev-dependencies]
+serde_json = "1.0.137"

rust/src/errors.rs

@@ -0,0 +1,29 @@
+use crate::exceptions::{ModuleNotPresent, NoSuchContainer};
+use pyo3::exceptions::PyValueError;
+use pyo3::PyErr;
+use thiserror::Error;
+
+#[derive(Debug, Error)]
+pub enum GrimpError {
+    #[error("Module {0} is not present in the graph.")]
+    ModuleNotPresent(String),
+
+    #[error("Container {0} does not exist.")]
+    NoSuchContainer(String),
+
+    #[error("Modules have shared descendants.")]
+    SharedDescendants,
+}
+
+pub type GrimpResult<T> = Result<T, GrimpError>;
+
+impl From<GrimpError> for PyErr {
+    fn from(value: GrimpError) -> Self {
+        // A default mapping from `GrimpError`s to python exceptions.
+        match value {
+            GrimpError::ModuleNotPresent(_) => ModuleNotPresent::new_err(value.to_string()),
+            GrimpError::NoSuchContainer(_) => NoSuchContainer::new_err(value.to_string()),
+            GrimpError::SharedDescendants => PyValueError::new_err(value.to_string()),
+        }
+    }
+}

rust/src/exceptions.rs

@@ -0,0 +1,4 @@
+use pyo3::create_exception;
+
+create_exception!(_rustgrimp, ModuleNotPresent, pyo3::exceptions::PyException);
+create_exception!(_rustgrimp, NoSuchContainer, pyo3::exceptions::PyException);

rust/src/graph/direct_import_queries.rs

@@ -0,0 +1,57 @@
+use crate::errors::{GrimpError, GrimpResult};
+use crate::graph::{
+    ExtendWithDescendants, Graph, ImportDetails, ModuleToken, EMPTY_IMPORT_DETAILS,
+    EMPTY_MODULE_TOKENS,
+};
+use rustc_hash::FxHashSet;
+
+impl Graph {
+    pub fn count_imports(&self) -> usize {
+        self.imports.values().map(|imports| imports.len()).sum()
+    }
+
+    pub fn direct_import_exists(
+        &self,
+        importer: ModuleToken,
+        imported: ModuleToken,
+        as_packages: bool,
+    ) -> GrimpResult<bool> {
+        let mut importer: FxHashSet<_> = importer.into();
+        let mut imported: FxHashSet<_> = imported.into();
+        if as_packages {
+            importer.extend_with_descendants(self);
+            imported.extend_with_descendants(self);
+            if !(&importer & &imported).is_empty() {
+                return Err(GrimpError::SharedDescendants);
+            }
+        }
+
+        let direct_imports = importer
+            .iter()
+            .flat_map(|module| self.imports.get(*module).unwrap().iter().cloned())
+            .collect::<FxHashSet<ModuleToken>>();
+
+        Ok(!(&direct_imports & &imported).is_empty())
+    }
+
+    pub fn modules_directly_imported_by(&self, importer: ModuleToken) -> &FxHashSet<ModuleToken> {
+        self.imports.get(importer).unwrap_or(&EMPTY_MODULE_TOKENS)
+    }
+
+    pub fn modules_that_directly_import(&self, imported: ModuleToken) -> &FxHashSet<ModuleToken> {
+        self.reverse_imports
+            .get(imported)
+            .unwrap_or(&EMPTY_MODULE_TOKENS)
+    }
+
+    pub fn get_import_details(
+        &self,
+        importer: ModuleToken,
+        imported: ModuleToken,
+    ) -> &FxHashSet<ImportDetails> {
+        match self.import_details.get(&(importer, imported)) {
+            Some(import_details) => import_details,
+            None => &EMPTY_IMPORT_DETAILS,
+        }
+    }
+}

rust/src/graph/graph_manipulation.rs

@@ -0,0 +1,195 @@
+use crate::graph::{Graph, ImportDetails, Module, ModuleIterator, ModuleToken, MODULE_NAMES};
+use rustc_hash::FxHashSet;
+use slotmap::secondary::Entry;
+
+impl Graph {
+    /// `foo.bar.baz => [foo.bar.baz, foo.bar, foo]`
+    pub(crate) fn module_name_to_self_and_ancestors(&self, name: &str) -> Vec<String> {
+        let mut names = vec![name.to_owned()];
+        while let Some(parent_name) = parent_name(names.last().unwrap()) {
+            names.push(parent_name);
+        }
+        names
+    }
+
+    pub fn get_or_add_module(&mut self, name: &str) -> &Module {
+        if let Some(module) = self.get_module_by_name(name) {
+            let module = self.modules.get_mut(module.token).unwrap();
+            module.is_invisible = false;
+            return module;
+        }
+
+        let mut ancestor_names = self.module_name_to_self_and_ancestors(name);
+
+        {
+            let mut interner = MODULE_NAMES.write().unwrap();
+            let mut parent: Option<ModuleToken> = None;
+            while let Some(name) = ancestor_names.pop() {
+                let name = interner.get_or_intern(name);
+                if let Some(module) = self.modules_by_name.get_by_left(&name) {
+                    parent = Some(*module)
+                } else {
+                    let module = self.modules.insert_with_key(|token| Module {
+                        token,
+                        name,
+                        is_invisible: !ancestor_names.is_empty(),
+                        is_squashed: false,
+                    });
+                    self.modules_by_name.insert(name, module);
+                    self.module_parents.insert(module, parent);
+                    self.module_children.insert(module, FxHashSet::default());
+                    self.imports.insert(module, FxHashSet::default());
+                    self.reverse_imports.insert(module, FxHashSet::default());
+                    if let Some(parent) = parent {
+                        self.module_children[parent].insert(module);
+                    }
+                    parent = Some(module)
+                }
+            }
+        }
+
+        self.get_module_by_name(name).unwrap()
+    }
+
+    pub fn get_or_add_squashed_module(&mut self, module: &str) -> &Module {
+        let module = self.get_or_add_module(module).token();
+        self.mark_module_squashed(module);
+        self.get_module(module).unwrap()
+    }
+
+    fn mark_module_squashed(&mut self, module: ModuleToken) {
+        let module = self.modules.get_mut(module).unwrap();
+        if !self.module_children[module.token].is_empty() {
+            panic!("cannot mark a module with children as squashed")
+        }
+        module.is_squashed = true;
+    }
+
+    pub fn remove_module(&mut self, module: ModuleToken) {
+        let module = self.get_module(module);
+        if module.is_none() {
+            return;
+        }
+        let module = module.unwrap().token();
+
+        // TODO(peter) Remove children automatically here, or raise an error?
+        if !self.module_children[module].is_empty() {
+            for child in self.module_children[module].clone() {
+                self.remove_module(child);
+            }
+        }
+
+        // Update hierarchy.
+        if let Some(parent) = self.module_parents[module] {
+            self.module_children[parent].remove(&module);
+        }
+        self.modules_by_name.remove_by_right(&module);
+        self.modules.remove(module);
+        self.module_parents.remove(module);
+        self.module_children.remove(module);
+
+        // Update imports.
+        for imported in self.modules_directly_imported_by(module).clone() {
+            self.remove_import(module, imported);
+        }
+        for importer in self.modules_that_directly_import(module).clone() {
+            self.remove_import(importer, module);
+        }
+        self.imports.remove(module);
+        self.reverse_imports.remove(module);
+    }
+
+    pub fn add_import(&mut self, importer: ModuleToken, imported: ModuleToken) {
+        self.imports
+            .entry(importer)
+            .unwrap()
+            .or_default()
+            .insert(imported);
+        self.reverse_imports
+            .entry(imported)
+            .unwrap()
+            .or_default()
+            .insert(importer);
+    }
+
+    pub fn add_detailed_import(
+        &mut self,
+        importer: ModuleToken,
+        imported: ModuleToken,
+        line_number: usize,
+        line_contents: &str,
+    ) {
+        self.imports
+            .entry(importer)
+            .unwrap()
+            .or_default()
+            .insert(imported);
+        self.reverse_imports
+            .entry(imported)
+            .unwrap()
+            .or_default()
+            .insert(importer);
+        self.import_details
+            .entry((importer, imported))
+            .or_default()
+            .insert(ImportDetails::new(line_number, line_contents.to_owned()));
+    }
+
+    pub fn remove_import(&mut self, importer: ModuleToken, imported: ModuleToken) {
+        match self.imports.entry(importer).unwrap() {
+            Entry::Occupied(mut entry) => {
+                entry.get_mut().remove(&imported);
+            }
+            Entry::Vacant(_) => {}
+        };
+        match self.reverse_imports.entry(imported).unwrap() {
+            Entry::Occupied(mut entry) => {
+                entry.get_mut().remove(&importer);
+            }
+            Entry::Vacant(_) => {}
+        };
+        self.import_details.remove(&(importer, imported));
+    }
+
+    pub fn squash_module(&mut self, module: ModuleToken) {
+        // Get descendants and their imports.
+        let descendants: FxHashSet<_> = self.get_module_descendants(module).tokens().collect();
+
+        let modules_imported_by_descendants: FxHashSet<_> = descendants
+            .iter()
+            .flat_map(|descendant| {
+                self.modules_directly_imported_by(*descendant)
+                    .iter()
+                    .cloned()
+            })
+            .collect();
+        let modules_that_import_descendants: FxHashSet<_> = descendants
+            .iter()
+            .flat_map(|descendant| {
+                self.modules_that_directly_import(*descendant)
+                    .iter()
+                    .cloned()
+            })
+            .collect();
+
+        // Remove any descendants.
+        for descendant in descendants {
+            self.remove_module(descendant);
+        }
+
+        // Add descendants and imports to parent module.
+        for imported in modules_imported_by_descendants {
+            self.add_import(module, imported);
+        }
+
+        for importer in modules_that_import_descendants {
+            self.add_import(importer, module);
+        }
+
+        self.mark_module_squashed(module);
+    }
+}
+
+fn parent_name(name: &str) -> Option<String> {
+    name.rsplit_once(".").map(|(base, _)| base.to_owned())
+}

rust/src/graph/hierarchy_queries.rs

@@ -0,0 +1,50 @@
+use crate::graph::{Graph, Module, ModuleToken, MODULE_NAMES};
+
+impl Graph {
+    pub fn get_module_by_name(&self, name: &str) -> Option<&Module> {
+        let interner = MODULE_NAMES.read().unwrap();
+        let name = interner.get(name)?;
+        match self.modules_by_name.get_by_left(&name) {
+            Some(token) => self.get_module(*token),
+            None => None,
+        }
+    }
+
+    pub fn get_module(&self, module: ModuleToken) -> Option<&Module> {
+        self.modules.get(module)
+    }
+
+    // TODO(peter) Guarantee order?
+    pub fn all_modules(&self) -> impl Iterator<Item = &Module> {
+        self.modules.values()
+    }
+
+    pub fn get_module_parent(&self, module: ModuleToken) -> Option<&Module> {
+        match self.module_parents.get(module) {
+            Some(parent) => parent.map(|parent| self.get_module(parent).unwrap()),
+            None => None,
+        }
+    }
+
+    pub fn get_module_children(&self, module: ModuleToken) -> impl Iterator<Item = &Module> {
+        let children = match self.module_children.get(module) {
+            Some(children) => children
+                .iter()
+                .map(|child| self.get_module(*child).unwrap())
+                .collect(),
+            None => Vec::new(),
+        };
+        children.into_iter()
+    }
+
+    /// Returns an iterator over the passed modules descendants.
+    ///
+    /// Parent modules will be yielded before their child modules.
+    pub fn get_module_descendants(&self, module: ModuleToken) -> impl Iterator<Item = &Module> {
+        let mut descendants = self.get_module_children(module).collect::<Vec<_>>();
+        for child in descendants.clone() {
+            descendants.extend(self.get_module_descendants(child.token).collect::<Vec<_>>())
+        }
+        descendants.into_iter()
+    }
+}