feat: better support for reducing Nat.rec (#3616)

closes #3022

With this commit, given the declaration
```
def foo : Nat → Nat
  | 0 => 2
  | n + 1 => foo n
```
when we unfold `foo (n+1)`, we now obtain `foo n` instead of `foo
(Nat.add n 0)`.

src/Lean/Expr.lean

@@ -2011,6 +2011,10 @@ private def natMulFn : Expr :=
   let nat := mkConst ``Nat
   mkApp4 (mkConst ``HMul.hMul [0, 0, 0]) nat nat nat (mkApp2 (mkConst ``instHMul [0]) nat (mkConst ``instMulNat))
 
+/-- Given `a : Nat`, returns `Nat.succ a` -/
+def mkNatSucc (a : Expr) : Expr :=
+  mkApp (mkConst ``Nat.succ) a
+
 /-- Given `a b : Nat`, returns `a + b` -/
 def mkNatAdd (a b : Expr) : Expr :=
   mkApp2 natAddFn a b

src/Lean/Meta/CtorRecognizer.lean

@@ -5,6 +5,7 @@ Authors: Leonardo de Moura
 -/
 prelude
 import Lean.Meta.LitValues
+import Lean.Meta.Offset
 
 namespace Lean.Meta
 
@@ -64,9 +65,17 @@ def constructorApp? (e : Expr) : MetaM (Option (ConstructorVal × Array Expr)) :
 
 /--
 Similar to `constructorApp?`, but on failure it puts `e` in WHNF and tries again.
+It also `isOffset?`
 -/
 def constructorApp'? (e : Expr) : MetaM (Option (ConstructorVal × Array Expr)) := do
-  if let some r ← constructorApp? e then
+  if let some (e, k) ← isOffset? e then
+    if k = 0 then
+      return none
+    else
+      let .ctorInfo val ← getConstInfo ``Nat.succ | return none
+      if k = 1 then return some (val, #[e])
+      else return some (val, #[mkNatAdd e (toExpr (k-1))])
+  else if let some r ← constructorApp? e then
     return some r
   else
     constructorApp? (← whnf e)

src/Lean/Meta/Offset.lean

@@ -67,7 +67,7 @@ private partial def getOffset (e : Expr) : MetaM (Expr × Nat) :=
 /--
 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
+partial def isOffset? (e : Expr) : OptionT MetaM (Expr × Nat) := do
   let add (a b : Expr) := do
     let v ← evalNat b
     let (s, k) ← getOffset a

src/Lean/Meta/WHNF.lean

@@ -8,6 +8,7 @@ import Lean.Structure
 import Lean.Util.Recognizers
 import Lean.Meta.GetUnfoldableConst
 import Lean.Meta.FunInfo
+import Lean.Meta.Offset
 import Lean.Meta.CtorRecognizer
 import Lean.Meta.Match.MatcherInfo
 import Lean.Meta.Match.MatchPatternAttr
@@ -161,6 +162,22 @@ private def toCtorWhenStructure (inductName : Name) (major : Expr) : MetaM Expr
 private def isWFRec (declName : Name) : Bool :=
   declName == ``Acc.rec || declName == ``WellFounded.rec
 
+/--
+Helper method for `reduceRec`.
+We use it to ensure we don't expose `Nat.add` when reducing `Nat.rec`.
+We we use the following trick, if `e` can be expressed as an offest `(a, k)` with `k > 0`,
+we create a new expression `Nat.succ e'` where `e'` is `a` for `k = 1`, or `a + (k-1)` for `k > 1`.
+See issue #3022
+-/
+private def cleanupNatOffsetMajor (e : Expr) : MetaM Expr := do
+  let some (e, k) ← isOffset? e | return e
+  if k = 0 then
+    return e
+  else if k = 1 then
+    return mkNatSucc e
+  else
+    return mkNatSucc (mkNatAdd e (toExpr (k - 1)))
+
 /-- Auxiliary function for reducing recursor applications. -/
 private def reduceRec (recVal : RecursorVal) (recLvls : List Level) (recArgs : Array Expr) (failK : Unit → MetaM α) (successK : Expr → MetaM α) : MetaM α :=
   let majorIdx := recVal.getMajorIdx
@@ -177,6 +194,7 @@ private def reduceRec (recVal : RecursorVal) (recLvls : List Level) (recArgs : A
     if recVal.k then
       major ← toCtorWhenK recVal major
     major := major.toCtorIfLit
+    major ← cleanupNatOffsetMajor major
     major ← toCtorWhenStructure recVal.getInduct major
     match getRecRuleFor recVal major with
     | some rule =>

tests/lean/445.lean.expected.out

@@ -1,15 +1,15 @@
 true
-(match Nat.add i 0 with
+(match i with
   | 0 => true
   | Nat.succ n => true) &&
-  f (Nat.add i 0)
+  f i
 if i < 5 then 0 else 1
-if i < 5 then 0 else g i (Nat.add j 0)
+if i < 5 then 0 else g i j
 i + 1
-i + h i (Nat.add j 0)
+i + h i j
 i + 1
 i + i * 2
-i + i * r i (Nat.add j 0)
 i + i * r i j
-let z := s i (Nat.add j 0);
+i + i * r i j
+let z := s i j;
 z + z

tests/lean/run/3022.lean

@@ -0,0 +1,13 @@
+def foo : Nat → Nat
+  | 0 => 2
+  | n + 1 => foo n
+
+example (n : Nat) : foo (n + 1) = 2 := by
+  unfold foo      -- should not produce `⊢ foo (Nat.add n 0) = 2`
+  guard_target =ₛ foo n = 2
+  sorry
+
+example (n : Nat) : foo (n + 1) = 2 := by
+  simp only [foo]      -- should not produce `⊢ foo (Nat.add n 0) = 2`
+  guard_target =ₛ foo n = 2
+  sorry

tests/lean/smartUnfolding.lean.expected.out

@@ -1,3 +1,3 @@
 x y : Nat
-h : Nat.add x 1 = Nat.add y 1
+h : x + 1 = y + 1
 ⊢ x = y

tests/lean/unfold1.lean.expected.out

@@ -1,3 +1,4 @@
+unfold1.lean:22:4-22:8: error: simp made no progress
 case succ
 x : Nat
 ih : isEven (2 * x) = true
@@ -8,4 +9,4 @@ ih : isEven (2 * x) = true
 case succ
 x : Nat
 ih : isEven (2 * x) = true
-⊢ isEven (Nat.add (2 * x) 0) = true
+⊢ isEven (2 * x) = true