feat: automatic attach introduction in well-founded recursion

src/Init.lean

@@ -38,3 +38,4 @@ import Init.Grind
 import Init.While
 import Init.Syntax
 import Init.Internal
+import Init.Attach

src/Init/Attach.lean

@@ -0,0 +1,83 @@
+/-
+Copyright (c) 2025 Lean FRO, LLC. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Joachim Breitner
+-/
+
+prelude
+
+import Init.WF
+import Init.Data.List.Attach
+import Init.Data.Array.Attach
+
+/-
+At some point I tried to use this rule to upgrade `wfParam` to `….unattach` to only
+have a single rule for each function below. But then there are `xs.attach.unattach` below that were
+sometimes hard to reliably get rid of.
+
+So for now I keep `wfParam`, and have duplicate rules below.
+Maybe later I can create an attribute that takes the `xs.unattach`-rule and then produce the
+alternative form from it (like Mathlib’s `assoc` attribute).
+
+theorem List.wfParam_to_attach (xs : List α) :
+    (wfParam xs) = xs.attach.unattach :=
+  List.unattach_attach.symm
+
+theorem Array.wfParam_to_attach (xs : Array α) :
+    (wfParam xs) = xs.attach.unattach :=
+  Array.unattach_attach.symm
+-/
+
+set_option linter.unusedVariables false in
+theorem List.map_wfParam (xs : List α) (f : α → β) :
+    (wfParam xs).map f = xs.attach.map (fun ⟨x, h⟩ => f (wfParam x)) := by
+  simp [wfParam]
+
+set_option linter.unusedVariables false in
+theorem List.map_unattach (P : α → Prop) (xs : List (Subtype P)) (f : α → β) :
+    xs.unattach.map f = xs.map (fun ⟨x, h⟩ => f (wfParam x)) := by
+  simp [wfParam]
+
+set_option linter.unusedVariables false in
+theorem List.foldl_wfParam (xs : List α) (f : β → α → β) (init : β):
+    (wfParam xs).foldl f init = xs.attach.foldl (fun s ⟨x, h⟩ => f s (wfParam x)) init := by
+  simp [wfParam]
+
+set_option linter.unusedVariables false in
+theorem List.foldl_unattach (P : α → Prop) (xs : List (Subtype P)) (f : β → α → β) (init : β):
+    xs.unattach.foldl f init = xs.foldl (fun s ⟨x, h⟩ => f s (wfParam x)) init := by
+  simp [wfParam]
+
+theorem List.unattach_foldl {p : α → Prop} {l : List { x // p x }}
+    {f : β → { x // p x } → β} {g : β → α → β} {hf : ∀ s x h, f s ⟨x, h⟩ = g s x} :
+    (l.foldl f) = l.unattach.foldl g := by
+  induction l with
+  | nil => simp
+  | cons a l ih =>
+    simp [foldl_cons, hf, List.unattach_cons]
+
+set_option linter.unusedVariables false in
+theorem List.filter_wfParam (xs : List α) (f : α → Bool) :
+    (wfParam xs).filter f = (xs.attach.filter (fun ⟨x, h⟩ => f (wfParam x))).unattach := by
+  simp [wfParam]
+
+set_option linter.unusedVariables false in
+theorem List.filter_unattach (P : α → Prop) (xs : List (Subtype P)) (f : α → Bool) :
+    xs.unattach.filter f = (xs.filter (fun ⟨x, h⟩ => f (wfParam x))).unattach := by
+  simp [wfParam]
+
+theorem List.reverse_wfParam (xs : List α) :
+    (wfParam xs).reverse = xs.attach.reverse.unattach := by simp [wfParam]
+
+theorem List.reverse_unattach (P : α → Prop) (xs : List (Subtype P)) :
+    xs.unattach.reverse = xs.reverse.unattach := by simp
+
+set_option linter.unusedVariables false in
+theorem Array.map_wfParam (xs : Array α) (f : α → β) :
+    (wfParam xs).map f = xs.attach.map (fun ⟨x, h⟩ => f (wfParam x)) := by
+  simp [wfParam]
+
+set_option linter.unusedVariables false in
+theorem Array.map_unattach (P : α → Prop) (xs : Array (Subtype P)) (f : α → β) :
+    xs.unattach.map f = xs.map (fun ⟨x, h⟩ => f (wfParam x)) := by
+  simp [wfParam]

src/Init/ByCases.lean

@@ -37,8 +37,16 @@ theorem apply_ite (f : α → β) (P : Prop) [Decidable P] (x y : α) :
     f (ite P x y) = ite P (f x) (f y) :=
   apply_dite f P (fun _ => x) (fun _ => y)
 
+/--
+`fun_h => …` is like `fun _ =>` but the binder name is based on `h`, not `x`.
+This is relevant when preparing well-founded recusion proof obligations, where the lemma below is
+rewritten with backwards, and the name shows up in the context (daggered, but still visible).
+-/
+local macro "fun_h" "=>" t:term : term => `(fun h => $t)
+
 /-- 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_h => a) (fun_h => 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/WF.lean

@@ -414,3 +414,10 @@ 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

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

@@ -8,9 +8,9 @@ import Lean.Meta.Tactic.Rewrite
 import Lean.Meta.Tactic.Split
 import Lean.Elab.PreDefinition.Basic
 import Lean.Elab.PreDefinition.Eqns
+import Lean.Elab.PreDefinition.WF.Ite
 import Lean.Meta.ArgsPacker.Basic
 import Init.Data.Array.Basic
-import Init.Internal.Order.Basic
 
 namespace Lean.Elab.WF
 open Meta
@@ -23,35 +23,33 @@ structure EqnInfo extends EqnInfoCore where
   argsPacker      : ArgsPacker
   deriving Inhabited
 
-private partial def deltaLHSUntilFix (mvarId : MVarId) : MetaM MVarId := mvarId.withContext do
-  let target ← mvarId.getType'
-  let some (_, lhs, _) := target.eq? | throwTacticEx `deltaLHSUntilFix mvarId "equality expected"
-  if lhs.isAppOf ``WellFounded.fix then
-    return mvarId
-  else if lhs.isAppOf ``Order.fix then
-    return mvarId
-  else
-    deltaLHSUntilFix (← deltaLHS mvarId)
-
 private def rwFixEq (mvarId : MVarId) : MetaM MVarId := mvarId.withContext do
   let target ← mvarId.getType'
   let some (_, lhs, rhs) := target.eq? | unreachable!
-  let h ←
-    if lhs.isAppOf ``WellFounded.fix then
-      pure <| mkAppN (mkConst ``WellFounded.fix_eq lhs.getAppFn.constLevels!) lhs.getAppArgs
-    else if lhs.isAppOf ``Order.fix then
-      let x := lhs.getAppArgs.back!
-      let args := lhs.getAppArgs.pop
-      mkAppM ``congrFun #[mkAppN (mkConst ``Order.fix_eq lhs.getAppFn.constLevels!) args, x]
-    else
-      throwTacticEx `rwFixEq mvarId "expected fixed-point application"
-  let some (_, _, lhsNew) := (← inferType h).eq? | unreachable!
+
+  -- lhs should be an application of the declNameNonrec, which unfolds to an
+  -- application of fix in one step
+  let some lhs' ← delta? lhs | throwError "rwFixEq: cannot delta-reduce {lhs}"
+  let_expr WellFounded.fix _α _C _r _hwf F x := lhs'
+    | throwTacticEx `rwFixEq mvarId "expected saturated fixed-point application in {lhs'}"
+  let h := mkAppN (mkConst ``WellFounded.fix_eq lhs'.getAppFn.constLevels!) lhs'.getAppArgs
+
+  -- We used to just rewrite with `fix_eq` and continue with whatever RHS that produces, but that
+  -- would include more copies of `fix` resulting in large and confusing terms.
+  -- Instead we manually construct the new term in terms of the current functions,
+  -- which should be headed by the `declNameNonRec`, and should be defeq to the expected type
+
+  -- if lhs == e x and lhs' == fix .., then lhsNew := e x = F x (fun y _ => e y)
+  let ftype := (← inferType (mkApp F x)).bindingDomain!
+  let f' ← forallBoundedTelescope ftype (some 2) fun ys _ => do
+    mkLambdaFVars ys (.app lhs.appFn! ys[0]!)
+  let lhsNew := mkApp2 F x f'
   let targetNew ← mkEq lhsNew rhs
   let mvarNew ← mkFreshExprSyntheticOpaqueMVar targetNew
   mvarId.assign (← mkEqTrans h mvarNew)
   return mvarNew.mvarId!
 
-private partial def mkProof (declName : Name) (type : Expr) : MetaM Expr := do
+private partial def mkProof (declName declNameNonRec : Name) (type : Expr) : MetaM Expr := do
   trace[Elab.definition.wf.eqns] "proving: {type}"
   withNewMCtxDepth do
     let main ← mkFreshExprSyntheticOpaqueMVar type
@@ -69,7 +67,9 @@ private partial def mkProof (declName : Name) (type : Expr) : MetaM Expr := do
       else if let some mvarId ← whnfReducibleLHS? mvarId then
         go mvarId
       else
-        let ctx ← Simp.mkContext (config := { dsimp := false })
+        let ctx ← Simp.mkContext
+          (config := { dsimp := false })
+          (simpTheorems  := #[(← getUnattachSimpTheorems)])
         match (← simpTargetStar mvarId ctx (simprocs := {})).1 with
         | TacticResultCNM.closed => return ()
         | TacticResultCNM.modified mvarId => go mvarId
@@ -83,7 +83,10 @@ private partial def mkProof (declName : Name) (type : Expr) : MetaM Expr := do
             -- LHS (introduced in 096e4eb), but it seems that code path was never used,
             -- so #3133 removed it again (and can be recovered from there if this was premature).
             throwError "failed to generate equational theorem for '{declName}'\n{MessageData.ofGoal mvarId}"
-    go (← rwFixEq (← deltaLHSUntilFix mvarId))
+
+    let mvarId ← if declName != declNameNonRec then deltaLHS mvarId else pure mvarId
+    let mvarId ← rwFixEq mvarId
+    go mvarId
     instantiateMVars main
 
 def mkEqns (declName : Name) (info : EqnInfo) : MetaM (Array Name) :=
@@ -101,7 +104,7 @@ def mkEqns (declName : Name) (info : EqnInfo) : MetaM (Array Name) :=
     trace[Elab.definition.wf.eqns] "{eqnTypes[i]}"
     let name := (Name.str baseName eqnThmSuffixBase).appendIndexAfter (i+1)
     thmNames := thmNames.push name
-    let value ← mkProof declName type
+    let value ← mkProof declName info.declNameNonRec type
     let (type, value) ← removeUnusedEqnHypotheses type value
     addDecl <| Declaration.thmDecl {
       name, type, value

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

@@ -5,26 +5,128 @@ Authors: Leonardo de Moura
 -/
 prelude
 import Lean.Meta.Transform
+import Lean.Meta.Match.MatcherApp.Basic
+import Lean.Elab.Tactic.Simp
+
+open Lean Meta
+
+namespace Lean.Elab.WF
+
+def getAttachSimpTheorems : MetaM SimpTheorems := do
+  let s : SimpTheorems := {}
+  let s ← s.addConst ``dite_eq_ite (inv := true)
+  -- let s ← s.addConst `List.wfParam_to_attach
+  -- let s ← s.addConst `Array.wfParam_to_attach
+  let s ← s.addConst `List.map_wfParam
+  let s ← s.addConst `List.map_unattach
+  let s ← s.addConst `List.filter_wfParam
+  let s ← s.addConst `List.filter_unattach
+  let s ← s.addConst `List.reverse_wfParam
+  let s ← s.addConst `List.reverse_unattach
+  let s ← s.addConst `List.foldl_wfParam
+  let s ← s.addConst `List.foldl_unattach
+  let s ← s.addConst `Array.map_wfParam
+  let s ← s.addConst `Array.map_unattach
+  pure s
+
+def getUnattachSimpTheorems : MetaM SimpTheorems := do
+  let s : SimpTheorems := {}
+  let s ← s.addConst ``List.unattach_attach
+  let s ← s.addConst ``List.map_subtype
+  let s ← s.addConst ``List.unattach_filter
+  let s ← s.addConst ``List.unattach_reverse
+  let s ← s.addConst `List.unattach_foldl
+  let s ← s.addConst ``Array.map_subtype
+  let s ← s.addConst ``Array.unattach_attach
+  pure s
+
+private def getSimpContext : MetaM Simp.Context := do
+  Simp.mkContext
+    (simpTheorems  := #[(← getAttachSimpTheorems)])
+    (congrTheorems := {})
+    (config        := { Simp.neutralConfig with dsimp := false })
+
+private def getCleanupSimpContext : MetaM Simp.Context := do
+  let s : SimpTheorems := {}
+  let s ← s.addConst ``List.unattach_attach
+  let s ← s.addConst ``Array.unattach_attach
+  let s ← s.addDeclToUnfold ``wfParam
+  Simp.mkContext
+    (simpTheorems  := #[s])
+    (congrTheorems := {})
+    (config        := { Simp.neutralConfig with dsimp := true })
+
+def isWfParam? (e : Expr) : Option Expr :=
+  if e.isAppOfArity ``wfParam 2 then
+    e.appArg!
+  else
+    none
+
+def mkWfParam (e : Expr) : MetaM Expr :=
+  mkAppM ``wfParam #[e]
+
+/-- `f (wfParam x) ==> wfParam (f x)` if `f` is a projection -/
+builtin_dsimproc paramProj (_) := fun e => do
+  if h : e.isApp then
+    let some a' := isWfParam? (e.appArg h) | return .continue
+    let f := e.getAppFn
+    unless f.isConst do return .continue
+    unless (← isProjectionFn f.constName!) do return .continue
+    let e' ← mkWfParam (.app (e.appFn h) a')
+    return .done e'
+  else
+    return .continue
+
+/-- `match (wfParam x) with con y => alt[y] ==> match x with con y => alt[wfParam y] -/
+builtin_dsimproc paramMatcher (_) := fun e => do
+  let some matcherApp ← matchMatcherApp? e (alsoCasesOn := true)  | return .continue
+  unless matcherApp.discrs.any (isWfParam? · |>.isSome) do return .continue
+  let discrs' := matcherApp.discrs.map (fun e => isWfParam? e |>.getD e)
+  let alts' ← matcherApp.alts.mapM fun alt =>
+    lambdaTelescope alt fun xs body => do
+      -- Annotate all xs with `wfParam`
+      let xs' ← xs.mapM (mkWfParam ·)
+      let body' := body.replaceFVars xs xs'
+      mkLambdaFVars xs body'
+  let matcherApp' := { matcherApp with discrs := discrs', alts := alts' }
+  return .continue <| matcherApp'.toExpr
 
-namespace Lean.Meta
 
-/--
-  Convert `ite` expressions in `e` to `dite`s.
-  It is useful to make this conversion in the `WF` module because the condition is often used in
-  termination proofs. -/
 def iteToDIte (e : Expr) : MetaM Expr := do
-  -- TODO: move this file to `Meta` if we decide to use it in other places.
-  let post (e : Expr) : MetaM TransformStep := do
-    if e.isAppOfArity ``ite 5 then
-      let f    := e.getAppFn
-      let args := e.getAppArgs
-      let c    := args[1]!
-      let h    ← mkFreshUserName `h
-      let args := args.set! 3 (Lean.mkLambda h BinderInfo.default c args[3]!)
-      let args := args.set! 4 (Lean.mkLambda h BinderInfo.default (mkNot c) args[4]!)
-      return .done <| mkAppN (mkConst ``dite f.constLevels!) args
-    else
-      return .done e
-  transform e (post := post)
-
-end Lean.Meta
+  lambdaTelescope e fun xs e => do
+    -- Annotate all xs with `wfParam`
+    let xs' ← xs.mapM mkWfParam
+    let e' := e.replaceFVars xs xs'
+
+    -- Now run the simplifier
+    let simprocs : Simprocs := {}
+    let simprocs ← simprocs.add ``paramProj (post := true)
+    let simprocs ← simprocs.add ``paramMatcher (post := false)
+    let (result, _) ← Meta.simp e' (← getSimpContext) (simprocs := #[simprocs])
+    let e' := result.expr
+
+    -- Remove markers
+    -- let (result, _) ← Meta.simp e' (← getCleanupSimpContext)
+    -- let e'' := result.expr
+    -- Simp, even with dsimp on, is not as thorough as `Core.transform`:
+    let e'' ← Core.transform e' fun e =>
+      e.withApp fun f as => do
+        if f.isConstOf ``wfParam then
+          if h : as.size ≥ 2 then
+            return .continue (mkAppN as[1] as[2:])
+        return .continue
+
+    trace[Elab.definition.wf] "Attach-introduction:{indentExpr e}\nto{indentExpr e'}\ncleaned up as{indentExpr e''}"
+
+    mkLambdaFVars xs e''
+
+/-
+run_elab do
+  let stx ← `(fun (n : Nat) => if n > 0 then 3 else 4)
+  let e ← Elab.Term.elabTerm stx .none
+  let e' ← iteToDIte e
+  logInfo m!"{e}\n{e'}"
+-/
+
+
+end Lean.Elab.WF