refactor: structural recursion: prove .eq_def directly

src/Init/Data/Dyadic/Basic.lean

@@ -9,6 +9,7 @@ prelude
 public import Init.Data.Rat.Lemmas
 import Init.Data.Int.Bitwise.Lemmas
 import Init.Data.Int.DivMod.Lemmas
+import Init.Hints
 
 /-!
 # The dyadic rationals

src/Lean/Elab/PreDefinition/Eqns.lean

@@ -101,10 +101,11 @@ partial def splitMatch? (mvarId : MVarId) (declNames : Array Name) : MetaM (Opti
                                        target declNames badCases then
       try
         Meta.Split.splitMatch mvarId e
-      catch _ =>
+      catch ex =>
+        trace[Elab.definition.eqns] "cannot split {e}\n{ex.toMessageData}"
         go (badCases.insert e)
     else
-      trace[Meta.Tactic.split] "did not find term to split\n{MessageData.ofGoal mvarId}"
+      trace[Elab.definition.eqns] "did not find term to split\n{MessageData.ofGoal mvarId}"
       return none
   go {}
 
@@ -288,7 +289,7 @@ public def deltaLHS (mvarId : MVarId) : MetaM MVarId := mvarId.withContext do
   let some lhs ← delta? lhs | throwTacticEx `deltaLHS mvarId "failed to delta reduce lhs"
   mvarId.replaceTargetDefEq (← mkEq lhs rhs)
 
-def deltaRHS? (mvarId : MVarId) (declName : Name) : MetaM (Option MVarId) := mvarId.withContext do
+public def deltaRHS? (mvarId : MVarId) (declName : Name) : MetaM (Option MVarId) := mvarId.withContext do
   let target ← mvarId.getType'
   let some (_, lhs, rhs) := target.eq? | return none
   let some rhs ← delta? rhs.consumeMData (· == declName) | return none
@@ -347,7 +348,7 @@ private def unfoldLHS (declName : Name) (mvarId : MVarId) : MetaM MVarId := mvar
     deltaLHS mvarId
 
 private partial def mkEqnProof (declName : Name) (type : Expr) (tryRefl : Bool) : MetaM Expr := do
-  trace[Elab.definition.eqns] "proving: {type}"
+  withTraceNode `Elab.definition.eqns (return m!"{exceptEmoji ·} proving:{indentExpr type}") do
   withNewMCtxDepth do
     let main ← mkFreshExprSyntheticOpaqueMVar type
     let (_, mvarId) ← main.mvarId!.intros
@@ -371,7 +372,7 @@ private partial def mkEqnProof (declName : Name) (type : Expr) (tryRefl : Bool)
   recursion and structural recursion can and should use this too.
   -/
   go (mvarId : MVarId) : MetaM Unit := do
-    trace[Elab.definition.eqns] "step\n{MessageData.ofGoal mvarId}"
+    withTraceNode `Elab.definition.eqns (return m!"{exceptEmoji ·} step:\n{MessageData.ofGoal mvarId}") do
     if (← tryURefl mvarId) then
       return ()
     else if (← tryContradiction mvarId) then

src/Lean/Elab/PreDefinition/PartialFixpoint/Eqns.lean

@@ -117,7 +117,7 @@ def getUnfoldFor? (declName : Name) : MetaM (Option Name) := do
 
 def getEqnsFor? (declName : Name) : MetaM (Option (Array Name)) := do
   if let some info := eqnInfoExt.find? (← getEnv) declName then
-    mkEqns declName info.declNames
+    mkEqns declName info.declNames (tryRefl := false)
   else
     return none
 

src/Lean/Elab/PreDefinition/Structural/Eqns.lean

@@ -16,6 +16,7 @@ import Lean.Meta.Tactic.Apply
 import Lean.Elab.PreDefinition.Basic
 import Lean.Elab.PreDefinition.Structural.Basic
 import Lean.Meta.Match.MatchEqs
+import Lean.Meta.Tactic.Rewrite
 
 namespace Lean.Elab
 open Meta
@@ -29,114 +30,115 @@ public structure EqnInfo extends EqnInfoCore where
   fixedParamPerms : FixedParamPerms
   deriving Inhabited
 
-private partial def mkProof (declName : Name) (type : Expr) : MetaM Expr := do
+/--
+Searches in the lhs of goal for a `.brecOn` application, possibly with extra arguments
+and under `PProd` projections. Returns the `.brecOn` application and the context
+`(fun x => (x).1.2.3 extraArgs = rhs)`.
+-/
+partial def findBRecOnLHS (goal : Expr) : MetaM (Expr × Expr) := do
+  let some (_, lhs, rhs) := goal.eq? | throwError "goal not an equality{indentExpr goal}"
+  go lhs fun brecOnApp x c =>
+    return (brecOnApp, ← mkLambdaFVars #[x] (← mkEq c rhs))
+where
+  go {α} (e : Expr) (k : Expr → Expr → Expr → MetaM α) : MetaM α := e.withApp fun f xs => do
+    if let .proj t n e := f then
+      return ← go e fun brecOnApp x c => k brecOnApp x (mkAppN (mkProj t n c) xs)
+    if let .const name _ := f then
+      if isBRecOnRecursor (← getEnv) name then
+        let arity ← forallTelescope (← inferType f) fun xs _ => return xs.size
+        if arity ≤ xs.size then
+          let brecOnApp := mkAppN f xs[:arity]
+          let extraArgs := xs[arity:]
+          return ← withLocalDeclD `x (← inferType brecOnApp) fun x =>
+            k brecOnApp x (mkAppN x extraArgs)
+    throwError "could not find `.brecOn` application in{indentExpr e}"
+
+/--
+Creates the proof of the unfolding theorem for `declName` with type `type`. It
+
+1. unfolds the function on the left to expose the `.brecOn` application
+2. rewrites that using the `.brecOn.eq` theorem, unrolling it once
+3. let-binds the last argument, which should be the `.brecOn.go` call of type `.below …`.
+   This way subsequent steps (which may involve `simp`) do not touch it and do
+   not break the definitional equality with the recursive calls on the RHS.
+4. repeatedly splits `match` statements (because on the left we have `match` statements with extra
+   `.below` arguments, and on the right we have the original `match` statements) until the goal
+   is solved using `rfl` or `contradiction`.
+-/
+partial def mkProof (declName : Name) (type : Expr) : MetaM Expr := do
   withTraceNode `Elab.definition.structural.eqns (return m!"{exceptEmoji ·} proving:{indentExpr type}") do
+    prependError m!"failed to generate equational theorem for `{.ofConstName declName}`" do
     withNewMCtxDepth do
       let main ← mkFreshExprSyntheticOpaqueMVar type
       let (_, mvarId) ← main.mvarId!.intros
       unless (← tryURefl mvarId) do -- catch easy cases
-        go1 mvarId
+        goUnfold (← deltaLHS mvarId)
       instantiateMVars main
 where
-  /--
-  Step 1: Split the function body into its cases, but keeping the LHS intact, because the
-  `.below`-added `match` statements and the `.rec` can quickly confuse `split`.
-  -/
-  go1 (mvarId : MVarId) : MetaM Unit := do
-    withTraceNode `Elab.definition.structural.eqns (return m!"{exceptEmoji ·} go1:\n{MessageData.ofGoal mvarId}") do
+  goUnfold (mvarId : MVarId) : MetaM Unit := do
+    withTraceNode `Elab.definition.structural.eqns (return m!"{exceptEmoji ·} goUnfold:\n{MessageData.ofGoal mvarId}") do
+    let mvarId' ← mvarId.withContext do
+      -- This should now be headed by `.brecOn`
+      let goal ← mvarId.getType
+      let (brecOnApp, context) ← findBRecOnLHS goal
+      let brecOnName := brecOnApp.getAppFn.constName!
+      let us := brecOnApp.getAppFn.constLevels!
+      let brecOnThmName := brecOnName.str "eq"
+      let brecOnAppArgs := brecOnApp.getAppArgs
+      unless (← hasConst brecOnThmName) do
+        throwError "no theorem `{brecOnThmName}`\n{MessageData.ofGoal mvarId}"
+      -- We don't just `← inferType eqThmApp` as that beta-reduces more than we want
+      let eqThmType ← inferType (mkConst brecOnThmName us)
+      let eqThmType ← instantiateForall eqThmType brecOnAppArgs
+      let some (_, _, rwRhs) := eqThmType.eq? | throwError "theorem `{brecOnThmName}` is not an equality\n{MessageData.ofGoal mvarId}"
+      let recArg := rwRhs.getAppArgs.back!
+      trace[Elab.definition.structural.eqns] "abstracting{inlineExpr recArg} from{indentExpr rwRhs}"
+      let mvarId2 ← mvarId.define `r (← inferType recArg) recArg
+      let (r, mvarId3) ← mvarId2.intro1P
+      let mvarId4 ← mvarId3.withContext do
+        let goal' := mkApp rwRhs.appFn! (mkFVar r)
+        let thm ← mkCongrArg context (mkAppN (mkConst brecOnThmName us) brecOnAppArgs)
+        mvarId3.replaceTargetEq (mkApp context goal') thm
+      pure mvarId4
+    go mvarId'
+
+  go (mvarId : MVarId) : MetaM Unit := do
+    withTraceNode `Elab.definition.structural.eqns (return m!"{exceptEmoji ·} step:\n{MessageData.ofGoal mvarId}") do
       if (← tryURefl mvarId) then
         trace[Elab.definition.structural.eqns] "tryURefl succeeded"
         return ()
       else if (← tryContradiction mvarId) then
         trace[Elab.definition.structural.eqns] "tryContadiction succeeded"
         return ()
+      else if let some mvarId ← whnfReducibleLHS? mvarId then
+        trace[Elab.definition.structural.eqns] "whnfReducibleLHS succeeded"
+        go mvarId
       else if let some mvarId ← simpMatch? mvarId then
         trace[Elab.definition.structural.eqns] "simpMatch? succeeded"
-        go1 mvarId
+        go mvarId
       else if let some mvarId ← simpIf? mvarId (useNewSemantics := true) then
         trace[Elab.definition.structural.eqns] "simpIf? succeeded"
-        go1 mvarId
+        go mvarId
       else
         let ctx ← Simp.mkContext
         match (← simpTargetStar mvarId ctx (simprocs := {})).1 with
         | TacticResultCNM.closed =>
           trace[Elab.definition.structural.eqns] "simpTargetStar closed the goal"
         | TacticResultCNM.modified mvarId =>
           trace[Elab.definition.structural.eqns] "simpTargetStar modified the goal"
-          go1 mvarId
+          go mvarId
         | TacticResultCNM.noChange =>
-          if let some mvarIds ← casesOnStuckLHS? mvarId then
+          if let some mvarId ← deltaRHS? mvarId declName then
+            trace[Elab.definition.structural.eqns] "deltaRHS? succeeded"
+            go mvarId
+          else if let some mvarIds ← casesOnStuckLHS? mvarId then
             trace[Elab.definition.structural.eqns] "casesOnStuckLHS? succeeded"
-            mvarIds.forM go1
+            mvarIds.forM go
           else if let some mvarIds ← splitTarget? mvarId (useNewSemantics := true) then
             trace[Elab.definition.structural.eqns] "splitTarget? succeeded"
-            mvarIds.forM go1
-          else
-            go2 (← deltaLHS mvarId)
-  /-- Step 2: Unfold the lhs to expose the recursor. -/
-  go2 (mvarId : MVarId) : MetaM Unit := do
-    withTraceNode `Elab.definition.structural.eqns (return m!"{exceptEmoji ·} go2:\n{MessageData.ofGoal mvarId}") do
-    if let some mvarId ← whnfReducibleLHS? mvarId then
-      go2 mvarId
-    else
-      go3 mvarId
-  /-- Step 3: Simplify the match and if statements on the left hand side, until we have rfl. -/
-  go3 (mvarId : MVarId) : MetaM Unit := do
-      withTraceNode `Elab.definition.structural.eqns (return m!"{exceptEmoji ·} go3:\n{MessageData.ofGoal mvarId}") do
-      if (← tryURefl mvarId) then
-        trace[Elab.definition.structural.eqns] "tryURefl succeeded"
-        return ()
-      else if (← tryContradiction mvarId) then
-        trace[Elab.definition.structural.eqns] "tryContadiction succeeded"
-        return ()
-      else if let some mvarId ← simpMatch? mvarId then
-        trace[Elab.definition.structural.eqns] "simpMatch? succeeded"
-        go3 mvarId
-      else if let some mvarId ← simpIf? mvarId (useNewSemantics := true) then
-        trace[Elab.definition.structural.eqns] "simpIf? succeeded"
-        go3 mvarId
-      else
-        let ctx ← Simp.mkContext
-        match (← simpTargetStar mvarId ctx (simprocs := {})).1 with
-        | TacticResultCNM.closed =>
-          trace[Elab.definition.structural.eqns] "simpTargetStar closed the goal"
-        | TacticResultCNM.modified mvarId =>
-          trace[Elab.definition.structural.eqns] "simpTargetStar modified the goal"
-          go3 mvarId
-        | TacticResultCNM.noChange =>
-          if let some mvarIds ← casesOnStuckLHS? mvarId then
-            trace[Elab.definition.structural.eqns] "casesOnStuckLHS? succeeded"
-            mvarIds.forM go3
+            mvarIds.forM go
           else
-            throwError "failed to generate equational theorem for `{.ofConstName declName}`\n{MessageData.ofGoal mvarId}"
-
-def mkEqns (info : EqnInfo) : MetaM (Array Name) :=
-  withOptions (tactic.hygienic.set · false) do
-  let eqnTypes ← withNewMCtxDepth <| lambdaTelescope (cleanupAnnotations := true) info.value fun xs body => do
-    let us := info.levelParams.map mkLevelParam
-    let target ← mkEq (mkAppN (Lean.mkConst info.declName us) xs) body
-    let goal ← mkFreshExprSyntheticOpaqueMVar target
-    mkEqnTypes info.declNames goal.mvarId!
-  let mut thmNames := #[]
-  for h : i in *...eqnTypes.size do
-    let type := eqnTypes[i]
-    trace[Elab.definition.structural.eqns] "eqnType {i+1}: {type}"
-    let name := mkEqLikeNameFor (← getEnv) info.declName s!"{eqnThmSuffixBasePrefix}{i+1}"
-    thmNames := thmNames.push name
-    -- determinism: `type` should be independent of the environment changes since `baseName` was
-    -- added
-    realizeConst info.declNames[0]! name (doRealize name type)
-  return thmNames
-where
-  doRealize name type := withOptions (tactic.hygienic.set · false) do
-    let value ← withoutExporting do mkProof info.declName type
-    let (type, value) ← removeUnusedEqnHypotheses type value
-    let type ← letToHave type
-    addDecl <| Declaration.thmDecl {
-      name, type, value
-      levelParams := info.levelParams
-    }
-    inferDefEqAttr name
+            throwError "no progress at goal\n{MessageData.ofGoal mvarId}"
 
 public builtin_initialize eqnInfoExt : MapDeclarationExtension EqnInfo ←
   mkMapDeclarationExtension (exportEntriesFn := fun env s _ =>
@@ -151,7 +153,7 @@ public def registerEqnsInfo (preDef : PreDefinition) (declNames : Array Name) (r
 
 def getEqnsFor? (declName : Name) : MetaM (Option (Array Name)) := do
   if let some info := eqnInfoExt.find? (← getEnv) declName then
-    mkEqns info
+    mkEqns declName info.declNames
   else
     return none
 
@@ -165,11 +167,10 @@ where
     lambdaTelescope info.value fun xs body => do
       let us := info.levelParams.map mkLevelParam
       let type ← mkEq (mkAppN (Lean.mkConst declName us) xs) body
-      let goal ← mkFreshExprSyntheticOpaqueMVar type
-      mkUnfoldProof declName goal.mvarId!
+      let value ← withoutExporting <| mkProof declName type
       let type ← mkForallFVars xs type
       let type ← letToHave type
-      let value ← mkLambdaFVars xs (← instantiateMVars goal)
+      let value ← mkLambdaFVars xs value
       addDecl <| Declaration.thmDecl {
         name, type, value
         levelParams := info.levelParams
@@ -187,6 +188,7 @@ def getStructuralRecArgPosImp? (declName : Name) : CoreM (Option Nat) := do
   let some info := eqnInfoExt.find? (← getEnv) declName | return none
   return some info.recArgPos
 
+
 builtin_initialize
   registerGetEqnsFn getEqnsFor?
   registerGetUnfoldEqnFn getUnfoldFor?