fix: always place module attribute generated items inside module (#5943)

# Description ## Problem Resolves #5875 (comment) ## Summary ## Additional Context ## Documentation Check one: - [x] No documentation needed. - [ ] Documentation included in this PR. - [ ] **[For Experimental Features]** Documentation to be submitted in a separate PR. # PR Checklist - [x] I have tested the changes locally. - [x] I have formatted the changes with [Prettier](https://prettier.io/) and/or `cargo fmt` on default settings. --------- Co-authored-by: jfecher <jake@aztecprotocol.com> Co-authored-by: Jake Fecher <jfecher11@gmail.com>
noir-lang · Sep 5, 2024 · 89ac6e0 · 89ac6e0
1 parent a297ec6
commit 89ac6e0
Show file tree

Hide file tree

Showing 3 changed files with 79 additions and 21 deletions.
diff --git a/compiler/noirc_frontend/src/elaborator/comptime.rs b/compiler/noirc_frontend/src/elaborator/comptime.rs
@@ -15,7 +15,7 @@ use crate::{
             },
             dc_mod,
         },
-        def_map::ModuleId,
+        def_map::{LocalModuleId, ModuleId},
         resolution::errors::ResolverError,
     },
     hir_def::expr::HirIdent,
@@ -30,6 +30,24 @@ use crate::{
 
 use super::{Elaborator, FunctionContext, ResolverMeta};
 
+#[derive(Debug, Copy, Clone)]
+struct AttributeContext {
+    // The file where generated items should be added
+    file: FileId,
+    // The module where generated items should be added
+    module: LocalModuleId,
+    // The file where the attribute is located
+    attribute_file: FileId,
+    // The module where the attribute is located
+    attribute_module: LocalModuleId,
+}
+
+impl AttributeContext {
+    fn new(file: FileId, module: LocalModuleId) -> Self {
+        Self { file, module, attribute_file: file, attribute_module: module }
+    }
+}
+
 impl<'context> Elaborator<'context> {
     /// Elaborate an expression from the middle of a comptime scope.
     /// When this happens we require additional information to know
@@ -89,15 +107,22 @@ impl<'context> Elaborator<'context> {
         }
     }
 
-    pub(super) fn run_comptime_attributes_on_item(
+    fn run_comptime_attributes_on_item(
         &mut self,
         attributes: &[SecondaryAttribute],
         item: Value,
         span: Span,
+        attribute_context: AttributeContext,
         generated_items: &mut CollectedItems,
     ) {
         for attribute in attributes {
-            self.run_comptime_attribute_on_item(attribute, &item, span, generated_items);
+            self.run_comptime_attribute_on_item(
+                attribute,
+                &item,
+                span,
+                attribute_context,
+                generated_items,
+            );
         }
     }
 
@@ -106,6 +131,7 @@ impl<'context> Elaborator<'context> {
         attribute: &SecondaryAttribute,
         item: &Value,
         span: Span,
+        attribute_context: AttributeContext,
         generated_items: &mut CollectedItems,
     ) {
         if let SecondaryAttribute::Custom(attribute) = attribute {
@@ -114,6 +140,7 @@ impl<'context> Elaborator<'context> {
                 item.clone(),
                 span,
                 attribute.contents_span,
+                attribute_context,
                 generated_items,
             ) {
                 self.errors.push(error);
@@ -127,8 +154,12 @@ impl<'context> Elaborator<'context> {
         item: Value,
         span: Span,
         attribute_span: Span,
+        attribute_context: AttributeContext,
         generated_items: &mut CollectedItems,
     ) -> Result<(), (CompilationError, FileId)> {
+        self.file = attribute_context.attribute_file;
+        self.local_module = attribute_context.attribute_module;
+
         let location = Location::new(attribute_span, self.file);
         let Some((function, arguments)) = Self::parse_attribute(attribute, location)? else {
             // Do not issue an error if the attribute is unknown
@@ -156,6 +187,9 @@ impl<'context> Elaborator<'context> {
             return Err((ResolverError::NonFunctionInAnnotation { span }.into(), self.file));
         };
 
+        self.file = attribute_context.file;
+        self.local_module = attribute_context.module;
+
         let mut interpreter = self.setup_interpreter();
         let mut arguments = Self::handle_attribute_arguments(
             &mut interpreter,
@@ -463,18 +497,28 @@ impl<'context> Elaborator<'context> {
             let attributes = &trait_.trait_def.attributes;
             let item = Value::TraitDefinition(*trait_id);
             let span = trait_.trait_def.span;
-            self.local_module = trait_.module_id;
-            self.file = trait_.file_id;
-            self.run_comptime_attributes_on_item(attributes, item, span, &mut generated_items);
+            let context = AttributeContext::new(trait_.file_id, trait_.module_id);
+            self.run_comptime_attributes_on_item(
+                attributes,
+                item,
+                span,
+                context,
+                &mut generated_items,
+            );
         }
 
         for (struct_id, struct_def) in types {
             let attributes = &struct_def.struct_def.attributes;
             let item = Value::StructDefinition(*struct_id);
             let span = struct_def.struct_def.span;
-            self.local_module = struct_def.module_id;
-            self.file = struct_def.file_id;
-            self.run_comptime_attributes_on_item(attributes, item, span, &mut generated_items);
+            let context = AttributeContext::new(struct_def.file_id, struct_def.module_id);
+            self.run_comptime_attributes_on_item(
+                attributes,
+                item,
+                span,
+                context,
+                &mut generated_items,
+            );
         }
 
         self.run_attributes_on_functions(functions, &mut generated_items);
@@ -496,10 +540,14 @@ impl<'context> Elaborator<'context> {
             let attribute = &module_attribute.attribute;
             let span = Span::default();
 
-            self.local_module = module_attribute.attribute_module_id;
-            self.file = module_attribute.attribute_file_id;
+            let context = AttributeContext {
+                file: module_attribute.file_id,
+                module: module_attribute.module_id,
+                attribute_file: module_attribute.attribute_file_id,
+                attribute_module: module_attribute.attribute_module_id,
+            };
 
-            self.run_comptime_attribute_on_item(attribute, &item, span, generated_items);
+            self.run_comptime_attribute_on_item(attribute, &item, span, context, generated_items);
         }
     }
 
@@ -509,15 +557,20 @@ impl<'context> Elaborator<'context> {
         generated_items: &mut CollectedItems,
     ) {
         for function_set in function_sets {
-            self.file = function_set.file_id;
             self.self_type = function_set.self_type.clone();
 
             for (local_module, function_id, function) in &function_set.functions {
-                self.local_module = *local_module;
+                let context = AttributeContext::new(function_set.file_id, *local_module);
                 let attributes = function.secondary_attributes();
                 let item = Value::FunctionDefinition(*function_id);
                 let span = function.span();
-                self.run_comptime_attributes_on_item(attributes, item, span, generated_items);
+                self.run_comptime_attributes_on_item(
+                    attributes,
+                    item,
+                    span,
+                    context,
+                    generated_items,
+                );
             }
         }
     }

diff --git a/test_programs/compile_success_empty/comptime_module/src/main.nr b/test_programs/compile_success_empty/comptime_module/src/main.nr
@@ -1,8 +1,8 @@
 #[outer_attribute]
 mod foo {
     #![some_attribute]
-    fn x() {}
-    fn y() {}
+    pub fn x() {}
+    pub fn y() {}
 }
 
 contract bar {}
@@ -13,7 +13,7 @@ mod another_module {}
 #[outer_attribute_func]
 mod yet_another_module {
     #![super::inner_attribute_func]
-    fn foo() {}
+    pub fn foo() {}
 }
 
 #[outer_attribute_separate_module]
@@ -25,14 +25,16 @@ fn increment_counter() {
     counter += 1;
 }
 
-fn outer_attribute_func(m: Module) {
+fn outer_attribute_func(m: Module) -> Quoted {
     assert_eq(m.name(), quote { yet_another_module });
     increment_counter();
+    quote { pub fn generated_outer_function() {} }
 }
 
-fn inner_attribute_func(m: Module) {
+fn inner_attribute_func(m: Module) -> Quoted {
     assert_eq(m.name(), quote { yet_another_module });
     increment_counter();
+    quote { pub fn generated_inner_function() {} }
 }
 
 fn outer_attribute_separate_module(m: Module) {
@@ -68,6 +70,9 @@ fn main() {
     }
 
     assert_eq(counter, 4);
+
+    yet_another_module::generated_outer_function();
+    yet_another_module::generated_inner_function();
 }
 
 // docs:start:as_module_example

diff --git a/tooling/lsp/src/requests/completion/builtins.rs b/tooling/lsp/src/requests/completion/builtins.rs
@@ -97,6 +97,7 @@ pub(super) fn keyword_builtin_type(keyword: &Keyword) -> Option<&'static str> {
         Keyword::Expr => Some("Expr"),
         Keyword::Field => Some("Field"),
         Keyword::FunctionDefinition => Some("FunctionDefinition"),
+        Keyword::Module => Some("Module"),
         Keyword::Quoted => Some("Quoted"),
         Keyword::StructDefinition => Some("StructDefinition"),
         Keyword::TraitConstraint => Some("TraitConstraint"),
@@ -128,7 +129,6 @@ pub(super) fn keyword_builtin_type(keyword: &Keyword) -> Option<&'static str> {
         | Keyword::In
         | Keyword::Let
         | Keyword::Mod
-        | Keyword::Module
         | Keyword::Mut
         | Keyword::Pub
         | Keyword::Return