From d0a957ba9bd743ae00959b0680e275a0a3992308 Mon Sep 17 00:00:00 2001
From: jfecher <jake@aztecprotocol.com>
Date: Fri, 2 Aug 2024 09:11:20 -0500
Subject: [PATCH] fix: Elaborate struct & trait annotations in the correct
 module (#5643)

# Description

## Problem\*

## Summary\*

Elaborates trait & struct annotations in the correct module.
We also save & restore type variables before elaborating items from the
interpreter. Otherwise this infers with checks that trait impls do not
overlap.

## Additional Context

This is necessary for `derive` to work in the stdlib. This will be
tested against regressions in the future as part of the PR adding
`derive` in the stdlib.

## 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.
---
 compiler/noirc_frontend/src/elaborator/mod.rs |  4 ++++
 .../src/hir/comptime/interpreter.rs           | 21 +++++++++++++------
 .../src/hir/comptime/interpreter/builtin.rs   |  8 +++----
 3 files changed, 22 insertions(+), 11 deletions(-)

diff --git a/compiler/noirc_frontend/src/elaborator/mod.rs b/compiler/noirc_frontend/src/elaborator/mod.rs
index 490e1b9c3c4..0b6233b3445 100644
--- a/compiler/noirc_frontend/src/elaborator/mod.rs
+++ b/compiler/noirc_frontend/src/elaborator/mod.rs
@@ -1625,6 +1625,8 @@ 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);
         }
 
@@ -1632,6 +1634,8 @@ impl<'context> Elaborator<'context> {
             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);
         }
 
diff --git a/compiler/noirc_frontend/src/hir/comptime/interpreter.rs b/compiler/noirc_frontend/src/hir/comptime/interpreter.rs
index a3d37fd76fc..8f3f1295cac 100644
--- a/compiler/noirc_frontend/src/hir/comptime/interpreter.rs
+++ b/compiler/noirc_frontend/src/hir/comptime/interpreter.rs
@@ -166,7 +166,7 @@ impl<'local, 'interner> Interpreter<'local, 'interner> {
             Some(body) => Ok(body),
             None => {
                 if matches!(&meta.function_body, FunctionBody::Unresolved(..)) {
-                    self.elaborator.elaborate_item_from_comptime(None, |elaborator| {
+                    self.elaborate_item(None, |elaborator| {
                         elaborator.elaborate_function(function);
                     });
 
@@ -179,6 +179,17 @@ impl<'local, 'interner> Interpreter<'local, 'interner> {
         }
     }
 
+    fn elaborate_item<T>(
+        &mut self,
+        function: Option<FuncId>,
+        f: impl FnOnce(&mut Elaborator) -> T,
+    ) -> T {
+        self.unbind_generics_from_previous_function();
+        let result = self.elaborator.elaborate_item_from_comptime(function, f);
+        self.rebind_generics_from_previous_function();
+        result
+    }
+
     fn call_special(
         &mut self,
         function: FuncId,
@@ -1236,11 +1247,9 @@ impl<'local, 'interner> Interpreter<'local, 'interner> {
                 let mut result = self.call_function(function_id, arguments, bindings, location)?;
                 if call.is_macro_call {
                     let expr = result.into_expression(self.elaborator.interner, location)?;
-                    let expr = self
-                        .elaborator
-                        .elaborate_item_from_comptime(self.current_function, |elab| {
-                            elab.elaborate_expression(expr).0
-                        });
+                    let expr = self.elaborate_item(self.current_function, |elaborator| {
+                        elaborator.elaborate_expression(expr).0
+                    });
                     result = self.evaluate(expr)?;
                 }
                 Ok(result)
diff --git a/compiler/noirc_frontend/src/hir/comptime/interpreter/builtin.rs b/compiler/noirc_frontend/src/hir/comptime/interpreter/builtin.rs
index 1904eaa02b9..645e4d707c0 100644
--- a/compiler/noirc_frontend/src/hir/comptime/interpreter/builtin.rs
+++ b/compiler/noirc_frontend/src/hir/comptime/interpreter/builtin.rs
@@ -403,8 +403,7 @@ fn quoted_as_trait_constraint(
     })?;
 
     let bound = interpreter
-        .elaborator
-        .elaborate_item_from_comptime(interpreter.current_function, |elaborator| {
+        .elaborate_item(interpreter.current_function, |elaborator| {
             elaborator.resolve_trait_bound(&trait_bound, Type::Unit)
         })
         .ok_or(InterpreterError::FailedToResolveTraitBound { trait_bound, location })?;
@@ -429,9 +428,8 @@ fn quoted_as_type(
         InterpreterError::FailedToParseMacro { error, tokens, rule, file: location.file }
     })?;
 
-    let typ = interpreter
-        .elaborator
-        .elaborate_item_from_comptime(interpreter.current_function, |elab| elab.resolve_type(typ));
+    let typ =
+        interpreter.elaborate_item(interpreter.current_function, |elab| elab.resolve_type(typ));
 
     Ok(Value::Type(typ))
 }