Add tests for ModuleTree

crates/common2/src/input.rs

@@ -27,6 +27,7 @@ pub struct InputIngot {
 
     /// A list of files which the current ingot contains.
     #[return_ref]
+    #[set(__set_files_impl)]
     pub files: BTreeSet<InputFile>,
 
     #[set(__set_root_file_impl)]
@@ -35,7 +36,7 @@ pub struct InputIngot {
 }
 impl InputIngot {
     pub fn new(
-        db: &mut dyn InputDb,
+        db: &dyn InputDb,
         path: &str,
         kind: IngotKind,
         version: Version,
@@ -60,6 +61,12 @@ impl InputIngot {
         self.__set_root_file_impl(db).to(Some(file));
     }
 
+    /// Set the list of files which the ingot contains.
+    /// All files must bee set before the ingot is used.
+    pub fn set_files(self, db: &mut dyn InputDb, files: BTreeSet<InputFile>) {
+        self.__set_files_impl(db).to(files);
+    }
+
     /// Returns the root file of the ingot.
     /// Panics if the root file is not set.
     pub fn root_file(&self, db: &dyn InputDb) -> InputFile {
@@ -109,6 +116,14 @@ pub struct IngotDependency {
     /// An ingot which the current ingot depends on.
     pub ingot: InputIngot,
 }
+impl IngotDependency {
+    pub fn new(name: &str, ingot: InputIngot) -> Self {
+        Self {
+            name: SmolStr::new(name),
+            ingot,
+        }
+    }
+}
 
 impl PartialOrd for IngotDependency {
     fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {

crates/hir/src/hir_def/item.rs

@@ -21,8 +21,8 @@ use crate::{
 };
 
 use super::{
-    ingot_module_tree_impl, AttrListId, Body, FnParamListId, GenericParamListId, IdentId,
-    IngotModuleTree, ItemTree, Partial, TypeId, WhereClauseId,
+    module_tree_impl, AttrListId, Body, FnParamListId, GenericParamListId, IdentId, ItemTree,
+    ModuleTree, Partial, TypeId, WhereClauseId,
 };
 
 #[derive(
@@ -71,11 +71,11 @@ impl TopLevelMod {
     }
 
     pub fn module_item_tree(self, db: &dyn HirDb) -> &ItemTree {
-        lower::module_item_tree_impl(db, self)
+        lower::item_tree_impl(db, self)
     }
 
-    pub fn ingot_module_tree(self, db: &dyn HirDb) -> &IngotModuleTree {
-        ingot_module_tree_impl(db, self.ingot(db))
+    pub fn ingot_module_tree(self, db: &dyn HirDb) -> &ModuleTree {
+        module_tree_impl(db, self.ingot(db))
     }
 }
 

crates/hir/src/hir_def/item_tree.rs

@@ -65,7 +65,7 @@ mod tests {
             }
         "#;
 
-        let (_, item_tree) = db.parse_source(text);
+        let item_tree = db.parse_source(text);
         let top_mod = item_tree.top_mod;
         assert_eq!(item_tree.dfs().count(), 8);
 

crates/hir/src/hir_def/module_tree.rs

@@ -55,17 +55,17 @@ use super::{IdentId, TopLevelMod};
 /// As a result, `baz` is represented as a "floating" node.
 /// In this case, the tree is actually a forest. But we don't need to care about it.
 #[derive(Debug, Clone, PartialEq, Eq)]
-pub struct IngotModuleTree {
+pub struct ModuleTree {
     pub(crate) root: ModuleTreeNodeId,
     pub(crate) module_tree: PrimaryMap<ModuleTreeNodeId, ModuleTreeNode>,
     pub(crate) mod_map: BTreeMap<TopLevelMod, ModuleTreeNodeId>,
 
     pub(crate) ingot: InputIngot,
 }
 
-impl IngotModuleTree {
+impl ModuleTree {
     /// Returns the tree node data of the given id.
-    pub fn tree_node_data(&self, id: ModuleTreeNodeId) -> &ModuleTreeNode {
+    pub fn node_data(&self, id: ModuleTreeNodeId) -> &ModuleTreeNode {
         &self.module_tree[id]
     }
 
@@ -74,11 +74,20 @@ impl IngotModuleTree {
         self.mod_map[&top_mod]
     }
 
+    /// Returns the tree node data of the given top level module.
+    pub fn tree_node_data(&self, top_mod: TopLevelMod) -> &ModuleTreeNode {
+        &self.module_tree[self.tree_node(top_mod)]
+    }
+
     /// Returns the root of the tree, which corresponds to the ingot root file.
     pub fn root(&self) -> ModuleTreeNodeId {
         self.root
     }
 
+    pub fn root_data(&self) -> &ModuleTreeNode {
+        self.node_data(self.root)
+    }
+
     /// Returns an iterator of all top level modules in this ingot.
     pub fn all_modules(&self) -> impl Iterator<Item = TopLevelMod> + '_ {
         self.mod_map.keys().copied()
@@ -89,8 +98,8 @@ impl IngotModuleTree {
 /// top level modules. This function only depends on an ingot structure and
 /// external ingot dependency, and not depends on file contents.
 #[salsa::tracked(return_ref)]
-pub fn ingot_module_tree_impl(db: &dyn HirDb, ingot: InputIngot) -> IngotModuleTree {
-    IngotModuleTreeBuilder::new(db, ingot).build()
+pub fn module_tree_impl(db: &dyn HirDb, ingot: InputIngot) -> ModuleTree {
+    ModuleTreeBuilder::new(db, ingot).build()
 }
 
 /// A top level module that is one-to-one mapped to a file.
@@ -103,18 +112,18 @@ pub struct ModuleTreeNode {
     /// 2. the module is a "floating" module.
     pub parent: Option<ModuleTreeNodeId>,
     /// A list of child top level module.
-    pub children: BTreeMap<IdentId, Vec<ModuleTreeNodeId>>,
+    pub children: Vec<ModuleTreeNodeId>,
 }
 
 impl ModuleTreeNode {
     fn new(top_mod: TopLevelMod) -> Self {
         Self {
             top_mod,
             parent: None,
-            children: BTreeMap::new(),
+            children: Vec::new(),
         }
     }
-    fn name(&self, db: &dyn HirDb) -> IdentId {
+    pub fn name(&self, db: &dyn HirDb) -> IdentId {
         self.top_mod.name(db)
     }
 }
@@ -124,15 +133,15 @@ impl ModuleTreeNode {
 pub struct ModuleTreeNodeId(u32);
 entity_impl!(ModuleTreeNodeId);
 
-struct IngotModuleTreeBuilder<'db> {
+struct ModuleTreeBuilder<'db> {
     db: &'db dyn HirDb,
     ingot: InputIngot,
     module_tree: PrimaryMap<ModuleTreeNodeId, ModuleTreeNode>,
     mod_map: BTreeMap<TopLevelMod, ModuleTreeNodeId>,
     path_map: BTreeMap<&'db Utf8Path, ModuleTreeNodeId>,
 }
 
-impl<'db> IngotModuleTreeBuilder<'db> {
+impl<'db> ModuleTreeBuilder<'db> {
     fn new(db: &'db dyn HirDb, ingot: InputIngot) -> Self {
         Self {
             db,
@@ -143,13 +152,13 @@ impl<'db> IngotModuleTreeBuilder<'db> {
         }
     }
 
-    fn build(mut self) -> IngotModuleTree {
+    fn build(mut self) -> ModuleTree {
         self.set_modules();
         self.build_tree();
 
-        let top_mod = map_file_to_mod_impl(self.db, self.ingot.root_file(self.db.upcast()));
-        let root = self.mod_map[&top_mod];
-        IngotModuleTree {
+        let root_mod = map_file_to_mod_impl(self.db, self.ingot.root_file(self.db.upcast()));
+        let root = self.mod_map[&root_mod];
+        ModuleTree {
             root,
             module_tree: self.module_tree,
             mod_map: self.mod_map,
@@ -213,13 +222,70 @@ impl<'db> IngotModuleTreeBuilder<'db> {
     }
 
     fn add_branch(&mut self, parent: ModuleTreeNodeId, child: ModuleTreeNodeId) {
-        let child_name = self.module_tree[child].name(self.db);
-        self.module_tree[parent]
-            .children
-            .entry(child_name)
-            .or_default()
-            .push(child);
+        self.module_tree[parent].children.push(child);
 
         self.module_tree[child].parent = Some(parent);
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use common::input::{IngotKind, Version};
+
+    use super::*;
+    use crate::{lower, test_db::TestDb};
+
+    #[test]
+    fn module_tree() {
+        let mut db = TestDb::default();
+
+        let local_ingot = InputIngot::new(
+            &mut db,
+            "/foo/fargo",
+            IngotKind::Local,
+            Version::new(0, 0, 1),
+            [].into(),
+        );
+        let local_root = InputFile::new(&db, local_ingot, "src/lib.fe".into(), "".into());
+        let mod1 = InputFile::new(&db, local_ingot, "src/mod1.fe".into(), "".into());
+        let mod2 = InputFile::new(&db, local_ingot, "src/mod2.fe".into(), "".into());
+        let foo = InputFile::new(&db, local_ingot, "src/mod1/foo.fe".into(), "".into());
+        let bar = InputFile::new(&db, local_ingot, "src/mod2/bar.fe".into(), "".into());
+        let baz = InputFile::new(&db, local_ingot, "src/mod2/baz.fe".into(), "".into());
+        let floating = InputFile::new(&db, local_ingot, "src/mod3/floating.fe".into(), "".into());
+        local_ingot.set_root_file(&mut db, local_root);
+        local_ingot.set_files(
+            &mut db,
+            [local_root, mod1, mod2, foo, bar, baz, floating].into(),
+        );
+
+        let local_root_mod = lower::map_file_to_mod(&db, local_root);
+        let mod1_mod = lower::map_file_to_mod(&db, mod1);
+        let mod2_mod = lower::map_file_to_mod(&db, mod2);
+        let foo_mod = lower::map_file_to_mod(&db, foo);
+        let bar_mod = lower::map_file_to_mod(&db, bar);
+        let baz_mod = lower::map_file_to_mod(&db, baz);
+
+        let local_tree = lower::module_tree(&db, local_ingot);
+        let root_node = local_tree.root_data();
+        assert_eq!(root_node.top_mod, local_root_mod);
+        assert_eq!(root_node.children.len(), 2);
+
+        for &child in &root_node.children {
+            if child == local_tree.tree_node(mod1_mod) {
+                let child = local_tree.node_data(child);
+                assert_eq!(child.parent, Some(local_tree.root()));
+                assert_eq!(child.children.len(), 1);
+                assert_eq!(child.children[0], local_tree.tree_node(foo_mod));
+            } else if child == local_tree.tree_node(mod2_mod) {
+                let child = local_tree.node_data(child);
+                assert_eq!(child.parent, Some(local_tree.root()));
+                assert_eq!(child.children.len(), 2);
+                assert_eq!(child.children[0], local_tree.tree_node(bar_mod));
+                assert_eq!(child.children[1], local_tree.tree_node(baz_mod));
+            } else {
+                panic!("unexpected child")
+            }
+        }
+    }
+}

crates/hir/src/lib.rs

@@ -1,9 +1,9 @@
 use common::{InputDb, Upcast};
-use hir_def::ingot_module_tree_impl;
+use hir_def::module_tree_impl;
 pub use lower::parse::ParseDiagnostic;
 
 use lower::{
-    map_file_to_mod_impl, module_item_tree_impl,
+    item_tree_impl, map_file_to_mod_impl,
     parse::{parse_file_impl, ParseDiagnosticAccumulator},
 };
 
@@ -49,8 +49,8 @@ pub struct Jar(
     /// Private tracked functions. These are not part of the public API, and
     /// thus, can't be accessed from outside of the crate without implementing
     /// [`LowerHirDb`] marker trait.
-    ingot_module_tree_impl,
-    module_item_tree_impl,
+    module_tree_impl,
+    item_tree_impl,
     map_file_to_mod_impl,
     parse_file_impl,
 );
@@ -89,7 +89,7 @@ mod test_db {
 
     use crate::{
         hir_def::{ItemKind, ItemTree, TopLevelMod},
-        lower::{map_file_to_mod, module_item_tree},
+        lower::{item_tree, map_file_to_mod},
         span::LazySpan,
         LowerHirDb, SpannedHirDb,
     };
@@ -125,23 +125,24 @@ mod test_db {
     }
 
     impl TestDb {
-        pub fn parse_source(&mut self, text: &str) -> (TopLevelMod, &ItemTree) {
+        pub fn parse_source(&mut self, text: &str) -> &ItemTree {
             let file = self.standalone_file(text);
             let top_mod = map_file_to_mod(self, file);
-            (top_mod, module_item_tree(self, top_mod))
+            item_tree(self, top_mod)
         }
 
         /// Parses the given source text and returns the first inner item in the
         /// file.
-        pub fn parse_source_to_first_item<T>(&mut self, text: &str) -> (TopLevelMod, T)
+        pub fn parse_source_to_first_item<T>(&mut self, text: &str) -> T
         where
             ItemKind: TryInto<T, Error = &'static str>,
         {
-            let (top_mod, tree) = self.parse_source(text);
-            (
-                top_mod,
-                tree.children(top_mod).next().unwrap().try_into().unwrap(),
-            )
+            let tree = self.parse_source(text);
+            tree.children(tree.top_mod)
+                .next()
+                .unwrap()
+                .try_into()
+                .unwrap()
         }
 
         pub fn text_at(&self, top_mod: TopLevelMod, span: &impl LazySpan) -> &str {
@@ -158,7 +159,7 @@ mod test_db {
             let ingot = InputIngot::new(self, path, kind, version, BTreeSet::default());
             let file = InputFile::new(self, ingot, "test_file.fe".into(), text.to_string());
             ingot.set_root_file(self, file);
-            ingot.set_files(self).to([file].into());
+            ingot.set_files(self, [file].into());
             file
         }
     }

crates/hir/src/lower/mod.rs

@@ -1,6 +1,6 @@
 use std::collections::{BTreeMap, BTreeSet};
 
-use common::InputFile;
+use common::{InputFile, InputIngot};
 use num_bigint::BigUint;
 use num_traits::Num;
 use parser::{
@@ -10,8 +10,8 @@ use parser::{
 
 use crate::{
     hir_def::{
-        IdentId, IntegerId, ItemKind, ItemTree, ItemTreeNode, LitKind, Partial, StringId,
-        TopLevelMod, TrackedItemId,
+        module_tree_impl, IdentId, IntegerId, ItemKind, ItemTree, ItemTreeNode, LitKind,
+        ModuleTree, Partial, StringId, TopLevelMod, TrackedItemId,
     },
     HirDb, LowerHirDb, ParseDiagnostic,
 };
@@ -43,8 +43,8 @@ pub fn map_file_to_mod(db: &dyn LowerHirDb, file: InputFile) -> TopLevelMod {
 }
 
 /// Returns the item tree of the given top-level module.
-pub fn module_item_tree(db: &dyn LowerHirDb, top_mod: TopLevelMod) -> &ItemTree {
-    module_item_tree_impl(db.upcast(), top_mod)
+pub fn item_tree(db: &dyn LowerHirDb, top_mod: TopLevelMod) -> &ItemTree {
+    item_tree_impl(db.upcast(), top_mod)
 }
 
 /// Returns the root node of the given top-level module.
@@ -65,6 +65,11 @@ pub fn parse_file(db: &dyn LowerHirDb, top_mod: TopLevelMod) -> GreenNode {
     parse_file_impl(db.upcast(), top_mod)
 }
 
+/// Returns the ingot module tree of the given ingot.
+pub fn module_tree(db: &dyn LowerHirDb, ingot: InputIngot) -> &ModuleTree {
+    module_tree_impl(db.upcast(), ingot)
+}
+
 #[salsa::tracked]
 pub(crate) fn map_file_to_mod_impl(db: &dyn HirDb, file: InputFile) -> TopLevelMod {
     let path = file.path(db.upcast());
@@ -75,7 +80,7 @@ pub(crate) fn map_file_to_mod_impl(db: &dyn HirDb, file: InputFile) -> TopLevelM
 }
 
 #[salsa::tracked(return_ref)]
-pub(crate) fn module_item_tree_impl(db: &dyn HirDb, top_mod: TopLevelMod) -> ItemTree {
+pub(crate) fn item_tree_impl(db: &dyn HirDb, top_mod: TopLevelMod) -> ItemTree {
     let ast = top_mod_ast(db, top_mod);
     let mut ctxt = FileLowerCtxt::new(db, top_mod);