feat: nested well-founded recursion via automatic preprocessing (#6744)

This PR extend the preprocessing of well-founded recursive definitions
to bring assumptions like `h✝ : x ∈ xs` into scope automatically.

This fixes #5471, and follows (roughly) the design written there.
See the module docs at `src/Lean/Elab/PreDefinition/WF/AutoAttach.lean`
for details on the implementation.

This only works for higher-order functions that have a suitable setup.
See for example section “Well-founded recursion preprocessing setup” in
`src/Init/Data/List/Attach.lean`.

This does not change the `decreasing_tactic`, so in some cases there is
still the need for a manual termination proof some cases. We expect a
better termination tactic in the near future.

Author: nomeata

src/Init/ByCases.lean

@@ -38,7 +38,8 @@ theorem apply_ite (f : α → β) (P : Prop) [Decidable P] (x y : α) :
   apply_dite f P (fun _ => x) (fun _ => y)
 
 /-- A `dite` whose results do not actually depend on the condition may be reduced to an `ite`. -/
-@[simp] theorem dite_eq_ite [Decidable P] : (dite P (fun _ => a) fun _ => b) = ite P a b := rfl
+@[simp] theorem dite_eq_ite [Decidable P] :
+  (dite P (fun _ => a) (fun _ => b)) = ite P a b := rfl
 
 @[deprecated "Use `ite_eq_right_iff`" (since := "2024-09-18")]
 theorem ite_some_none_eq_none [Decidable P] :

src/Init/Control/Basic.lean

@@ -5,6 +5,7 @@ Author: Leonardo de Moura, Sebastian Ullrich
 -/
 prelude
 import Init.Core
+import Init.BinderNameHint
 
 universe u v w
 
@@ -35,6 +36,12 @@ instance (priority := 500) instForInOfForIn' [ForIn' m ρ α d] : ForIn m ρ α
   simp [h]
   rfl
 
+@[wf_preprocess] theorem forIn_eq_forin' [d : Membership α ρ] [ForIn' m ρ α d] {β} [Monad m]
+    (x : ρ) (b : β) (f : (a : α) → β → m (ForInStep β)) :
+    forIn x b f = forIn' x b (fun x h => binderNameHint x f <| binderNameHint h () <| f x) := by
+  simp [binderNameHint]
+  rfl -- very strange why `simp` did not close it
+
 /-- Extract the value from a `ForInStep`, ignoring whether it is `done` or `yield`. -/
 def ForInStep.value (x : ForInStep α) : α :=
   match x with

src/Init/Data/Array/Attach.lean

@@ -650,6 +650,16 @@ and simplifies these to the function directly taking the value.
   rw [List.filterMap_subtype]
   simp [hf]
 
+
+@[simp] theorem flatMap_subtype {p : α → Prop} {l : Array { x // p x }}
+    {f : { x // p x } → Array β} {g : α → Array β} (hf : ∀ x h, f ⟨x, h⟩ = g x) :
+    (l.flatMap f) = l.unattach.flatMap g := by
+  cases l
+  simp only [size_toArray, List.flatMap_toArray, List.unattach_toArray, List.length_unattach,
+    mk.injEq]
+  rw [List.flatMap_subtype]
+  simp [hf]
+
 @[simp] theorem findSome?_subtype {p : α → Prop} {l : Array { x // p x }}
     {f : { x // p x } → Option β} {g : α → Option β} (hf : ∀ x h, f ⟨x, h⟩ = g x) :
     l.findSome? f = l.unattach.findSome? g := by
@@ -695,4 +705,67 @@ and simplifies these to the function directly taking the value.
     (Array.mkArray n x).unattach = Array.mkArray n x.1 := by
   simp [unattach]
 
+/-! ### Well-founded recursion preprocessing setup -/
+
+@[wf_preprocess] theorem Array.map_wfParam (xs : Array α) (f : α → β) :
+    (wfParam xs).map f = xs.attach.unattach.map f := by
+  simp [wfParam]
+
+@[wf_preprocess] theorem Array.map_unattach (P : α → Prop) (xs : Array (Subtype P)) (f : α → β) :
+    xs.unattach.map f = xs.map fun ⟨x, h⟩ =>
+      binderNameHint x f <| binderNameHint h () <| f (wfParam x) := by
+  simp [wfParam]
+
+@[wf_preprocess] theorem foldl_wfParam (xs : Array α) (f : β → α → β) (x : β) :
+    (wfParam xs).foldl f x = xs.attach.unattach.foldl f x := by
+  simp [wfParam]
+
+@[wf_preprocess] theorem foldl_unattach (P : α → Prop) (xs : Array (Subtype P)) (f : β → α → β) (x : β):
+    xs.unattach.foldl f x = xs.foldl (fun s ⟨x, h⟩ =>
+      binderNameHint s f <| binderNameHint x (f s) <| binderNameHint h () <| f s (wfParam x)) x := by
+  simp [wfParam]
+
+@[wf_preprocess] theorem foldr_wfParam (xs : Array α) (f : α → β → β) (x : β) :
+    (wfParam xs).foldr f x = xs.attach.unattach.foldr f x := by
+  simp [wfParam]
+
+@[wf_preprocess] theorem foldr_unattach (P : α → Prop) (xs : Array (Subtype P)) (f : α → β → β) (x : β):
+    xs.unattach.foldr f x = xs.foldr (fun ⟨x, h⟩ s =>
+      binderNameHint x f <| binderNameHint s (f x) <| binderNameHint h () <| f (wfParam x) s) x := by
+  simp [wfParam]
+
+@[wf_preprocess] theorem filter_wfParam (xs : Array α) (f : α → Bool) :
+    (wfParam xs).filter f = xs.attach.unattach.filter f:= by
+  simp [wfParam]
+
+@[wf_preprocess] theorem filter_unattach (P : α → Prop) (xs : Array (Subtype P)) (f : α → Bool) :
+    xs.unattach.filter f = (xs.filter (fun ⟨x, h⟩ =>
+      binderNameHint x f <| binderNameHint h () <| f (wfParam x))).unattach := by
+  simp [wfParam]
+
+@[wf_preprocess] theorem reverse_wfParam (xs : Array α) :
+    (wfParam xs).reverse = xs.attach.unattach.reverse := by simp [wfParam]
+
+@[wf_preprocess] theorem reverse_unattach (P : α → Prop) (xs : Array (Subtype P)) :
+    xs.unattach.reverse = xs.reverse.unattach := by simp
+
+@[wf_preprocess] theorem filterMap_wfParam (xs : Array α) (f : α → Option β) :
+    (wfParam xs).filterMap f = xs.attach.unattach.filterMap f := by
+  simp [wfParam]
+
+@[wf_preprocess] theorem filterMap_unattach (P : α → Prop) (xs : Array (Subtype P)) (f : α → Option β) :
+    xs.unattach.filterMap f = xs.filterMap fun ⟨x, h⟩ =>
+      binderNameHint x f <| binderNameHint h () <| f (wfParam x) := by
+  simp [wfParam]
+
+@[wf_preprocess] theorem flatMap_wfParam (xs : Array α) (f : α → Array β) :
+    (wfParam xs).flatMap f = xs.attach.unattach.flatMap f := by
+  simp [wfParam]
+
+@[wf_preprocess] theorem flatMap_unattach (P : α → Prop) (xs : Array (Subtype P)) (f : α → Array β) :
+    xs.unattach.flatMap f = xs.flatMap fun ⟨x, h⟩ =>
+      binderNameHint x f <| binderNameHint h () <| f (wfParam x) := by
+  simp [wfParam]
+
+
 end Array

src/Init/Data/List/Attach.lean

@@ -6,6 +6,7 @@ Authors: Mario Carneiro
 prelude
 import Init.Data.List.Count
 import Init.Data.Subtype
+import Init.BinderNameHint
 
 namespace List
 
@@ -796,4 +797,66 @@ and simplifies these to the function directly taking the value.
     (List.replicate n x).unattach = List.replicate n x.1 := by
   simp [unattach, -map_subtype]
 
+/-! ### Well-founded recursion preprocessing setup -/
+
+@[wf_preprocess] theorem map_wfParam (xs : List α) (f : α → β) :
+    (wfParam xs).map f = xs.attach.unattach.map f := by
+  simp [wfParam]
+
+@[wf_preprocess] theorem map_unattach (P : α → Prop) (xs : List (Subtype P)) (f : α → β) :
+    xs.unattach.map f = xs.map fun ⟨x, h⟩ =>
+      binderNameHint x f <| binderNameHint h () <| f (wfParam x) := by
+  simp [wfParam]
+
+@[wf_preprocess] theorem foldl_wfParam (xs : List α) (f : β → α → β) (x : β) :
+    (wfParam xs).foldl f x = xs.attach.unattach.foldl f x := by
+  simp [wfParam]
+
+@[wf_preprocess] theorem foldl_unattach (P : α → Prop) (xs : List (Subtype P)) (f : β → α → β) (x : β):
+    xs.unattach.foldl f x = xs.foldl (fun s ⟨x, h⟩ =>
+      binderNameHint s f <| binderNameHint x (f s) <| binderNameHint h () <| f s (wfParam x)) x := by
+  simp [wfParam]
+
+@[wf_preprocess] theorem foldr_wfParam (xs : List α) (f : α → β → β) (x : β) :
+    (wfParam xs).foldr f x = xs.attach.unattach.foldr f x := by
+  simp [wfParam]
+
+@[wf_preprocess] theorem foldr_unattach (P : α → Prop) (xs : List (Subtype P)) (f : α → β → β) (x : β):
+    xs.unattach.foldr f x = xs.foldr (fun ⟨x, h⟩ s =>
+      binderNameHint x f <| binderNameHint s (f x) <| binderNameHint h () <| f (wfParam x) s) x := by
+  simp [wfParam]
+
+@[wf_preprocess] theorem filter_wfParam (xs : List α) (f : α → Bool) :
+    (wfParam xs).filter f = xs.attach.unattach.filter f:= by
+  simp [wfParam]
+
+@[wf_preprocess] theorem filter_unattach (P : α → Prop) (xs : List (Subtype P)) (f : α → Bool) :
+    xs.unattach.filter f = (xs.filter (fun ⟨x, h⟩ =>
+      binderNameHint x f <| binderNameHint h () <| f (wfParam x))).unattach := by
+  simp [wfParam]
+
+@[wf_preprocess] theorem reverse_wfParam (xs : List α) :
+    (wfParam xs).reverse = xs.attach.unattach.reverse := by simp [wfParam]
+
+@[wf_preprocess] theorem reverse_unattach (P : α → Prop) (xs : List (Subtype P)) :
+    xs.unattach.reverse = xs.reverse.unattach := by simp
+
+@[wf_preprocess] theorem filterMap_wfParam (xs : List α) (f : α → Option β) :
+    (wfParam xs).filterMap f = xs.attach.unattach.filterMap f := by
+  simp [wfParam]
+
+@[wf_preprocess] theorem filterMap_unattach (P : α → Prop) (xs : List (Subtype P)) (f : α → Option β) :
+    xs.unattach.filterMap f = xs.filterMap fun ⟨x, h⟩ =>
+      binderNameHint x f <| binderNameHint h () <| f (wfParam x) := by
+  simp [wfParam]
+
+@[wf_preprocess] theorem flatMap_wfParam (xs : List α) (f : α → List β) :
+    (wfParam xs).flatMap f = xs.attach.unattach.flatMap f := by
+  simp [wfParam]
+
+@[wf_preprocess] theorem flatMap_unattach (P : α → Prop) (xs : List (Subtype P)) (f : α → List β) :
+    xs.unattach.flatMap f = xs.flatMap fun ⟨x, h⟩ =>
+      binderNameHint x f <| binderNameHint h () <| f (wfParam x) := by
+  simp [wfParam]
+
 end List

src/Init/WF.lean

@@ -5,6 +5,7 @@ Author: Leonardo de Moura
 -/
 prelude
 import Init.SizeOf
+import Init.BinderNameHint
 import Init.Data.Nat.Basic
 
 universe u v
@@ -414,3 +415,18 @@ theorem mkSkipLeft {α : Type u} {β : Type v} {b₁ b₂ : β} {s : β → β 
 end
 
 end PSigma
+
+/--
+The `wfParam` gadget is used internally during the construction of recursive functions by
+wellfounded recursion, to keep track of the parameter for which the automatic introduction
+of `List.attach` (or similar) is plausible.
+-/
+def wfParam {α : Sort u} (a : α) : α := a
+
+/--
+Reverse direction of `dite_eq_ite`. Used by the well-founded definition preprocessor to extend the
+context of a termination proof inside `if-then-else` with the condition.
+-/
+@[wf_preprocess] theorem ite_eq_dite [Decidable P] :
+    ite P a b = (dite P (fun h => binderNameHint h () a) (fun h => binderNameHint h () b)) := by
+  rfl

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

@@ -0,0 +1,36 @@
+/-
+Copyright (c) 2021 Microsoft Corporation. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Leonardo de Moura
+-/
+prelude
+import Lean.Meta.Transform
+import Lean.Elab.RecAppSyntax
+
+namespace Lean.Elab.WF
+open Meta
+
+/--
+Preprocesses the expressions to improve the effectiveness of `wfRecursion`.
+
+* Floats out the RecApp markers.
+  Example:
+  ```
+  def f : Nat → Nat
+    | 0 => 1
+    | i+1 => (f x) i
+  ```
+
+Unlike `Lean.Elab.Structural.preprocess`, do _not_ beta-reduce, as it could
+remove `let_fun`-lambdas that contain explicit termination proofs.
+-/
+def floatRecApp (e : Expr) : CoreM Expr :=
+  Core.transform e
+    (post := fun e => do
+      if e.isApp && e.getAppFn.isMData then
+        let .mdata m f := e.getAppFn | unreachable!
+        if m.isRecApp then
+          return .done (.mdata m (f.beta e.getAppArgs))
+      return .continue)
+
+end Lean.Elab.WF

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

@@ -8,12 +8,11 @@ import Lean.Elab.PreDefinition.Basic
 import Lean.Elab.PreDefinition.TerminationMeasure
 import Lean.Elab.PreDefinition.Mutual
 import Lean.Elab.PreDefinition.WF.PackMutual
-import Lean.Elab.PreDefinition.WF.Preprocess
+import Lean.Elab.PreDefinition.WF.FloatRecApp
 import Lean.Elab.PreDefinition.WF.Rel
 import Lean.Elab.PreDefinition.WF.Fix
 import Lean.Elab.PreDefinition.WF.Unfold
-import Lean.Elab.PreDefinition.WF.Ite
-import Lean.Elab.PreDefinition.WF.AutoAttach
+import Lean.Elab.PreDefinition.WF.Preprocess
 import Lean.Elab.PreDefinition.WF.GuessLex
 
 namespace Lean.Elab
@@ -23,8 +22,8 @@ open Meta
 def wfRecursion (preDefs : Array PreDefinition) (termMeasure?s : Array (Option TerminationMeasure)) : TermElabM Unit := do
   let termMeasures? := termMeasure?s.mapM id -- Either all or none, checked by `elabTerminationByHints`
   let preDefs ← preDefs.mapM fun preDef =>
-    return { preDef with value := (← preprocess preDef.value) }
-  let (fixedPrefixSize, argsPacker, unaryPreDef) ← withoutModifyingEnv do
+    return { preDef with value := (← floatRecApp preDef.value) }
+  let (fixedPrefixSize, argsPacker, unaryPreDef, wfPreprocessProofs) ← withoutModifyingEnv do
     for preDef in preDefs do
       addAsAxiom preDef
     let fixedPrefixSize ← Mutual.getFixedPrefix preDefs
@@ -34,8 +33,10 @@ def wfRecursion (preDefs : Array PreDefinition) (termMeasure?s : Array (Option T
       if varNames.isEmpty then
         throwError "well-founded recursion cannot be used, '{preDef.declName}' does not take any (non-fixed) arguments"
     let argsPacker := { varNamess }
-    let preDefsDIte ← preDefs.mapM fun preDef => return { preDef with value := (← iteToDIte preDef.value) }
-    return (fixedPrefixSize, argsPacker, ← packMutual fixedPrefixSize argsPacker preDefsDIte)
+    let (preDefsAttached, wfPreprocessProofs) ← Array.unzip <$> preDefs.mapM fun preDef => do
+      let result ← preprocess preDef.value
+      return ({preDef with value := result.expr}, result)
+    return (fixedPrefixSize, argsPacker, ← packMutual fixedPrefixSize argsPacker preDefsAttached, wfPreprocessProofs)
 
   let wf : TerminationMeasures ← do
     if let some tms := termMeasures? then pure tms else
@@ -62,10 +63,10 @@ def wfRecursion (preDefs : Array PreDefinition) (termMeasure?s : Array (Option T
   Mutual.addPreDefsFromUnary preDefs preDefsNonrec preDefNonRec
   let preDefs ← Mutual.cleanPreDefs preDefs
   registerEqnsInfo preDefs preDefNonRec.declName fixedPrefixSize argsPacker
-  for preDef in preDefs do
+  for preDef in preDefs, wfPreprocessProof in wfPreprocessProofs do
     unless preDef.kind.isTheorem do
       unless (← isProp preDef.type) do
-        WF.mkUnfoldEq preDef preDefNonRec.declName
+        WF.mkUnfoldEq preDef preDefNonRec.declName wfPreprocessProof
   Mutual.addPreDefAttributes preDefs
 
 builtin_initialize registerTraceClass `Elab.definition.wf