Reuse BTreeSets in module resolver

Author: charliermarsh

crates/ruff/src/commands/analyze_graph.rs

@@ -169,8 +169,8 @@ pub(crate) fn analyze_graph(
 
     // Generate the import map.
     let import_map = match args.direction {
-        Direction::Dependencies => ImportMap::from_iter(imports),
-        Direction::Dependents => ImportMap::reverse(imports),
+        Direction::Dependencies => ImportMap::dependencies(imports),
+        Direction::Dependents => ImportMap::dependents(imports),
     };
 
     // Print to JSON.

crates/ruff_graph/src/lib.rs

@@ -91,17 +91,24 @@ impl ModuleImports {
 
 #[derive(Debug, Default)]
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
-pub struct ImportMap(BTreeMap<SystemPathBuf, ModuleImports>);
+pub struct ImportMap(BTreeMap<SystemPathBuf, BTreeSet<SystemPathBuf>>);
 
 impl ImportMap {
-    /// Insert a module's imports into the map.
-    pub fn insert(&mut self, path: SystemPathBuf, imports: ModuleImports) {
-        self.0.insert(path, imports);
+    /// Create an [`ImportMap`] of file to its dependencies.
+    ///
+    /// Assumes that the input is a collection of unique file paths and their imports.
+    pub fn dependencies(imports: impl IntoIterator<Item = (SystemPathBuf, ModuleImports)>) -> Self {
+        let mut map = ImportMap::default();
+        for (path, imports) in imports {
+            map.0.insert(path, imports.0);
+        }
+        map
     }
 
-    /// Reverse the [`ImportMap`], e.g., to convert from dependencies to dependents.
-    #[must_use]
-    pub fn reverse(imports: impl IntoIterator<Item = (SystemPathBuf, ModuleImports)>) -> Self {
+    /// Create an [`ImportMap`] of file to its dependents.
+    ///
+    /// Assumes that the input is a collection of unique file paths and their imports.
+    pub fn dependents(imports: impl IntoIterator<Item = (SystemPathBuf, ModuleImports)>) -> Self {
         let mut reverse = ImportMap::default();
         for (path, imports) in imports {
             for import in imports.0 {
@@ -112,13 +119,3 @@ impl ImportMap {
         reverse
     }
 }
-
-impl FromIterator<(SystemPathBuf, ModuleImports)> for ImportMap {
-    fn from_iter<I: IntoIterator<Item = (SystemPathBuf, ModuleImports)>>(iter: I) -> Self {
-        let mut map = ImportMap::default();
-        for (path, imports) in iter {
-            map.0.entry(path).or_default().0.extend(imports.0);
-        }
-        map
-    }
-}