4 files changed, 131 insertions, 15 deletions

diff --git a/crates/libanalysis/src/imp.rs b/crates/libanalysis/src/imp.rs
index 6a6e1e933..f1d72da15 100644
--- a/crates/libanalysis/src/imp.rs
+++ b/crates/libanalysis/src/imp.rs
@@ -5,7 +5,7 @@ use std::{
     },
     fmt,
     time::Instant,
-    collections::HashMap,
+    collections::{HashMap, HashSet, VecDeque},
     panic,
 };
 
@@ -23,7 +23,7 @@ use {
     module_map::Problem,
     symbol_index::FileSymbols,
     module_map::{ModuleMap, ChangeKind},
-    JobToken,
+    JobToken, CrateGraph, CrateId,
 };
 
 #[derive(Debug)]
@@ -37,7 +37,6 @@ impl AnalysisHostImpl {
             data: Arc::new(WorldData::default()),
         }
     }
-
     pub fn analysis(
         &self,
         file_resolver: Arc<dyn FileResolver>,
@@ -48,7 +47,6 @@ impl AnalysisHostImpl {
             data: self.data.clone(),
         }
     }
-
     pub fn change_files(&mut self, changes: &mut dyn Iterator<Item=(FileId, Option<String>)>) {
         let data = self.data_mut();
         for (file_id, text) in changes {
@@ -71,7 +69,15 @@ impl AnalysisHostImpl {
             }
         }
     }
-
+    pub fn set_crate_graph(&mut self, graph: CrateGraph) {
+        let mut visited = HashSet::new();
+        for &file_id in graph.crate_roots.values() {
+            if !visited.insert(file_id) {
+                panic!("duplicate crate root: {:?}", file_id);
+            }
+        }
+        self.data_mut().crate_graph = graph;
+    }
     fn data_mut(&mut self) -> &mut WorldData {
         Arc::make_mut(&mut self.data)
     }
@@ -145,6 +151,33 @@ impl AnalysisImpl {
             .collect()
     }
 
+    pub fn crate_root(&self, id: FileId) -> Vec<CrateId> {
+        let module_map = &self.data.module_map;
+        let crate_graph = &self.data.crate_graph;
+        let mut res = Vec::new();
+        let mut work = VecDeque::new();
+        work.push_back(id);
+        let mut visited = HashSet::new();
+        while let Some(id) = work.pop_front() {
+            if let Some(crate_id) = crate_graph.crate_id_for_crate_root(id) {
+                res.push(crate_id);
+                continue;
+            }
+            let mid = module_map.file2module(id);
+            let parents = module_map
+                .parent_module_ids(
+                    mid,
+                    &*self.file_resolver,
+                    &|file_id| self.file_syntax(file_id),
+                )
+                .into_iter()
+                .map(|id| module_map.module2file(id))
+                .filter(|&id| visited.insert(id));
+            work.extend(parents);
+        }
+        res
+    }
+
     pub fn approximately_resolve_symbol(
         &self,
         id: FileId,
@@ -295,6 +328,7 @@ impl AnalysisImpl {
 
 #[derive(Clone, Default, Debug)]
 struct WorldData {
+    crate_graph: CrateGraph,
     file_map: HashMap<FileId, Arc<FileData>>,
     module_map: ModuleMap,
 }
@@ -356,3 +390,12 @@ impl SourceChange {
         }
     }
 }
+
+impl CrateGraph {
+    fn crate_id_for_crate_root(&self, file_id: FileId) -> Option<CrateId> {
+        let (&crate_id, _) = self.crate_roots
+            .iter()
+            .find(|(_crate_id, &root_id)| root_id == file_id)?;
+        Some(crate_id)
+    }
+}
diff --git a/crates/libanalysis/src/lib.rs b/crates/libanalysis/src/lib.rs
index a59fd1c09..041abbb89 100644
--- a/crates/libanalysis/src/lib.rs
+++ b/crates/libanalysis/src/lib.rs
@@ -15,7 +15,10 @@ mod module_map;
 mod imp;
 mod job;
 
-use std::sync::Arc;
+use std::{
+    sync::Arc,
+    collections::HashMap,
+};
 
 use relative_path::{RelativePath, RelativePathBuf};
 use libsyntax2::{File, TextRange, TextUnit, AtomEdit};
@@ -30,6 +33,14 @@ pub use job::{JobToken, JobHandle};
 #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
 pub struct FileId(pub u32);
 
+#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
+pub struct CrateId(pub u32);
+
+#[derive(Debug, Clone, Default)]
+pub struct CrateGraph {
+    pub crate_roots: HashMap<CrateId, FileId>,
+}
+
 pub trait FileResolver: Send + Sync + 'static {
     fn file_stem(&self, id: FileId) -> String;
     fn resolve(&self, id: FileId, path: &RelativePath) -> Option<FileId>;
@@ -53,6 +64,9 @@ impl AnalysisHost {
     pub fn change_files(&mut self, mut changes: impl Iterator<Item=(FileId, Option<String>)>) {
         self.imp.change_files(&mut changes)
     }
+    pub fn set_crate_graph(&mut self, graph: CrateGraph) {
+        self.imp.set_crate_graph(graph)
+    }
 }
 
 #[derive(Debug)]
@@ -168,6 +182,9 @@ impl Analysis {
     pub fn parent_module(&self, file_id: FileId) -> Vec<(FileId, FileSymbol)> {
         self.imp.parent_module(file_id)
     }
+    pub fn crate_root(&self, file_id: FileId) -> Vec<CrateId> {
+        self.imp.crate_root(file_id)
+    }
     pub fn runnables(&self, file_id: FileId) -> Vec<Runnable> {
         let file = self.file_syntax(file_id);
         libeditor::runnables(&file)
diff --git a/crates/libanalysis/src/module_map.rs b/crates/libanalysis/src/module_map.rs
index 47c0018e1..38935807d 100644
--- a/crates/libanalysis/src/module_map.rs
+++ b/crates/libanalysis/src/module_map.rs
@@ -91,16 +91,38 @@ impl ModuleMap {
         file_resolver: &FileResolver,
         syntax_provider: &SyntaxProvider,
     ) -> Vec<(ModuleId, SmolStr, SyntaxNode)> {
-        let links = self.links(file_resolver, syntax_provider);
-        let res = links
+        let mut res = Vec::new();
+        self.for_each_parent_link(m, file_resolver, syntax_provider, |link| {
+            res.push(
+                (link.owner, link.name().clone(), link.syntax.clone())
+            )
+        });
+        res
+    }
+
+    pub fn parent_module_ids(
+        &self,
+        m: ModuleId,
+        file_resolver: &FileResolver,
+        syntax_provider: &SyntaxProvider,
+    ) -> Vec<ModuleId> {
+        let mut res = Vec::new();
+        self.for_each_parent_link(m, file_resolver, syntax_provider, |link| res.push(link.owner));
+        res
+    }
+
+    fn for_each_parent_link(
+        &self,
+        m: ModuleId,
+        file_resolver: &FileResolver,
+        syntax_provider: &SyntaxProvider,
+        f: impl FnMut(&Link)
+    ) {
+        self.links(file_resolver, syntax_provider)
             .links
             .iter()
             .filter(move |link| link.points_to.iter().any(|&it| it == m))
-            .map(|link| {
-                (link.owner, link.name().clone(), link.syntax.clone())
-            })
-            .collect();
-        res
+            .for_each(f)
     }
 
     pub fn problems(
diff --git a/crates/libanalysis/tests/tests.rs b/crates/libanalysis/tests/tests.rs
index 887e13fa3..c098c8e8c 100644
--- a/crates/libanalysis/tests/tests.rs
+++ b/crates/libanalysis/tests/tests.rs
@@ -2,10 +2,13 @@ extern crate libanalysis;
 extern crate relative_path;
 extern crate test_utils;
 
-use std::path::{Path};
+use std::{
+    collections::HashMap,
+    path::{Path},
+};
 
 use relative_path::RelativePath;
-use libanalysis::{AnalysisHost, FileId, FileResolver, JobHandle};
+use libanalysis::{AnalysisHost, FileId, FileResolver, JobHandle, CrateGraph, CrateId};
 use test_utils::assert_eq_dbg;
 
 struct FileMap(&'static [(u32, &'static str)]);
@@ -112,3 +115,34 @@ fn test_resolve_parent_module() {
         &symbols,
     );
 }
+
+#[test]
+fn test_resolve_crate_root() {
+    let mut world = AnalysisHost::new();
+    world.change_file(FileId(1), Some("mod foo;".to_string()));
+    world.change_file(FileId(2), Some("".to_string()));
+
+    let snap = world.analysis(FileMap(&[
+        (1, "/lib.rs"),
+        (2, "/foo.rs"),
+    ]));
+    assert!(snap.crate_root(FileId(2)).is_empty());
+
+    let crate_graph = CrateGraph {
+        crate_roots: {
+            let mut m = HashMap::new();
+            m.insert(CrateId(1), FileId(1));
+            m
+        },
+    };
+    world.set_crate_graph(crate_graph);
+
+    let snap = world.analysis(FileMap(&[
+        (1, "/lib.rs"),
+        (2, "/foo.rs"),
+    ]));
+    assert_eq!(
+        snap.crate_root(FileId(2)),
+        vec![CrateId(1)],
+    );
+}