chore: Rework CrateGraph to only have one root crate (#2391)

crates/nargo/src/lib.rs

@@ -17,7 +17,7 @@ pub mod workspace;
 use std::collections::BTreeMap;
 
 use fm::FileManager;
-use noirc_driver::{add_dep, prepare_crate};
+use noirc_driver::{add_dep, prepare_crate, prepare_dependency};
 use noirc_frontend::{
     graph::{CrateGraph, CrateId, CrateName},
     hir::Context,
@@ -34,7 +34,7 @@ pub fn prepare_dependencies(
     for (dep_name, dep) in dependencies.iter() {
         match dep {
             Dependency::Remote { package } | Dependency::Local { package } => {
-                let crate_id = prepare_crate(context, &package.entry_path);
+                let crate_id = prepare_dependency(context, &package.entry_path);
                 add_dep(context, parent_crate, crate_id, dep_name.clone());
                 prepare_dependencies(context, crate_id, &package.dependencies);
             }

crates/noirc_driver/src/lib.rs

@@ -56,13 +56,20 @@ pub fn compile_file(
     compile_main(context, crate_id, &CompileOptions::default())
 }
 
-/// Adds the file from the file system at `Path` to the crate graph
+/// Adds the file from the file system at `Path` to the crate graph as a root file
 pub fn prepare_crate(context: &mut Context, file_name: &Path) -> CrateId {
     let root_file_id = context.file_manager.add_file(file_name).unwrap();
 
     context.crate_graph.add_crate_root(root_file_id)
 }
 
+// Adds the file from the file system at `Path` to the crate graph
+pub fn prepare_dependency(context: &mut Context, file_name: &Path) -> CrateId {
+    let root_file_id = context.file_manager.add_file(file_name).unwrap();
+
+    context.crate_graph.add_crate(root_file_id)
+}
+
 /// Adds a edge in the crate graph for two crates
 pub fn add_dep(
     context: &mut Context,

crates/noirc_frontend/src/graph/mod.rs

@@ -12,18 +12,24 @@ use smol_str::SmolStr;
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
 pub enum CrateId {
+    Root(usize),
     Crate(usize),
     Stdlib(usize),
+    Dummy,
 }
 
 impl CrateId {
     pub fn dummy_id() -> CrateId {
-        CrateId::Crate(std::usize::MAX)
+        CrateId::Dummy
     }
 
     pub fn is_stdlib(&self) -> bool {
         matches!(self, CrateId::Stdlib(_))
     }
+
+    pub fn is_root(&self) -> bool {
+        matches!(self, CrateId::Root(_))
+    }
 }
 
 #[derive(Debug, Clone, PartialEq, Eq, Hash, Ord, PartialOrd)]
@@ -117,29 +123,73 @@ impl Dependency {
 }
 
 impl CrateGraph {
+    pub fn root_crate_id(&self) -> &CrateId {
+        self.arena
+            .keys()
+            .find(|crate_id| crate_id.is_root())
+            .expect("ICE: A root crate should exist in the CrateGraph")
+    }
+
     pub fn add_crate_root(&mut self, file_id: FileId) -> CrateId {
-        let mut roots_with_file_id =
-            self.arena.iter().filter(|(_, crate_data)| crate_data.root_file_id == file_id);
+        for (crate_id, crate_data) in self.arena.iter() {
+            if crate_id.is_root() {
+                panic!("ICE: Cannot add two crate roots to a graph - use `add_crate` instead");
+            }
 
-        let next_file_id = roots_with_file_id.next();
-        if let Some(file_id) = next_file_id {
-            return *file_id.0;
+            if crate_data.root_file_id == file_id {
+                panic!("ICE: This FileId was already added to the CrateGraph")
+            }
         }
 
         let data = CrateData { root_file_id: file_id, dependencies: Vec::new() };
-        let crate_id = CrateId::Crate(self.arena.len());
+        let crate_id = CrateId::Root(self.arena.len());
         let prev = self.arena.insert(crate_id, data);
         assert!(prev.is_none());
         crate_id
     }
 
+    pub fn add_crate(&mut self, file_id: FileId) -> CrateId {
+        let mut crates_with_file_id = self
+            .arena
+            .iter()
+            .filter(|(_, crate_data)| crate_data.root_file_id == file_id)
+            .peekable();
+
+        let matching_id = crates_with_file_id.next();
+        if crates_with_file_id.peek().is_some() {
+            panic!("ICE: Found multiple crates with the same FileId");
+        }
+
+        match matching_id {
+            Some((crate_id @ CrateId::Crate(_), _)) => *crate_id,
+            Some((CrateId::Root(_), _)) => {
+                panic!("ICE: Tried to re-add the root crate as a regular crate")
+            }
+            Some((CrateId::Stdlib(_), _)) => {
+                panic!("ICE: Tried to re-add the stdlib crate as a regular crate")
+            }
+            Some((CrateId::Dummy, _)) => {
+                panic!("ICE: A dummy CrateId should not exist in the CrateGraph")
+            }
+            None => {
+                let data = CrateData { root_file_id: file_id, dependencies: Vec::new() };
+                let crate_id = CrateId::Crate(self.arena.len());
+                let prev = self.arena.insert(crate_id, data);
+                assert!(prev.is_none());
+                crate_id
+            }
+        }
+    }
+
     pub fn add_stdlib(&mut self, file_id: FileId) -> CrateId {
-        let mut roots_with_file_id =
-            self.arena.iter().filter(|(_, crate_data)| crate_data.root_file_id == file_id);
+        for (crate_id, crate_data) in self.arena.iter() {
+            if crate_id.is_stdlib() {
+                panic!("ICE: Cannot add two stdlib crates to a graph - use `add_crate` instead");
+            }
 
-        let next_file_id = roots_with_file_id.next();
-        if let Some(file_id) = next_file_id {
-            return *file_id.0;
+            if crate_data.root_file_id == file_id {
+                panic!("ICE: This FileId was already added to the CrateGraph")
+            }
         }
 
         let data = CrateData { root_file_id: file_id, dependencies: Vec::new() };
@@ -263,8 +313,8 @@ mod tests {
 
         let mut graph = CrateGraph::default();
         let crate1 = graph.add_crate_root(file_ids[0]);
-        let crate2 = graph.add_crate_root(file_ids[1]);
-        let crate3 = graph.add_crate_root(file_ids[2]);
+        let crate2 = graph.add_crate(file_ids[1]);
+        let crate3 = graph.add_crate(file_ids[2]);
 
         assert!(graph.add_dep(crate1, "crate2".parse().unwrap(), crate2).is_ok());
         assert!(graph.add_dep(crate2, "crate3".parse().unwrap(), crate3).is_ok());
@@ -279,8 +329,8 @@ mod tests {
         let file_id_2 = file_ids[2];
         let mut graph = CrateGraph::default();
         let crate1 = graph.add_crate_root(file_id_0);
-        let crate2 = graph.add_crate_root(file_id_1);
-        let crate3 = graph.add_crate_root(file_id_2);
+        let crate2 = graph.add_crate(file_id_1);
+        let crate3 = graph.add_crate(file_id_2);
         assert!(graph.add_dep(crate1, "crate2".parse().unwrap(), crate2).is_ok());
         assert!(graph.add_dep(crate2, "crate3".parse().unwrap(), crate3).is_ok());
     }
@@ -292,11 +342,68 @@ mod tests {
         let file_id_2 = file_ids[2];
         let mut graph = CrateGraph::default();
         let _crate1 = graph.add_crate_root(file_id_0);
-        let _crate2 = graph.add_crate_root(file_id_1);
+        let _crate2 = graph.add_crate(file_id_1);
 
         // Adding the same file, so the crate should be the same.
-        let crate3 = graph.add_crate_root(file_id_2);
-        let crate3_2 = graph.add_crate_root(file_id_2);
+        let crate3 = graph.add_crate(file_id_2);
+        let crate3_2 = graph.add_crate(file_id_2);
         assert_eq!(crate3, crate3_2);
     }
+
+    #[test]
+    #[should_panic = "ICE: Cannot add two crate roots to a graph - use `add_crate` instead"]
+    fn panics_if_adding_two_roots() {
+        let file_ids = dummy_file_ids(2);
+        let mut graph = CrateGraph::default();
+        let _ = graph.add_crate_root(file_ids[0]);
+        let _ = graph.add_crate_root(file_ids[1]);
+    }
+
+    #[test]
+    #[should_panic = "ICE: This FileId was already added to the CrateGraph"]
+    fn panics_if_adding_existing_file_as_root() {
+        let file_ids = dummy_file_ids(1);
+        let mut graph = CrateGraph::default();
+        let file_id_0 = file_ids[0];
+        let _ = graph.add_crate(file_id_0);
+        let _ = graph.add_crate_root(file_id_0);
+    }
+
+    #[test]
+    #[should_panic = "ICE: Cannot add two stdlib crates to a graph - use `add_crate` instead"]
+    fn panics_if_adding_two_stdlib() {
+        let file_ids = dummy_file_ids(2);
+        let mut graph = CrateGraph::default();
+        let _ = graph.add_stdlib(file_ids[0]);
+        let _ = graph.add_stdlib(file_ids[1]);
+    }
+
+    #[test]
+    #[should_panic = "ICE: This FileId was already added to the CrateGraph"]
+    fn panics_if_adding_existing_file_as_stdlib() {
+        let file_ids = dummy_file_ids(1);
+        let mut graph = CrateGraph::default();
+        let file_id_0 = file_ids[0];
+        let _ = graph.add_crate(file_id_0);
+        let _ = graph.add_stdlib(file_id_0);
+    }
+
+    #[test]
+    #[should_panic = "ICE: Tried to re-add the root crate as a regular crate"]
+    fn panics_if_adding_root_as_regular() {
+        let file_ids = dummy_file_ids(1);
+        let mut graph = CrateGraph::default();
+        let file_id_0 = file_ids[0];
+        let _ = graph.add_crate_root(file_id_0);
+        let _ = graph.add_crate(file_id_0);
+    }
+    #[test]
+    #[should_panic = "ICE: Tried to re-add the stdlib crate as a regular crate"]
+    fn panics_if_adding_stdlib_as_regular() {
+        let file_ids = dummy_file_ids(1);
+        let mut graph = CrateGraph::default();
+        let file_id_0 = file_ids[0];
+        let _ = graph.add_stdlib(file_id_0);
+        let _ = graph.add_crate(file_id_0);
+    }
 }

crates/noirc_frontend/src/hir/mod.rs

@@ -59,6 +59,10 @@ impl Context {
         self.crate_graph.iter_keys()
     }
 
+    pub fn root_crate_id(&self) -> &CrateId {
+        self.crate_graph.root_crate_id()
+    }
+
     // TODO: Decide if we actually need `function_name` and `fully_qualified_function_name`
     pub fn function_name(&self, id: &FuncId) -> &str {
         self.def_interner.function_name(id)

crates/wasm/src/compile.rs

@@ -3,8 +3,8 @@ use fm::FileManager;
 use gloo_utils::format::JsValueSerdeExt;
 use log::debug;
 use noirc_driver::{
-    check_crate, compile_contracts, compile_no_check, prepare_crate, propagate_dep, CompileOptions,
-    CompiledContract,
+    check_crate, compile_contracts, compile_no_check, prepare_crate, prepare_dependency,
+    propagate_dep, CompileOptions, CompiledContract,
 };
 use noirc_frontend::{graph::CrateGraph, hir::Context};
 use serde::{Deserialize, Serialize};
@@ -60,7 +60,7 @@ impl Default for WASMCompileOptions {
 
 fn add_noir_lib(context: &mut Context, crate_name: &str) {
     let path_to_lib = Path::new(&crate_name).join("lib.nr");
-    let library_crate = prepare_crate(context, &path_to_lib);
+    let library_crate = prepare_dependency(context, &path_to_lib);
 
     propagate_dep(context, library_crate, &crate_name.parse().unwrap());
 }