Prevent accidential reduction of open fuel

Author: nomeata

src/Init/WF.lean

@@ -446,12 +446,27 @@ 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 (h x + 1) x (Nat.lt_add_one _)
+  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
@@ -464,8 +479,10 @@ protected theorem Nat.fix.go_congr (x : α) (fuel₁ fuel₂ : Nat) (h₁ : h x
       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) :=
-  congrArg (F x) (funext fun _ => funext fun _ => Nat.fix.go_congr ..)
+    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
 

tests/lean/run/wfrec-nat.lean

@@ -0,0 +1,71 @@
+/-!
+Tests around the special case of well-founded recursion on Nat.
+-/
+
+namespace T1
+
+def foo : List α → Nat
+  | [] => 0
+  | _::xs => 1 + (foo xs)
+termination_by xs => xs.length
+
+-- Closed terms should evaluate
+
+example : foo ([] : List Unit) = 0 := rfl
+example : foo ([] : List Unit) = 0 := by decide
+example : foo ([] : List Unit) = 0 := by decide +kernel
+example : foo [1,2,3,4,5] = 5 := rfl
+example : foo [1,2,3,4,5] = 5 := by decide
+example : foo [1,2,3,4,5] = 5 := by decide +kernel
+
+-- Open terms should not (these wouldn't even without the provisions with `WellFounded.Nat.eager`,
+-- the fuel does not line up)
+
+example : foo (x::xs) = 1 + foo xs := by (fail_if_success rfl); simp [foo]
+example : foo (x::y::z::xs) = 1+ (1+(1+ foo xs)) := by (fail_if_success rfl); simp [foo]
+
+end T1
+
+-- Variant where the fuel does not line up
+
+namespace T2
+def foo : List α → Nat
+  | [] => 0
+  | _::xs => 1 + (foo xs)
+termination_by xs => 2 * xs.length
+
+example : foo ([] : List Unit) = 0 := rfl
+example : foo ([] : List Unit) = 0 := by decide
+example : foo ([] : List Unit) = 0 := by decide +kernel
+example : foo [1,2,3,4,5] = 5 := rfl
+example : foo [1,2,3,4,5] = 5 := by decide
+example : foo [1,2,3,4,5] = 5 := by decide +kernel
+
+-- Open terms should not (these wouldn't even without the provisions, the fuel does not line up)
+
+example : foo (x::xs) = 1 + foo xs := by (fail_if_success rfl); simp [foo]
+example : foo (x::y::z::xs) = 1+ (1 + ( 1+ foo xs)) := by (fail_if_success rfl); simp [foo]
+
+end T2
+
+-- Idiom to switch to `WellFounded.fix`
+
+namespace T3
+def foo : List α → Nat
+  | [] => 0
+  | _::xs => 1 + (foo xs)
+termination_by xs => (xs.length, 0)
+
+example : foo ([] : List Unit) = 0 := by (fail_if_success rfl); simp [foo]
+example : foo ([] : List Unit) = 0 := by (fail_if_success decide); simp [foo]
+example : foo ([] : List Unit) = 0 := by (fail_if_success decide +kernel); simp [foo]
+example : foo [1,2,3,4,5] = 5 := by (fail_if_success rfl); simp [foo]
+example : foo [1,2,3,4,5] = 5 := by (fail_if_success decide); simp [foo]
+example : foo [1,2,3,4,5] = 5 := by (fail_if_success decide +kernel); simp [foo]
+
+-- Open terms should not (these wouldn't even without the provisions, the fuel does not line up)
+
+example : foo (x::xs) = 1 + foo xs := by (fail_if_success rfl); simp [foo]
+example : foo (x::y::z::xs) = 1+ (1 + ( 1+ foo xs)) := by (fail_if_success rfl); simp [foo]
+
+end T3