diff --git a/compiler/noirc_frontend/src/attribute_order.rs b/compiler/noirc_frontend/src/attribute_order.rs new file mode 100644 index 00000000000..f4a45f7ff33 --- /dev/null +++ b/compiler/noirc_frontend/src/attribute_order.rs @@ -0,0 +1,93 @@ +use fm::FileId; +use noirc_errors::Span; +use petgraph::prelude::{DiGraph, NodeIndex}; +use rustc_hash::FxHashMap as HashMap; + +use crate::{ + ast::Expression, + hir::{comptime::Value, def_map::LocalModuleId}, + node_interner::FuncId, +}; + +#[derive(Debug)] +pub struct AttributeGraph { + default_stage: NodeIndex, + + order: DiGraph, + + indices: HashMap, + + modified_functions: std::collections::HashSet, +} + +#[derive(Debug, Copy, Clone)] +pub(crate) struct AttributeContext { + // The file where generated items should be added + pub(crate) file: FileId, + // The module where generated items should be added + pub(crate) module: LocalModuleId, + // The file where the attribute is located + pub(crate) attribute_file: FileId, + // The module where the attribute is located + pub(crate) attribute_module: LocalModuleId, +} + +pub(crate) type CollectedAttributes = Vec<(FuncId, Value, Vec, AttributeContext, Span)>; + +impl AttributeContext { + pub(crate) fn new(file: FileId, module: LocalModuleId) -> Self { + Self { file, module, attribute_file: file, attribute_module: module } + } +} + +impl Default for AttributeGraph { + fn default() -> Self { + let mut order = DiGraph::default(); + let mut indices = HashMap::default(); + + let default_stage = order.add_node(FuncId::dummy_id()); + indices.insert(FuncId::dummy_id(), default_stage); + + Self { default_stage, order, indices, modified_functions: Default::default() } + } +} + +impl AttributeGraph { + pub fn get_or_insert(&mut self, attr: FuncId) -> NodeIndex { + if let Some(index) = self.indices.get(&attr) { + return *index; + } + + let index = self.order.add_node(attr); + self.indices.insert(attr, index); + index + } + + pub fn add_ordering_constraint(&mut self, run_first: FuncId, run_second: FuncId) { + let first_index = self.get_or_insert(run_first); + let second_index = self.get_or_insert(run_second); + + // Just for debugging + if run_first != FuncId::dummy_id() { + self.modified_functions.insert(run_first); + self.modified_functions.insert(run_second); + } + + self.order.update_edge(second_index, first_index, 1.0); + } + + /// The default ordering of an attribute: run in the default stage + pub fn run_in_default_stage(&mut self, attr: FuncId) { + let index = self.get_or_insert(attr); + self.order.update_edge(self.default_stage, index, 1.0); + } + + pub(crate) fn sort_attributes_by_run_order(&self, attributes: &mut CollectedAttributes) { + let topological_sort = petgraph::algo::toposort(&self.order, None).unwrap(); + + let ordering: HashMap = + topological_sort.into_iter().map(|index| (self.order[index], index.index())).collect(); + + attributes.sort_by_key(|(f, ..)| ordering[f]); + } +} diff --git a/compiler/noirc_frontend/src/elaborator/comptime.rs b/compiler/noirc_frontend/src/elaborator/comptime.rs index 65cb6072c62..920191e4173 100644 --- a/compiler/noirc_frontend/src/elaborator/comptime.rs +++ b/compiler/noirc_frontend/src/elaborator/comptime.rs @@ -6,6 +6,7 @@ use noirc_errors::{Location, Span}; use crate::{ ast::{Documented, Expression, ExpressionKind}, + attribute_order::{AttributeContext, CollectedAttributes}, hir::{ comptime::{Interpreter, InterpreterError, Value}, def_collector::{ @@ -15,7 +16,7 @@ use crate::{ }, dc_mod, }, - def_map::{LocalModuleId, ModuleId}, + def_map::ModuleId, resolution::errors::ResolverError, }, hir_def::expr::{HirExpression, HirIdent}, @@ -28,24 +29,6 @@ use crate::{ use super::{Elaborator, FunctionContext, ResolverMeta}; -#[derive(Debug, Copy, Clone)] -struct AttributeContext { - // The file where generated items should be added - file: FileId, - // The module where generated items should be added - module: LocalModuleId, - // The file where the attribute is located - attribute_file: FileId, - // The module where the attribute is located - attribute_module: LocalModuleId, -} - -impl AttributeContext { - fn new(file: FileId, module: LocalModuleId) -> Self { - Self { file, module, attribute_file: file, attribute_module: module } - } -} - impl<'context> Elaborator<'context> { /// Elaborate an expression from the middle of a comptime scope. /// When this happens we require additional information to know @@ -131,42 +114,42 @@ impl<'context> Elaborator<'context> { } } - fn run_comptime_attributes_on_item( + fn collect_comptime_attributes_on_item( &mut self, attributes: &[SecondaryAttribute], item: Value, span: Span, attribute_context: AttributeContext, - generated_items: &mut CollectedItems, + attributes_to_run: &mut CollectedAttributes, ) { for attribute in attributes { - self.run_comptime_attribute_on_item( + self.collect_comptime_attribute_on_item( attribute, &item, span, attribute_context, - generated_items, + attributes_to_run, ); } } - fn run_comptime_attribute_on_item( + fn collect_comptime_attribute_on_item( &mut self, attribute: &SecondaryAttribute, item: &Value, span: Span, attribute_context: AttributeContext, - generated_items: &mut CollectedItems, + attributes_to_run: &mut CollectedAttributes, ) { if let SecondaryAttribute::Meta(attribute) = attribute { self.elaborate_in_comptime_context(|this| { - if let Err(error) = this.run_comptime_attribute_name_on_item( + if let Err(error) = this.collect_comptime_attribute_name_on_item( &attribute.contents, item.clone(), span, attribute.contents_span, attribute_context, - generated_items, + attributes_to_run, ) { this.errors.push(error); } @@ -174,14 +157,15 @@ impl<'context> Elaborator<'context> { } } - fn run_comptime_attribute_name_on_item( + /// Resolve an attribute to the function it refers to and add it to `attributes_to_run` + fn collect_comptime_attribute_name_on_item( &mut self, attribute: &str, item: Value, span: Span, attribute_span: Span, attribute_context: AttributeContext, - generated_items: &mut CollectedItems, + attributes_to_run: &mut CollectedAttributes, ) -> Result<(), (CompilationError, FileId)> { self.file = attribute_context.attribute_file; self.local_module = attribute_context.attribute_module; @@ -233,6 +217,19 @@ impl<'context> Elaborator<'context> { return Err((ResolverError::NonFunctionInAnnotation { span }.into(), self.file)); }; + attributes_to_run.push((function, item, arguments, attribute_context, span)); + Ok(()) + } + + fn run_attribute( + &mut self, + attribute_context: AttributeContext, + function: FuncId, + arguments: Vec, + item: Value, + location: Location, + generated_items: &mut CollectedItems, + ) -> Result<(), (CompilationError, FileId)> { self.file = attribute_context.file; self.local_module = attribute_context.module; @@ -244,10 +241,7 @@ impl<'context> Elaborator<'context> { arguments, location, ) - .map_err(|error| { - let file = error.get_location().file; - (error.into(), file) - })?; + .map_err(|error| error.into_compilation_error_pair())?; arguments.insert(0, (item, location)); @@ -537,19 +531,19 @@ impl<'context> Elaborator<'context> { functions: &[UnresolvedFunctions], module_attributes: &[ModuleAttribute], ) -> CollectedItems { - let mut generated_items = CollectedItems::default(); + let mut attributes_to_run = Vec::new(); for (trait_id, trait_) in traits { let attributes = &trait_.trait_def.attributes; let item = Value::TraitDefinition(*trait_id); let span = trait_.trait_def.span; let context = AttributeContext::new(trait_.file_id, trait_.module_id); - self.run_comptime_attributes_on_item( + self.collect_comptime_attributes_on_item( attributes, item, span, context, - &mut generated_items, + &mut attributes_to_run, ); } @@ -558,26 +552,38 @@ impl<'context> Elaborator<'context> { let item = Value::StructDefinition(*struct_id); let span = struct_def.struct_def.span; let context = AttributeContext::new(struct_def.file_id, struct_def.module_id); - self.run_comptime_attributes_on_item( + self.collect_comptime_attributes_on_item( attributes, item, span, context, - &mut generated_items, + &mut attributes_to_run, ); } - self.run_attributes_on_functions(functions, &mut generated_items); + self.collect_attributes_on_functions(functions, &mut attributes_to_run); + self.collect_attributes_on_modules(module_attributes, &mut attributes_to_run); - self.run_attributes_on_modules(module_attributes, &mut generated_items); + self.interner.attribute_order.sort_attributes_by_run_order(&mut attributes_to_run); + + // run + let mut generated_items = CollectedItems::default(); + for (attribute, item, args, context, span) in attributes_to_run { + let location = Location::new(span, context.attribute_file); + if let Err(error) = + self.run_attribute(context, attribute, args, item, location, &mut generated_items) + { + self.errors.push(error); + } + } generated_items } - fn run_attributes_on_modules( + fn collect_attributes_on_modules( &mut self, module_attributes: &[ModuleAttribute], - generated_items: &mut CollectedItems, + attributes_to_run: &mut CollectedAttributes, ) { for module_attribute in module_attributes { let local_id = module_attribute.module_id; @@ -593,14 +599,20 @@ impl<'context> Elaborator<'context> { attribute_module: module_attribute.attribute_module_id, }; - self.run_comptime_attribute_on_item(attribute, &item, span, context, generated_items); + self.collect_comptime_attribute_on_item( + attribute, + &item, + span, + context, + attributes_to_run, + ); } } - fn run_attributes_on_functions( + fn collect_attributes_on_functions( &mut self, function_sets: &[UnresolvedFunctions], - generated_items: &mut CollectedItems, + attributes_to_run: &mut CollectedAttributes, ) { for function_set in function_sets { self.self_type = function_set.self_type.clone(); @@ -610,12 +622,12 @@ impl<'context> Elaborator<'context> { let attributes = function.secondary_attributes(); let item = Value::FunctionDefinition(*function_id); let span = function.span(); - self.run_comptime_attributes_on_item( + self.collect_comptime_attributes_on_item( attributes, item, span, context, - generated_items, + attributes_to_run, ); } } @@ -650,4 +662,50 @@ impl<'context> Elaborator<'context> { _ => false, } } + + pub(super) fn register_attribute_order( + &mut self, + id: FuncId, + attributes: &[SecondaryAttribute], + ) { + let mut has_order = false; + + for attribute in attributes { + let (name, run_before) = match attribute { + SecondaryAttribute::RunBefore(name) => (name, true), + SecondaryAttribute::RunAfter(name) => (name, false), + _ => continue, + }; + + // Parse a path from #[run_before(path)] + let Some(path) = Parser::for_str(name).parse_path_no_turbofish() else { + todo!("function should be a path") + }; + + let definition_id = self.resolve_variable(path).id; + + match self.interner.definition(definition_id).kind { + DefinitionKind::Function(attribute_arg) => { + if attribute_arg == id { + todo!("Attribute cannot be run before or after itself"); + } + + has_order = true; + if run_before { + self.interner.attribute_order.add_ordering_constraint(attribute_arg, id); + } else { + self.interner.attribute_order.add_ordering_constraint(id, attribute_arg); + } + } + _ => { + todo!("path doesn't refer to a function") + } + } + } + + // If no ordering was specified, set it to the default stage. + if !has_order { + self.interner.attribute_order.run_in_default_stage(id); + } + } } diff --git a/compiler/noirc_frontend/src/elaborator/mod.rs b/compiler/noirc_frontend/src/elaborator/mod.rs index aef0771c486..e87accb14db 100644 --- a/compiler/noirc_frontend/src/elaborator/mod.rs +++ b/compiler/noirc_frontend/src/elaborator/mod.rs @@ -820,6 +820,8 @@ impl<'context> Elaborator<'context> { None }; + self.register_attribute_order(func_id, func.secondary_attributes()); + let attributes = func.secondary_attributes().iter(); let attributes = attributes.filter_map(|secondary_attribute| secondary_attribute.as_custom()); diff --git a/compiler/noirc_frontend/src/elaborator/patterns.rs b/compiler/noirc_frontend/src/elaborator/patterns.rs index d55011f98a1..5226a7402c9 100644 --- a/compiler/noirc_frontend/src/elaborator/patterns.rs +++ b/compiler/noirc_frontend/src/elaborator/patterns.rs @@ -498,7 +498,7 @@ impl<'context> Elaborator<'context> { (id, typ) } - fn resolve_variable(&mut self, path: Path) -> HirIdent { + pub fn resolve_variable(&mut self, path: Path) -> HirIdent { if let Some(trait_path_resolution) = self.resolve_trait_generic_path(&path) { if let Some(error) = trait_path_resolution.error { self.push_err(error); diff --git a/compiler/noirc_frontend/src/elaborator/types.rs b/compiler/noirc_frontend/src/elaborator/types.rs index 8ffbd15bdab..7051c5ad184 100644 --- a/compiler/noirc_frontend/src/elaborator/types.rs +++ b/compiler/noirc_frontend/src/elaborator/types.rs @@ -566,7 +566,7 @@ impl<'context> Elaborator<'context> { let path_resolution = self.resolve_path(path.clone()).ok()?; let ModuleDefId::FunctionId(func_id) = path_resolution.module_def_id else { return None }; - let meta = self.interner.function_meta(&func_id); + let meta = self.interner.try_function_meta(&func_id)?; let the_trait = self.interner.get_trait(meta.trait_id?); let method = the_trait.find_method(path.last_name())?; let constraint = the_trait.as_constraint(path.span); diff --git a/compiler/noirc_frontend/src/lexer/token.rs b/compiler/noirc_frontend/src/lexer/token.rs index 8f05832d26d..f6d81a860ec 100644 --- a/compiler/noirc_frontend/src/lexer/token.rs +++ b/compiler/noirc_frontend/src/lexer/token.rs @@ -832,6 +832,8 @@ impl Attribute { ["varargs"] => Attribute::Secondary(SecondaryAttribute::Varargs), ["use_callers_scope"] => Attribute::Secondary(SecondaryAttribute::UseCallersScope), ["allow", tag] => Attribute::Secondary(SecondaryAttribute::Allow(tag.to_string())), + ["run_before", f] => Attribute::Secondary(SecondaryAttribute::RunBefore(f.to_string())), + ["run_after", f] => Attribute::Secondary(SecondaryAttribute::RunAfter(f.to_string())), tokens => { tokens.iter().try_for_each(|token| validate(token))?; Attribute::Secondary(SecondaryAttribute::Meta(CustomAttribute { @@ -972,6 +974,12 @@ pub enum SecondaryAttribute { /// Allow chosen warnings to happen so they are silenced. Allow(String), + + /// Specifies this attribute should run before the attribute specified as an argument + RunBefore(String), + + /// Specifies this attribute should run after the attribute specified as an argument + RunAfter(String), } impl SecondaryAttribute { @@ -997,6 +1005,8 @@ impl SecondaryAttribute { SecondaryAttribute::Varargs => Some("varargs".to_string()), SecondaryAttribute::UseCallersScope => Some("use_callers_scope".to_string()), SecondaryAttribute::Allow(_) => Some("allow".to_string()), + SecondaryAttribute::RunBefore(_) => Some("run_before".to_string()), + SecondaryAttribute::RunAfter(_) => Some("run_after".to_string()), } } @@ -1012,27 +1022,29 @@ impl SecondaryAttribute { } pub(crate) fn contents(&self) -> String { - match self { - SecondaryAttribute::Deprecated(None) => "deprecated".to_string(), - SecondaryAttribute::Deprecated(Some(ref note)) => { - format!("deprecated({note:?})") - } - SecondaryAttribute::Tag(ref attribute) => format!("'{}", attribute.contents), - SecondaryAttribute::Meta(ref attribute) => attribute.contents.to_string(), - SecondaryAttribute::ContractLibraryMethod => "contract_library_method".to_string(), - SecondaryAttribute::Export => "export".to_string(), - SecondaryAttribute::Field(ref k) => format!("field({k})"), - SecondaryAttribute::Abi(ref k) => format!("abi({k})"), - SecondaryAttribute::Varargs => "varargs".to_string(), - SecondaryAttribute::UseCallersScope => "use_callers_scope".to_string(), - SecondaryAttribute::Allow(ref k) => format!("allow({k})"), - } + self.to_string() } } impl fmt::Display for SecondaryAttribute { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - write!(f, "#[{}]", self.contents()) + match self { + SecondaryAttribute::Deprecated(None) => write!(f, "deprecated"), + SecondaryAttribute::Deprecated(Some(ref note)) => { + write!(f, "deprecated({note:?})") + } + SecondaryAttribute::Tag(ref attribute) => write!(f, "'{}", attribute.contents), + SecondaryAttribute::Meta(ref attribute) => write!(f, "{}", attribute.contents), + SecondaryAttribute::ContractLibraryMethod => write!(f, "contract_library_method"), + SecondaryAttribute::Export => write!(f, "export"), + 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"), + SecondaryAttribute::Allow(ref k) => write!(f, "allow({k})"), + SecondaryAttribute::RunBefore(attr) => write!(f, "#[run_before({attr})]"), + SecondaryAttribute::RunAfter(attr) => write!(f, "#[run_after({attr})]"), + } } } @@ -1081,6 +1093,8 @@ impl AsRef for SecondaryAttribute { SecondaryAttribute::Meta(attribute) => &attribute.contents, SecondaryAttribute::Field(string) | SecondaryAttribute::Abi(string) + | SecondaryAttribute::RunBefore(string) + | SecondaryAttribute::RunAfter(string) | SecondaryAttribute::Allow(string) => string, SecondaryAttribute::ContractLibraryMethod => "", SecondaryAttribute::Export => "", diff --git a/compiler/noirc_frontend/src/lib.rs b/compiler/noirc_frontend/src/lib.rs index 9d98b125e32..1cbb2057fd0 100644 --- a/compiler/noirc_frontend/src/lib.rs +++ b/compiler/noirc_frontend/src/lib.rs @@ -11,6 +11,7 @@ #![warn(clippy::semicolon_if_nothing_returned)] pub mod ast; +pub mod attribute_order; pub mod debug; pub mod elaborator; pub mod graph; diff --git a/compiler/noirc_frontend/src/node_interner.rs b/compiler/noirc_frontend/src/node_interner.rs index ca7e0c6aa59..eba685d8f4f 100644 --- a/compiler/noirc_frontend/src/node_interner.rs +++ b/compiler/noirc_frontend/src/node_interner.rs @@ -16,6 +16,7 @@ use rustc_hash::FxHashMap as HashMap; use crate::ast::{ ExpressionKind, Ident, LValue, Pattern, StatementKind, UnaryOp, UnresolvedTypeData, }; +use crate::attribute_order::AttributeGraph; use crate::graph::CrateId; use crate::hir::comptime; use crate::hir::def_collector::dc_crate::CompilationError; @@ -274,6 +275,8 @@ pub struct NodeInterner { /// Captures the documentation comments for each module, struct, trait, function, etc. pub(crate) doc_comments: HashMap>, + + pub(crate) attribute_order: AttributeGraph, } /// A dependency in the dependency graph may be a type or a definition. @@ -681,6 +684,7 @@ impl Default for NodeInterner { trait_impl_associated_types: HashMap::default(), usage_tracker: UsageTracker::default(), doc_comments: HashMap::default(), + attribute_order: AttributeGraph::default(), } } } diff --git a/noir_stdlib/src/meta/mod.nr b/noir_stdlib/src/meta/mod.nr index ff662b878ec..383ca9d73ed 100644 --- a/noir_stdlib/src/meta/mod.nr +++ b/noir_stdlib/src/meta/mod.nr @@ -28,6 +28,10 @@ pub comptime fn unquote(code: Quoted) -> Quoted { pub comptime fn type_of(x: T) -> Type {} // docs:end:type_of +/// Stub function used so that other attributes can specify if they +/// want to run before or after any attributes without a stage annotation. +pub comptime fn default_stage() {} + // docs:start:derive_example // These are needed for the unconstrained hashmap we're using to store derive functions use crate::collections::umap::UHashMap; @@ -47,6 +51,7 @@ comptime mut global HANDLERS: UHashMap Quoted { // docs:end:derive let mut result = quote {}; @@ -66,6 +71,7 @@ pub comptime fn derive(s: StructDefinition, traits: [TraitDefinition]) -> Quoted // docs:start:derive_via // To register a handler for a trait, just add it to our handlers map // docs:start:derive_via_signature +#[run_before(derive)] pub comptime fn derive_via(t: TraitDefinition, f: DeriveFunction) { // docs:end:derive_via_signature HANDLERS.insert(t, f);