feat(metaprogramming): Add #[use_callers_scope] (#6050)

# Description

## Problem\*

We have many methods now such as `Quoted::as_type` or
`Quoted::as_trait_constraint` which resolve items in the current scope.
It is useful to change the scope these are resolved in but somewhat
impractical to do what we did with `Expression::resolve` and have them
all take module parameters.

## Summary\*

I've added `#[use_callers_scope]` which can be added to a comptime
function to change the scope any builtins resolve in to that of the
caller. This also works in helper functions and closures.

## Additional Context



## Documentation\*

Check one:
- [ ] No documentation needed.
- [x] 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/comptime.rs

@@ -102,6 +102,9 @@ impl<'context> Elaborator<'context> {
         elaborator.function_context.push(FunctionContext::default());
         elaborator.scopes.start_function();
 
+        elaborator.local_module = self.local_module;
+        elaborator.file = self.file;
+
         setup(&mut elaborator);
 
         elaborator.populate_scope_from_comptime_scopes();
@@ -316,7 +319,7 @@ impl<'context> Elaborator<'context> {
         // If the function is varargs, push the type of the last slice element N times
         // to account for N extra arguments.
         let modifiers = interpreter.elaborator.interner.function_modifiers(&function);
-        let is_varargs = modifiers.attributes.is_varargs();
+        let is_varargs = modifiers.attributes.has_varargs();
         let varargs_type = if is_varargs { parameters.pop() } else { None };
 
         let varargs_elem_type = varargs_type.as_ref().and_then(|t| t.slice_element_type());

compiler/noirc_frontend/src/elaborator/scope.rs

@@ -38,11 +38,20 @@ impl<'context> Elaborator<'context> {
         })
     }
 
-    pub(super) fn module_id(&self) -> ModuleId {
+    pub fn module_id(&self) -> ModuleId {
         assert_ne!(self.local_module, LocalModuleId::dummy_id(), "local_module is unset");
         ModuleId { krate: self.crate_id, local_id: self.local_module }
     }
 
+    pub fn replace_module(&mut self, new_module: ModuleId) -> ModuleId {
+        assert_ne!(new_module.local_id, LocalModuleId::dummy_id(), "local_module is unset");
+        let current_module = self.module_id();
+
+        self.crate_id = new_module.krate;
+        self.local_module = new_module.local_id;
+        current_module
+    }
+
     pub(super) fn resolve_path_or_error(
         &mut self,
         path: Path,

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

@@ -133,9 +133,14 @@ impl<'local, 'interner> Interpreter<'local, 'interner> {
             return self.call_special(function, arguments, return_type, location);
         }
 
-        // Wait until after call_special to set the current function so that builtin functions like
-        // `.as_type()` still call the resolver in the caller's scope.
-        let old_function = self.current_function.replace(function);
+        // Don't change the current function scope if we're in a #[use_callers_scope] function.
+        // This will affect where `Expression::resolve`, `Quoted::as_type`, and similar functions resolve.
+        let mut old_function = self.current_function;
+        let modifiers = self.elaborator.interner.function_modifiers(&function);
+        if !modifiers.attributes.has_use_callers_scope() {
+            self.current_function = Some(function);
+        }
+
         let result = self.call_user_defined_function(function, arguments, location);
         self.current_function = old_function;
         result
@@ -242,6 +247,26 @@ impl<'local, 'interner> Interpreter<'local, 'interner> {
     }
 
     fn call_closure(
+        &mut self,
+        closure: HirLambda,
+        environment: Vec<Value>,
+        arguments: Vec<(Value, Location)>,
+        function_scope: Option<FuncId>,
+        module_scope: ModuleId,
+        call_location: Location,
+    ) -> IResult<Value> {
+        // Set the closure's scope to that of the function it was originally evaluated in
+        let old_module = self.elaborator.replace_module(module_scope);
+        let old_function = std::mem::replace(&mut self.current_function, function_scope);
+
+        let result = self.call_closure_inner(closure, environment, arguments, call_location);
+
+        self.current_function = old_function;
+        self.elaborator.replace_module(old_module);
+        result
+    }
+
+    fn call_closure_inner(
         &mut self,
         closure: HirLambda,
         environment: Vec<Value>,
@@ -1276,7 +1301,9 @@ impl<'local, 'interner> Interpreter<'local, 'interner> {
                 }
                 Ok(result)
             }
-            Value::Closure(closure, env, _) => self.call_closure(closure, env, arguments, location),
+            Value::Closure(closure, env, _, function_scope, module_scope) => {
+                self.call_closure(closure, env, arguments, function_scope, module_scope, location)
+            }
             value => {
                 let typ = value.get_type().into_owned();
                 Err(InterpreterError::NonFunctionCalled { typ, location })
@@ -1458,7 +1485,8 @@ impl<'local, 'interner> Interpreter<'local, 'interner> {
             try_vecmap(&lambda.captures, |capture| self.lookup_id(capture.ident.id, location))?;
 
         let typ = self.elaborator.interner.id_type(id).follow_bindings();
-        Ok(Value::Closure(lambda, environment, typ))
+        let module = self.elaborator.module_id();
+        Ok(Value::Closure(lambda, environment, typ, self.current_function, module))
     }
 
     fn evaluate_quote(&mut self, mut tokens: Tokens, expr_id: ExprId) -> IResult<Value> {

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

@@ -51,7 +51,11 @@ pub enum Value {
     FormatString(Rc<String>, Type),
     CtString(Rc<String>),
     Function(FuncId, Type, Rc<TypeBindings>),
-    Closure(HirLambda, Vec<Value>, Type),
+
+    // Closures also store their original scope (function & module)
+    // in case they use functions such as `Quoted::as_type` which require them.
+    Closure(HirLambda, Vec<Value>, Type, Option<FuncId>, ModuleId),
+
     Tuple(Vec<Value>),
     Struct(HashMap<Rc<String>, Value>, Type),
     Pointer(Shared<Value>, /* auto_deref */ bool),
@@ -125,7 +129,7 @@ impl Value {
             }
             Value::FormatString(_, typ) => return Cow::Borrowed(typ),
             Value::Function(_, typ, _) => return Cow::Borrowed(typ),
-            Value::Closure(_, _, typ) => return Cow::Borrowed(typ),
+            Value::Closure(_, _, typ, ..) => return Cow::Borrowed(typ),
             Value::Tuple(fields) => {
                 Type::Tuple(vecmap(fields, |field| field.get_type().into_owned()))
             }
@@ -221,11 +225,6 @@ impl Value {
                 interner.store_instantiation_bindings(expr_id, unwrap_rc(bindings));
                 ExpressionKind::Resolved(expr_id)
             }
-            Value::Closure(_lambda, _env, _typ) => {
-                // TODO: How should a closure's environment be inlined?
-                let item = "Returning closures from a comptime fn".into();
-                return Err(InterpreterError::Unimplemented { item, location });
-            }
             Value::Tuple(fields) => {
                 let fields = try_vecmap(fields, |field| field.into_expression(interner, location))?;
                 ExpressionKind::Tuple(fields)
@@ -293,6 +292,7 @@ impl Value {
             | Value::Zeroed(_)
             | Value::Type(_)
             | Value::UnresolvedType(_)
+            | Value::Closure(..)
             | Value::ModuleDefinition(_) => {
                 let typ = self.get_type().into_owned();
                 let value = self.display(interner).to_string();
@@ -370,11 +370,6 @@ impl Value {
                 interner.store_instantiation_bindings(expr_id, unwrap_rc(bindings));
                 return Ok(expr_id);
             }
-            Value::Closure(_lambda, _env, _typ) => {
-                // TODO: How should a closure's environment be inlined?
-                let item = "Returning closures from a comptime fn".into();
-                return Err(InterpreterError::Unimplemented { item, location });
-            }
             Value::Tuple(fields) => {
                 let fields =
                     try_vecmap(fields, |field| field.into_hir_expression(interner, location))?;
@@ -427,6 +422,7 @@ impl Value {
             | Value::Zeroed(_)
             | Value::Type(_)
             | Value::UnresolvedType(_)
+            | Value::Closure(..)
             | Value::ModuleDefinition(_) => {
                 let typ = self.get_type().into_owned();
                 let value = self.display(interner).to_string();
@@ -601,7 +597,7 @@ impl<'value, 'interner> Display for ValuePrinter<'value, 'interner> {
             Value::CtString(value) => write!(f, "{value}"),
             Value::FormatString(value, _) => write!(f, "{value}"),
             Value::Function(..) => write!(f, "(function)"),
-            Value::Closure(_, _, _) => write!(f, "(closure)"),
+            Value::Closure(..) => write!(f, "(closure)"),
             Value::Tuple(fields) => {
                 let fields = vecmap(fields, |field| field.display(self.interner).to_string());
                 write!(f, "({})", fields.join(", "))

compiler/noirc_frontend/src/lexer/token.rs

@@ -698,9 +698,13 @@ impl Attributes {
         self.function.as_ref().map_or(false, |func_attribute| func_attribute.is_no_predicates())
     }
 
-    pub fn is_varargs(&self) -> bool {
+    pub fn has_varargs(&self) -> bool {
         self.secondary.iter().any(|attr| matches!(attr, SecondaryAttribute::Varargs))
     }
+
+    pub fn has_use_callers_scope(&self) -> bool {
+        self.secondary.iter().any(|attr| matches!(attr, SecondaryAttribute::UseCallersScope))
+    }
 }
 
 /// An Attribute can be either a Primary Attribute or a Secondary Attribute
@@ -799,6 +803,7 @@ impl Attribute {
                 ))
             }
             ["varargs"] => Attribute::Secondary(SecondaryAttribute::Varargs),
+            ["use_callers_scope"] => Attribute::Secondary(SecondaryAttribute::UseCallersScope),
             tokens => {
                 tokens.iter().try_for_each(|token| validate(token))?;
                 Attribute::Secondary(SecondaryAttribute::Custom(CustomAttribute {
@@ -915,6 +920,11 @@ pub enum SecondaryAttribute {
 
     /// A variable-argument comptime function.
     Varargs,
+
+    /// Treat any metaprogramming functions within this one as resolving
+    /// within the scope of the calling function/module rather than this one.
+    /// This affects functions such as `Expression::resolve` or `Quoted::as_type`.
+    UseCallersScope,
 }
 
 impl SecondaryAttribute {
@@ -937,6 +947,7 @@ impl SecondaryAttribute {
             SecondaryAttribute::Custom(custom) => custom.name(),
             SecondaryAttribute::Abi(_) => Some("abi".to_string()),
             SecondaryAttribute::Varargs => Some("varargs".to_string()),
+            SecondaryAttribute::UseCallersScope => Some("use_callers_scope".to_string()),
         }
     }
 }
@@ -954,6 +965,7 @@ impl fmt::Display for SecondaryAttribute {
             SecondaryAttribute::Field(ref k) => write!(f, "#[field({k})]"),
             SecondaryAttribute::Abi(ref k) => write!(f, "#[abi({k})]"),
             SecondaryAttribute::Varargs => write!(f, "#[varargs]"),
+            SecondaryAttribute::UseCallersScope => write!(f, "#[use_callers_scope]"),
         }
     }
 }
@@ -1003,6 +1015,7 @@ impl AsRef<str> for SecondaryAttribute {
             SecondaryAttribute::ContractLibraryMethod => "",
             SecondaryAttribute::Export => "",
             SecondaryAttribute::Varargs => "",
+            SecondaryAttribute::UseCallersScope => "",
         }
     }
 }

docs/docs/noir/concepts/comptime.md

@@ -238,6 +238,17 @@ The following is an incomplete list of some `comptime` types along with some use
 There are many more functions available by exploring the `std::meta` module and its submodules.
 Using these methods is the key to writing powerful metaprogramming libraries.
 
+## `#[use_callers_scope]`
+
+Since certain functions such as `Quoted::as_type`, `Expression::as_type`, or `Quoted::as_trait_constraint` will attempt
+to resolve their contents in a particular scope - it can be useful to change the scope they resolve in. By default
+these functions will resolve in the current function's scope which is usually the attribute function they are called in.
+If you're working on a library however, this may be a completely different module or crate to the item you're trying to
+use the attribute on. If you want to be able to use `Quoted::as_type` to refer to types local to the caller's scope for
+example, you can annotate your attribute function with `#[use_callers_scope]`. This will ensure your attribute, and any
+closures it uses, can refer to anything in the caller's scope. `#[use_callers_scope]` also works recursively. So if both
+your attribute function and a helper function it calls use it, then they can both refer to the same original caller.
+
 ---
 
 # Example: Derive

test_programs/compile_success_empty/use_callers_scope/Nargo.toml

@@ -0,0 +1,7 @@
+[package]
+name = "use_callers_scope"
+type = "bin"
+authors = [""]
+compiler_version = ">=0.34.0"
+
+[dependencies]

test_programs/compile_success_empty/use_callers_scope/src/main.nr

@@ -0,0 +1,34 @@
+#[bar::struct_attr]
+struct Foo {}
+
+struct Bar {}
+
+#[bar::fn_attr]
+fn main() {}
+
+mod bar {
+    #[use_callers_scope]
+    pub comptime fn struct_attr(_: StructDefinition) {
+        let _ = quote { Bar }.as_type();
+    }
+
+    #[use_callers_scope]
+    pub comptime fn fn_attr(_: FunctionDefinition) {
+        let _ = quote { Bar }.as_type();
+        let _ = nested();
+
+        // Ensure closures can still access Bar even
+        // though `map` separates them from `fn_attr`.
+        let _ = &[1, 2, 3].map(
+            |_| {
+            quote { Bar }.as_type()
+        }
+        );
+    }
+
+    // use_callers_scope should also work nested
+    #[use_callers_scope]
+    comptime fn nested() -> Type {
+        quote { Bar }.as_type()
+    }
+}