From 09c267f0d75f7d02f6977aabbb226b2c5ff982f3 Mon Sep 17 00:00:00 2001
From: Charlie Marsh <charlie.r.marsh@gmail.com>
Date: Sat, 21 Sep 2024 14:23:32 -0400
Subject: [PATCH] Reuse BTreeSets in module resolver

---
 crates/ruff/src/commands/analyze_graph.rs |  4 ++--
 crates/ruff_graph/src/lib.rs              | 29 ++++++++++-------------
 2 files changed, 15 insertions(+), 18 deletions(-)

diff --git a/crates/ruff/src/commands/analyze_graph.rs b/crates/ruff/src/commands/analyze_graph.rs
index 4b5ce6ee994ce..58df61a36d551 100644
--- a/crates/ruff/src/commands/analyze_graph.rs
+++ b/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.
diff --git a/crates/ruff_graph/src/lib.rs b/crates/ruff_graph/src/lib.rs
index f11f03ffe5894..6df130757987d 100644
--- a/crates/ruff_graph/src/lib.rs
+++ b/crates/ruff_graph/src/lib.rs
@@ -94,14 +94,21 @@ impl ModuleImports {
 pub struct ImportMap(BTreeMap<SystemPathBuf, ModuleImports>);
 
 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);
+        }
+        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
-    }
-}