feat: dedicated fix operator for well-founded recursion on Nat (#7965)

This PR lets recursive functions defined by well-founded recursion use a
different `fix` function when the termination measure is of type `Nat`.
This fix-point operator use structural recursion on “fuel”, initialized
by the given measure, and is thus reasonable to reduce, e.g. in `by
decide` proofs.

Extra provisions are in place that the fixpoint operator only starts
reducing when the fuel is fully known, to prevent “accidential” defeqs
when the remaining fuel for the recursive calls match the initial fuel
for that recursive argument.

To opt-out, the idiom `termination_by (n,0)` can be used.

We still use `@[irreducible]` as the default for such recursive
definitions, to avoid unexpected `defeq` lemmas. Making these functions
`@[semireducible]` by default showed performance regressions in lean.
When the measure is of type `Nat`, the system will accept an explicit
`@[semireducible]` without the usual warning.

Fixes #5234. Fixes: #11181.

Author: nomeata

src/Init/WF.lean

@@ -439,6 +439,53 @@ end
 
 end PSigma
 
+namespace WellFounded
+
+variable {α : Sort u}
+variable {motive : α → Sort v}
+variable (h : α → Nat)
+variable (F : (x : α) → ((y : α) → InvImage (· < ·) h y x → motive y) → motive x)
+
+/-- Helper gadget that prevents reduction of `Nat.eager n` unless `n` evalutes to a ground term. -/
+def Nat.eager (n : Nat) : Nat :=
+  if Nat.beq n n = true then n else n
+
+theorem Nat.eager_eq (n : Nat) : Nat.eager n = n := ite_self n
+
+/--
+A well-founded fixpoint operator specialized for `Nat`-valued measures. Given a measure `h`, it expects
+its higher order function argument `F` to invoke its argument only on values `y` that are smaller
+than `x` with regard to `h`.
+
+In contrast to to `WellFounded.fix`, this fixpoint operator reduces on closed terms. (More precisely:
+when `h x` evalutes to a ground value)
+
+-/
+def Nat.fix : (x : α) → motive x :=
+  let rec go : ∀ (fuel : Nat) (x : α), (h x < fuel) → motive x :=
+    Nat.rec
+      (fun _ hfuel => (Nat.not_succ_le_zero _ hfuel).elim)
+      (fun _ ih x hfuel => F x (fun y hy => ih y (Nat.lt_of_lt_of_le hy (Nat.le_of_lt_add_one hfuel))))
+  fun x => go (Nat.eager (h x + 1)) x (Nat.eager_eq _ ▸ Nat.lt_add_one _)
+
+protected theorem Nat.fix.go_congr (x : α) (fuel₁ fuel₂ : Nat) (h₁ : h x < fuel₁) (h₂ : h x < fuel₂) :
+    Nat.fix.go h F fuel₁ x h₁ = Nat.fix.go h F fuel₂ x h₂ := by
+  induction fuel₁ generalizing x fuel₂ with
+  | zero => contradiction
+  | succ fuel₁ ih =>
+    cases fuel₂ with
+    | zero => contradiction
+    | succ fuel₂ =>
+      exact congrArg (F x) (funext fun y => funext fun hy => ih y fuel₂ _ _ )
+
+theorem Nat.fix_eq (x : α) :
+    Nat.fix h F x = F x (fun y _ => Nat.fix h F y) := by
+  unfold Nat.fix
+  simp [Nat.eager_eq]
+  exact congrArg (F x) (funext fun _ => funext fun _ => Nat.fix.go_congr ..)
+
+end WellFounded
+
 /--
 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

src/Lean/Compiler/InlineAttrs.lean

@@ -37,6 +37,7 @@ private def isValidMacroInline (declName : Name) : CoreM Bool := do
   let isRec (declName' : Name) : Bool :=
     isBRecOnRecursor env declName' ||
     declName' == ``WellFounded.fix ||
+    declName' == ``WellFounded.Nat.fix ||
     declName' == declName ++ `_unary -- Auxiliary declaration created by `WF` module
   if Option.isSome <| info.value.find? fun e => e.isConst && isRec e.constName! then
     -- It contains a `brecOn` or `WellFounded.fix` application. So, it should be recursvie

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

@@ -237,21 +237,32 @@ def solveDecreasingGoals (funNames : Array Name) (argsPacker : ArgsPacker) (decr
             Term.reportUnsolvedGoals remainingGoals
   instantiateMVars value
 
+def isNatLtWF (wfRel : Expr) : MetaM (Option Expr) := do
+  match_expr wfRel with
+  | invImage _ β f wfRelβ =>
+    unless (← isDefEq β (mkConst ``Nat)) do return none
+    unless (← isDefEq wfRelβ (mkConst ``Nat.lt_wfRel)) do return none
+    return f
+  | _ => return none
+
 def mkFix (preDef : PreDefinition) (prefixArgs : Array Expr) (argsPacker : ArgsPacker)
     (wfRel : Expr) (funNames : Array Name) (decrTactics : Array (Option DecreasingBy)) :
     TermElabM Expr := do
   let type ← instantiateForall preDef.type prefixArgs
   let (wfFix, varName) ← forallBoundedTelescope type (some 1) fun x type => do
     let x := x[0]!
+    let varName ← x.fvarId!.getUserName -- See comment below.
     let α ← inferType x
     let u ← getLevel α
     let v ← getLevel type
     let motive ← mkLambdaFVars #[x] type
-    let rel := mkProj ``WellFoundedRelation 0 wfRel
-    let wf  := mkProj ``WellFoundedRelation 1 wfRel
-    let wf ← mkAppM `Lean.opaqueId #[wf]
-    let varName ← x.fvarId!.getUserName -- See comment below.
-    return (mkApp4 (mkConst ``WellFounded.fix [u, v]) α motive rel wf, varName)
+    if let some measure ← isNatLtWF wfRel then
+      return (mkApp3 (mkConst `WellFounded.Nat.fix [u, v]) α motive measure, varName)
+    else
+      let rel := mkProj ``WellFoundedRelation 0 wfRel
+      let wf  := mkProj ``WellFoundedRelation 1 wfRel
+      let wf ← mkAppM `Lean.opaqueId #[wf]
+      return (mkApp4 (mkConst ``WellFounded.fix [u, v]) α motive rel wf, varName)
   forallBoundedTelescope (← whnf (← inferType wfFix)).bindingDomain! (some 2) fun xs _ => do
     let x   := xs[0]!
     -- Remark: we rename `x` here to make sure we preserve the variable name in the

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

@@ -53,19 +53,21 @@ def wfRecursion (docCtx : LocalContext × LocalInstances) (preDefs : Array PreDe
     -- No termination_by here, so use GuessLex to infer one
     guessLex preDefs unaryPreDefProcessed fixedParamPerms argsPacker
 
-  -- Warn about likely unwanted reducibility attributes
-  preDefs.forM fun preDef =>
-    preDef.modifiers.attrs.forM fun a => do
-      if a.name = `reducible || a.name = `semireducible then
-        logWarningAt a.stx s!"marking functions defined by well-founded recursion as `{a.name}` is not effective"
-
   let preDefNonRec ← forallBoundedTelescope unaryPreDef.type fixedParamPerms.numFixed fun fixedArgs type => do
     let type ← whnfForall type
     unless type.isForall do
       throwError "wfRecursion: expected unary function type: {type}"
     let packedArgType := type.bindingDomain!
     elabWFRel (preDefs.map (·.declName)) unaryPreDef.declName fixedParamPerms fixedArgs argsPacker packedArgType wf fun wfRel => do
       trace[Elab.definition.wf] "wfRel: {wfRel}"
+      let useNatRec := (← isNatLtWF wfRel).isSome
+      -- Warn about likely unwanted reducibility attributes
+      unless useNatRec do
+        preDefs.forM fun preDef =>
+          preDef.modifiers.attrs.forM fun a => do
+            if a.name = `reducible || a.name = `semireducible then
+              logWarningAt a.stx s!"marking functions defined by well-founded recursion as `{a.name}` is not effective"
+
       let (value, envNew) ← withoutModifyingEnv' do
         addAsAxiom unaryPreDef
         let value ← mkFix unaryPreDefProcessed fixedArgs argsPacker wfRel (preDefs.map (·.declName)) (preDefs.map (·.termination.decreasingBy?))

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

@@ -36,9 +36,14 @@ def rwFixEq (mvarId : MVarId) : MetaM MVarId := mvarId.withContext do
   -- 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
+  let h ← match_expr lhs' with
+    | WellFounded.fix _α _C _r _hwf _F _x =>
+      pure <| mkAppN (mkConst ``WellFounded.fix_eq lhs'.getAppFn.constLevels!) lhs'.getAppArgs
+    | WellFounded.Nat.fix _α _C _motive _F _x =>
+      pure <| mkAppN (mkConst ``WellFounded.Nat.fix_eq lhs'.getAppFn.constLevels!) lhs'.getAppArgs
+    | _ => throwTacticEx `rwFixEq mvarId m!"expected saturated fixed-point application in {lhs'}"
+  let F := lhs'.appFn!.appArg!
+  let x := lhs'.appArg!
 
   -- 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.

src/Lean/Meta/Tactic/FunInd.lean

@@ -927,76 +927,84 @@ where doRealize (inductName : Name) := do
   -- to make sure that `target` indeed the last parameter
   let e := info.value
   let e ← lambdaTelescope e fun params body => do
-    if body.isAppOfArity ``WellFounded.fix 5 then
+    if body.isAppOfArity ``WellFounded.fix 5 || body.isAppOfArity ``WellFounded.Nat.fix 4 then
       forallBoundedTelescope (← inferType body) (some 1) fun xs _ => do
         unless xs.size = 1 do
           throwError "functional induction: Failed to eta-expand{indentExpr e}"
         mkLambdaFVars (params ++ xs) (mkAppN body xs)
     else
       pure e
   let (e', paramMask) ← lambdaTelescope e fun params funBody => MatcherApp.withUserNames params varNames do
-    match_expr funBody with
-    | [email protected] α _motive rel wf body target =>
-      unless params.back! == target do
-        throwError "functional induction: expected the target as last parameter{indentExpr e}"
-      let fixedParamPerms := params.pop
-      let motiveType ←
-        if unfolding then
-          withLocalDeclD `r (← instantiateForall info.type params) fun r =>
-            mkForallFVars #[target, r] (.sort 0)
+    unless funBody.isApp && funBody.appFn!.isApp do
+      throwError "functional induction: unexpected body {funBody}"
+    let body := funBody.appFn!.appArg!
+    let target := funBody.appArg!
+    unless params.back! == target do
+      throwError "functional induction: expected the target as last parameter{indentExpr e}"
+    let fixedParamPerms := params.pop
+    let motiveType ←
+      if unfolding then
+        withLocalDeclD `r (← instantiateForall info.type params) fun r =>
+          mkForallFVars #[target, r] (.sort 0)
+      else
+        mkForallFVars #[target] (.sort 0)
+    withLocalDeclD `motive motiveType fun motive => do
+      let fn := mkAppN (← mkConstWithLevelParams name) fixedParamPerms
+      let isRecCall : Expr → Option Expr := fun e =>
+        e.withApp fun f xs =>
+          if f.isFVarOf motive.fvarId! && xs.size > 0 then
+          mkApp fn xs[0]!
         else
-          mkForallFVars #[target] (.sort 0)
-      withLocalDeclD `motive motiveType fun motive => do
-        let fn := mkAppN (← mkConstWithLevelParams name) fixedParamPerms
-        let isRecCall : Expr → Option Expr := fun e =>
-          e.withApp fun f xs =>
-            if f.isFVarOf motive.fvarId! && xs.size > 0 then
-            mkApp fn xs[0]!
-          else
-            none
+          none
 
-        let motiveArg ←
-          if unfolding then
-            let motiveArg := mkApp2 motive target (mkAppN (← mkConstWithLevelParams name) params)
-            mkLambdaFVars #[target] motiveArg
+      let motiveArg ←
+        if unfolding then
+          let motiveArg := mkApp2 motive target (mkAppN (← mkConstWithLevelParams name) params)
+          mkLambdaFVars #[target] motiveArg
+        else
+          pure motive
+
+      let e' ← match_expr funBody with
+        | [email protected] α _motive rel wf _body _target =>
+          let e' := .const ``WellFounded.fix [fix.constLevels![0]!, levelZero]
+          pure <| mkApp4 e' α motiveArg rel wf
+        | [email protected] α _motive measure _body _target =>
+          let e' := .const `WellFounded.Nat.fix [fix.constLevels![0]!, levelZero]
+          pure <| mkApp3 e' α motiveArg measure
+        | _ =>
+          if funBody.isAppOf ``WellFounded.fix || funBody.isAppOf `WellFounded.Nat.Fix then
+            throwError "Function {name} defined via WellFounded.fix with unexpected arity {funBody.getAppNumArgs}:{indentExpr funBody}"
           else
-            pure motive
-        let e' := .const ``WellFounded.fix [fix.constLevels![0]!, levelZero]
-        let e' := mkApp4 e' α motiveArg rel wf
-        check e'
-        let (body', mvars) ← M2.run do
-          forallTelescope (← inferType e').bindingDomain! fun xs goal => do
-            if xs.size ≠ 2 then
-              throwError "expected recursor argument to take 2 parameters, got {xs}" else
-            let targets : Array Expr := xs[*...1]
-            let genIH := xs[1]!
-            let extraParams := xs[2...*]
-            -- open body with the same arg
-            let body ← instantiateLambda body targets
-            lambdaTelescope1 body fun oldIH body => do
-              let body ← instantiateLambda body extraParams
-              let body' ← withRewrittenMotiveArg goal (rwFun #[name]) fun goal => do
-                buildInductionBody #[oldIH, genIH.fvarId!] #[] goal oldIH genIH.fvarId! isRecCall body
-              if body'.containsFVar oldIH then
-                throwError m!"Did not fully eliminate `{mkFVar oldIH}` from induction principle body:{indentExpr body}"
-              mkLambdaFVars (targets.push genIH) (← mkLambdaFVars extraParams body')
-        let e' := mkApp2 e' body' target
-        let e' ← mkLambdaFVars #[target] e'
-        let e' ← abstractIndependentMVars mvars (← motive.fvarId!.getDecl).index e'
-        let e' ← mkLambdaFVars #[motive] e'
-
-        -- We used to pass (usedOnly := false) below in the hope that the types of the
-        -- induction principle match the type of the function better.
-        -- But this leads to avoidable parameters that make functional induction strictly less
-        -- useful (e.g. when the unused parameter mentions bound variables in the users' goal)
-        let (paramMask, e') ← mkLambdaFVarsMasked fixedParamPerms e'
-        let e' ← instantiateMVars e'
-        return (e', paramMask)
-    | _ =>
-      if funBody.isAppOf ``WellFounded.fix then
-        throwError "Function `{name}` defined via `{.ofConstName ``WellFounded.fix}` with unexpected arity {funBody.getAppNumArgs}:{indentExpr funBody}"
-      else
-        throwError "Function `{name}` not defined via `{.ofConstName ``WellFounded.fix}`:{indentExpr funBody}"
+            throwError "Function {name} not defined via WellFounded.fix:{indentExpr funBody}"
+      check e'
+      let (body', mvars) ← M2.run do
+        forallTelescope (← inferType e').bindingDomain! fun xs goal => do
+          if xs.size ≠ 2 then
+            throwError "expected recursor argument to take 2 parameters, got {xs}" else
+          let targets : Array Expr := xs[*...1]
+          let genIH := xs[1]!
+          let extraParams := xs[2...*]
+          -- open body with the same arg
+          let body ← instantiateLambda body targets
+          lambdaTelescope1 body fun oldIH body => do
+            let body ← instantiateLambda body extraParams
+            let body' ← withRewrittenMotiveArg goal (rwFun #[name]) fun goal => do
+              buildInductionBody #[oldIH, genIH.fvarId!] #[] goal oldIH genIH.fvarId! isRecCall body
+            if body'.containsFVar oldIH then
+              throwError m!"Did not fully eliminate `{mkFVar oldIH}` from induction principle body:{indentExpr body}"
+            mkLambdaFVars (targets.push genIH) (← mkLambdaFVars extraParams body')
+      let e' := mkApp2 e' body' target
+      let e' ← mkLambdaFVars #[target] e'
+      let e' ← abstractIndependentMVars mvars (← motive.fvarId!.getDecl).index e'
+      let e' ← mkLambdaFVars #[motive] e'
+
+      -- We used to pass (usedOnly := false) below in the hope that the types of the
+      -- induction principle match the type of the function better.
+      -- But this leads to avoidable parameters that make functional induction strictly less
+      -- useful (e.g. when the unused parameter mentions bound variables in the users' goal)
+      let (paramMask, e') ← mkLambdaFVarsMasked fixedParamPerms e'
+      let e' ← instantiateMVars e'
+      return (e', paramMask)
 
   unless (← isTypeCorrect e') do
     logError m!"failed to derive a type-correct induction principle:{indentExpr e'}"

src/Std/Sat/AIG/Basic.lean

@@ -485,6 +485,7 @@ where
       let lval := go lhs.gate decls assign (by omega) h2
       let rval := go rhs.gate decls assign (by omega) h2
       xor lval lhs.invert && xor rval rhs.invert
+  termination_by (x, 0) -- Don't allow reduction, we have large concrete gate entries
 
 /--
 Denotation of an `AIG` at a specific `Entrypoint`.

stage0/src/stdlib_flags.h

@@ -1,5 +1,7 @@
 #include "util/options.h"
 
+// please update stage0
+
 namespace lean {
 options get_default_options() {
     options opts;

tests/lean/run/1026.lean

@@ -16,7 +16,7 @@ info: foo.eq_def (n : Nat) :
     if n = 0 then 0
     else
       have x := n - 1;
-      have this := foo._proof_4;
+      have this := foo._proof_3;
       foo x
 -/
 #guard_msgs in
@@ -28,7 +28,7 @@ info: foo.eq_unfold :
     if n = 0 then 0
     else
       have x := n - 1;
-      have this := foo._proof_4;
+      have this := foo._proof_3;
       foo x
 -/
 #guard_msgs in

tests/lean/run/4928.lean

@@ -46,7 +46,7 @@ end
 /--
 error: Failed: `fail` tactic was invoked
 x : List Nat
-⊢ (invImage (fun x => PSum.casesOn x (fun x => x.length) fun x => x.length) sizeOfWFRel).1 (PSum.inr x.tail)
+⊢ InvImage (fun x1 x2 => x1 < x2) (fun x => PSum.casesOn x (fun x => x.length) fun x => x.length) (PSum.inr x.tail)
     (PSum.inl x)
 -/
 #guard_msgs in