feat: Guess lexicographic order for well-founded recursion (leanprove…

…r#2874)

This improves Lean’s capabilities to guess the termination measure for
well-founded
recursion, by also trying lexicographic orders.  For example:

    def ackermann (n m : Nat) := match n, m with
      | 0, m => m + 1
      | .succ n, 0 => ackermann n 1
      | .succ n, .succ m => ackermann n (ackermann (n + 1) m)

now just works.

The module docstring of `Lean.Elab.PreDefinition.WF.GuessLex` tells the
technical story.
Fixes leanprover#2837

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

@@ -12,6 +12,7 @@ import Lean.Elab.PreDefinition.WF.Rel
 import Lean.Elab.PreDefinition.WF.Fix
 import Lean.Elab.PreDefinition.WF.Eqns
 import Lean.Elab.PreDefinition.WF.Ite
+import Lean.Elab.PreDefinition.WF.GuessLex
 
 namespace Lean.Elab
 open WF
@@ -89,10 +90,17 @@ def wfRecursion (preDefs : Array PreDefinition) (wf? : Option TerminationWF) (de
     let preDefsDIte ← preDefs.mapM fun preDef => return { preDef with value := (← iteToDIte preDef.value) }
     let unaryPreDefs ← packDomain fixedPrefixSize preDefsDIte
     return (← packMutual fixedPrefixSize preDefs unaryPreDefs, fixedPrefixSize)
+
+  let wf ←
+    if let .some wf := wf? then
+      pure wf
+    else
+      guessLex preDefs unaryPreDef fixedPrefixSize decrTactic?
+
   let preDefNonRec ← forallBoundedTelescope unaryPreDef.type fixedPrefixSize fun prefixArgs type => do
     let type ← whnfForall type
     let packedArgType := type.bindingDomain!
-    elabWFRel preDefs unaryPreDef.declName fixedPrefixSize packedArgType wf? fun wfRel => do
+    elabWFRel preDefs unaryPreDef.declName fixedPrefixSize packedArgType wf fun wfRel => do
       trace[Elab.definition.wf] "wfRel: {wfRel}"
       let (value, envNew) ← withoutModifyingEnv' do
         addAsAxiom unaryPreDef

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

@@ -53,67 +53,21 @@ private partial def unpackUnary (preDef : PreDefinition) (prefixSize : Nat) (mva
       mvarId.rename fvarId varNames.back
   go 0 mvarId fvarId
 
-def getNumCandidateArgs (fixedPrefixSize : Nat) (preDefs : Array PreDefinition) : MetaM (Array Nat) := do
-  preDefs.mapM fun preDef =>
-    lambdaTelescope preDef.value fun xs _ =>
-      return xs.size - fixedPrefixSize
-
-/--
-  Given a predefinition with value `fun (x_₁ ... xₙ) (y_₁ : α₁)... (yₘ : αₘ) => ...`,
-  where `n = fixedPrefixSize`, return an array `A` s.t. `i ∈ A` iff `sizeOf yᵢ` reduces to a literal.
-  This is the case for types such as `Prop`, `Type u`, etc.
-  This arguments should not be considered when guessing a well-founded relation.
-  See `generateCombinations?`
--/
-def getForbiddenByTrivialSizeOf (fixedPrefixSize : Nat) (preDef : PreDefinition) : MetaM (Array Nat) :=
-  lambdaTelescope preDef.value fun xs _ => do
-    let mut result := #[]
-    for x in xs[fixedPrefixSize:], i in [:xs.size] do
-      try
-        let sizeOf ← whnfD (← mkAppM ``sizeOf #[x])
-        if sizeOf.isLit then
-         result := result.push i
-      catch _ =>
-        result := result.push i
-    return result
-
-def generateCombinations? (forbiddenArgs : Array (Array Nat)) (numArgs : Array Nat) (threshold : Nat := 32) : Option (Array (Array Nat)) :=
-  go 0 #[] |>.run #[] |>.2
-where
-  isForbidden (fidx : Nat) (argIdx : Nat) : Bool :=
-    if h : fidx < forbiddenArgs.size then
-       forbiddenArgs.get ⟨fidx, h⟩ |>.contains argIdx
-    else
-      false
-
-  go (fidx : Nat) : OptionT (ReaderT (Array Nat) (StateM (Array (Array Nat)))) Unit := do
-    if h : fidx < numArgs.size then
-      let n := numArgs.get ⟨fidx, h⟩
-      for argIdx in [:n] do
-        unless isForbidden fidx argIdx do
-          withReader (·.push argIdx) (go (fidx + 1))
-    else
-      modify (·.push (← read))
-      if (← get).size > threshold then
-        failure
-termination_by _ fidx => numArgs.size - fidx
-
-def elabWFRel (preDefs : Array PreDefinition) (unaryPreDefName : Name) (fixedPrefixSize : Nat) (argType : Expr) (wf? : Option TerminationWF) (k : Expr → TermElabM α) : TermElabM α := do
+def elabWFRel (preDefs : Array PreDefinition) (unaryPreDefName : Name) (fixedPrefixSize : Nat)
+    (argType : Expr) (wf : TerminationWF) (k : Expr → TermElabM α) : TermElabM α := do
   let α := argType
   let u ← getLevel α
   let expectedType := mkApp (mkConst ``WellFoundedRelation [u]) α
   trace[Elab.definition.wf] "elabWFRel start: {(← mkFreshTypeMVar).mvarId!}"
-  match wf? with
-  | some (TerminationWF.core wfStx) => withDeclName unaryPreDefName do
+  withDeclName unaryPreDefName do
+  match wf with
+  | TerminationWF.core wfStx =>
       let wfRel ← instantiateMVars (← withSynthesize <| elabTermEnsuringType wfStx expectedType)
       let pendingMVarIds ← getMVars wfRel
       discard <| logUnassignedUsingErrorInfos pendingMVarIds
       k wfRel
-  | some (TerminationWF.ext elements) => go expectedType elements
-  | none => guess expectedType
-where
-  go (expectedType : Expr) (elements : Array TerminationByElement) : TermElabM α :=
-    withDeclName unaryPreDefName <| withRef (getRefFromElems elements) do
+  | TerminationWF.ext elements =>
+    withRef (getRefFromElems elements) do
       let mainMVarId := (← mkFreshExprSyntheticOpaqueMVar expectedType).mvarId!
       let [fMVarId, wfRelMVarId, _] ← mainMVarId.apply (← mkConstWithFreshMVarLevels ``invImage) | throwError "failed to apply 'invImage'"
       let (d, fMVarId) ← fMVarId.intro1
@@ -127,38 +81,4 @@ where
       wfRelMVarId.assign wfRelVal
       k (← instantiateMVars (mkMVar mainMVarId))
 
-  generateElements (numArgs : Array Nat) (argCombination : Array Nat) : TermElabM (Array TerminationByElement) := do
-    let mut result := #[]
-    let var ← `(x)
-    let hole ← `(_)
-    for preDef in preDefs, numArg in numArgs, argIdx in argCombination, i in [:preDefs.size] do
-      let mut vars := #[var]
-      for _ in [:numArg - argIdx - 1] do
-        vars := vars.push hole
-      -- TODO: improve this.
-      -- The following trick allows a function `f` in a mutual block to invoke `g` appearing before it with the input argument.
-      -- We should compute the "right" order (if there is one) in the future.
-      let body ← if preDefs.size > 1 then `((sizeOf x, $(quote i))) else `(sizeOf x)
-      result := result.push {
-        ref := preDef.ref
-        declName := preDef.declName
-        vars := vars
-        body := body
-        implicit := false
-      }
-    return result
-
-  guess (expectedType : Expr) : TermElabM α := do
-    -- TODO: add support for lex
-    let numArgs ← getNumCandidateArgs fixedPrefixSize preDefs
-    -- TODO: include in `forbiddenArgs` arguments that are fixed but are not in the fixed prefix
-    let forbiddenArgs ← preDefs.mapM fun preDef => getForbiddenByTrivialSizeOf fixedPrefixSize preDef
-    -- TODO: add option to control the maximum number of cases to try
-    if let some combs := generateCombinations? forbiddenArgs numArgs then
-      for comb in combs do
-        let elements ← generateElements numArgs comb
-        if let some r ← observing? (go expectedType elements) then
-          return r
-    throwError "failed to prove termination, use `termination_by` to specify a well-founded relation"
-
 end Lean.Elab.WF

tests/lean/guessLexFailures.lean

@@ -0,0 +1,34 @@
+/-!
+A few cases where guessing the lexicographic order fails, and
+where we want to keep tabs on the output.
+-/
+
+def nonTerminating : Nat → Nat
+  | 0 => 0
+  | n => nonTerminating (.succ n)
+
+-- Saying decreasing_by forces Lean to use structural recursion, which gives a different
+-- error message
+def nonTerminating2 : Nat → Nat
+  | 0 => 0
+  | n => nonTerminating2 (.succ n)
+decreasing_by decreasing_tactic
+
+
+-- The GuessLex code does not like eta-contracted motives in `casesOn`.
+-- At the time of writing, the error message is swallowed
+-- When guessing the lexicographic order becomes more verbose this will improve.
+def FinPlus1 n := Fin (n + 1)
+def badCasesOn (n : Nat) : Fin (n + 1) :=
+   Nat.casesOn (motive := FinPlus1) n (⟨0,Nat.zero_lt_succ _⟩) (fun n => Fin.succ (badCasesOn n))
+decreasing_by decreasing_tactic
+-- termination_by badCasesOn n => n
+
+
+-- Like above, but now with a `casesOn` alternative with insufficient lambdas
+def Fin_succ_comp (f : (n : Nat) → Fin (n + 1)) : (n : Nat) → Fin (n + 2) := fun n => Fin.succ (f n)
+def badCasesOn2 (n : Nat) : Fin (n + 1) :=
+   Nat.casesOn (motive := fun n => Fin (n + 1)) n (⟨0,Nat.zero_lt_succ _⟩)
+      (Fin_succ_comp (fun n => badCasesOn2 n))
+decreasing_by decreasing_tactic
+-- termination_by badCasesOn2 n => n

tests/lean/guessLexFailures.lean.expected.out

@@ -0,0 +1,13 @@
+guessLexFailures.lean:8:9-8:33: error: fail to show termination for
+  nonTerminating
+with errors
+argument #1 was not used for structural recursion
+  failed to eliminate recursive application
+    nonTerminating (Nat.succ n)
+
+structural recursion cannot be used
+
+failed to prove termination, use `termination_by` to specify a well-founded relation
+guessLexFailures.lean:12:0-15:31: error: failed to prove termination, use `termination_by` to specify a well-founded relation
+guessLexFailures.lean:22:0-24:31: error: failed to prove termination, use `termination_by` to specify a well-founded relation
+guessLexFailures.lean:30:0-33:31: error: failed to prove termination, use `termination_by` to specify a well-founded relation

tests/lean/run/guessLex.lean

@@ -0,0 +1,109 @@
+/-!
+This files tests Lean's ability to guess the right lexicographic order.
+
+TODO: Once lean spits out the guessed order (probably guarded by an
+option), turn this on and check the output.
+-/
+
+def ackermann (n m : Nat) := match n, m with
+  | 0, m => m + 1
+  | .succ n, 0 => ackermann n 1
+  | .succ n, .succ m => ackermann n (ackermann (n + 1) m)
+
+def ackermann2 (n m : Nat) := match n, m with
+  | m, 0 => m + 1
+  | 0, .succ n => ackermann2 1 n
+  | .succ m, .succ n => ackermann2 (ackermann2 m (n + 1)) n
+
+def ackermannList (n m : List Unit) := match n, m with
+  | [], m => () :: m
+  | ()::n, [] => ackermannList n [()]
+  | ()::n, ()::m => ackermannList n (ackermannList (()::n) m)
+
+def foo2 : Nat → Nat → Nat
+  | .succ n, 1 => foo2 n 1
+  | .succ n, 2 => foo2 (.succ n) 1
+  | n,       3 => foo2 (.succ n) 2
+  | .succ n, 4 => foo2 (if n > 10 then n else .succ n) 3
+  | n,       5 => foo2 (n - 1) 4
+  | n, .succ m => foo2 n m
+  | _, _ => 0
+
+mutual
+def even : Nat → Bool
+  | 0 => true
+  | .succ n => not (odd n)
+def odd : Nat → Bool
+  | 0 => false
+  | .succ n => not (even n)
+end
+
+mutual
+def evenWithFixed (m : String) : Nat → Bool
+  | 0 => true
+  | .succ n => not (oddWithFixed m n)
+def oddWithFixed (m : String) : Nat → Bool
+  | 0 => false
+  | .succ n => not (evenWithFixed m n)
+end
+
+def ping (n : Nat) := pong n
+   where pong : Nat → Nat
+  | 0 => 0
+  | .succ n => ping n
+
+def hasForbiddenArg (n : Nat) (_h : n = n) (m : Nat) : Nat :=
+  match n, m with
+  | 0, 0 => 0
+  | .succ m, n => hasForbiddenArg m rfl n
+  | m, .succ n => hasForbiddenArg (.succ m) rfl n
+
+/-!
+Example from “Finding Lexicographic Orders for Termination Proofs in
+Isabelle/HOL” by Lukas Bulwahn, Alexander Krauss, and Tobias Nipkow,
+10.1007/978-3-540-74591-4_5
+-/
+def blowup : Nat → Nat → Nat → Nat → Nat → Nat → Nat → Nat → Nat
+  | 0, 0, 0, 0, 0, 0, 0, 0 => 0
+  | 0, 0, 0, 0, 0, 0, 0, .succ i => .succ (blowup i i i i i i i i)
+  | 0, 0, 0, 0, 0, 0, .succ h, i => .succ (blowup h h h h h h h i)
+  | 0, 0, 0, 0, 0, .succ g, h, i => .succ (blowup g g g g g g h i)
+  | 0, 0, 0, 0, .succ f, g, h, i => .succ (blowup f f f f f g h i)
+  | 0, 0, 0, .succ e, f, g, h, i => .succ (blowup e e e e f g h i)
+  | 0, 0, .succ d, e, f, g, h, i => .succ (blowup d d d e f g h i)
+  | 0, .succ c, d, e, f, g, h, i => .succ (blowup c c d e f g h i)
+  | .succ b, c, d, e, f, g, h, i => .succ (blowup b c d e f g h i)
+
+-- Let’s try to confuse the lexicographic guessing function's
+-- unpacking of packed n-ary arguments
+def confuseLex1 : Nat → @PSigma Nat (fun _ => Nat) → Nat
+  | 0, _p => 0
+  | .succ n, ⟨x,y⟩ => confuseLex1 n ⟨x,y⟩
+
+def confuseLex2 : @PSigma Nat (fun _ => Nat) → Nat
+  | ⟨_y,0⟩ => 0
+  | ⟨y,.succ n⟩ => confuseLex2 ⟨y,n⟩
+
+def dependent : (n : Nat) → (m : Fin n) → Nat
+ | 0, i => Fin.elim0 i
+ | .succ 0, 0 => 0
+ | .succ (.succ n), 0 => dependent (.succ n) ⟨n, n.lt_succ_self⟩
+ | .succ (.succ n), ⟨.succ m, h⟩ =>
+  dependent (.succ (.succ n)) ⟨m, Nat.lt_of_le_of_lt (Nat.le_succ _) h⟩
+
+
+-- An example based on a real world problem, condensed by Leo
+inductive Expr where
+  | add (a b : Expr)
+  | val (n : Nat)
+
+mutual
+def eval (a : Expr) : Nat :=
+  match a with
+  | .add x y => eval_add (x, y)
+  | .val n => n
+
+def eval_add (a : Expr × Expr) : Nat :=
+  match a with
+  | (x, y) => eval x + eval y
+end

tests/lean/run/guessLexTricky.lean

@@ -0,0 +1,51 @@
+/-!
+A “tricky” example from “Finding Lexicographic Orders for Termination Proofs in
+Isabelle/HOL” by Lukas Bulwahn, Alexander Krauss, and Tobias Nipkow,
+10.1007/978-3-540-74591-4_5
+
+At the time of writing, Lean is able to find the lexicographic order
+just fine, but only if the tactic is powerful enough. In partiuclar,
+the default `decreasing_tactic` can only handle lexicographic descend when either
+the left gets smaller, or the left stays equal and the right gets smaller.
+But here we need to allow the general form, where the left is ≤ and the right
+gets smaller. This needs a backtracking proof search, it seems, which we build here
+(`search_lex`).
+-/
+
+macro_rules | `(tactic| decreasing_trivial) =>
+              `(tactic| apply Nat.le_refl)
+macro_rules | `(tactic| decreasing_trivial) =>
+              `(tactic| apply Nat.succ_lt_succ; decreasing_trivial)
+macro_rules | `(tactic| decreasing_trivial) =>
+              `(tactic| apply Nat.sub_le)
+macro_rules | `(tactic| decreasing_trivial) =>
+              `(tactic| apply Nat.div_le_self)
+
+syntax "search_lex " tacticSeq : tactic
+
+macro_rules | `(tactic|search_lex $ts:tacticSeq) => `(tactic| (
+    solve
+    | apply Prod.Lex.right'
+      · $ts
+      · search_lex $ts
+    | apply Prod.Lex.left
+      · $ts
+    | $ts
+    ))
+
+-- set_option trace.Elab.definition.wf true in
+mutual
+def prod (x y z : Nat) : Nat :=
+  if y % 2 = 0 then eprod x y z else oprod x y z
+def oprod (x y z : Nat) := eprod x (y - 1) (z + x)
+def eprod (x y z : Nat) := if y = 0 then z else prod (2 * x) (y / 2) z
+end
+-- termination_by
+--   prod x y z => (y, 2)
+--   oprod x y z => (y, 1)
+--   eprod x y z => (y, 0)
+decreasing_by
+  simp_wf
+  search_lex solve
+    | decreasing_trivial
+    | apply Nat.bitwise_rec_lemma; assumption

tests/lean/run/guessLexTricky2.lean

@@ -0,0 +1,23 @@
+/-!
+Another “tricky” example from “Finding Lexicographic Orders for Termination Proofs in
+Isabelle/HOL” by Lukas Bulwahn, Alexander Krauss, and Tobias Nipkow,
+10.1007/978-3-540-74591-4_5.
+
+Works out of the box!
+-/
+
+mutual
+def pedal : Nat → Nat → Nat → Nat
+  | 0, _m, c => c
+  | _n, 0, c => c
+  | n+1, m+1, c =>
+    if n < m
+    then coast n m (c + m + 1)
+    else pedal n (m + 1) (c + m + 1)
+
+def coast : Nat → Nat → Nat → Nat
+  | n, m , c =>
+  if n < m
+  then coast n (m - 1) (c + n)
+  else pedal n m (c + n)
+end