Track referrers of modules

Cargo.lock

@@ -161,6 +161,7 @@ dependencies = [
  "once_cell",
  "ordered-float",
  "salsa",
+ "smallvec",
  "smol_str",
  "syntax",
  "url",

crates/ide/Cargo.toml

@@ -16,6 +16,7 @@ nix-interop = { path = "../nix-interop" }
 once_cell = "1.17.0"
 ordered-float = "3.4.0"
 salsa = "0.17.0-pre.2"
+smallvec = { version = "1.10.0", features = ["const_generics", "union"] }
 smol_str = "0.1.23"
 syntax = { path = "../syntax" }
 url = "2.3.1"

crates/ide/src/base.rs

@@ -5,10 +5,10 @@ use std::path::{Path, PathBuf};
 use std::sync::Arc;
 use syntax::{TextRange, TextSize};
 
-#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
 pub struct FileId(pub u32);
 
-#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
 pub struct SourceRootId(pub u32);
 
 /// An path in the virtual filesystem.
@@ -126,7 +126,7 @@ impl FileSet {
         &self.paths[&file]
     }
 
-    pub fn iter(&self) -> impl Iterator<Item = (FileId, &'_ VfsPath)> + '_ {
+    pub fn iter(&self) -> impl Iterator<Item = (FileId, &'_ VfsPath)> + ExactSizeIterator + '_ {
         self.paths.iter().map(|(&file, path)| (file, path))
     }
 }
@@ -157,7 +157,7 @@ impl SourceRoot {
         self.file_set.path_for_file(file)
     }
 
-    pub fn iter(&self) -> impl Iterator<Item = (FileId, &'_ VfsPath)> + '_ {
+    pub fn files(&self) -> impl Iterator<Item = (FileId, &'_ VfsPath)> + ExactSizeIterator + '_ {
         self.file_set.iter()
     }
 
@@ -281,7 +281,7 @@ impl Change {
         if let Some(roots) = self.roots {
             u32::try_from(roots.len()).expect("Length overflow");
             for (sid, root) in (0u32..).map(SourceRootId).zip(roots) {
-                for (fid, _) in root.iter() {
+                for (fid, _) in root.files() {
                     db.set_file_source_root_with_durability(fid, sid, Durability::HIGH);
                 }
                 db.set_source_root_with_durability(sid, Arc::new(root), Durability::HIGH);

crates/ide/src/def/mod.rs

@@ -11,6 +11,7 @@ use crate::{Diagnostic, FileId, SourceRootId, VfsPath};
 use la_arena::{Arena, ArenaMap, Idx};
 use nix_interop::DEFAULT_IMPORT_FILE;
 use ordered_float::OrderedFloat;
+use smallvec::SmallVec;
 use smol_str::SmolStr;
 use std::collections::{HashMap, HashSet};
 use std::ops;
@@ -42,8 +43,17 @@ pub trait DefDatabase: SourceDatabase {
     #[salsa::invoke(Module::module_references_query)]
     fn module_references(&self, file_id: FileId) -> Arc<HashSet<FileId>>;
 
+    fn source_root_referrer_graph(
+        &self,
+        sid: SourceRootId,
+    ) -> Arc<HashMap<FileId, ModuleReferrers>>;
+
     fn source_root_closure(&self, id: SourceRootId) -> Arc<HashSet<FileId>>;
 
+    // The result is not wrapped in Arc. Typically, the number of referrers is just 1 or 0.
+    // And also this method is not call so often.
+    fn module_referrers(&self, file_id: FileId) -> ModuleReferrers;
+
     #[salsa::invoke(Path::resolve_path_query)]
     fn resolve_path(&self, path: Path) -> Option<VfsPath>;
 
@@ -84,6 +94,40 @@ fn source_map(db: &dyn DefDatabase, file_id: FileId) -> Arc<ModuleSourceMap> {
     db.module_with_source_map(file_id).1
 }
 
+pub type ModuleReferrers = SmallVec<[FileId; 2]>;
+
+fn source_root_referrer_graph(
+    db: &dyn DefDatabase,
+    sid: SourceRootId,
+) -> Arc<HashMap<FileId, ModuleReferrers>> {
+    // Assert our inline threshould costs no extra memory.
+    const _: [(); std::mem::size_of::<Vec<FileId>>()] =
+        [(); std::mem::size_of::<ModuleReferrers>()];
+
+    let source_root = db.source_root(sid);
+    let mut graph = HashMap::<FileId, ModuleReferrers>::with_capacity(source_root.files().len());
+    for (file, _) in source_root.files() {
+        for &referrer in &*db.module_references(file) {
+            // This never duplicates, since `module_references` returns a `HashSet`.
+            graph.entry(referrer).or_default().push(file);
+        }
+    }
+
+    // Keep the order deterministic.
+    for referees in graph.values_mut() {
+        referees.sort();
+    }
+
+    graph.shrink_to_fit();
+    Arc::new(graph)
+}
+
+fn module_referrers(db: &dyn DefDatabase, file_id: FileId) -> ModuleReferrers {
+    let sid = db.file_source_root(file_id);
+    let graph = db.source_root_referrer_graph(sid);
+    graph.get(&file_id).cloned().unwrap_or_default()
+}
+
 fn source_root_closure(db: &dyn DefDatabase, id: SourceRootId) -> Arc<HashSet<FileId>> {
     let entry = match db.source_root(id).entry() {
         Some(file) => file,

crates/ide/src/def/tests.rs

@@ -2,6 +2,7 @@ use super::DefDatabase;
 use crate::tests::TestDB;
 use crate::{FlakeInfo, ModuleKind, SourceDatabase, VfsPath};
 use expect_test::expect;
+use itertools::Itertools;
 use std::collections::{HashMap, HashSet};
 
 #[test]
@@ -64,6 +65,64 @@ baz/../../bar.nix + ../default.nix
     }
 }
 
+#[test]
+fn source_root_referrer_graph() {
+    let (db, f) = TestDB::from_fixture(
+        "
+#- /default.nix
+./foo/bar.nix
+
+#- /foo/bar.nix
+baz/../../bar.nix + ../default.nix
+
+#- /bar.nix
+./.
+
+#- /single.nix
+42
+
+#- /mutual1.nix
+./mutual2.nix
+
+#- /mutual2.nix
+./mutual1.nix
+    ",
+    )
+    .unwrap();
+
+    let sid = db.file_source_root(f["/default.nix"]);
+    let source_root = db.source_root(sid);
+    let graph = db.source_root_referrer_graph(sid);
+
+    let got = source_root
+        .files()
+        .sorted_by_key(|&(f, _)| f)
+        .map(|(referee, _)| {
+            let referrers = graph.get(&referee).cloned().unwrap_or_default();
+            let referee = source_root.path_for_file(referee).display();
+            let referrers = referrers
+                .iter()
+                .map(|&f| source_root.path_for_file(f).display())
+                .join(", ");
+            format!("{referee} <- [{referrers}]\n")
+        })
+        .collect::<String>();
+    expect![[r#"
+        /default.nix <- [/foo/bar.nix, /bar.nix]
+        /foo/bar.nix <- [/default.nix]
+        /bar.nix <- [/foo/bar.nix]
+        /single.nix <- []
+        /mutual1.nix <- [/mutual2.nix]
+        /mutual2.nix <- [/mutual1.nix]
+    "#]]
+    .assert_eq(&got);
+
+    assert_eq!(
+        db.module_referrers(f["/default.nix"]).into_vec(),
+        vec![f["/foo/bar.nix"], f["/bar.nix"]],
+    );
+}
+
 #[test]
 fn source_root_closure() {
     let (db, f) = TestDB::from_fixture(