feat: add FunctionDef::has_named_attribute

compiler/noirc_frontend/src/elaborator/comptime.rs

@@ -168,7 +168,7 @@ impl<'context> Elaborator<'context> {
     /// Parses an attribute in the form of a function call (e.g. `#[foo(a b, c d)]`) into
     /// the function and quoted arguments called (e.g. `("foo", vec![(a b, location), (c d, location)])`)
     #[allow(clippy::type_complexity)]
-    fn parse_attribute(
+    pub(crate) fn parse_attribute(
         annotation: &str,
         file: FileId,
     ) -> Result<Option<(Expression, Vec<Expression>)>, (CompilationError, FileId)> {

compiler/noirc_frontend/src/elaborator/mod.rs

@@ -816,6 +816,11 @@ impl<'context> Elaborator<'context> {
             None
         };
 
+        let attributes = func.secondary_attributes().iter();
+        let attributes =
+            attributes.filter_map(|secondary_attribute| secondary_attribute.as_custom());
+        let attributes = attributes.map(|str| str.to_string()).collect();
+
         let meta = FuncMeta {
             name: name_ident,
             kind: func.kind,
@@ -839,6 +844,7 @@ impl<'context> Elaborator<'context> {
             function_body: FunctionBody::Unresolved(func.kind, body, func.def.span),
             self_type: self.self_type.clone(),
             source_file: self.file,
+            custom_attributes: attributes,
         };
 
         self.interner.push_fn_meta(meta, func_id);

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

@@ -25,6 +25,7 @@ use crate::{
         FunctionReturnType, IntegerBitSize, LValue, Literal, Statement, StatementKind, UnaryOp,
         UnresolvedType, UnresolvedTypeData, Visibility,
     },
+    elaborator::Elaborator,
     hir::comptime::{
         errors::IResult,
         value::{ExprValue, TypedExpr},
@@ -94,6 +95,9 @@ impl<'local, 'context> Interpreter<'local, 'context> {
             "expr_resolve" => expr_resolve(self, arguments, location),
             "is_unconstrained" => Ok(Value::Bool(true)),
             "function_def_body" => function_def_body(interner, arguments, location),
+            "function_def_has_named_attribute" => {
+                function_def_has_named_attribute(interner, arguments, location)
+            }
             "function_def_name" => function_def_name(interner, arguments, location),
             "function_def_parameters" => function_def_parameters(interner, arguments, location),
             "function_def_return_type" => function_def_return_type(interner, arguments, location),
@@ -1412,36 +1416,53 @@ where
     option(return_type, option_value)
 }
 
-// fn resolve(self) -> TypedExpr
+// fn resolve(self, in_function: Option<FunctionDefinition>) -> TypedExpr
 fn expr_resolve(
     interpreter: &mut Interpreter,
     arguments: Vec<(Value, Location)>,
     location: Location,
 ) -> IResult<Value> {
-    let self_argument = check_one_argument(arguments, location)?;
+    let (self_argument, func) = check_two_arguments(arguments, location)?;
     let self_argument_location = self_argument.1;
     let expr_value = get_expr(interpreter.elaborator.interner, self_argument)?;
     let expr_value = unwrap_expr_value(interpreter.elaborator.interner, expr_value);
 
-    let value =
-        interpreter.elaborate_item(interpreter.current_function, |elaborator| match expr_value {
-            ExprValue::Expression(expression_kind) => {
-                let expr = Expression { kind: expression_kind, span: self_argument_location.span };
-                let (expr_id, _) = elaborator.elaborate_expression(expr);
-                Value::TypedExpr(TypedExpr::ExprId(expr_id))
-            }
-            ExprValue::Statement(statement_kind) => {
-                let statement =
-                    Statement { kind: statement_kind, span: self_argument_location.span };
-                let (stmt_id, _) = elaborator.elaborate_statement(statement);
-                Value::TypedExpr(TypedExpr::StmtId(stmt_id))
-            }
-            ExprValue::LValue(lvalue) => {
-                let expr = lvalue.as_expression();
-                let (expr_id, _) = elaborator.elaborate_expression(expr);
-                Value::TypedExpr(TypedExpr::ExprId(expr_id))
-            }
-        });
+    let Value::Struct(fields, _) = func.0 else {
+        panic!("Expected second argument to be a struct");
+    };
+
+    let is_some = fields.get(&Rc::new("_is_some".to_string())).unwrap();
+    let Value::Bool(is_some) = is_some else {
+        panic!("Expected is_some to be a boolean");
+    };
+
+    let function_to_resolve_in = if *is_some {
+        let value = fields.get(&Rc::new("_value".to_string())).unwrap();
+        let Value::FunctionDefinition(func_id) = value else {
+            panic!("Expected option value to be a FunctionDefinition");
+        };
+        Some(*func_id)
+    } else {
+        interpreter.current_function
+    };
+
+    let value = interpreter.elaborate_item(function_to_resolve_in, |elaborator| match expr_value {
+        ExprValue::Expression(expression_kind) => {
+            let expr = Expression { kind: expression_kind, span: self_argument_location.span };
+            let (expr_id, _) = elaborator.elaborate_expression(expr);
+            Value::TypedExpr(TypedExpr::ExprId(expr_id))
+        }
+        ExprValue::Statement(statement_kind) => {
+            let statement = Statement { kind: statement_kind, span: self_argument_location.span };
+            let (stmt_id, _) = elaborator.elaborate_statement(statement);
+            Value::TypedExpr(TypedExpr::StmtId(stmt_id))
+        }
+        ExprValue::LValue(lvalue) => {
+            let expr = lvalue.as_expression();
+            let (expr_id, _) = elaborator.elaborate_expression(expr);
+            Value::TypedExpr(TypedExpr::ExprId(expr_id))
+        }
+    });
 
     Ok(value)
 }
@@ -1487,6 +1508,41 @@ fn function_def_body(
     }
 }
 
+// fn has_named_attribute(self, name: Quoted) -> bool
+fn function_def_has_named_attribute(
+    interner: &NodeInterner,
+    arguments: Vec<(Value, Location)>,
+    location: Location,
+) -> IResult<Value> {
+    let (self_argument, name) = check_two_arguments(arguments, location)?;
+    let func_id = get_function_def(self_argument)?;
+    let name = get_quoted(name)?;
+    let func_meta = interner.function_meta(&func_id);
+    let attributes = &func_meta.custom_attributes;
+    if attributes.is_empty() {
+        return Ok(Value::Bool(false));
+    };
+
+    let name = name.iter().map(|token| token.to_string()).collect::<Vec<_>>().join("");
+
+    for attribute in attributes {
+        let parse_result = Elaborator::parse_attribute(attribute, location.file);
+        let Ok(Some((function, _arguments))) = parse_result else {
+            continue;
+        };
+
+        let ExpressionKind::Variable(path) = function.kind else {
+            continue;
+        };
+
+        if path.last_name() == name {
+            return Ok(Value::Bool(true));
+        }
+    }
+
+    Ok(Value::Bool(false))
+}
+
 // fn name(self) -> Quoted
 fn function_def_name(
     interner: &NodeInterner,

compiler/noirc_frontend/src/hir_def/function.rs

@@ -164,6 +164,9 @@ pub struct FuncMeta {
     /// If this function is from an impl (trait or regular impl), this
     /// is the object type of the impl. Otherwise this is None.
     pub self_type: Option<Type>,
+
+    /// Custom attributes attached to this function.
+    pub custom_attributes: Vec<String>,
 }
 
 #[derive(Debug, Clone)]

compiler/noirc_frontend/src/lexer/token.rs

@@ -870,6 +870,16 @@ pub enum SecondaryAttribute {
     Varargs,
 }
 
+impl SecondaryAttribute {
+    pub(crate) fn as_custom(&self) -> Option<&str> {
+        if let Self::Custom(str) = self {
+            Some(str)
+        } else {
+            None
+        }
+    }
+}
+
 impl fmt::Display for SecondaryAttribute {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         match self {

docs/docs/noir/standard_library/meta/expr.md

@@ -190,4 +190,12 @@ Returns this expression as a `Quoted` value. It's the same as `quote { $self }`.
 
 #include_code resolve noir_stdlib/src/meta/expr.nr rust
 
-Resolves and type-checks this expression and returns the result as a `TypedExpr`. If any names used by this expression are not in scope or if there are any type errors, this will give compiler errors as if the expression was written directly into the current `comptime` function.
+Resolves and type-checks this expression and returns the result as a `TypedExpr`. 
+
+The `in_function` argument specifies where the expression is resolved:
+- If it's `none`, the expression is resolved in the function where `resolve` was called
+- If it's `some`, the expression is resolved in the given function
+
+If any names used by this expression are not in scope or if there are any type errors, 
+this will give compiler errors as if the expression was written directly into 
+the current `comptime` function.

docs/docs/noir/standard_library/meta/function_def.md

@@ -15,6 +15,12 @@ Returns the body of the function as an expression. This is only valid
 on functions in the current crate which have not yet been resolved.
 This means any functions called at compile-time are invalid targets for this method.
 
+### has_named_attribute
+
+#include_code has_named_attribute noir_stdlib/src/meta/function_def.nr rust
+
+Returns true if this function has a custom attribute with the given name.
+
 ### name
 
 #include_code name noir_stdlib/src/meta/function_def.nr rust

noir_stdlib/src/meta/expr.nr

@@ -154,7 +154,7 @@ impl Expr {
 
     #[builtin(expr_resolve)]
     // docs:start:resolve
-    fn resolve(self) -> TypedExpr {}
+    fn resolve(self, in_function: Option<FunctionDefinition>) -> TypedExpr {}
     // docs:end:resolve
 }
 

noir_stdlib/src/meta/function_def.nr

@@ -4,6 +4,11 @@ impl FunctionDefinition {
     fn body(self) -> Expr {}
     // docs:end:body
 
+    #[builtin(function_def_has_named_attribute)]
+    // docs:start:has_named_attribute
+    fn has_named_attribute(self, name: Quoted) -> bool {}
+    // docs:end:has_named_attribute
+
     #[builtin(function_def_name)]
     // docs:start:name
     fn name(self) -> Quoted {}

test_programs/compile_success_empty/comptime_function_definition/src/main.nr

@@ -32,6 +32,8 @@ comptime fn function_attr(f: FunctionDefinition) {
 
     // Check FunctionDefinition::name
     assert_eq(f.name(), quote { foo });
+
+    assert(f.has_named_attribute(quote { function_attr }));
 }
 
 #[mutate_add_one]

test_programs/compile_success_empty/inject_context_attribute/src/main.nr

@@ -6,7 +6,7 @@ struct Context {
 fn foo(x: Field) {
     if true {
         // 20 + 1 => 21
-        bar(qux(x + 1));
+        bar(qux(x + 1) + zero());
     } else {
         assert(false);
     }
@@ -24,30 +24,40 @@ fn qux(x: Field) -> Field {
     x * 2
 }
 
+fn zero() -> Field {
+    0
+}
+
 fn inject_context(f: FunctionDefinition) {
     // Add a `_context: Context` parameter to the function
     let parameters = f.parameters();
     let parameters = parameters.push_front((quote { _context }, quote { Context }.as_type()));
     f.set_parameters(parameters);
 
     // Create a new body where every function call has `_context` added to the list of arguments.
-    let body = f.body().modify(mapping_function);
+    let body = f.body().modify(|expr| mapping_function(expr, f));
     f.set_body(body);
 }
 
-fn mapping_function(expr: Expr) -> Option<Expr> {
-    expr.as_function_call().map(
+fn mapping_function(expr: Expr, f: FunctionDefinition) -> Option<Expr> {
+    expr.as_function_call().and_then(
         |func_call: (Expr, [Expr])| {
         let (name, arguments) = func_call;
-        let arguments = arguments.push_front(quote { _context }.as_expr().unwrap());
-        let arguments = arguments.map(|arg: Expr| arg.quoted()).join(quote { , });
-        quote { $name($arguments) }.as_expr().unwrap()
-    }
+        name.resolve(Option::some(f)).as_function_definition().and_then(
+            |function_definition: FunctionDefinition| {
+            if function_definition.has_named_attribute(quote { inject_context }) {
+                let arguments = arguments.push_front(quote { _context }.as_expr().unwrap());
+                let arguments = arguments.map(|arg: Expr| arg.quoted()).join(quote { , });
+                Option::some(quote { $name($arguments) }.as_expr().unwrap())
+            } else {
+                Option::none()
+            }
+        }
+        )}
     )
 }
 
 fn main() {
     let context = Context { value: 42 };
     foo(context, 20);
 }
-

test_programs/noir_test_success/comptime_expr/src/main.nr

@@ -582,7 +582,7 @@ mod tests {
         comptime
         {
             let expr = quote { times_two }.as_expr().unwrap();
-            let func = expr.resolve().as_function_definition().unwrap();
+            let func = expr.resolve(Option::none()).as_function_definition().unwrap();
             assert_eq(func.name(), quote { times_two });
             assert_eq(func.parameters().len(), 1);
         }