From 409c6cac4c20c32d512d2976b23f26ffd4800a03 Mon Sep 17 00:00:00 2001
From: Joachim Breitner <mail@joachim-breitner.de>
Date: Wed, 24 Jan 2024 09:37:16 +0100
Subject: [PATCH] fix: predefinition preprocessing: float .mdata out of
 non-unary applications (#3204)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Recursive predefinitions contains “rec app” markers as mdata in the
predefinitions,
but sometimes these get in the way of termination checking, when you
have
```
  [mdata (fun x => f)] arg
```

Therefore, the `preprocess` pass floats them out of applications
(originally
only for structural recursion, since #2818 also for well-founded
recursion).

But the code was incomplete: Because `Meta.transform` calls `post` on `f
x y` only
once (and not also on `f x`) one has to float out of nested applications
as well.

A consequence of this can be that in a recursive proof, `rw [foo]` does
not work
although `rw [foo _ _]` does.

Also adding the testcase where @david-christiansen and I stumbled over
this


(Maybe the two preprocess modules can be combined, now that #2973 is
landed, will try that
in a follow-up).
---
 .../PreDefinition/Structural/Preprocess.lean  | 12 ++---
 src/Lean/Elab/PreDefinition/WF/Main.lean      |  3 +-
 .../Elab/PreDefinition/WF/Preprocess.lean     | 18 ++++---
 tests/lean/run/2810.lean                      | 49 +++++++++++++++++++
 tests/lean/run/issue3204.lean                 | 17 +++++++
 5 files changed, 84 insertions(+), 15 deletions(-)
 create mode 100644 tests/lean/run/issue3204.lean

diff --git a/src/Lean/Elab/PreDefinition/Structural/Preprocess.lean b/src/Lean/Elab/PreDefinition/Structural/Preprocess.lean
index 7354102412e..24c6f17a1d2 100644
--- a/src/Lean/Elab/PreDefinition/Structural/Preprocess.lean
+++ b/src/Lean/Elab/PreDefinition/Structural/Preprocess.lean
@@ -41,14 +41,12 @@ def preprocess (e : Expr) (recFnName : Name) : CoreM Expr :=
         return .visit e.headBeta
       else
         return .continue)
-    (post := fun e =>
-      match e with
-      | .app (.mdata m f) a =>
+    (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 (.app f a))
-        else
-          return .done e
-      | _ => return .done e)
+          return .done (.mdata m (f.beta e.getAppArgs))
+      return .continue)
 
 
 end Lean.Elab.Structural
diff --git a/src/Lean/Elab/PreDefinition/WF/Main.lean b/src/Lean/Elab/PreDefinition/WF/Main.lean
index 6f69214c67c..879bf44a77f 100644
--- a/src/Lean/Elab/PreDefinition/WF/Main.lean
+++ b/src/Lean/Elab/PreDefinition/WF/Main.lean
@@ -81,7 +81,8 @@ private def isOnlyOneUnaryDef (preDefs : Array PreDefinition) (fixedPrefixSize :
     return false
 
 def wfRecursion (preDefs : Array PreDefinition) : TermElabM Unit := do
-  let preDefs ← preDefs.mapM fun preDef => return { preDef with value := (← preprocess preDef.value) }
+  let preDefs ← preDefs.mapM fun preDef =>
+    return { preDef with value := (← preprocess preDef.value) }
   let (unaryPreDef, fixedPrefixSize) ← withoutModifyingEnv do
     for preDef in preDefs do
       addAsAxiom preDef
diff --git a/src/Lean/Elab/PreDefinition/WF/Preprocess.lean b/src/Lean/Elab/PreDefinition/WF/Preprocess.lean
index afcd6f76233..54242a884bd 100644
--- a/src/Lean/Elab/PreDefinition/WF/Preprocess.lean
+++ b/src/Lean/Elab/PreDefinition/WF/Preprocess.lean
@@ -9,6 +9,12 @@ import Lean.Elab.RecAppSyntax
 namespace Lean.Elab.WF
 open Meta
 
+private def shouldBetaReduce (e : Expr) (recFnNames : Array Name) : Bool :=
+  if e.isHeadBetaTarget then
+    e.getAppFn.find? (fun e => recFnNames.any (e.isConstOf ·)) |>.isSome
+  else
+    false
+
 /--
 Preprocesses the expessions to improve the effectiveness of `wfRecursion`.
 
@@ -25,13 +31,11 @@ remove `let_fun`-lambdas that contain explicit termination proofs.
 -/
 def preprocess (e : Expr) : CoreM Expr :=
   Core.transform e
-    (post := fun e =>
-      match e with
-      | .app (.mdata m f) a =>
+    (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 (.app f a))
-        else
-          return .done e
-      | _ => return .done e)
+          return .done (.mdata m (f.beta e.getAppArgs))
+      return .continue)
 
 end Lean.Elab.WF
diff --git a/tests/lean/run/2810.lean b/tests/lean/run/2810.lean
index 0be282ac59b..d57b1f01573 100644
--- a/tests/lean/run/2810.lean
+++ b/tests/lean/run/2810.lean
@@ -3,11 +3,15 @@ Test that parentheses don't get in the way of structural recursion.
 https://github.com/leanprover/lean4/issues/2810
 -/
 
+namespace Unary
+
 def f (n : Nat) : Nat :=
   match n with
   | 0 => 0
   | n + 1 => (f) n
 
+-- TODO: How can we assert that this was compiled structurally?
+
 -- with beta-reduction
 def f2 (n : Nat) : Nat :=
   match n with
@@ -39,3 +43,48 @@ theorem f_zero' (n : Nat) : f n = 0 := by
   | .succ n  =>
     unfold f
     simp only [f_zero']
+
+end Unary
+
+namespace Binary
+
+def f (n m : Nat) : Nat :=
+  match n with
+  | 0 => 0
+  | n + 1 => (f) n m
+
+-- TODO: How can we assert that this was compiled structurally?
+
+-- with beta-reduction
+def f2 (n m : Nat) : Nat :=
+  match n with
+  | 0 => 0
+  | n + 1 => (fun n' => (f2) n' m) n
+
+-- structural recursion
+def f3 (n m : Nat) : Nat :=
+  match n with
+  | 0 => 0
+  | n + 1 => (f3) n m
+termination_by n
+
+-- Same with rewrite
+
+theorem f_zero (n m : Nat) : f n m = 0 := by
+  match n with
+  | .zero => rfl
+  | .succ n  =>
+    unfold f
+    rewrite [f_zero]
+    rfl
+
+-- Same with simp
+
+theorem f_zero' (n m : Nat) : f n m = 0 := by
+  match n with
+  | .zero => rfl
+  | .succ n  =>
+    unfold f
+    simp only [f_zero']
+
+end Binary
diff --git a/tests/lean/run/issue3204.lean b/tests/lean/run/issue3204.lean
new file mode 100644
index 00000000000..9a0309d801e
--- /dev/null
+++ b/tests/lean/run/issue3204.lean
@@ -0,0 +1,17 @@
+def zero_out (arr : Array Nat) (i : Nat) : Array Nat :=
+  if h : i < arr.size then
+    zero_out (arr.set ⟨i, h⟩ 0) (i + 1)
+  else
+    arr
+termination_by arr.size - i
+decreasing_by simp_wf; apply Nat.sub_succ_lt_self _ _ h
+
+-- set_option trace.Elab.definition true
+theorem size_zero_out (arr : Array Nat) (i : Nat) : (zero_out arr i).size = arr.size := by
+  unfold zero_out
+  split
+  · rw [size_zero_out]
+    rw [Array.size_set]
+  · rfl
+termination_by arr.size - i
+decreasing_by simp_wf; apply Nat.sub_succ_lt_self; assumption