From 903493799dd943193aa019e7df042b1e656203a7 Mon Sep 17 00:00:00 2001
From: Joe Hendrix <jhendrix@whoisjoe.info>
Date: Fri, 5 Jan 2024 10:08:26 -0800
Subject: [PATCH] fix: reduceNat? match terms with free or meta variables
 (#3139)

This removes checks in `Lean.Meta.reduceNat?` that caused it to fail on
terms it could handle because they contain meta variables in arguments.
This lead to those operations being reduced using their equational
definitions and slow performance on large patterns:

```
set_option profiler true
set_option profiler.threshold 1

def testMod (x:Nat) :=
  match x with
  | 128 % 1024 => true
  | _ => false
-- elaboration took 3.02ms

def testMul (x:Nat) :=
  match x with
  | 128 * 1 => true
  | _ => false
-- type checking took 11.1ms
-- compilation of testMul.match_1 took 313ms
-- compilation of testMul took 65.7ms
-- elaboration took 58.9ms
```

Performance is slower on `testMul` than `testMod` because `whnf` ends up
evaluateing `128 * 1` using Peano arithmetic while `128 % 1024` is able
to avoid that treatment since `128 < 1024`.
---
 src/Lean/Meta/WHNF.lean                 | 39 ++++++++++++-------------
 tests/elabissues/reduceNatWithMeta.lean | 17 +++++++++++
 2 files changed, 36 insertions(+), 20 deletions(-)
 create mode 100644 tests/elabissues/reduceNatWithMeta.lean

diff --git a/src/Lean/Meta/WHNF.lean b/src/Lean/Meta/WHNF.lean
index f721a1d4aa..95d2f9df3d 100644
--- a/src/Lean/Meta/WHNF.lean
+++ b/src/Lean/Meta/WHNF.lean
@@ -850,27 +850,26 @@ def reduceBinNatPred (f : Nat → Nat → Bool) (a b : Expr) : MetaM (Option Exp
   return toExpr <| f a b
 
 def reduceNat? (e : Expr) : MetaM (Option Expr) :=
-  if e.hasFVar || e.hasMVar then
-    return none
-  else match e with
-    | .app (.const fn _) a                =>
-      if fn == ``Nat.succ then
-        reduceUnaryNatOp Nat.succ a
-      else
-        return none
-    | .app (.app (.const fn _) a1) a2 =>
-      if fn == ``Nat.add then reduceBinNatOp Nat.add a1 a2
-      else if fn == ``Nat.sub then reduceBinNatOp Nat.sub a1 a2
-      else if fn == ``Nat.mul then reduceBinNatOp Nat.mul a1 a2
-      else if fn == ``Nat.div then reduceBinNatOp Nat.div a1 a2
-      else if fn == ``Nat.mod then reduceBinNatOp Nat.mod a1 a2
-      else if fn == ``Nat.pow then reduceBinNatOp Nat.pow a1 a2
-      else if fn == ``Nat.gcd then reduceBinNatOp Nat.gcd a1 a2
-      else if fn == ``Nat.beq then reduceBinNatPred Nat.beq a1 a2
-      else if fn == ``Nat.ble then reduceBinNatPred Nat.ble a1 a2
-      else return none
-    | _ =>
+  match e with
+  | .app (.const fn _) a =>
+    if fn == ``Nat.succ then
+      reduceUnaryNatOp Nat.succ a
+    else
       return none
+  | .app (.app (.const fn _) a1) a2 =>
+    match fn with
+    | ``Nat.add => reduceBinNatOp Nat.add a1 a2
+    | ``Nat.sub => reduceBinNatOp Nat.sub a1 a2
+    | ``Nat.mul => reduceBinNatOp Nat.mul a1 a2
+    | ``Nat.div => reduceBinNatOp Nat.div a1 a2
+    | ``Nat.mod => reduceBinNatOp Nat.mod a1 a2
+    | ``Nat.pow => reduceBinNatOp Nat.pow a1 a2
+    | ``Nat.gcd => reduceBinNatOp Nat.gcd a1 a2
+    | ``Nat.beq => reduceBinNatPred Nat.beq a1 a2
+    | ``Nat.ble => reduceBinNatPred Nat.ble a1 a2
+    | _ => return none
+  | _ =>
+    return none
 
 
 @[inline] private def useWHNFCache (e : Expr) : MetaM Bool := do
diff --git a/tests/elabissues/reduceNatWithMeta.lean b/tests/elabissues/reduceNatWithMeta.lean
new file mode 100644
index 0000000000..a3b1032acf
--- /dev/null
+++ b/tests/elabissues/reduceNatWithMeta.lean
@@ -0,0 +1,17 @@
+-- This validates that Lean is able to simplify patterns containing operations
+-- on ground natural literals.
+--
+-- This is a regression test for #3139.
+set_option maxHeartbeats 1000
+
+-- This fails without the fix and maxHeartbeats 1000.
+def testZeroAdd (x:Nat) :=
+  match x with
+  | 0 + 128 => true
+  | _ => false
+
+-- This succeeds in all cases
+def testAddZero (x:Nat) :=
+  match x with
+  | 128 + 0 => true
+  | _ => false