feat: add StructDefinition::add_attribute and has_named_attribute

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

@@ -9,8 +9,8 @@ use builtin_helpers::{
     check_one_argument, check_three_arguments, check_two_arguments, get_expr, get_field,
     get_format_string, get_function_def, get_module, get_quoted, get_slice, get_struct,
     get_trait_constraint, get_trait_def, get_trait_impl, get_tuple, get_type, get_typed_expr,
-    get_u32, get_unresolved_type, hir_pattern_to_tokens, mutate_func_meta_type, parse,
-    replace_func_meta_parameters, replace_func_meta_return_type,
+    get_u32, get_unresolved_type, has_named_attribute, hir_pattern_to_tokens,
+    mutate_func_meta_type, parse, replace_func_meta_parameters, replace_func_meta_return_type,
 };
 use chumsky::{chain::Chain, prelude::choice, Parser};
 use im::Vector;
@@ -25,7 +25,6 @@ use crate::{
         FunctionReturnType, IntegerBitSize, LValue, Literal, Statement, StatementKind, UnaryOp,
         UnresolvedType, UnresolvedTypeData, Visibility,
     },
-    elaborator::Elaborator,
     hir::comptime::{
         errors::IResult,
         value::{ExprValue, TypedExpr},
@@ -132,9 +131,13 @@ impl<'local, 'context> Interpreter<'local, 'context> {
             "slice_push_front" => slice_push_front(interner, arguments, location),
             "slice_remove" => slice_remove(interner, arguments, location, call_stack),
             "str_as_bytes" => str_as_bytes(interner, arguments, location),
+            "struct_def_add_attribute" => struct_def_add_attribute(self, arguments, location),
             "struct_def_as_type" => struct_def_as_type(interner, arguments, location),
             "struct_def_fields" => struct_def_fields(interner, arguments, location),
             "struct_def_generics" => struct_def_generics(interner, arguments, location),
+            "struct_def_has_named_attribute" => {
+                struct_def_has_named_attribute(interner, arguments, location)
+            }
             "struct_def_set_fields" => struct_def_set_fields(interner, arguments, location),
             "to_le_radix" => to_le_radix(arguments, return_type, location),
             "trait_constraint_eq" => trait_constraint_eq(interner, arguments, location),
@@ -265,6 +268,50 @@ fn str_as_bytes(
     Ok(Value::Array(bytes, byte_array_type))
 }
 
+// fn add_attribute<let N: u32>(self, attribute: str<N>)
+fn struct_def_add_attribute(
+    interpreter: &mut Interpreter,
+    arguments: Vec<(Value, Location)>,
+    location: Location,
+) -> IResult<Value> {
+    let (self_argument, attribute) = check_two_arguments(arguments, location)?;
+    let attribute_location = attribute.1;
+    let attribute = get_str(interpreter.elaborator.interner, attribute)?;
+
+    let mut tokens = Lexer::lex(&format!("#[{}]", attribute)).0 .0;
+    if let Some(Token::EOF) = tokens.last().map(|token| token.token()) {
+        tokens.pop();
+    }
+    if tokens.len() != 1 {
+        return Err(InterpreterError::InvalidAttribute {
+            attribute: attribute.to_string(),
+            location: attribute_location,
+        });
+    }
+
+    let token = tokens.into_iter().next().unwrap().into_token();
+    let Token::Attribute(attribute) = token else {
+        return Err(InterpreterError::InvalidAttribute {
+            attribute: attribute.to_string(),
+            location: attribute_location,
+        });
+    };
+
+    let Attribute::Secondary(attribute) = attribute else {
+        return Err(InterpreterError::InvalidAttribute {
+            attribute: attribute.to_string(),
+            location: attribute_location,
+        });
+    };
+
+    let struct_id = get_struct(self_argument)?;
+    interpreter.elaborator.interner.update_struct_attributes(struct_id, |attributes| {
+        attributes.push(attribute.clone());
+    });
+
+    Ok(Value::Unit)
+}
+
 /// fn as_type(self) -> Type
 fn struct_def_as_type(
     interner: &NodeInterner,
@@ -302,6 +349,25 @@ fn struct_def_generics(
     Ok(Value::Slice(generics.collect(), typ))
 }
 
+// fn has_named_attribute(self, name: Quoted) -> bool
+fn struct_def_has_named_attribute(
+    interner: &NodeInterner,
+    arguments: Vec<(Value, Location)>,
+    location: Location,
+) -> IResult<Value> {
+    let (self_argument, name) = check_two_arguments(arguments, location)?;
+    let struct_id = get_struct(self_argument)?;
+
+    let name = get_quoted(name)?;
+    let name = name.iter().map(|token| token.to_string()).collect::<Vec<_>>().join("");
+
+    let attributes = interner.struct_attributes(&struct_id);
+    let attributes = attributes.iter().filter_map(|attribute| attribute.as_custom());
+    let attributes = attributes.map(|attribute| &attribute.contents);
+
+    Ok(Value::Bool(has_named_attribute(&name, attributes, location)))
+}
+
 /// fn fields(self) -> [(Quoted, Type)]
 /// Returns (name, type) pairs of each field of this StructDefinition
 fn struct_def_fields(
@@ -1688,7 +1754,7 @@ fn function_def_add_attribute(
         });
     }
 
-    let token = tokens[0].token();
+    let token = tokens.into_iter().next().unwrap().into_token();
     let Token::Attribute(attribute) = token else {
         return Err(InterpreterError::InvalidAttribute {
             attribute: attribute.to_string(),
@@ -1701,7 +1767,7 @@ fn function_def_add_attribute(
 
     let function_modifiers = interpreter.elaborator.interner.function_modifiers_mut(&func_id);
 
-    match attribute {
+    match &attribute {
         Attribute::Function(attribute) => {
             function_modifiers.attributes.function = Some(attribute.clone());
         }
@@ -1712,7 +1778,7 @@ fn function_def_add_attribute(
 
     if let Attribute::Secondary(SecondaryAttribute::Custom(attribute)) = attribute {
         let func_meta = interpreter.elaborator.interner.function_meta_mut(&func_id);
-        func_meta.custom_attributes.push(attribute.clone());
+        func_meta.custom_attributes.push(attribute);
     }
 
     Ok(Value::Unit)
@@ -1742,31 +1808,15 @@ fn function_def_has_named_attribute(
 ) -> 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 = get_quoted(name)?;
     let name = name.iter().map(|token| token.to_string()).collect::<Vec<_>>().join("");
 
-    for attribute in attributes {
-        let parse_result = Elaborator::parse_attribute(&attribute.contents, location);
-        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));
-        }
-    }
+    let attributes = &func_meta.custom_attributes;
+    let attributes = attributes.iter().map(|attribute| &attribute.contents);
 
-    Ok(Value::Bool(false))
+    Ok(Value::Bool(has_named_attribute(&name, attributes, location)))
 }
 
 // fn name(self) -> Quoted
@@ -1968,27 +2018,13 @@ fn module_has_named_attribute(
     let (self_argument, name) = check_two_arguments(arguments, location)?;
     let module_id = get_module(self_argument)?;
     let module_data = interpreter.elaborator.get_module(module_id);
-    let name = get_quoted(name)?;
 
+    let name = get_quoted(name)?;
     let name = name.iter().map(|token| token.to_string()).collect::<Vec<_>>().join("");
 
     let attributes = module_data.outer_attributes.iter().chain(&module_data.inner_attributes);
-    for attribute in attributes {
-        let parse_result = Elaborator::parse_attribute(attribute, location);
-        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))
+    Ok(Value::Bool(has_named_attribute(&name, attributes, location)))
 }
 
 // fn is_contract(self) -> bool

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

@@ -8,6 +8,7 @@ use crate::{
         BlockExpression, ExpressionKind, IntegerBitSize, LValue, Signedness, StatementKind,
         UnresolvedTypeData,
     },
+    elaborator::Elaborator,
     hir::{
         comptime::{
             errors::IResult,
@@ -437,3 +438,26 @@ pub(super) fn block_expression_to_value(block_expr: BlockExpression) -> Value {
 
     Value::Slice(statements, typ)
 }
+
+pub(super) fn has_named_attribute<'a>(
+    name: &'a str,
+    attributes: impl Iterator<Item = &'a String>,
+    location: Location,
+) -> bool {
+    for attribute in attributes {
+        let parse_result = Elaborator::parse_attribute(attribute, location);
+        let Ok(Some((function, _arguments))) = parse_result else {
+            continue;
+        };
+
+        let ExpressionKind::Variable(path) = function.kind else {
+            continue;
+        };
+
+        if path.last_name() == name {
+            return true;
+        }
+    }
+
+    false
+}

docs/docs/noir/standard_library/meta/struct_def.md

@@ -7,6 +7,12 @@ This type corresponds to `struct Name { field1: Type1, ... }` items in the sourc
 
 ## Methods
 
+### add_attribute
+
+#include_code add_attribute noir_stdlib/src/meta/struct_def.nr rust
+
+Adds an attribute to the struct.
+
 ### as_type
 
 #include_code as_type noir_stdlib/src/meta/struct_def.nr rust
@@ -44,6 +50,12 @@ comptime fn example(foo: StructDefinition) {
 
 Returns each field of this struct as a pair of (field name, field type).
 
+### has_named_attribute
+
+#include_code has_named_attribute noir_stdlib/src/meta/struct_def.nr rust
+
+Returns true if this struct has a custom attribute with the given name.
+
 ### set_fields
 
 #include_code set_fields noir_stdlib/src/meta/struct_def.nr rust

noir_stdlib/src/meta/struct_def.nr

@@ -1,11 +1,21 @@
 impl StructDefinition {
+    #[builtin(struct_def_add_attribute)]
+    // docs:start:add_attribute
+    fn add_attribute<let N: u32>(self, attribute: str<N>) {}
+    // docs:end:add_attribute
+
     /// Return a syntactic version of this struct definition as a type.
     /// For example, `as_type(quote { type Foo<A, B> { ... } })` would return `Foo<A, B>`
     #[builtin(struct_def_as_type)]
 // docs:start:as_type
     fn as_type(self) -> Type {}
     // docs:end:as_type
 
+    #[builtin(struct_def_has_named_attribute)]
+    // docs:start:has_named_attribute
+    fn has_named_attribute(self, name: Quoted) -> bool {}
+    // docs:end:has_named_attribute
+
     /// Return each generic on this struct.
     #[builtin(struct_def_generics)]
 // docs:start:generics

test_programs/compile_success_empty/attributes_struct/src/main.nr

@@ -6,3 +6,16 @@ struct SomeStruct {
 }
 
 fn main() {}
+
+// Check that add_attribute and has_named_attribute work well
+
+#[add_attribute]
+struct Foo {
+
+}
+
+fn add_attribute(s: StructDefinition) {
+    assert(!s.has_named_attribute(quote { foo }));
+    s.add_attribute("foo");
+    assert(s.has_named_attribute(quote { foo }));
+}