3 files changed, 228 insertions, 104 deletions

diff --git a/crates/hir_def/src/nameres/collector.rs b/crates/hir_def/src/nameres/collector.rs
index 3896be25d..e76d039b8 100644
--- a/crates/hir_def/src/nameres/collector.rs
+++ b/crates/hir_def/src/nameres/collector.rs
@@ -20,7 +20,7 @@ use rustc_hash::{FxHashMap, FxHashSet};
 use syntax::ast;
 
 use crate::{
-    attr::{AttrId, Attrs},
+    attr::{Attr, AttrId, AttrInput, Attrs},
     builtin_attr,
     db::DefDatabase,
     derive_macro_as_call_id,
@@ -94,14 +94,16 @@ pub(super) fn collect_defs(
         unresolved_imports: Vec::new(),
         resolved_imports: Vec::new(),
 
-        unexpanded_macros: Vec::new(),
+        unresolved_macros: Vec::new(),
         mod_dirs: FxHashMap::default(),
         cfg_options,
         proc_macros,
         exports_proc_macros: false,
         from_glob_import: Default::default(),
-        ignore_attrs_on: FxHashSet::default(),
-        derive_helpers_in_scope: FxHashMap::default(),
+        skip_attrs: Default::default(),
+        derive_helpers_in_scope: Default::default(),
+        registered_attrs: Default::default(),
+        registered_tools: Default::default(),
     };
     match block {
         Some(block) => {
@@ -237,7 +239,7 @@ struct DefCollector<'a> {
     glob_imports: FxHashMap<LocalModuleId, Vec<(LocalModuleId, Visibility)>>,
     unresolved_imports: Vec<ImportDirective>,
     resolved_imports: Vec<ImportDirective>,
-    unexpanded_macros: Vec<MacroDirective>,
+    unresolved_macros: Vec<MacroDirective>,
     mod_dirs: FxHashMap<LocalModuleId, ModDir>,
     cfg_options: &'a CfgOptions,
     /// List of procedural macros defined by this crate. This is read from the dynamic library
@@ -247,10 +249,20 @@ struct DefCollector<'a> {
     proc_macros: Vec<(Name, ProcMacroExpander)>,
     exports_proc_macros: bool,
     from_glob_import: PerNsGlobImports,
-    ignore_attrs_on: FxHashSet<InFile<ModItem>>,
+    /// If we fail to resolve an attribute on a `ModItem`, we fall back to ignoring the attribute.
+    /// This map is used to skip all attributes up to and including the one that failed to resolve,
+    /// in order to not expand them twice.
+    ///
+    /// This also stores the attributes to skip when we resolve derive helpers and non-macro
+    /// non-builtin attributes in general.
+    skip_attrs: FxHashMap<InFile<ModItem>, AttrId>,
     /// Tracks which custom derives are in scope for an item, to allow resolution of derive helper
     /// attributes.
     derive_helpers_in_scope: FxHashMap<AstId<ast::Item>, Vec<Name>>,
+    /// Custom attributes registered with `#![register_attr]`.
+    registered_attrs: Vec<String>,
+    /// Custom tool modules registered with `#![register_tool]`.
+    registered_tools: Vec<String>,
 }
 
 impl DefCollector<'_> {
@@ -259,11 +271,39 @@ impl DefCollector<'_> {
         let item_tree = self.db.file_item_tree(file_id.into());
         let module_id = self.def_map.root;
         self.def_map.modules[module_id].origin = ModuleOrigin::CrateRoot { definition: file_id };
-        if item_tree
-            .top_level_attrs(self.db, self.def_map.krate)
-            .cfg()
-            .map_or(true, |cfg| self.cfg_options.check(&cfg) != Some(false))
-        {
+
+        let attrs = item_tree.top_level_attrs(self.db, self.def_map.krate);
+        if attrs.cfg().map_or(true, |cfg| self.cfg_options.check(&cfg) != Some(false)) {
+            // Process other crate-level attributes.
+            for attr in &*attrs {
+                let attr_name = match attr.path.as_ident() {
+                    Some(name) => name,
+                    None => continue,
+                };
+
+                let registered_name = if *attr_name == hir_expand::name![register_attr]
+                    || *attr_name == hir_expand::name![register_tool]
+                {
+                    match &attr.input {
+                        Some(AttrInput::TokenTree(subtree)) => match &*subtree.token_trees {
+                            [tt::TokenTree::Leaf(tt::Leaf::Ident(name))] => name.as_name(),
+                            _ => continue,
+                        },
+                        _ => continue,
+                    }
+                } else {
+                    continue;
+                };
+
+                if *attr_name == hir_expand::name![register_attr] {
+                    self.registered_attrs.push(registered_name.to_string());
+                    cov_mark::hit!(register_attr);
+                } else {
+                    self.registered_tools.push(registered_name.to_string());
+                    cov_mark::hit!(register_tool);
+                }
+            }
+
             ModCollector {
                 def_collector: &mut *self,
                 macro_depth: 0,
@@ -319,7 +359,7 @@ impl DefCollector<'_> {
                 }
             }
 
-            if self.reseed_with_unresolved_attributes() == ReachedFixedPoint::Yes {
+            if self.reseed_with_unresolved_attribute() == ReachedFixedPoint::Yes {
                 break;
             }
         }
@@ -362,27 +402,23 @@ impl DefCollector<'_> {
     }
 
     /// When the fixed-point loop reaches a stable state, we might still have some unresolved
-    /// attributes (or unexpanded attribute proc macros) left over. This takes them, and feeds the
-    /// item they're applied to back into name resolution.
+    /// attributes (or unexpanded attribute proc macros) left over. This takes one of them, and
+    /// feeds the item it's applied to back into name resolution.
     ///
     /// This effectively ignores the fact that the macro is there and just treats the items as
     /// normal code.
     ///
     /// This improves UX when proc macros are turned off or don't work, and replicates the behavior
     /// before we supported proc. attribute macros.
-    fn reseed_with_unresolved_attributes(&mut self) -> ReachedFixedPoint {
+    fn reseed_with_unresolved_attribute(&mut self) -> ReachedFixedPoint {
         cov_mark::hit!(unresolved_attribute_fallback);
 
-        let mut added_items = false;
-        let unexpanded_macros = std::mem::replace(&mut self.unexpanded_macros, Vec::new());
-        for directive in &unexpanded_macros {
-            if let MacroDirectiveKind::Attr { ast_id, mod_item, .. } = &directive.kind {
-                // Make sure to only add such items once.
-                if !self.ignore_attrs_on.insert(ast_id.ast_id.with_value(*mod_item)) {
-                    continue;
-                }
+        let mut unresolved_macros = std::mem::replace(&mut self.unresolved_macros, Vec::new());
+        let pos = unresolved_macros.iter().position(|directive| {
+            if let MacroDirectiveKind::Attr { ast_id, mod_item, attr } = &directive.kind {
+                self.skip_attrs.insert(ast_id.ast_id.with_value(*mod_item), *attr);
 
-                let file_id = self.def_map[directive.module_id].definition_source(self.db).file_id;
+                let file_id = ast_id.ast_id.file_id;
                 let item_tree = self.db.file_item_tree(file_id);
                 let mod_dir = self.mod_dirs[&directive.module_id].clone();
                 ModCollector {
@@ -394,14 +430,20 @@ impl DefCollector<'_> {
                     mod_dir,
                 }
                 .collect(&[*mod_item]);
-                added_items = true;
+                true
+            } else {
+                false
             }
+        });
+
+        if let Some(pos) = pos {
+            unresolved_macros.remove(pos);
         }
 
         // The collection above might add new unresolved macros (eg. derives), so merge the lists.
-        self.unexpanded_macros.extend(unexpanded_macros);
+        self.unresolved_macros.extend(unresolved_macros);
 
-        if added_items {
+        if pos.is_some() {
             // Continue name resolution with the new data.
             ReachedFixedPoint::No
         } else {
@@ -873,7 +915,7 @@ impl DefCollector<'_> {
     }
 
     fn resolve_macros(&mut self) -> ReachedFixedPoint {
-        let mut macros = std::mem::replace(&mut self.unexpanded_macros, Vec::new());
+        let mut macros = std::mem::replace(&mut self.unresolved_macros, Vec::new());
         let mut resolved = Vec::new();
         let mut res = ReachedFixedPoint::Yes;
         macros.retain(|directive| {
@@ -922,14 +964,43 @@ impl DefCollector<'_> {
                         Err(UnresolvedMacro { .. }) => (),
                     }
                 }
-                MacroDirectiveKind::Attr { .. } => {
-                    // not yet :)
+                MacroDirectiveKind::Attr { ast_id, mod_item, attr } => {
+                    if let Some(ident) = ast_id.path.as_ident() {
+                        if let Some(helpers) = self.derive_helpers_in_scope.get(&ast_id.ast_id) {
+                            if helpers.contains(ident) {
+                                cov_mark::hit!(resolved_derive_helper);
+
+                                // Resolved to derive helper. Collect the item's attributes again,
+                                // starting after the derive helper.
+                                let file_id = ast_id.ast_id.file_id;
+                                let item_tree = self.db.file_item_tree(file_id);
+                                let mod_dir = self.mod_dirs[&directive.module_id].clone();
+                                self.skip_attrs.insert(InFile::new(file_id, *mod_item), *attr);
+                                ModCollector {
+                                    def_collector: &mut *self,
+                                    macro_depth: directive.depth,
+                                    module_id: directive.module_id,
+                                    file_id,
+                                    item_tree: &item_tree,
+                                    mod_dir,
+                                }
+                                .collect(&[*mod_item]);
+
+                                // Remove the original directive since we resolved it.
+                                return false;
+                            }
+                        }
+                    }
+
+                    // Not resolved to a derive helper, so try to resolve as a macro.
+                    // FIXME: not yet :)
                 }
             }
 
             true
         });
-        self.unexpanded_macros = macros;
+        // Attribute resolution can add unresolved macro invocations, so concatenate the lists.
+        self.unresolved_macros.extend(macros);
 
         for (module_id, macro_call_id, depth) in resolved {
             self.collect_macro_expansion(module_id, macro_call_id, depth);
@@ -1000,7 +1071,7 @@ impl DefCollector<'_> {
     fn finish(mut self) -> DefMap {
         // Emit diagnostics for all remaining unexpanded macros.
 
-        for directive in &self.unexpanded_macros {
+        for directive in &self.unresolved_macros {
             match &directive.kind {
                 MacroDirectiveKind::FnLike { ast_id, fragment } => match macro_call_as_call_id(
                     ast_id,
@@ -1102,7 +1173,7 @@ impl ModCollector<'_, '_> {
 
         // Prelude module is always considered to be `#[macro_use]`.
         if let Some(prelude_module) = self.def_collector.def_map.prelude {
-            if prelude_module.krate != self.def_collector.def_map.krate {
+            if prelude_module.krate != krate {
                 cov_mark::hit!(prelude_is_macro_use);
                 self.def_collector.import_all_macros_exported(self.module_id, prelude_module.krate);
             }
@@ -1203,11 +1274,6 @@ impl ModCollector<'_, '_> {
                 ModItem::Struct(id) => {
                     let it = &self.item_tree[id];
 
-                    // FIXME: check attrs to see if this is an attribute macro invocation;
-                    // in which case we don't add the invocation, just a single attribute
-                    // macro invocation
-                    self.collect_derives(&attrs, it.ast_id.upcast());
-
                     def = Some(DefData {
                         id: StructLoc { container: module, id: ItemTreeId::new(self.file_id, id) }
                             .intern(self.def_collector.db)
@@ -1220,11 +1286,6 @@ impl ModCollector<'_, '_> {
                 ModItem::Union(id) => {
                     let it = &self.item_tree[id];
 
-                    // FIXME: check attrs to see if this is an attribute macro invocation;
-                    // in which case we don't add the invocation, just a single attribute
-                    // macro invocation
-                    self.collect_derives(&attrs, it.ast_id.upcast());
-
                     def = Some(DefData {
                         id: UnionLoc { container: module, id: ItemTreeId::new(self.file_id, id) }
                             .intern(self.def_collector.db)
@@ -1237,11 +1298,6 @@ impl ModCollector<'_, '_> {
                 ModItem::Enum(id) => {
                     let it = &self.item_tree[id];
 
-                    // FIXME: check attrs to see if this is an attribute macro invocation;
-                    // in which case we don't add the invocation, just a single attribute
-                    // macro invocation
-                    self.collect_derives(&attrs, it.ast_id.upcast());
-
                     def = Some(DefData {
                         id: EnumLoc { container: module, id: ItemTreeId::new(self.file_id, id) }
                             .intern(self.def_collector.db)
@@ -1453,76 +1509,103 @@ impl ModCollector<'_, '_> {
     ///
     /// Returns `Err` when some attributes could not be resolved to builtins and have been
     /// registered as unresolved.
+    ///
+    /// If `ignore_up_to` is `Some`, attributes precending and including that attribute will be
+    /// assumed to be resolved already.
     fn resolve_attributes(&mut self, attrs: &Attrs, mod_item: ModItem) -> Result<(), ()> {
-        fn is_builtin_attr(path: &ModPath) -> bool {
-            if path.kind == PathKind::Plain {
-                if let Some(tool_module) = path.segments().first() {
-                    let tool_module = tool_module.to_string();
-                    if builtin_attr::TOOL_MODULES.iter().any(|m| tool_module == *m) {
-                        return true;
-                    }
-                }
-
-                if let Some(name) = path.as_ident() {
-                    let name = name.to_string();
-                    if builtin_attr::INERT_ATTRIBUTES
-                        .iter()
-                        .chain(builtin_attr::EXTRA_ATTRIBUTES)
-                        .any(|attr| name == *attr)
-                    {
-                        return true;
-                    }
-                }
+        let mut ignore_up_to =
+            self.def_collector.skip_attrs.get(&InFile::new(self.file_id, mod_item)).copied();
+        for attr in attrs.iter().skip_while(|attr| match ignore_up_to {
+            Some(id) if attr.id == id => {
+                ignore_up_to = None;
+                true
             }
-
-            false
-        }
-
-        // We failed to resolve an attribute on this item earlier, and are falling back to treating
-        // the item as-is.
-        if self.def_collector.ignore_attrs_on.contains(&InFile::new(self.file_id, mod_item)) {
-            return Ok(());
-        }
-
-        match attrs.iter().find(|attr| !is_builtin_attr(&attr.path)) {
-            Some(non_builtin_attr) => {
-                log::debug!("non-builtin attribute {}", non_builtin_attr.path);
+            Some(_) => true,
+            None => false,
+        }) {
+            if attr.path.as_ident() == Some(&hir_expand::name![derive]) {
+                self.collect_derive(attr, mod_item);
+            } else if self.is_builtin_or_registered_attr(&attr.path) {
+                continue;
+            } else {
+                log::debug!("non-builtin attribute {}", attr.path);
 
                 let ast_id = AstIdWithPath::new(
                     self.file_id,
                     mod_item.ast_id(self.item_tree),
-                    non_builtin_attr.path.as_ref().clone(),
+                    attr.path.as_ref().clone(),
                 );
-                self.def_collector.unexpanded_macros.push(MacroDirective {
+                self.def_collector.unresolved_macros.push(MacroDirective {
                     module_id: self.module_id,
                     depth: self.macro_depth + 1,
-                    kind: MacroDirectiveKind::Attr { ast_id, attr: non_builtin_attr.id, mod_item },
+                    kind: MacroDirectiveKind::Attr { ast_id, attr: attr.id, mod_item },
                 });
 
-                Err(())
+                return Err(());
             }
-            None => Ok(()),
         }
+
+        Ok(())
     }
 
-    fn collect_derives(&mut self, attrs: &Attrs, ast_id: FileAstId<ast::Item>) {
-        for derive in attrs.by_key("derive").attrs() {
-            match derive.parse_derive() {
-                Some(derive_macros) => {
-                    for path in derive_macros {
-                        let ast_id = AstIdWithPath::new(self.file_id, ast_id, path);
-                        self.def_collector.unexpanded_macros.push(MacroDirective {
-                            module_id: self.module_id,
-                            depth: self.macro_depth + 1,
-                            kind: MacroDirectiveKind::Derive { ast_id, derive_attr: derive.id },
-                        });
-                    }
+    fn is_builtin_or_registered_attr(&self, path: &ModPath) -> bool {
+        if path.kind == PathKind::Plain {
+            if let Some(tool_module) = path.segments().first() {
+                let tool_module = tool_module.to_string();
+                if builtin_attr::TOOL_MODULES
+                    .iter()
+                    .copied()
+                    .chain(self.def_collector.registered_tools.iter().map(|s| &**s))
+                    .any(|m| tool_module == *m)
+                {
+                    return true;
                 }
-                None => {
-                    // FIXME: diagnose
-                    log::debug!("malformed derive: {:?}", derive);
+            }
+
+            if let Some(name) = path.as_ident() {
+                let name = name.to_string();
+                if builtin_attr::INERT_ATTRIBUTES
+                    .iter()
+                    .chain(builtin_attr::EXTRA_ATTRIBUTES)
+                    .copied()
+                    .chain(self.def_collector.registered_attrs.iter().map(|s| &**s))
+                    .any(|attr| name == *attr)
+                {
+                    return true;
+                }
+            }
+        }
+
+        false
+    }
+
+    fn collect_derive(&mut self, attr: &Attr, mod_item: ModItem) {
+        let ast_id: FileAstId<ast::Item> = match mod_item {
+            ModItem::Struct(it) => self.item_tree[it].ast_id.upcast(),
+            ModItem::Union(it) => self.item_tree[it].ast_id.upcast(),
+            ModItem::Enum(it) => self.item_tree[it].ast_id.upcast(),
+            _ => {
+                // Cannot use derive on this item.
+                // FIXME: diagnose
+                return;
+            }
+        };
+
+        match attr.parse_derive() {
+            Some(derive_macros) => {
+                for path in derive_macros {
+                    let ast_id = AstIdWithPath::new(self.file_id, ast_id, path);
+                    self.def_collector.unresolved_macros.push(MacroDirective {
+                        module_id: self.module_id,
+                        depth: self.macro_depth + 1,
+                        kind: MacroDirectiveKind::Derive { ast_id, derive_attr: attr.id },
+                    });
                 }
             }
+            None => {
+                // FIXME: diagnose
+                log::debug!("malformed derive: {:?}", attr);
+            }
         }
     }
 
@@ -1686,7 +1769,7 @@ impl ModCollector<'_, '_> {
             ast_id.path.kind = PathKind::Super(0);
         }
 
-        self.def_collector.unexpanded_macros.push(MacroDirective {
+        self.def_collector.unresolved_macros.push(MacroDirective {
             module_id: self.module_id,
             depth: self.macro_depth + 1,
             kind: MacroDirectiveKind::FnLike { ast_id, fragment: mac.fragment },
@@ -1731,14 +1814,16 @@ mod tests {
             glob_imports: FxHashMap::default(),
             unresolved_imports: Vec::new(),
             resolved_imports: Vec::new(),
-            unexpanded_macros: Vec::new(),
+            unresolved_macros: Vec::new(),
             mod_dirs: FxHashMap::default(),
             cfg_options: &CfgOptions::default(),
             proc_macros: Default::default(),
             exports_proc_macros: false,
             from_glob_import: Default::default(),
-            ignore_attrs_on: FxHashSet::default(),
-            derive_helpers_in_scope: FxHashMap::default(),
+            skip_attrs: Default::default(),
+            derive_helpers_in_scope: Default::default(),
+            registered_attrs: Default::default(),
+            registered_tools: Default::default(),
         };
         collector.seed_with_top_level();
         collector.collect();
diff --git a/crates/hir_def/src/nameres/tests/diagnostics.rs b/crates/hir_def/src/nameres/tests/diagnostics.rs
index 543975e07..75147d973 100644
--- a/crates/hir_def/src/nameres/tests/diagnostics.rs
+++ b/crates/hir_def/src/nameres/tests/diagnostics.rs
@@ -237,3 +237,20 @@ fn good_out_dir_diagnostic() {
         "#,
     );
 }
+
+#[test]
+fn register_attr_and_tool() {
+    cov_mark::check!(register_attr);
+    cov_mark::check!(register_tool);
+    check_no_diagnostics(
+        r#"
+#![register_tool(tool)]
+#![register_attr(attr)]
+
+#[tool::path]
+#[attr]
+struct S;
+        "#,
+    );
+    // NB: we don't currently emit diagnostics here
+}
diff --git a/crates/hir_def/src/nameres/tests/macros.rs b/crates/hir_def/src/nameres/tests/macros.rs
index 6eb5f97be..04de107f5 100644
--- a/crates/hir_def/src/nameres/tests/macros.rs
+++ b/crates/hir_def/src/nameres/tests/macros.rs
@@ -736,6 +736,28 @@ fn unresolved_attributes_fall_back_track_per_file_moditems() {
 }
 
 #[test]
+fn resolves_derive_helper() {
+    cov_mark::check!(resolved_derive_helper);
+    check(
+        r#"
+//- /main.rs crate:main deps:proc
+#[derive(proc::Derive)]
+#[helper]
+#[unresolved]
+struct S;
+
+//- /proc.rs crate:proc
+#[proc_macro_derive(Derive, attributes(helper))]
+fn derive() {}
+        "#,
+        expect![[r#"
+            crate
+            S: t v
+        "#]],
+    )
+}
+
+#[test]
 fn macro_expansion_overflow() {
     cov_mark::check!(macro_expansion_overflow);
     check(