From 578175298514ecd4490652a854ad75928414dfbc Mon Sep 17 00:00:00 2001
From: Gabriel Ebner <gebner@gebner.org>
Date: Tue, 7 Mar 2023 23:35:32 +0000
Subject: [PATCH] fix: offset unification with a+a+1

Fixes #2136
---
 src/Lean/Meta/Offset.lean | 53 ++++++++++++++++++---------------------
 tests/lean/run/2136.lean  | 14 +++++++++++
 2 files changed, 38 insertions(+), 29 deletions(-)
 create mode 100644 tests/lean/run/2136.lean

diff --git a/src/Lean/Meta/Offset.lean b/src/Lean/Meta/Offset.lean
index bed0af22c9..21ade50c12 100644
--- a/src/Lean/Meta/Offset.lean
+++ b/src/Lean/Meta/Offset.lean
@@ -60,48 +60,43 @@ where
         failure
     | _ => failure
 
-/-- Quick function for converting `e` into `s + k` s.t. `e` is definitionally equal to `Nat.add s k`. -/
-private partial def getOffsetAux (e : Expr) (top : Bool) : OptionT MetaM (Expr × Nat) := do
+mutual
+
+/--
+Quick function for converting `e` into `s + k` s.t. `e` is definitionally equal to `Nat.add s k`.
+This function always succeeds in finding such `s` and `k`
+(as a last resort it returns `e` and `0`).
+-/
+private partial def getOffset (e : Expr) : MetaM (Expr × Nat) :=
+  return (← isOffset? e).getD (e, 0)
+
+/--
+Similar to `getOffset` but returns `none` if the expression is not syntactically an offset.
+-/
+private partial def isOffset? (e : Expr) : OptionT MetaM (Expr × Nat) := do
   match e with
   | .app _ a => do
     let f := e.getAppFn
     match f with
-    | .mvar .. => withInstantiatedMVars e (getOffsetAux · top)
+    | .mvar .. => withInstantiatedMVars e isOffset?
     | .const c _ =>
       let nargs := e.getAppNumArgs
       if c == ``Nat.succ && nargs == 1 then
-        let (s, k) ← getOffsetAux a false
+        let (s, k) ← getOffset a
         pure (s, k+1)
       else if c == ``Nat.add && nargs == 2 then
         let v ← evalNat (e.getArg! 1)
-        let (s, k) ← getOffsetAux (e.getArg! 0) false
+        let (s, k) ← getOffset (e.getArg! 0)
         pure (s, k+v)
       else if (c == ``Add.add && nargs == 4) || (c == ``HAdd.hAdd && nargs == 6) then
-        getOffsetAux (← unfoldProjInst? e) false
+        isOffset? (← unfoldProjInst? e)
       else
-        guard (!top); return (e, 0)
-    | _ => guard (!top); return (e, 0)
-  | _ => guard (!top); return (e, 0)
-
-private def getOffset (e : Expr) : OptionT MetaM (Expr × Nat) :=
-  getOffsetAux e true
-
-private partial def isOffset (e : Expr) : OptionT MetaM (Expr × Nat) := do
-  match e with
-  | .app .. =>
-    let f := e.getAppFn
-    match f with
-    | .mvar ..   => withInstantiatedMVars e isOffset
-    | .const c _ =>
-      let nargs := e.getAppNumArgs
-      guard ((c == ``Nat.succ && nargs == 1) ||
-             (c == ``Nat.add && nargs == 2) ||
-             (c == ``Add.add && nargs == 4) ||
-             (c == ``HAdd.hAdd && nargs == 6))
-      getOffset e
+        failure
     | _ => failure
   | _ => failure
 
+end
+
 private def isNatZero (e : Expr) : MetaM Bool := do
   match (← evalNat e) with
   | some v => return v == 0
@@ -128,9 +123,9 @@ def isDefEqOffset (s t : Expr) : MetaM LBool := do
   if !(← getConfig).offsetCnstrs then
     return LBool.undef
   else
-    match (← isOffset s) with
+    match (← isOffset? s) with
     | some (s, k₁) =>
-      match (← isOffset t) with
+      match (← isOffset? t) with
       | some (t, k₂) => -- s+k₁ =?= t+k₂
         if k₁ == k₂ then
           isDefEq s t
@@ -150,7 +145,7 @@ def isDefEqOffset (s t : Expr) : MetaM LBool := do
     | none =>
       match (← evalNat s) with
       | some v₁ =>
-        match (← isOffset t) with
+        match (← isOffset? t) with
         | some (t, k₂) => -- v₁ =?= t+k₂
           if v₁ ≥ k₂ then
             isDefEq (mkNatLit <| v₁ - k₂) t
diff --git a/tests/lean/run/2136.lean b/tests/lean/run/2136.lean
new file mode 100644
index 0000000000..666eff05dd
--- /dev/null
+++ b/tests/lean/run/2136.lean
@@ -0,0 +1,14 @@
+opaque bar : Nat → Nat
+theorem foo : bar (Nat.succ a) = 0 := sorry
+example : bar (a + 1) = 0 := by with_reducible exact foo -- ok
+example : bar (a + a + 1) = 0 := by with_reducible exact foo -- should also work
+
+def factorial : Nat → Nat
+  | 0 => 1
+  | n + 1 => (n + 1) * factorial n
+
+example (a : Nat) : factorial (a + 1) = (a + 1) * factorial a := by
+  rw [factorial] -- works
+
+example (a b : Nat) : factorial ((a + b) + 1) = ((a + b) + 1) * factorial (a + b) := by
+  rw [factorial] -- should also work