chore: encapsulate Index within LocalModuleId

compiler/noirc_arena/src/lib.rs

@@ -9,22 +9,11 @@ use std::fmt;
 pub struct Index(usize);
 
 impl Index {
-    #[cfg(test)]
-    pub fn test_new(index: usize) -> Index {
-        Self(index)
-    }
-
     /// Return a dummy index (max value internally).
     /// This should be avoided over `Option<Index>` if possible.
     pub fn dummy() -> Self {
         Self(usize::MAX)
     }
-
-    /// Return the zeroed index. This is unsafe since we don't know
-    /// if this is a valid index for any particular map yet.
-    pub fn unsafe_zeroed() -> Self {
-        Self(0)
-    }
 }
 
 impl fmt::Display for Index {

compiler/noirc_frontend/src/elaborator/mod.rs

@@ -1544,12 +1544,12 @@ impl<'context> Elaborator<'context> {
 
     pub fn get_module(&self, module: ModuleId) -> &ModuleData {
         let message = "A crate should always be present for a given crate id";
-        &self.def_maps.get(&module.krate).expect(message).modules[module.local_id.0]
+        &self.def_maps.get(&module.krate).expect(message)[module.local_id]
     }
 
     fn get_module_mut(def_maps: &mut DefMaps, module: ModuleId) -> &mut ModuleData {
         let message = "A crate should always be present for a given crate id";
-        &mut def_maps.get_mut(&module.krate).expect(message).modules[module.local_id.0]
+        &mut def_maps.get_mut(&module.krate).expect(message)[module.local_id]
     }
 
     fn declare_methods_on_struct(

compiler/noirc_frontend/src/hir/comptime/tests.rs

@@ -1,10 +1,9 @@
 #![cfg(test)]
 
-use std::collections::{BTreeMap, HashMap};
+use std::collections::HashMap;
 use std::path::PathBuf;
 
 use fm::{FileId, FileManager};
-use noirc_arena::Index;
 use noirc_errors::Location;
 
 use super::Interpreter;
@@ -13,7 +12,7 @@ use super::value::Value;
 use crate::elaborator::{Elaborator, ElaboratorOptions};
 use crate::hir::def_collector::dc_crate::{CompilationError, DefCollector};
 use crate::hir::def_collector::dc_mod::collect_defs;
-use crate::hir::def_map::{CrateDefMap, LocalModuleId, ModuleData};
+use crate::hir::def_map::{CrateDefMap, ModuleData};
 use crate::hir::{Context, ParsedFiles};
 use crate::node_interner::FuncId;
 use crate::parse_program;
@@ -27,19 +26,15 @@ pub(crate) fn with_interpreter<T>(
 ) -> T {
     let file = FileId::default();
 
-    // Can't use Index::test_new here for some reason, even with #[cfg(test)].
-    let module_id = LocalModuleId(Index::unsafe_zeroed());
-    let mut modules = noirc_arena::Arena::default();
     let location = Location::new(Default::default(), file);
-    let root = LocalModuleId(modules.insert(ModuleData::new(
+    let root_module = ModuleData::new(
         None,
         location,
         Vec::new(),
         Vec::new(),
         false, // is contract
         false, // is struct
-    )));
-    assert_eq!(root, module_id);
+    );
 
     let file_manager = FileManager::new(&PathBuf::new());
     let parsed_files = ParsedFiles::new();
@@ -52,10 +47,11 @@ pub(crate) fn with_interpreter<T>(
     assert_eq!(errors.len(), 0);
     let ast = module.into_sorted();
 
-    let def_map = CrateDefMap { root: module_id, modules, krate, extern_prelude: BTreeMap::new() };
+    let def_map = CrateDefMap::new(krate, root_module);
+    let root_module_id = def_map.root();
     let mut collector = DefCollector::new(def_map);
 
-    collect_defs(&mut collector, ast, FileId::dummy(), module_id, krate, &mut context);
+    collect_defs(&mut collector, ast, FileId::dummy(), root_module_id, krate, &mut context);
     context.def_maps.insert(krate, collector.def_map);
 
     let main = context.get_main_function(&krate).expect("Expected 'main' function");

compiler/noirc_frontend/src/hir/def_collector/dc_crate.rs

@@ -35,6 +35,7 @@ use iter_extended::vecmap;
 use rustc_hash::FxHashMap as HashMap;
 use std::collections::BTreeMap;
 use std::fmt::Write;
+use std::ops::IndexMut;
 use std::path::PathBuf;
 use std::vec;
 
@@ -313,7 +314,7 @@ impl DefCollector {
         options: FrontendOptions,
     ) -> Vec<CompilationError> {
         let mut errors: Vec<CompilationError> = vec![];
-        let crate_id = def_map.krate;
+        let crate_id = def_map.krate();
 
         // Recursively resolve the dependencies
         //
@@ -328,7 +329,7 @@ impl DefCollector {
             let dep_def_map =
                 context.def_map(&dep.crate_id).expect("ice: def map was just created");
 
-            let dep_def_root = dep_def_map.root;
+            let dep_def_root = dep_def_map.root();
             let module_id = ModuleId { krate: dep.crate_id, local_id: dep_def_root };
             // Add this crate as a dependency by linking it's root module
             def_map.extern_prelude.insert(dep.as_name(), module_id);
@@ -346,7 +347,7 @@ impl DefCollector {
         // At this point, all dependencies are resolved and type checked.
         //
         // It is now possible to collect all of the definitions of this crate.
-        let crate_root = def_map.root;
+        let crate_root = def_map.root();
         let mut def_collector = DefCollector::new(def_map);
 
         let module_id = ModuleId { krate: crate_id, local_id: crate_root };
@@ -367,13 +368,14 @@ impl DefCollector {
             context,
         ));
 
-        let submodules = vecmap(def_collector.def_map.modules().iter(), |(index, _)| index);
+        let submodules =
+            vecmap(def_collector.def_map.modules().iter(), |(index, _)| LocalModuleId::new(index));
         // Add the current crate to the collection of DefMaps
         context.def_maps.insert(crate_id, def_collector.def_map);
 
         inject_prelude(crate_id, context, crate_root, &mut def_collector.imports);
         for submodule in submodules {
-            inject_prelude(crate_id, context, LocalModuleId(submodule), &mut def_collector.imports);
+            inject_prelude(crate_id, context, submodule, &mut def_collector.imports);
         }
 
         // Resolve unresolved imports collected from the crate, one by one.
@@ -417,7 +419,7 @@ impl DefCollector {
                         }
                         let visibility = visibility.min(item_visibility);
 
-                        let result = current_def_map.modules[local_module_id.0].import(
+                        let result = current_def_map.index_mut(local_module_id).import(
                             name.clone(),
                             visibility,
                             module_def_id,

compiler/noirc_frontend/src/hir/def_collector/dc_mod.rs

@@ -383,7 +383,7 @@ impl ModCollector<'_> {
             context.def_interner.set_doc_comments(ReferenceId::Alias(type_alias_id), doc_comments);
 
             // Add the type alias to scope so its path can be looked up later
-            let result = self.def_collector.def_map.modules[self.module_id.0].declare_type_alias(
+            let result = self.def_collector.def_map[self.module_id].declare_type_alias(
                 name.clone(),
                 visibility,
                 type_alias_id,
@@ -461,7 +461,7 @@ impl ModCollector<'_> {
 
             // Add the trait to scope so its path can be looked up later
             let visibility = trait_definition.visibility;
-            let result = self.def_collector.def_map.modules[self.module_id.0].declare_trait(
+            let result = self.def_collector.def_map[self.module_id].declare_trait(
                 name.clone(),
                 visibility,
                 trait_id,
@@ -542,9 +542,11 @@ impl ModCollector<'_> {
                             );
                         }
 
-                        match self.def_collector.def_map.modules[trait_id.0.local_id.0]
-                            .declare_function(name.clone(), ItemVisibility::Public, func_id)
-                        {
+                        match self.def_collector.def_map[trait_id.0.local_id].declare_function(
+                            name.clone(),
+                            ItemVisibility::Public,
+                            func_id,
+                        ) {
                             Ok(()) => {
                                 if let Some(body) = body {
                                     let impl_method =
@@ -585,8 +587,8 @@ impl ModCollector<'_> {
                             false,
                         );
 
-                        if let Err((first_def, second_def)) = self.def_collector.def_map.modules
-                            [trait_id.0.local_id.0]
+                        if let Err((first_def, second_def)) = self.def_collector.def_map
+                            [trait_id.0.local_id]
                             .declare_global(name.clone(), ItemVisibility::Public, global_id)
                         {
                             let error = DefCollectorErrorKind::Duplicate {
@@ -610,9 +612,8 @@ impl ModCollector<'_> {
                         }
                     }
                     TraitItem::Type { name } => {
-                        if let Err((first_def, second_def)) = self.def_collector.def_map.modules
-                            [trait_id.0.local_id.0]
-                            .declare_type_alias(
+                        if let Err((first_def, second_def)) =
+                            self.def_collector.def_map[trait_id.0.local_id].declare_type_alias(
                                 name.clone(),
                                 ItemVisibility::Public,
                                 TypeAliasId::dummy_id(),
@@ -914,13 +915,14 @@ fn push_child_module(
         is_type,
     );
 
-    let module_id = def_map.modules.insert(new_module);
-    let modules = &mut def_map.modules;
+    let krate = def_map.krate();
+    let module_id = def_map.insert_module(new_module);
+    let parent_module = &mut def_map[parent];
 
     // Update the parent module to reference the child
-    modules[parent.0].children.insert(mod_name.clone(), LocalModuleId(module_id));
+    parent_module.children.insert(mod_name.clone(), module_id);
 
-    let mod_id = ModuleId { krate: def_map.krate, local_id: LocalModuleId(module_id) };
+    let mod_id = ModuleId { krate, local_id: module_id };
 
     // Add this child module into the scope of the parent module as a module definition
     // module definitions are definitions which can only exist at the module level.
@@ -931,7 +933,7 @@ fn push_child_module(
     // the struct name.
     if add_to_parent_scope {
         if let Err((first_def, second_def)) =
-            modules[parent.0].declare_child_module(mod_name.to_owned(), visibility, mod_id)
+            parent_module.declare_child_module(mod_name.to_owned(), visibility, mod_id)
         {
             let err = DefCollectorErrorKind::Duplicate {
                 typ: DuplicateType::Module,
@@ -952,7 +954,7 @@ fn push_child_module(
         );
 
         if interner.is_in_lsp_mode() {
-            let parent_module_id = ModuleId { krate: def_map.krate, local_id: parent };
+            let parent_module_id = ModuleId { krate: def_map.krate(), local_id: parent };
             interner.register_module(
                 mod_id,
                 location,
@@ -982,7 +984,7 @@ pub fn collect_function(
         }
     }
 
-    let module_data = &mut def_map.modules[module.local_id.0];
+    let module_data = &mut def_map[module.local_id];
 
     let is_test = function.def.attributes.is_test_function();
     let is_entry_point_function = if module_data.is_contract {
@@ -1009,7 +1011,7 @@ pub fn collect_function(
     interner.set_doc_comments(ReferenceId::Function(func_id), doc_comments);
 
     // Add function to scope/ns of the module
-    let result = def_map.modules[module.local_id.0].declare_function(name, visibility, func_id);
+    let result = module_data.declare_function(name, visibility, func_id);
     if let Err((first_def, second_def)) = result {
         let error = DefCollectorErrorKind::Duplicate {
             typ: DuplicateType::Function,
@@ -1083,7 +1085,7 @@ pub fn collect_struct(
 
     // Add the struct to scope so its path can be looked up later
     let visibility = unresolved.struct_def.visibility;
-    let result = def_map.modules[module_id.0].declare_type(name.clone(), visibility, id);
+    let result = def_map[module_id].declare_type(name.clone(), visibility, id);
 
     let parent_module_id = ModuleId { krate, local_id: module_id };
 
@@ -1174,7 +1176,7 @@ pub fn collect_enum(
 
     // Add the enum to scope so its path can be looked up later
     let visibility = unresolved.enum_def.visibility;
-    let result = def_map.modules[module_id.0].declare_type(name.clone(), visibility, id);
+    let result = def_map[module_id].declare_type(name.clone(), visibility, id);
 
     let parent_module_id = ModuleId { krate, local_id: module_id };
 
@@ -1400,7 +1402,7 @@ pub(crate) fn collect_global(
     );
 
     // Add the statement to the scope so its path can be looked up later
-    let result = def_map.modules[module_id.0].declare_global(name.clone(), visibility, global_id);
+    let result = def_map[module_id].declare_global(name.clone(), visibility, global_id);
 
     // Globals marked as ABI don't have to be used.
     if !is_abi {