feat: LSP auto-import completion

compiler/noirc_frontend/src/ast/expression.rs

@@ -8,7 +8,7 @@ use crate::ast::{
 use crate::hir::def_collector::errors::DefCollectorErrorKind;
 use crate::macros_api::StructId;
 use crate::node_interner::{ExprId, QuotedTypeId};
-use crate::token::{Attributes, Token, Tokens};
+use crate::token::{Attributes, FunctionAttribute, Token, Tokens};
 use crate::{Kind, Type};
 use acvm::{acir::AcirField, FieldElement};
 use iter_extended::vecmap;
@@ -478,6 +478,20 @@ pub struct FunctionDefinition {
     pub return_visibility: Visibility,
 }
 
+impl FunctionDefinition {
+    pub fn is_private(&self) -> bool {
+        self.visibility == ItemVisibility::Private
+    }
+
+    pub fn is_test(&self) -> bool {
+        if let Some(attribute) = &self.attributes.function {
+            matches!(attribute, FunctionAttribute::Test(..))
+        } else {
+            false
+        }
+    }
+}
+
 #[derive(Debug, PartialEq, Eq, Clone)]
 pub struct Param {
     pub visibility: Visibility,

compiler/noirc_frontend/src/elaborator/comptime.rs

@@ -267,6 +267,14 @@ impl<'context> Elaborator<'context> {
                 let id = self.interner.push_empty_fn();
                 let module = self.module_id();
                 self.interner.push_function(id, &function.def, module, location);
+
+                if self.interner.is_in_lsp_mode()
+                    && !function.def.is_test()
+                    && !function.def.is_private()
+                {
+                    self.interner.register_function(id, &function.def);
+                }
+
                 let functions = vec![(self.local_module, id, function)];
                 generated_items.functions.push(UnresolvedFunctions {
                     file_id: self.file,

compiler/noirc_frontend/src/elaborator/mod.rs

@@ -1297,6 +1297,12 @@ impl<'context> Elaborator<'context> {
         self.current_item = Some(DependencyId::Global(global_id));
         let let_stmt = global.stmt_def;
 
+        let name = if self.interner.is_in_lsp_mode() {
+            Some(let_stmt.pattern.name_ident().to_string())
+        } else {
+            None
+        };
+
         if !self.in_contract()
             && let_stmt.attributes.iter().any(|attr| matches!(attr, SecondaryAttribute::Abi(_)))
         {
@@ -1319,8 +1325,9 @@ impl<'context> Elaborator<'context> {
             self.elaborate_comptime_global(global_id);
         }
 
-        self.interner
-            .add_definition_location(ReferenceId::Global(global_id), Some(self.module_id()));
+        if let Some(name) = name {
+            self.interner.register_global(global_id, name, self.module_id());
+        }
 
         self.local_module = old_module;
         self.file = old_file;

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

@@ -14,7 +14,7 @@
     TypeImpl,
 };
 use crate::macros_api::NodeInterner;
-use crate::node_interner::{ModuleAttributes, ReferenceId};
+use crate::node_interner::ModuleAttributes;
 use crate::{
     graph::CrateId,
     hir::def_collector::dc_crate::{UnresolvedStruct, UnresolvedTrait},
@@ -232,6 +232,13 @@
             let location = Location::new(function.span(), self.file_id);
             context.def_interner.push_function(func_id, &function.def, module, location);
 
+            if context.def_interner.is_in_lsp_mode()
+                && !function.def.is_test()
+                && !function.def.is_private()
+            {
+                context.def_interner.register_function(func_id, &function.def);
+            }
+
             // Now link this func_id to a crate level map with the noir function and the module id
             // Encountering a NoirFunction, we retrieve it's module_data to get the namespace
             // Once we have lowered it to a HirFunction, we retrieve it's Id from the DefInterner
@@ -307,8 +314,8 @@
             };
 
             // Add the struct to scope so its path can be looked up later
-            let result =
-                self.def_collector.def_map.modules[self.module_id.0].declare_struct(name, id);
+            let result = self.def_collector.def_map.modules[self.module_id.0]
+                .declare_struct(name.clone(), id);
 
             if let Err((first_def, second_def)) = result {
                 let error = DefCollectorErrorKind::Duplicate {
@@ -322,10 +329,10 @@
             // And store the TypeId -> StructType mapping somewhere it is reachable
             self.def_collector.items.types.insert(id, unresolved);
 
-            context.def_interner.add_definition_location(
-                ReferenceId::Struct(id),
-                Some(ModuleId { krate, local_id: self.module_id }),
-            );
+            if context.def_interner.is_in_lsp_mode() {
+                let parent_module_id = ModuleId { krate, local_id: self.module_id };
+                context.def_interner.register_struct(id, name.to_string(), parent_module_id);
+            }
         }
         definition_errors
     }
@@ -383,7 +390,7 @@
 
             // 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(name, type_alias_id);
+                .declare_type_alias(name.clone(), type_alias_id);
 
             if let Err((first_def, second_def)) = result {
                 let err = DefCollectorErrorKind::Duplicate {
@@ -396,10 +403,11 @@
 
             self.def_collector.items.type_aliases.insert(type_alias_id, unresolved);
 
-            context.def_interner.add_definition_location(
-                ReferenceId::Alias(type_alias_id),
-                Some(ModuleId { krate, local_id: self.module_id }),
-            );
+            if context.def_interner.is_in_lsp_mode() {
+                let parent_module_id = ModuleId { krate, local_id: self.module_id };
+                let name = name.to_string();
+                context.def_interner.register_type_alias(type_alias_id, name, parent_module_id);
+            }
         }
         errors
     }
@@ -432,8 +440,8 @@
             };
 
             // Add the trait to scope so its path can be looked up later
-            let result =
-                self.def_collector.def_map.modules[self.module_id.0].declare_trait(name, trait_id);
+            let result = self.def_collector.def_map.modules[self.module_id.0]
+                .declare_trait(name.clone(), trait_id);
 
             if let Err((first_def, second_def)) = result {
                 let error = DefCollectorErrorKind::Duplicate {
@@ -537,7 +545,7 @@
                         }
                     }
                     TraitItem::Type { name } => {
                         // TODO(nickysn or alexvitkov): implement context.def_interner.push_empty_type_alias and get an id, instead of using TypeAliasId::dummy_id()
                         if let Err((first_def, second_def)) = self.def_collector.def_map.modules
                             [trait_id.0.local_id.0]
                             .declare_type_alias(name.clone(), TypeAliasId::dummy_id())
@@ -567,10 +575,10 @@
             };
             context.def_interner.push_empty_trait(trait_id, &unresolved, resolved_generics);
 
-            context.def_interner.add_definition_location(
-                ReferenceId::Trait(trait_id),
-                Some(ModuleId { krate, local_id: self.module_id }),
-            );
+            if context.def_interner.is_in_lsp_mode() {
+                let parent_module_id = ModuleId { krate, local_id: self.module_id };
+                context.def_interner.register_trait(trait_id, name.to_string(), parent_module_id);
+            }
 
             self.def_collector.items.traits.insert(trait_id, unresolved);
         }
@@ -723,7 +731,7 @@
         //   if it's an inline module, or the first char of a the file if it's an external module.
         // - `location` will always point to the token "foo" in `mod foo` regardless of whether
         //   it's inline or external.
         // Eventually the location put in `ModuleData` is used for codelenses about `contract`s,
         // so we keep using `location` so that it continues to work as usual.
         let location = Location::new(mod_name.span(), mod_location.file);
         let new_module = ModuleData::new(parent, location, is_contract);
@@ -766,6 +774,10 @@
                     parent: self.module_id,
                 },
             );
+
+            if context.def_interner.is_in_lsp_mode() {
+                context.def_interner.register_module(mod_id, mod_name.0.contents.clone());
+            }
         }
 
         Ok(mod_id)

compiler/noirc_frontend/src/hir/def_map/module_def.rs

@@ -3,7 +3,7 @@ use crate::node_interner::{FuncId, GlobalId, StructId, TraitId, TypeAliasId};
 use super::ModuleId;
 
 /// A generic ID that references either a module, function, type, interface or global
-#[derive(Debug, Copy, Clone, PartialEq, Eq)]
+#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
 pub enum ModuleDefId {
     ModuleId(ModuleId),
     FunctionId(FuncId),

compiler/noirc_frontend/src/locations.rs

@@ -1,9 +1,10 @@
 use fm::FileId;
 use noirc_errors::Location;
 use rangemap::RangeMap;
-use rustc_hash::FxHashMap;
+use rustc_hash::FxHashMap as HashMap;
 
 use crate::{
+    ast::{FunctionDefinition, ItemVisibility},
     hir::def_map::{ModuleDefId, ModuleId},
     macros_api::{NodeInterner, StructId},
     node_interner::{DefinitionId, FuncId, GlobalId, ReferenceId, TraitId, TypeAliasId},
@@ -12,7 +13,7 @@ use petgraph::prelude::NodeIndex as PetGraphIndex;
 
 #[derive(Debug, Default)]
 pub(crate) struct LocationIndices {
-    map_file_to_range: FxHashMap<FileId, RangeMap<u32, PetGraphIndex>>,
+    map_file_to_range: HashMap<FileId, RangeMap<u32, PetGraphIndex>>,
 }
 
 impl LocationIndices {
@@ -275,4 +276,75 @@ impl NodeInterner {
             .neighbors_directed(reference_index, petgraph::Direction::Outgoing)
             .next()
     }
+
+    pub(crate) fn register_module(&mut self, id: ModuleId, name: String) {
+        self.register_name_for_auto_import(name, ModuleDefId::ModuleId(id), ItemVisibility::Public);
+    }
+
+    pub(crate) fn register_global(
+        &mut self,
+        id: GlobalId,
+        name: String,
+        parent_module_id: ModuleId,
+    ) {
+        self.add_definition_location(ReferenceId::Global(id), Some(parent_module_id));
+
+        let visibility = ItemVisibility::Public;
+        self.register_name_for_auto_import(name, ModuleDefId::GlobalId(id), visibility);
+    }
+
+    pub(crate) fn register_struct(
+        &mut self,
+        id: StructId,
+        name: String,
+        parent_module_id: ModuleId,
+    ) {
+        self.add_definition_location(ReferenceId::Struct(id), Some(parent_module_id));
+
+        let visibility = ItemVisibility::Public;
+        self.register_name_for_auto_import(name, ModuleDefId::TypeId(id), visibility);
+    }
+
+    pub(crate) fn register_trait(&mut self, id: TraitId, name: String, parent_module_id: ModuleId) {
+        self.add_definition_location(ReferenceId::Trait(id), Some(parent_module_id));
+
+        self.register_name_for_auto_import(name, ModuleDefId::TraitId(id), ItemVisibility::Public);
+    }
+
+    pub(crate) fn register_type_alias(
+        &mut self,
+        id: TypeAliasId,
+        name: String,
+        parent_module_id: ModuleId,
+    ) {
+        self.add_definition_location(ReferenceId::Alias(id), Some(parent_module_id));
+
+        let visibility = ItemVisibility::Public;
+        self.register_name_for_auto_import(name, ModuleDefId::TypeAliasId(id), visibility);
+    }
+
+    pub(crate) fn register_function(&mut self, id: FuncId, func_def: &FunctionDefinition) {
+        self.register_name_for_auto_import(
+            func_def.name.0.contents.clone(),
+            ModuleDefId::FunctionId(id),
+            func_def.visibility,
+        );
+    }
+
+    fn register_name_for_auto_import(
+        &mut self,
+        name: String,
+        module_def_id: ModuleDefId,
+        visibility: ItemVisibility,
+    ) {
+        if !self.lsp_mode {
+            return;
+        }
+
+        self.auto_import_names.entry(name).or_default().push((module_def_id, visibility));
+    }
+
+    pub fn get_auto_import_names(&self) -> &HashMap<String, Vec<(ModuleDefId, ItemVisibility)>> {
+        &self.auto_import_names
+    }
 }

compiler/noirc_frontend/src/node_interner.rs

@@ -19,6 +19,7 @@
 use crate::hir::def_collector::dc_crate::CompilationError;
 use crate::hir::def_collector::dc_crate::{UnresolvedStruct, UnresolvedTrait, UnresolvedTypeAlias};
 use crate::hir::def_map::{LocalModuleId, ModuleId};
+use crate::macros_api::ModuleDefId;
 use crate::macros_api::UnaryOp;
 use crate::QuotedType;
 
@@ -194,13 +195,13 @@
     /// Stores the [Location] of a [Type] reference
     pub(crate) type_ref_locations: Vec<(Type, Location)>,
 
     /// In Noir's metaprogramming, a noir type has the type `Type`. When these are spliced
     /// into `quoted` expressions, we preserve the original type by assigning it a unique id
     /// and creating a `Token::QuotedType(id)` from this id. We cannot create a token holding
     /// the actual type since types do not implement Send or Sync.
     quoted_types: noirc_arena::Arena<Type>,
 
     /// Determins whether to run in LSP mode. In LSP mode references are tracked.
     pub(crate) lsp_mode: bool,
 
     /// Store the location of the references in the graph.
@@ -231,6 +232,11 @@
     // (ReferenceId::Reference and ReferenceId::Local aren't included here)
     pub(crate) reference_modules: HashMap<ReferenceId, ModuleId>,
 
+    // All names (and their definitions) that can be offered for auto_import.
+    // These include top-level functions, global variables and types, but excludes
+    // impl and trait-impl methods.
+    pub(crate) auto_import_names: HashMap<String, Vec<(ModuleDefId, ItemVisibility)>>,
+
     /// Each value currently in scope in the comptime interpreter.
     /// Each element of the Vec represents a scope with every scope together making
     /// up all currently visible definitions. The first scope is always the global scope.
@@ -602,6 +608,7 @@
             reference_graph: petgraph::graph::DiGraph::new(),
             reference_graph_indices: HashMap::default(),
             reference_modules: HashMap::default(),
+            auto_import_names: HashMap::default(),
             comptime_scopes: vec![HashMap::default()],
         }
     }
@@ -910,6 +917,7 @@
         // This needs to be done after pushing the definition since it will reference the
         // location that was stored
         self.add_definition_location(ReferenceId::Function(id), Some(module));
+
         definition_id
     }