fix: make grind_lint see custom grind attrs

src/Lean/Elab/Tactic/Grind/Lint.lean

@@ -31,8 +31,14 @@ builtin_initialize muteExt : SimplePersistentEnvExtension Name NameSet ←
 
 open Command Meta Grind
 
+def isEMatchTheorem (declName : Name) : CoreM Bool := do
+  for ext in (← Grind.getRegisteredExtensions) do
+    if (← ext.isEMatchTheorem declName) then
+      return true
+  return false
+
 def checkEMatchTheorem (declName : Name) : CoreM Unit := do
-  unless (← Grind.grindExt.isEMatchTheorem declName) do
+  unless (← isEMatchTheorem declName) do
     throwError "`{declName}` is not marked with the `@[grind]` attribute for theorem instantiation"
 
 @[builtin_command_elab Lean.Grind.grindLintSkip]
@@ -159,29 +165,30 @@ def nameEndsWithSuffix (name suff : Name) : Bool :=
 def getTheorems (prefixes? : Option (Array Name)) (inModule : Bool) : CoreM (List Name) := do
   let skip := skipExt.getState (← getEnv)
   let skipSuffixes := skipSuffixExt.getState (← getEnv)
-  let origins := (← Grind.grindExt.getEMatchTheorems).getOrigins
   let env ← getEnv
-  return origins.filterMap fun origin => Id.run do
-    let .decl declName := origin | return none
-    if skip.contains declName then return none
-    -- Check if declName's last component ends with any of the skip suffixes
-    if skipSuffixes.any fun suff => nameEndsWithSuffix declName suff then return none
-    let some prefixes := prefixes? | return some declName
-    if inModule then
-      let some modIdx := env.getModuleIdxFor? declName | return none
-      let modName := env.header.moduleNames[modIdx]!
-      if prefixes.any fun pre => pre.isPrefixOf modName then
-        return some declName
-      else
-        return none
-    else
-      let keep := prefixes.any fun pre =>
-        if pre == `_root_ then
-          declName.components.length == 1
+  let mut decls := #[]
+  for ext in (← Grind.getRegisteredExtensions) do
+    for origin in (← ext.getEMatchTheorems).getOrigins do
+      let .decl declName := origin | continue
+      if skip.contains declName then continue
+      -- Check if declName's last component ends with any of the skip suffixes
+      if skipSuffixes.any fun suff => nameEndsWithSuffix declName suff then continue
+      if let some prefixes := prefixes? then
+        if inModule then
+          let some modIdx := env.getModuleIdxFor? declName | continue
+          let modName := env.header.moduleNames[modIdx]!
+          unless prefixes.any fun pre => pre.isPrefixOf modName do
+            continue
         else
-          pre.isPrefixOf declName
-      unless keep do return none
-      return some declName
+          let keep := prefixes.any fun pre =>
+            if pre == `_root_ then
+              declName.components.length == 1
+            else
+              pre.isPrefixOf declName
+          unless keep do continue
+      unless decls.contains declName do
+        decls := decls.push declName
+  return decls.toList
 
 @[builtin_command_elab Lean.Grind.grindLintCheck]
 def elabGrindLintCheck : CommandElab := fun stx => liftTermElabM <| withTheReader Core.Context (fun c => { c with maxHeartbeats := 0 }) do

src/Lean/Meta/Tactic/Grind/Attr.lean

@@ -237,6 +237,15 @@ def getExtension? (attrName : Name) : CoreM (Option Extension) := do
     recordExtraModUseFromDecl (isMeta := true) ext.ext.name
   return ext?
 
+/-- Returns the registered grind attribute extensions. -/
+def getRegisteredExtensions : CoreM (Array Extension) := do
+  let map ← extensionMapRef.get
+  let mut exts := #[]
+  for (_, ext) in map.toArray do
+    recordExtraModUseFromDecl (isMeta := true) ext.ext.name
+    exts := exts.push ext
+  return exts
+
 def registerAttr (attrName : Name) (ref : Name := by exact decl_name%) : IO Extension := do
   let ext ← mkExtension ref
   mkGrindAttr attrName (minIndexable := false) (showInfo := true) (ext := ext) (ref := ref)

tests/elab/grind_lint_custom_attr.lean

@@ -0,0 +1,13 @@
+import Lean.Elab.Tactic.Grind.Lint
+import Lean.Meta.Tactic.Grind.RegisterCommand
+
+register_grind_attr foo
+
+@[foo =] theorem foo_add_zero (x : Nat) : x + 0 = x := by
+  simp
+
+#grind_lint skip foo_add_zero
+
+/-- error: `foo_add_zero` is already in the `#grind_lint` skip set -/
+#guard_msgs in
+#grind_lint skip foo_add_zero