From b59a29e5b246121a4d81e4894a4b10f5df4dd5cf Mon Sep 17 00:00:00 2001 From: jfecher Date: Fri, 12 Jul 2024 11:08:24 -0500 Subject: [PATCH] feat: Allow comptime attributes on traits & functions (#5496) # Description ## Problem\* Resolves https://github.com/noir-lang/noir/issues/5477 ## Summary\* Adds: - The ability to run `comptime` attribute functions on traits & functions in the program - The `TraitDefinition` type - The `FunctionDefinition` type - The `Module` type - the only one of the new types which you still can't run attributes on. See: https://github.com/noir-lang/noir/issues/5495. Running these on modules is a bit more difficult since modules don't have an entry in `CollectedItems` to run them on. So I'm delaying this for a later PR. ## Additional Context ## Documentation\* Check one: - [ ] No documentation needed. - [ ] Documentation included in this PR. - [x] **[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/ast/traits.rs | 2 + compiler/noirc_frontend/src/elaborator/mod.rs | 50 ++++++++++++++----- .../noirc_frontend/src/elaborator/traits.rs | 13 ++++- .../noirc_frontend/src/hir/comptime/value.rs | 24 +++++++-- compiler/noirc_frontend/src/hir_def/types.rs | 6 +++ compiler/noirc_frontend/src/lexer/token.rs | 13 ++++- .../src/parser/parser/traits.rs | 18 +++++-- .../noirc_frontend/src/parser/parser/types.rs | 20 ++++++++ .../function_attribute/Nargo.toml | 7 +++ .../function_attribute/src/main.nr | 18 +++++++ .../trait_attribute/Nargo.toml | 7 +++ .../trait_attribute/src/main.nr | 19 +++++++ tooling/nargo_cli/build.rs | 4 +- 13 files changed, 176 insertions(+), 25 deletions(-) create mode 100644 test_programs/compile_success_empty/function_attribute/Nargo.toml create mode 100644 test_programs/compile_success_empty/function_attribute/src/main.nr create mode 100644 test_programs/compile_success_empty/trait_attribute/Nargo.toml create mode 100644 test_programs/compile_success_empty/trait_attribute/src/main.nr diff --git a/compiler/noirc_frontend/src/ast/traits.rs b/compiler/noirc_frontend/src/ast/traits.rs index b1b14e3f657..064ea8ba9d9 100644 --- a/compiler/noirc_frontend/src/ast/traits.rs +++ b/compiler/noirc_frontend/src/ast/traits.rs @@ -7,6 +7,7 @@ use crate::ast::{ BlockExpression, Expression, FunctionReturnType, Ident, NoirFunction, Path, UnresolvedGenerics, UnresolvedType, }; +use crate::macros_api::SecondaryAttribute; use crate::node_interner::TraitId; /// AST node for trait definitions: @@ -18,6 +19,7 @@ pub struct NoirTrait { pub where_clause: Vec, pub span: Span, pub items: Vec, + pub attributes: Vec, } /// Any declaration inside the body of a trait that a user is required to diff --git a/compiler/noirc_frontend/src/elaborator/mod.rs b/compiler/noirc_frontend/src/elaborator/mod.rs index f0d3e3cae1f..b0fd8f79557 100644 --- a/compiler/noirc_frontend/src/elaborator/mod.rs +++ b/compiler/noirc_frontend/src/elaborator/mod.rs @@ -252,11 +252,11 @@ impl<'context> Elaborator<'context> { } // Must resolve structs before we resolve globals. - let generated_items = self.collect_struct_definitions(items.types); + let mut generated_items = self.collect_struct_definitions(items.types); self.define_function_metas(&mut items.functions, &mut items.impls, &mut items.trait_impls); - self.collect_traits(items.traits); + self.collect_traits(items.traits, &mut generated_items); // Before we resolve any function symbols we must go through our impls and // re-collect the methods within into their proper module. This cannot be @@ -278,6 +278,10 @@ impl<'context> Elaborator<'context> { self.elaborate_global(global); } + // We have to run any comptime attributes on functions before the function is elaborated + // since the generated items are checked beforehand as well. + self.run_attributes_on_functions(&items.functions, &mut generated_items); + // After everything is collected, we can elaborate our generated items. // It may be better to inline these within `items` entirely since elaborating them // all here means any globals will not see these. Inlining them completely within `items` @@ -1242,7 +1246,8 @@ impl<'context> Elaborator<'context> { .add_definition_location(ReferenceId::StructMember(type_id, field_index), None); } - self.run_comptime_attributes_on_struct(attributes, type_id, span, &mut generated_items); + let item = Value::StructDefinition(type_id); + self.run_comptime_attributes_on_item(&attributes, item, span, &mut generated_items); } // Check whether the struct fields have nested slices @@ -1268,17 +1273,17 @@ impl<'context> Elaborator<'context> { generated_items } - fn run_comptime_attributes_on_struct( + fn run_comptime_attributes_on_item( &mut self, - attributes: Vec, - struct_id: StructId, + attributes: &[SecondaryAttribute], + item: Value, span: Span, generated_items: &mut CollectedItems, ) { for attribute in attributes { if let SecondaryAttribute::Custom(name) = attribute { if let Err(error) = - self.run_comptime_attribute_on_struct(name, struct_id, span, generated_items) + self.run_comptime_attribute_on_item(name, item.clone(), span, generated_items) { self.errors.push(error); } @@ -1286,16 +1291,16 @@ impl<'context> Elaborator<'context> { } } - fn run_comptime_attribute_on_struct( + fn run_comptime_attribute_on_item( &mut self, - attribute: String, - struct_id: StructId, + attribute: &str, + item: Value, span: Span, generated_items: &mut CollectedItems, ) -> Result<(), (CompilationError, FileId)> { let location = Location::new(span, self.file); - let (function_name, mut arguments) = - Self::parse_attribute(&attribute, location).unwrap_or((attribute, Vec::new())); + let (function_name, mut arguments) = Self::parse_attribute(attribute, location) + .unwrap_or_else(|| (attribute.to_string(), Vec::new())); let id = self .lookup_global(Path::from_single(function_name, span)) @@ -1307,7 +1312,7 @@ impl<'context> Elaborator<'context> { }; self.handle_varargs_attribute(function, &mut arguments, location); - arguments.insert(0, (Value::StructDefinition(struct_id), location)); + arguments.insert(0, (item, location)); let mut interpreter_errors = vec![]; let mut interpreter = self.setup_interpreter(&mut interpreter_errors); @@ -1741,4 +1746,23 @@ impl<'context> Elaborator<'context> { )); } } + + fn run_attributes_on_functions( + &mut self, + function_sets: &[UnresolvedFunctions], + 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 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); + } + } + } } diff --git a/compiler/noirc_frontend/src/elaborator/traits.rs b/compiler/noirc_frontend/src/elaborator/traits.rs index 4cd20820c56..0627121f836 100644 --- a/compiler/noirc_frontend/src/elaborator/traits.rs +++ b/compiler/noirc_frontend/src/elaborator/traits.rs @@ -7,7 +7,7 @@ use crate::{ ast::{ FunctionKind, TraitItem, UnresolvedGeneric, UnresolvedGenerics, UnresolvedTraitConstraint, }, - hir::def_collector::dc_crate::UnresolvedTrait, + hir::def_collector::dc_crate::{CollectedItems, UnresolvedTrait}, hir_def::traits::{TraitConstant, TraitFunction, TraitType}, macros_api::{ BlockExpression, FunctionDefinition, FunctionReturnType, Ident, ItemVisibility, @@ -21,7 +21,11 @@ use crate::{ use super::Elaborator; impl<'context> Elaborator<'context> { - pub fn collect_traits(&mut self, traits: BTreeMap) { + pub fn collect_traits( + &mut self, + traits: BTreeMap, + generated_items: &mut CollectedItems, + ) { for (trait_id, unresolved_trait) in traits { self.recover_generics(|this| { let resolved_generics = this.interner.get_trait(trait_id).generics.clone(); @@ -41,6 +45,11 @@ impl<'context> Elaborator<'context> { this.interner.update_trait(trait_id, |trait_def| { trait_def.set_methods(methods); }); + + let attributes = &unresolved_trait.trait_def.attributes; + let item = crate::hir::comptime::Value::TraitDefinition(trait_id); + let span = unresolved_trait.trait_def.span; + this.run_comptime_attributes_on_item(attributes, item, span, generated_items); }); // This check needs to be after the trait's methods are set since diff --git a/compiler/noirc_frontend/src/hir/comptime/value.rs b/compiler/noirc_frontend/src/hir/comptime/value.rs index 9e15b73324f..46d1fc23973 100644 --- a/compiler/noirc_frontend/src/hir/comptime/value.rs +++ b/compiler/noirc_frontend/src/hir/comptime/value.rs @@ -8,12 +8,13 @@ use noirc_errors::Location; use crate::{ ast::{ArrayLiteral, ConstructorExpression, Ident, IntegerBitSize, Signedness}, + hir::def_map::ModuleId, hir_def::expr::{HirArrayLiteral, HirConstructorExpression, HirIdent, HirLambda, ImplKind}, macros_api::{ Expression, ExpressionKind, HirExpression, HirLiteral, Literal, NodeInterner, Path, StructId, }, - node_interner::{ExprId, FuncId}, + node_interner::{ExprId, FuncId, TraitId}, parser::{self, NoirParser, TopLevelStatement}, token::{SpannedToken, Token, Tokens}, QuotedType, Shared, Type, TypeBindings, @@ -45,6 +46,9 @@ pub enum Value { Slice(Vector, Type), Code(Rc), StructDefinition(StructId), + TraitDefinition(TraitId), + FunctionDefinition(FuncId), + ModuleDefinition(ModuleId), } impl Value { @@ -79,6 +83,9 @@ impl Value { let element = element.borrow().get_type().into_owned(); Type::MutableReference(Box::new(element)) } + Value::TraitDefinition(_) => Type::Quoted(QuotedType::TraitDefinition), + Value::FunctionDefinition(_) => Type::Quoted(QuotedType::FunctionDefinition), + Value::ModuleDefinition(_) => Type::Quoted(QuotedType::Module), }) } @@ -192,7 +199,11 @@ impl Value { } }; } - Value::Pointer(_) | Value::StructDefinition(_) => { + Value::Pointer(_) + | Value::StructDefinition(_) + | Value::TraitDefinition(_) + | Value::FunctionDefinition(_) + | Value::ModuleDefinition(_) => { return Err(InterpreterError::CannotInlineMacro { value: self, location }) } }; @@ -298,7 +309,11 @@ impl Value { HirExpression::Literal(HirLiteral::Slice(HirArrayLiteral::Standard(elements))) } Value::Code(block) => HirExpression::Unquote(unwrap_rc(block)), - Value::Pointer(_) | Value::StructDefinition(_) => { + Value::Pointer(_) + | Value::StructDefinition(_) + | Value::TraitDefinition(_) + | Value::FunctionDefinition(_) + | Value::ModuleDefinition(_) => { return Err(InterpreterError::CannotInlineMacro { value: self, location }) } }; @@ -402,6 +417,9 @@ impl Display for Value { write!(f, " }}") } Value::StructDefinition(_) => write!(f, "(struct definition)"), + Value::TraitDefinition(_) => write!(f, "(trait definition)"), + Value::FunctionDefinition(_) => write!(f, "(function definition)"), + Value::ModuleDefinition(_) => write!(f, "(module)"), } } } diff --git a/compiler/noirc_frontend/src/hir_def/types.rs b/compiler/noirc_frontend/src/hir_def/types.rs index 6919f697d27..91a374a0787 100644 --- a/compiler/noirc_frontend/src/hir_def/types.rs +++ b/compiler/noirc_frontend/src/hir_def/types.rs @@ -142,6 +142,9 @@ pub enum QuotedType { TopLevelItem, Type, StructDefinition, + TraitDefinition, + FunctionDefinition, + Module, } /// A list of TypeVariableIds to bind to a type. Storing the @@ -677,6 +680,9 @@ impl std::fmt::Display for QuotedType { QuotedType::TopLevelItem => write!(f, "TopLevelItem"), QuotedType::Type => write!(f, "Type"), QuotedType::StructDefinition => write!(f, "StructDefinition"), + QuotedType::TraitDefinition => write!(f, "TraitDefinition"), + QuotedType::FunctionDefinition => write!(f, "FunctionDefinition"), + QuotedType::Module => write!(f, "Module"), } } } diff --git a/compiler/noirc_frontend/src/lexer/token.rs b/compiler/noirc_frontend/src/lexer/token.rs index 41de13fb17e..59ae3d8fd7a 100644 --- a/compiler/noirc_frontend/src/lexer/token.rs +++ b/compiler/noirc_frontend/src/lexer/token.rs @@ -895,12 +895,14 @@ pub enum Keyword { Fn, For, FormatString, + FunctionDefinition, Global, If, Impl, In, Let, Mod, + Module, Mut, Pub, Quoted, @@ -908,12 +910,13 @@ pub enum Keyword { ReturnData, String, Struct, + StructDefinition, Super, TopLevelItem, Trait, + TraitDefinition, Type, TypeType, - StructDefinition, Unchecked, Unconstrained, Use, @@ -943,12 +946,14 @@ impl fmt::Display for Keyword { Keyword::Fn => write!(f, "fn"), Keyword::For => write!(f, "for"), Keyword::FormatString => write!(f, "fmtstr"), + Keyword::FunctionDefinition => write!(f, "FunctionDefinition"), Keyword::Global => write!(f, "global"), Keyword::If => write!(f, "if"), Keyword::Impl => write!(f, "impl"), Keyword::In => write!(f, "in"), Keyword::Let => write!(f, "let"), Keyword::Mod => write!(f, "mod"), + Keyword::Module => write!(f, "Module"), Keyword::Mut => write!(f, "mut"), Keyword::Pub => write!(f, "pub"), Keyword::Quoted => write!(f, "Quoted"), @@ -956,12 +961,13 @@ impl fmt::Display for Keyword { Keyword::ReturnData => write!(f, "return_data"), Keyword::String => write!(f, "str"), Keyword::Struct => write!(f, "struct"), + Keyword::StructDefinition => write!(f, "StructDefinition"), Keyword::Super => write!(f, "super"), Keyword::TopLevelItem => write!(f, "TopLevelItem"), Keyword::Trait => write!(f, "trait"), + Keyword::TraitDefinition => write!(f, "TraitDefinition"), Keyword::Type => write!(f, "type"), Keyword::TypeType => write!(f, "Type"), - Keyword::StructDefinition => write!(f, "StructDefinition"), Keyword::Unchecked => write!(f, "unchecked"), Keyword::Unconstrained => write!(f, "unconstrained"), Keyword::Use => write!(f, "use"), @@ -994,12 +1000,14 @@ impl Keyword { "fn" => Keyword::Fn, "for" => Keyword::For, "fmtstr" => Keyword::FormatString, + "FunctionDefinition" => Keyword::FunctionDefinition, "global" => Keyword::Global, "if" => Keyword::If, "impl" => Keyword::Impl, "in" => Keyword::In, "let" => Keyword::Let, "mod" => Keyword::Mod, + "Module" => Keyword::Module, "mut" => Keyword::Mut, "pub" => Keyword::Pub, "Quoted" => Keyword::Quoted, @@ -1010,6 +1018,7 @@ impl Keyword { "super" => Keyword::Super, "TopLevelItem" => Keyword::TopLevelItem, "trait" => Keyword::Trait, + "TraitDefinition" => Keyword::TraitDefinition, "type" => Keyword::Type, "Type" => Keyword::TypeType, "StructDefinition" => Keyword::StructDefinition, diff --git a/compiler/noirc_frontend/src/parser/parser/traits.rs b/compiler/noirc_frontend/src/parser/parser/traits.rs index 1aec57c8e41..e64de584da4 100644 --- a/compiler/noirc_frontend/src/parser/parser/traits.rs +++ b/compiler/noirc_frontend/src/parser/parser/traits.rs @@ -1,5 +1,6 @@ use chumsky::prelude::*; +use super::attributes::{attributes, validate_secondary_attributes}; use super::function::function_return_type; use super::{block, expression, fresh_statement, function, function_declaration_parameters}; @@ -18,15 +19,24 @@ use crate::{ use super::{generic_type_args, parse_type, path, primitives::ident}; pub(super) fn trait_definition() -> impl NoirParser { - keyword(Keyword::Trait) - .ignore_then(ident()) + attributes() + .then_ignore(keyword(Keyword::Trait)) + .then(ident()) .then(function::generics()) .then(where_clause()) .then_ignore(just(Token::LeftBrace)) .then(trait_body()) .then_ignore(just(Token::RightBrace)) - .map_with_span(|(((name, generics), where_clause), items), span| { - TopLevelStatement::Trait(NoirTrait { name, generics, where_clause, span, items }) + .validate(|((((attributes, name), generics), where_clause), items), span, emit| { + let attributes = validate_secondary_attributes(attributes, span, emit); + TopLevelStatement::Trait(NoirTrait { + name, + generics, + where_clause, + span, + items, + attributes, + }) }) } diff --git a/compiler/noirc_frontend/src/parser/parser/types.rs b/compiler/noirc_frontend/src/parser/parser/types.rs index 6e8064c6423..7b69b309bae 100644 --- a/compiler/noirc_frontend/src/parser/parser/types.rs +++ b/compiler/noirc_frontend/src/parser/parser/types.rs @@ -26,6 +26,9 @@ pub(super) fn parse_type_inner<'a>( string_type(), expr_type(), struct_definition_type(), + trait_definition_type(), + function_definition_type(), + module_type(), top_level_item_type(), type_of_quoted_types(), quoted_type(), @@ -79,6 +82,23 @@ pub(super) fn struct_definition_type() -> impl NoirParser { }) } +pub(super) fn trait_definition_type() -> impl NoirParser { + keyword(Keyword::TraitDefinition).map_with_span(|_, span| { + UnresolvedTypeData::Quoted(QuotedType::TraitDefinition).with_span(span) + }) +} + +pub(super) fn function_definition_type() -> impl NoirParser { + keyword(Keyword::FunctionDefinition).map_with_span(|_, span| { + UnresolvedTypeData::Quoted(QuotedType::FunctionDefinition).with_span(span) + }) +} + +pub(super) fn module_type() -> impl NoirParser { + keyword(Keyword::Module) + .map_with_span(|_, span| UnresolvedTypeData::Quoted(QuotedType::Module).with_span(span)) +} + /// This is the type `TopLevelItem` - the type of a quoted statement in the top level. /// E.g. a type definition, trait definition, trait impl, function, etc. fn top_level_item_type() -> impl NoirParser { diff --git a/test_programs/compile_success_empty/function_attribute/Nargo.toml b/test_programs/compile_success_empty/function_attribute/Nargo.toml new file mode 100644 index 00000000000..94b5c5da6a8 --- /dev/null +++ b/test_programs/compile_success_empty/function_attribute/Nargo.toml @@ -0,0 +1,7 @@ +[package] +name = "function_attribute" +type = "bin" +authors = [""] +compiler_version = ">=0.31.0" + +[dependencies] \ No newline at end of file diff --git a/test_programs/compile_success_empty/function_attribute/src/main.nr b/test_programs/compile_success_empty/function_attribute/src/main.nr new file mode 100644 index 00000000000..ec22b730d3f --- /dev/null +++ b/test_programs/compile_success_empty/function_attribute/src/main.nr @@ -0,0 +1,18 @@ +#[function_attr] +fn foo() {} + +struct Foo {} + +comptime fn function_attr(_f: FunctionDefinition) -> Quoted { + quote { + impl Default for Foo { + fn default() -> Foo { + Foo {} + } + } + } +} + +fn main() { + let _ = Foo::default(); +} diff --git a/test_programs/compile_success_empty/trait_attribute/Nargo.toml b/test_programs/compile_success_empty/trait_attribute/Nargo.toml new file mode 100644 index 00000000000..c72fe5e3e89 --- /dev/null +++ b/test_programs/compile_success_empty/trait_attribute/Nargo.toml @@ -0,0 +1,7 @@ +[package] +name = "trait_attribute" +type = "bin" +authors = [""] +compiler_version = ">=0.31.0" + +[dependencies] diff --git a/test_programs/compile_success_empty/trait_attribute/src/main.nr b/test_programs/compile_success_empty/trait_attribute/src/main.nr new file mode 100644 index 00000000000..87f4893e3e5 --- /dev/null +++ b/test_programs/compile_success_empty/trait_attribute/src/main.nr @@ -0,0 +1,19 @@ +#[trait_attr] +trait Foo { + fn foo(self) -> Self; +} + +comptime fn trait_attr(_t: TraitDefinition) -> Quoted { + quote { + impl Foo for Field { + fn foo(self) -> Self { + self + 1 + } + } + } +} + +fn main() { + assert_eq(1.foo(), 2); + assert_eq(10.foo(), 11); +} diff --git a/tooling/nargo_cli/build.rs b/tooling/nargo_cli/build.rs index 5e255da135d..a8c02a0ecd1 100644 --- a/tooling/nargo_cli/build.rs +++ b/tooling/nargo_cli/build.rs @@ -61,7 +61,7 @@ const IGNORED_BRILLIG_TESTS: [&str; 11] = [ /// Certain features are only available in the elaborator. /// We skip these tests for non-elaborator code since they are not /// expected to work there. This can be removed once the old code is removed. -const IGNORED_NEW_FEATURE_TESTS: [&str; 11] = [ +const IGNORED_NEW_FEATURE_TESTS: [&str; 13] = [ "macros", "wildcard_type", "type_definition_annotation", @@ -71,6 +71,8 @@ const IGNORED_NEW_FEATURE_TESTS: [&str; 11] = [ "comptime_slice_methods", "unary_operator_overloading", "unquote_multiple_items_from_annotation", + "function_attribute", + "trait_attribute", "regression_5428", "unquote", ];