feat: add support for casesOn applications to structural and well-founded recursion modules

RELEASES.md

@@ -1,6 +1,26 @@
 Unreleased
 ---------
 
+* Add support for `casesOn` applications to structural and well-founded recursion modules.
+  This feature is useful when writing definitions using tactics. Example:
+  ```lean
+  inductive Foo where
+    | a | b | c
+    | pair: Foo × Foo → Foo
+
+  def Foo.deq (a b : Foo) : Decidable (a = b) := by
+    cases a <;> cases b
+    any_goals apply isFalse Foo.noConfusion
+    any_goals apply isTrue rfl
+    case pair a b =>
+      let (a₁, a₂) := a
+      let (b₁, b₂) := b
+      exact match deq a₁ b₁, deq a₂ b₂ with
+      | isTrue h₁, isTrue h₂ => isTrue (by rw [h₁,h₂])
+      | isFalse h₁, _ => isFalse (fun h => by cases h; cases (h₁ rfl))
+      | _, isFalse h₂ => isFalse (fun h => by cases h; cases (h₂ rfl))
+  ```
+
 * `Option` is again a monad. The auxiliary type `OptionM` has been removed. See [Zulip thread](https://leanprover.zulipchat.com/#narrow/stream/270676-lean4/topic/Do.20we.20still.20need.20OptionM.3F/near/279761084).
 
 * Improve `split` tactic. It used to fail on `match` expressions of the form `match h : e with ...` where `e` is not a free variable.

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

@@ -4,6 +4,7 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Leonardo de Moura
 -/
 import Lean.Util.HasConstCache
+import Lean.Meta.CasesOn
 import Lean.Meta.Match.Match
 import Lean.Elab.RecAppSyntax
 import Lean.Elab.PreDefinition.Basic
@@ -142,8 +143,7 @@ private partial def replaceRecApps (recFnName : Name) (recArgInfo : RecArgInfo)
             return mkAppN f fArgs
           else
             return mkAppN (← loop below f) (← args.mapM (loop below))
-      let matcherApp? ← matchMatcherApp? e
-      match matcherApp? with
+      match (← matchMatcherApp? e) with
       | some matcherApp =>
         if !recArgHasLooseBVarsAt recFnName recArgInfo.recArgPos e then
           processApp e
@@ -177,6 +177,21 @@ private partial def replaceRecApps (recFnName : Name) (recArgInfo : RecArgInfo)
                 let belowForAlt := xs[numParams - 1]
                 mkLambdaFVars xs (← loop belowForAlt altBody)
             pure { matcherApp with alts := altsNew }.toExpr
+      | none =>
+      match (← toCasesOnApp? e) with
+      | some casesOnApp =>
+        if !recArgHasLooseBVarsAt recFnName recArgInfo.recArgPos e then
+          processApp e
+        else if let some casesOnApp ← casesOnApp.addArg? below (checkIfRefined := true) then
+          let altsNew ← (Array.zip casesOnApp.alts casesOnApp.altNumParams).mapM fun (alt, numParams) =>
+            lambdaTelescope alt fun xs altBody => do
+              unless xs.size >= numParams do
+                throwError "unexpected `casesOn` application alternative{indentExpr alt}\nat application{indentExpr e}"
+              let belowForAlt := xs[numParams]
+              mkLambdaFVars xs (← loop belowForAlt altBody)
+          return { casesOnApp with alts := altsNew }.toExpr
+        else
+          processApp e
       | none => processApp e
     | e => ensureNoRecFn recFnName e
   loop below e |>.run' {}

src/Lean/Elab/PreDefinition/WF/Fix.lean

@@ -4,6 +4,7 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Leonardo de Moura
 -/
 import Lean.Util.HasConstCache
+import Lean.Meta.CasesOn
 import Lean.Meta.Match.Match
 import Lean.Meta.Tactic.Simp.Main
 import Lean.Meta.Tactic.Cleanup
@@ -30,7 +31,9 @@ private def mkDecreasingProof (decreasingProp : Expr) (decrTactic? : Option Synt
   | some decrTactic => Term.runTactic mvarId decrTactic
   instantiateMVars mvar
 
-private partial def replaceRecApps (recFnName : Name) (fixedPrefixSize : Nat) (decrTactic? : Option Syntax) (F : Expr) (e : Expr) : TermElabM Expr :=
+private partial def replaceRecApps (recFnName : Name) (fixedPrefixSize : Nat) (decrTactic? : Option Syntax) (F : Expr) (e : Expr) : TermElabM Expr := do
+  trace[Elab.definition.wf] "replaceRecApps:{indentExpr e}"
+  trace[Elab.definition.wf] "{F} : {← inferType F}"
   loop F e |>.run' {}
 where
   processRec (F : Expr) (e : Expr) : StateRefT (HasConstCache recFnName) TermElabM Expr := do
@@ -73,8 +76,7 @@ where
     | Expr.proj n i e _  => return mkProj n i (← loop F e)
     | Expr.const .. => if e.isConstOf recFnName then processRec F e else return e
     | Expr.app .. =>
-      let matcherApp? ← matchMatcherApp? e
-      match matcherApp? with
+      match (← matchMatcherApp? e) with
       | some matcherApp =>
         if !Structural.recArgHasLooseBVarsAt recFnName fixedPrefixSize e then
           processApp F e
@@ -90,6 +92,23 @@ where
                 let FAlt := xs[numParams - 1]
                 mkLambdaFVars xs (← loop FAlt altBody)
             return { matcherApp with alts := altsNew, discrs := (← matcherApp.discrs.mapM (loop F)) }.toExpr
+      | none =>
+      match (← toCasesOnApp? e) with
+      | some casesOnApp =>
+        if !Structural.recArgHasLooseBVarsAt recFnName fixedPrefixSize e then
+          processApp F e
+        else if let some casesOnApp ← casesOnApp.addArg? F (checkIfRefined := true) then
+          let altsNew ← (Array.zip casesOnApp.alts casesOnApp.altNumParams).mapM fun (alt, numParams) =>
+            lambdaTelescope alt fun xs altBody => do
+              unless xs.size >= numParams do
+                throwError "unexpected `casesOn` application alternative{indentExpr alt}\nat application{indentExpr e}"
+              let FAlt := xs[numParams]
+              mkLambdaFVars xs (← loop FAlt altBody)
+          return { casesOnApp with
+                   alts      := altsNew
+                   remaining := (← casesOnApp.remaining.mapM (loop F)) }.toExpr
+        else
+          processApp F e
       | none => processApp F e
     | e => ensureNoRecFn recFnName e
 

src/Lean/Meta.lean

@@ -39,3 +39,4 @@ import Lean.Meta.Structure
 import Lean.Meta.Constructions
 import Lean.Meta.CongrTheorems
 import Lean.Meta.Eqns
+import Lean.Meta.CasesOn

src/Lean/Meta/CasesOn.lean

@@ -0,0 +1,118 @@
+/-
+Copyright (c) 2022 Microsoft Corporation. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Leonardo de Moura
+-/
+import Lean.Meta.KAbstract
+import Lean.Meta.Check
+
+namespace Lean.Meta
+
+structure CasesOnApp where
+  declName     : Name
+  us           : List Level
+  params       : Array Expr
+  motive       : Expr
+  indices      : Array Expr
+  major        : Expr
+  alts         : Array Expr
+  altNumParams : Array Nat
+  remaining    : Array Expr
+  /-- `true` if the `casesOn` can only eliminate into `Prop` -/
+  propOnly     : Bool
+
+/-- Return `some c` if `e` is a `casesOn` application. -/
+def toCasesOnApp? (e : Expr) : MetaM (Option CasesOnApp) := do
+  let f := e.getAppFn
+  let .const declName us _ := f | return none
+  unless isCasesOnRecursor (← getEnv) declName do return none
+  let indName := declName.getPrefix
+  let .inductInfo info ← getConstInfo indName | return none
+  let args := e.getAppArgs
+  unless args.size >= info.numParams + 1 /- motive -/ + info.numIndices + 1 /- major -/ + info.numCtors do return none
+  let params    := args[:info.numParams]
+  let motive    := args[info.numParams]
+  let indices   := args[info.numParams + 1 : info.numParams + 1 + info.numIndices]
+  let major     := args[info.numParams + 1 + info.numIndices]
+  let alts      := args[info.numParams + 1 + info.numIndices + 1 : info.numParams + 1 + info.numIndices + 1 + info.numCtors]
+  let remaining := args[info.numParams + 1 + info.numIndices + 1 + info.numCtors :]
+  let propOnly  := info.levelParams.length == us.length
+  let mut altNumParams := #[]
+  for ctor in info.ctors do
+    let .ctorInfo ctorInfo ← getConstInfo ctor | unreachable!
+    altNumParams := altNumParams.push ctorInfo.numFields
+  return some { declName, us, params, motive, indices, major, alts, remaining, propOnly, altNumParams }
+
+/-- Convert `c` back to `Expr` -/
+def CasesOnApp.toExpr (c : CasesOnApp) : Expr :=
+  mkAppN (mkAppN (mkApp (mkAppN (mkApp (mkAppN (mkConst c.declName c.us) c.params) c.motive) c.indices) c.major) c.alts) c.remaining
+
+/--
+  Given a `casesOn` application `c` of the form
+  ```
+  casesOn As (fun is x => motive[i, xs]) is major  (fun ys_1 => (alt_1 : motive (C_1[ys_1])) ... (fun ys_n => (alt_n : motive (C_n[ys_n]) remaining
+  ```
+  and an expression `e : B[is, major]`, construct the term
+  ```
+  casesOn As (fun is x => B[is, x] → motive[i, xs]) is major (fun ys_1 (y : B[C_1[ys_1]]) => (alt_1 : motive (C_1[ys_1])) ... (fun ys_n (y : B[C_n[ys_n]]) => (alt_n : motive (C_n[ys_n]) e remaining
+  ```
+  We use `kabstract` to abstract the `is` and `major` from `B[is, major]`.
+-/
+def CasesOnApp.addArg (c : CasesOnApp) (arg : Expr) (checkIfRefined : Bool := false) : MetaM CasesOnApp := do
+  lambdaTelescope c.motive fun motiveArgs motiveBody => do
+    unless motiveArgs.size == c.indices.size + 1 do
+      throwError "failed to add argument to `casesOn` application, motive must be lambda expression with #{c.indices.size + 1} binders"
+    let argType ← inferType arg
+    let discrs := c.indices ++ #[c.major]
+    let mut argTypeAbst := argType
+    for motiveArg in motiveArgs.reverse, discr in discrs.reverse do
+      argTypeAbst := (← kabstract argTypeAbst discr).instantiate1 motiveArg
+    let motiveBody ← mkArrow argTypeAbst motiveBody
+    let us ← if c.propOnly then pure c.us else pure ((← getLevel motiveBody) :: c.us.tail!)
+    let motive ← mkLambdaFVars motiveArgs motiveBody
+    let remaining := #[arg] ++ c.remaining
+    let aux := mkAppN (mkConst c.declName us) c.params
+    let aux := mkApp aux motive
+    let aux := mkAppN aux discrs
+    trace[Meta.debug] "params: {c.params}\nmotive: {motive}\ndiscrs: {discrs}\nindices:{c.indices}\nmajor:{c.major}"
+    check aux
+    unless (← isTypeCorrect aux) do
+      throwError "failed to add argument to `casesOn` application, type error when constructing the new motive{indentExpr aux}"
+    let auxType ← inferType aux
+    let alts ← updateAlts argType auxType
+    return { c with us, motive, alts, remaining }
+where
+  updateAlts (argType : Expr) (auxType : Expr) : MetaM (Array Expr) := do
+    let indName := c.declName.getPrefix
+    let mut auxType := auxType
+    let mut altsNew := #[]
+    let mut refined := false
+    for alt in c.alts, numParams in c.altNumParams do
+      auxType ← whnfD auxType
+      match auxType with
+      | .forallE n d b _ =>
+        let (altNew, refinedAt) ← forallBoundedTelescope d (some numParams) fun xs d => do
+          forallBoundedTelescope d (some 1) fun x d => do
+            let alt := alt.beta xs
+            let alt ← mkLambdaFVars x alt -- x is the new argument we are adding to the alternative
+            if checkIfRefined then
+              return (← mkLambdaFVars xs alt, !(← isDefEq argType (← inferType x[0])))
+            else
+              return (← mkLambdaFVars xs alt, true)
+        if refinedAt then
+          refined := true
+        auxType := b.instantiate1 altNew
+        altsNew := altsNew.push altNew
+      | _ => throwError "unexpected type at `casesOnAddArg`"
+    unless refined do
+      throwError "failed to add argument to `casesOn` application, argument type was not refined by `casesOn`"
+    return altsNew
+
+/-- Similar `CasesOnApp.addArg`, but returns `none` on failure. -/
+def CasesOnApp.addArg? (c : CasesOnApp) (arg : Expr) (checkIfRefined : Bool := false) : MetaM (Option CasesOnApp) :=
+  try
+    return some (← c.addArg arg checkIfRefined)
+  catch _ =>
+    return none
+
+end Lean.Meta

tests/lean/run/casesRec.lean

@@ -0,0 +1,51 @@
+namespace Ex1
+def f (x : Nat) : Nat := by
+  cases x with
+  | zero => exact 1
+  | succ x' =>
+    apply Nat.mul 2
+    exact f x'
+
+#eval f 10
+
+example : f x.succ = 2 * f x := rfl
+end Ex1
+
+namespace Ex2
+inductive Foo where
+  | mk : List Foo → Foo
+
+mutual
+def g (x : Foo) : Nat := by
+  cases x with
+  | mk xs => exact gs xs
+
+def gs (xs : List Foo) : Nat := by
+  cases xs with
+  | nil => exact 1
+  | cons x xs =>
+    apply Nat.add
+    exact g x
+    exact gs xs
+end
+end Ex2
+
+namespace Ex3
+
+inductive Foo where
+  | a | b | c
+  | pair: Foo × Foo → Foo
+
+def Foo.deq (a b : Foo) : Decidable (a = b) := by
+  cases a <;> cases b
+  any_goals apply isFalse Foo.noConfusion
+  any_goals apply isTrue rfl
+  case pair a b =>
+    let (a₁, a₂) := a
+    let (b₁, b₂) := b
+    exact match deq a₁ b₁, deq a₂ b₂ with
+    | isTrue h₁, isTrue h₂ => isTrue (by rw [h₁,h₂])
+    | isFalse h₁, _ => isFalse (fun h => by cases h; cases (h₁ rfl))
+    | _, isFalse h₂ => isFalse (fun h => by cases h; cases (h₂ rfl))
+
+end Ex3