From 35697973771b8dbbbb6b3872a40cc604764d36ae Mon Sep 17 00:00:00 2001
From: Joachim Breitner <mail@joachim-breitner.de>
Date: Tue, 21 Jan 2025 11:16:42 +0100
Subject: [PATCH] feat: functional cases theorem for non-recursive functions
 (#6261)

This PR adds `foo.fun_cases`, an automatically generated theorem that
splits the goal according to the branching structure of `foo`, much like
the Functional Induction Principle, but for all functions (not just
recursive ones), and without providing inductive hypotheses.

The design isn't quite final yet as to which function parameters should
become targets of the motive, and which parameters of the theorem, but
the current version is already proven to be useful, so start with this
and iterate later.
---
 src/Lean/Meta/Tactic/FunInd.lean | 77 +++++++++++++++++++++++++++++++-
 tests/lean/run/fun_cases.lean    | 32 +++++++++++++
 2 files changed, 108 insertions(+), 1 deletion(-)
 create mode 100644 tests/lean/run/fun_cases.lean

diff --git a/src/Lean/Meta/Tactic/FunInd.lean b/src/Lean/Meta/Tactic/FunInd.lean
index 6cabf6eb1e..ef1a3d45d9 100644
--- a/src/Lean/Meta/Tactic/FunInd.lean
+++ b/src/Lean/Meta/Tactic/FunInd.lean
@@ -1085,6 +1085,63 @@ def deriveInductionStructural (names : Array Name) (numFixed : Nat) : MetaM Unit
   if names.size > 1 then
     projectMutualInduct names inductName
 
+
+/--
+For non-recursive (and recursive functions) functions we derive a “functional case splitting theorem”. This is very similar
+than the functional induction theorem. It splits the goal, but does not give you inductive hyptheses.
+
+For these, it is not really clear which parameters should be “targets” of the motive, as there is
+no “fixed prefix” to guide this decision. All? None? Some?
+
+We tried none, but that did not work well. Right now it's all parameters, and it seems to work well.
+In the future, we might post-process the theorem (or run the code below iteratively) and remove
+targets that are unchanged in each case, so simplify applying the lemma when these “fixed” parameters
+are not variables, to avoid having to generalize them.
+-/
+def deriveCases (name : Name) : MetaM Unit := do
+  mapError (f := (m!"Cannot derive functional cases principle (please report this issue)\n{indentD ·}")) do
+    let info ← getConstInfo name
+    let value ←
+      if let some unfoldEqnName ← getUnfoldEqnFor? (nonRec := false) name then
+        let eqInfo ← getConstInfo unfoldEqnName
+        forallTelescope eqInfo.type fun xs body => do
+          let some (_, _, rhs) := body.eq?
+            | throwError "Type of {unfoldEqnName} not an equality: {body}"
+          mkLambdaFVars xs rhs
+      else if let some value := info.value? then
+        pure value
+      else
+        throwError "'{name}' does not have an unfold theorem nor a value"
+    let motiveType ← lambdaTelescope value fun xs _body => do
+      mkForallFVars xs (.sort 0)
+    let e' ← withLocalDeclD `motive motiveType fun motive => do
+      lambdaTelescope value fun xs body => do
+        let (e',mvars) ← M2.run do
+          let goal := mkAppN motive xs
+          -- We bring an unused FVars into scope to pass as `oldIH` and `newIH`. These do not appear anywhere
+          -- so `buildInductionBody` should just do the right thing
+          withLocalDeclD `fakeIH (mkConst ``Unit) fun fakeIH =>
+            let isRecCall := fun _ => none
+            buildInductionBody #[fakeIH.fvarId!] goal fakeIH.fvarId! fakeIH.fvarId! isRecCall body
+        let e' ← mkLambdaFVars xs e'
+        let e' ← abstractIndependentMVars mvars (← motive.fvarId!.getDecl).index e'
+        let e' ← mkLambdaFVars #[motive] e'
+        pure e'
+
+    unless (← isTypeCorrect e') do
+      logError m!"constructed functional cases principle is not type correct:{indentExpr e'}"
+      check e'
+
+    let eTyp ← inferType e'
+    let eTyp ← elimOptParam eTyp
+    -- logInfo m!"eTyp: {eTyp}"
+    let params := (collectLevelParams {} eTyp).params
+    -- Prune unused level parameters, preserving the original order
+    let us := info.levelParams.filter (params.contains ·)
+
+    addDecl <| Declaration.thmDecl
+      { name := info.name ++ `fun_cases, levelParams := us, type := eTyp, value := e' }
+
 /--
 Given a recursively defined function `foo`, derives `foo.induct`. See the module doc for details.
 -/
@@ -1097,7 +1154,7 @@ def deriveInduction (name : Name) : MetaM Unit := do
     else if let some eqnInfo := Structural.eqnInfoExt.find? (← getEnv) name then
       deriveInductionStructural eqnInfo.declNames eqnInfo.numFixed
     else
-      throwError "{name} is not defined by structural or well-founded recursion"
+      throwError "constant '{name}' is not structurally or well-founded recursive"
 
 def isFunInductName (env : Environment) (name : Name) : Bool := Id.run do
   let .str p s := name | return false
@@ -1117,14 +1174,32 @@ def isFunInductName (env : Environment) (name : Name) : Bool := Id.run do
     return false
   | _ => return false
 
+
+def isFunCasesName (env : Environment) (name : Name) : Bool := Id.run do
+  let .str p s := name | return false
+  match s with
+  | "fun_cases" =>
+    if (WF.eqnInfoExt.find? env p).isSome then return true
+    if (Structural.eqnInfoExt.find? env p).isSome then return true
+    if let some ci := env.find? p then
+      if ci.hasValue then
+        return true
+    return false
+  | _ => return false
+
 builtin_initialize
   registerReservedNamePredicate isFunInductName
+  registerReservedNamePredicate isFunCasesName
 
   registerReservedNameAction fun name => do
     if isFunInductName (← getEnv) name then
       let .str p _ := name | return false
       MetaM.run' <| deriveInduction p
       return true
+    if isFunCasesName (← getEnv) name then
+      let .str p _ := name | return false
+      MetaM.run' <| deriveCases p
+      return true
     return false
 
 end Lean.Tactic.FunInd
diff --git a/tests/lean/run/fun_cases.lean b/tests/lean/run/fun_cases.lean
new file mode 100644
index 0000000000..37967767ef
--- /dev/null
+++ b/tests/lean/run/fun_cases.lean
@@ -0,0 +1,32 @@
+/--
+info: Option.map.fun_cases.{u_1, u_2} (motive : {α : Type u_1} → {β : Type u_2} → (α → β) → Option α → Prop)
+  (case1 : ∀ {α : Type u_1} {β : Type u_2} (f : α → β) (x : α), motive f (some x))
+  (case2 : ∀ {α : Type u_1} {β : Type u_2} (f : α → β), motive f none) {α : Type u_1} {β : Type u_2} (f : α → β)
+  (x✝ : Option α) : motive f x✝
+-/
+#guard_msgs in
+#check Option.map.fun_cases
+
+example (x : Option Nat) (f : Nat → Nat) : (x.map f).isSome = x.isSome := by
+  cases f, x using Option.map.fun_cases
+  case case1 x => simp [-Option.isSome_map']
+  case case2 =>   simp [-Option.isSome_map']
+
+/--
+info: List.map.fun_cases.{u, v} (motive : {α : Type u} → {β : Type v} → (α → β) → List α → Prop)
+  (case1 : ∀ {α : Type u} {β : Type v} (f : α → β), motive f [])
+  (case2 : ∀ {α : Type u} {β : Type v} (f : α → β) (a : α) (as : List α), motive f (a :: as)) {α : Type u} {β : Type v}
+  (f : α → β) (x✝ : List α) : motive f x✝
+-/
+#guard_msgs in
+#check List.map.fun_cases
+
+/--
+info: List.find?.fun_cases.{u} (motive : {α : Type u} → (α → Bool) → List α → Prop)
+  (case1 : ∀ {α : Type u} (p : α → Bool), motive p [])
+  (case2 : ∀ {α : Type u} (p : α → Bool) (a : α) (as : List α), p a = true → motive p (a :: as))
+  (case3 : ∀ {α : Type u} (p : α → Bool) (a : α) (as : List α), p a = false → motive p (a :: as)) {α : Type u}
+  (p : α → Bool) (x✝ : List α) : motive p x✝
+-/
+#guard_msgs in
+#check List.find?.fun_cases