feat: add [grind cases] and [grind cases eager] attributes

src/Init/Grind/Tactics.lean

@@ -14,10 +14,12 @@ syntax grindEqRhs  := atomic("=" "_")
 syntax grindEqBwd  := atomic("←" "=")
 syntax grindBwd    := "←"
 syntax grindFwd    := "→"
+syntax grindCases  := &"cases"
+syntax grindCasesEager := atomic(&"cases" &"eager")
 
-syntax grindThmMod := grindEqBoth <|> grindEqRhs <|> grindEq <|> grindEqBwd <|> grindBwd <|> grindFwd
+syntax grindMod := grindEqBoth <|> grindEqRhs <|> grindEq <|> grindEqBwd <|> grindBwd <|> grindFwd <|> grindCasesEager <|> grindCases
 
-syntax (name := grind) "grind" (grindThmMod)? : attr
+syntax (name := grind) "grind" (grindMod)? : attr
 
 end Lean.Parser.Attr
 
@@ -65,7 +67,7 @@ namespace Lean.Parser.Tactic
 -/
 
 syntax grindErase := "-" ident
-syntax grindLemma := (Attr.grindThmMod)? ident
+syntax grindLemma := (Attr.grindMod)? ident
 syntax grindParam := grindErase <|> grindLemma
 
 syntax (name := grind)

src/Lean/Elab/Tactic/Grind.lean

@@ -51,41 +51,52 @@ def elabGrindParams (params : Grind.Params) (ps :  TSyntaxArray ``Parser.Tactic.
     match p with
     | `(Parser.Tactic.grindParam| - $id:ident) =>
       let declName ← realizeGlobalConstNoOverloadWithInfo id
-      if (← isInductivePredicate declName) then
-        throwErrorAt p "NIY"
+      if (← Grind.isCasesAttrCandidate declName false) then
+        params := { params with casesTypes := (← params.casesTypes.eraseDecl declName) }
       else
         params := { params with ematch := (← params.ematch.eraseDecl declName) }
-    | `(Parser.Tactic.grindParam| $[$mod?:grindThmMod]? $id:ident) =>
+    | `(Parser.Tactic.grindParam| $[$mod?:grindMod]? $id:ident) =>
       let declName ← realizeGlobalConstNoOverloadWithInfo id
-      let kind ← if let some mod := mod? then Grind.getTheoremKindCore mod else pure .default
-      if (← isInductivePredicate declName) then
-        throwErrorAt p "NIY"
-      else
-        let info ← getConstInfo declName
-        match info with
-        | .thmInfo _ =>
-          if kind == .eqBoth then
-            params := { params with extra := params.extra.push (← Grind.mkEMatchTheoremForDecl declName .eqLhs) }
-            params := { params with extra := params.extra.push (← Grind.mkEMatchTheoremForDecl declName .eqRhs) }
-          else
-            params := { params with extra := params.extra.push (← Grind.mkEMatchTheoremForDecl declName kind) }
-        | .defnInfo _ =>
-          if (← isReducible declName) then
-            throwErrorAt p "`{declName}` is a reducible definition, `grind` automatically unfolds them"
-          if kind != .eqLhs && kind != .default then
-            throwErrorAt p "invalid `grind` parameter, `{declName}` is a definition, the only acceptable (and redundant) modifier is '='"
-          let some thms ← Grind.mkEMatchEqTheoremsForDef? declName
-            | throwErrorAt p "failed to genereate equation theorems for `{declName}`"
-          params := { params with extra := params.extra ++ thms.toPArray' }
-        | _ =>
-          throwErrorAt p "invalid `grind` parameter, `{declName}` is not a theorem, definition, or inductive type"
+      let kind ← if let some mod := mod? then Grind.getAttrKindCore mod else pure .infer
+      match kind with
+      | .ematch kind =>
+        params ← withRef p <| addEMatchTheorem params declName kind
+      | .cases eager =>
+        withRef p <| Grind.validateCasesAttr declName eager
+        params := { params with casesTypes := params.casesTypes.insert declName eager }
+      | .infer =>
+        if (← Grind.isCasesAttrCandidate declName false) then
+          params := { params with casesTypes := params.casesTypes.insert declName false }
+        else
+          params ← withRef p <| addEMatchTheorem params declName .default
     | _ => throwError "unexpected `grind` parameter{indentD p}"
   return params
+where
+  addEMatchTheorem (params : Grind.Params) (declName : Name) (kind : Grind.TheoremKind) : MetaM Grind.Params := do
+    let info ← getConstInfo declName
+    match info with
+    | .thmInfo _ =>
+      if kind == .eqBoth then
+        let params := { params with extra := params.extra.push (← Grind.mkEMatchTheoremForDecl declName .eqLhs) }
+        return { params with extra := params.extra.push (← Grind.mkEMatchTheoremForDecl declName .eqRhs) }
+      else
+        return { params with extra := params.extra.push (← Grind.mkEMatchTheoremForDecl declName kind) }
+    | .defnInfo _ =>
+      if (← isReducible declName) then
+        throwError "`{declName}` is a reducible definition, `grind` automatically unfolds them"
+      if kind != .eqLhs && kind != .default then
+        throwError "invalid `grind` parameter, `{declName}` is a definition, the only acceptable (and redundant) modifier is '='"
+      let some thms ← Grind.mkEMatchEqTheoremsForDef? declName
+        | throwError "failed to genereate equation theorems for `{declName}`"
+      return { params with extra := params.extra ++ thms.toPArray' }
+    | _ =>
+      throwError "invalid `grind` parameter, `{declName}` is not a theorem, definition, or inductive type"
 
 def mkGrindParams (config : Grind.Config) (only : Bool) (ps :  TSyntaxArray ``Parser.Tactic.grindParam) : MetaM Grind.Params := do
   let params ← Grind.mkParams config
   let ematch ← if only then pure {} else Grind.getEMatchTheorems
-  let params := { params with ematch }
+  let casesTypes ← if only then pure {} else Grind.getCasesTypes
+  let params := { params with ematch, casesTypes }
   elabGrindParams params ps
 
 def grind

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

@@ -4,12 +4,11 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Leonardo de Moura
 -/
 prelude
-import Lean.Meta.Tactic.Simp.Attr
-import Lean.Meta.Tactic.Simp.Simproc
+import Lean.Meta.Tactic.Grind.EMatchTheorem
+import Lean.Meta.Tactic.Grind.Cases
 
 namespace Lean.Meta.Grind
-open Simp
-
+--- TODO: delete
 builtin_initialize grindCasesExt : SimpleScopedEnvExtension Name NameSet ←
   registerSimpleScopedEnvExtension {
     initial        := {}
@@ -59,5 +58,68 @@ builtin_initialize
       validateGrindCasesAttr declName
       grindCasesExt.add declName attrKind
   }
+--- END of TODO: detele
+
+inductive AttrKind where
+  | ematch (k : TheoremKind)
+  | cases (eager : Bool)
+  | infer
+
+/-- Return theorem kind for `stx` of the form `Attr.grindThmMod` -/
+def getAttrKindCore (stx : Syntax) : CoreM AttrKind := do
+  match stx with
+  | `(Parser.Attr.grindMod| =) => return .ematch .eqLhs
+  | `(Parser.Attr.grindMod| →) => return .ematch .fwd
+  | `(Parser.Attr.grindMod| ←) => return .ematch .bwd
+  | `(Parser.Attr.grindMod| =_) => return .ematch .eqRhs
+  | `(Parser.Attr.grindMod| _=_) => return .ematch .eqBoth
+  | `(Parser.Attr.grindMod| ←=) => return .ematch .eqBwd
+  | `(Parser.Attr.grindMod| cases) => return .cases false
+  | `(Parser.Attr.grindMod| cases eager) => return .cases true
+  | _ => throwError "unexpected `grind` theorem kind: `{stx}`"
+
+/-- Return theorem kind for `stx` of the form `(Attr.grindMod)?` -/
+def getAttrKindFromOpt (stx : Syntax) : CoreM AttrKind := do
+  if stx[1].isNone then
+    return .infer
+  else
+    getAttrKindCore stx[1][0]
+
+builtin_initialize
+  registerBuiltinAttribute {
+    name := `grind
+    descr :=
+      "The `[grind]` attribute is used to annotate declarations.\
+      \
+      When applied to an equational theorem, `[grind =]`, `[grind =_]`, or `[grind _=_]`\
+      will mark the theorem for use in heuristic instantiations by the `grind` tactic,
+      using respectively the left-hand side, the right-hand side, or both sides of the theorem.\
+      When applied to a function, `[grind =]` automatically annotates the equational theorems associated with that function.\
+      When applied to a theorem `[grind ←]` will instantiate the theorem whenever it encounters the conclusion of the theorem
+      (that is, it will use the theorem for backwards reasoning).\
+      When applied to a theorem `[grind →]` will instantiate the theorem whenever it encounters sufficiently many of the propositional hypotheses
+      (that is, it will use the theorem for forwards reasoning).\
+      \
+      The attribute `[grind]` by itself will effectively try `[grind ←]` (if the conclusion is sufficient for instantiation) and then `[grind →]`.\
+      \
+      The `grind` tactic utilizes annotated theorems to add instances of matching patterns into the local context during proof search.\
+      For example, if a theorem `@[grind =] theorem foo_idempotent : foo (foo x) = foo x` is annotated,\
+      `grind` will add an instance of this theorem to the local context whenever it encounters the pattern `foo (foo x)`."
+    applicationTime := .afterCompilation
+    add := fun declName stx attrKind => MetaM.run' do
+      match (← getAttrKindFromOpt stx) with
+      | .ematch k => addEMatchAttr declName attrKind k
+      | .cases eager => addCasesAttr declName eager attrKind
+      | .infer =>
+        if (← isCasesAttrCandidate declName false) then
+          addCasesAttr declName false attrKind
+        else
+          addEMatchAttr declName attrKind .default
+    erase := fun declName => MetaM.run' do
+      if (← isCasesAttrCandidate declName false) then
+        eraseCasesAttr declName
+      else
+        eraseEMatchAttr declName
+  }
 
 end Lean.Meta.Grind

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

@@ -7,6 +7,79 @@ prelude
 import Lean.Meta.Tactic.Cases
 
 namespace Lean.Meta.Grind
+
+/-- Types that `grind` will case-split on. -/
+structure CasesTypes where
+  casesMap : PHashMap Name Bool := {}
+  deriving Inhabited
+
+structure CasesEntry where
+  declName : Name
+  eager : Bool
+  deriving Inhabited
+
+/-- Returns `true` if `s` contains a `declName`. -/
+def CasesTypes.contains (s : CasesTypes) (declName : Name) : Bool :=
+  s.casesMap.contains declName
+
+/-- Removes the given declaration from `s`. -/
+def CasesTypes.erase (s : CasesTypes) (declName : Name) : CasesTypes :=
+  { s with casesMap := s.casesMap.erase declName }
+
+def CasesTypes.insert (s : CasesTypes) (declName : Name) (eager : Bool) : CasesTypes :=
+  { s with casesMap := s.casesMap.insert declName eager }
+
+def CasesTypes.find? (s : CasesTypes) (declName : Name) : Option Bool :=
+  s.casesMap.find? declName
+
+builtin_initialize casesExt : SimpleScopedEnvExtension CasesEntry CasesTypes ←
+  registerSimpleScopedEnvExtension {
+    initial        := {}
+    addEntry       := fun s {declName, eager} => s.insert declName eager
+  }
+
+def getCasesTypes : CoreM CasesTypes :=
+  return casesExt.getState (← getEnv)
+
+private def getAlias? (value : Expr) : MetaM (Option Name) :=
+  lambdaTelescope value fun _ body => do
+    if let .const declName _ := body.getAppFn' then
+      return some declName
+    else
+      return none
+
+partial def isCasesAttrCandidate (declName : Name) (eager : Bool) : CoreM Bool := do
+  match (← getConstInfo declName) with
+  | .inductInfo info => return !info.isRec || !eager
+  | .defnInfo info =>
+    let some declName ← getAlias? info.value |>.run' {} {}
+      | return false
+    isCasesAttrCandidate declName eager
+  | _ => return false
+
+def validateCasesAttr (declName : Name) (eager : Bool) : CoreM Unit := do
+  unless (← isCasesAttrCandidate declName eager) do
+    if eager then
+      throwError "invalid `[grind cases eager]`, `{declName}` is not a non-recursive inductive datatype or an alias for one"
+    else
+      throwError "invalid `[grind cases]`, `{declName}` is not an inductive datatype or an alias for one"
+
+def addCasesAttr (declName : Name) (eager : Bool) (attrKind : AttributeKind) : CoreM Unit := do
+  validateCasesAttr declName eager
+  casesExt.add { declName, eager } attrKind
+
+def CasesTypes.eraseDecl (s : CasesTypes) (declName : Name) : CoreM CasesTypes := do
+  if s.contains declName then
+    return s.erase declName
+  else
+    throwError "`{declName}` is not marked with the `[grind]` attribute"
+
+def eraseCasesAttr (declName : Name) : CoreM Unit := do
+  let s := casesExt.getState (← getEnv)
+  let s ← s.eraseDecl declName
+  modifyEnv fun env => casesExt.modifyState env fun _ => s
+
+
 /--
 The `grind` tactic includes an auxiliary `cases` tactic that is not intended for direct use by users.
 This method implements it.