fix: remove special handling of numerals in DiscrTree

Fixes #2867
by removing special handling of `OfNat.ofNat` numerals from `DiscrTree`.

The special handling was added in 2b8e55c
to fix a bug represented in the test `tests/lean/run/bugNatLitDiscrTree.lean`,
but with #2916 fixed, that test case passes without the special handling.

Some special handling of `Nat.zero` and `Nat.succ` remains.

src/Lean/Meta/DiscrTree.lean

@@ -198,30 +198,27 @@ private partial def isNumeral (e : Expr) : Bool :=
       else if fName == ``Nat.zero && e.getAppNumArgs == 0 then true
       else false
 
-private partial def toNatLit? (e : Expr) : Option Literal :=
-  if isNumeral e then
-    if let some n := loop e then
-      some (.natVal n)
-    else
-      none
+/-- Returns `some k` if `e` consists of `Nat.zero` followed by `k` applications of `Nat.succ`.-/
+private partial def natConst? (e : Expr) : Option Nat :=
+  match_expr e with
+  | Nat.zero => some 0
+  | Nat.succ e' => (· + 1) <$> natConst? e'
+  | _ => none
+
+/-- If `e` consists of `Nat.zero` followed by `k` applications of `Nat.succ`,
+  returns `OfNat.ofNat (nat_lit k)`. Otherwise, returns `e` unchanged.
+
+  Remark:
+    We need to ensure that `Nat.zero` and `0` (i.e. `OfNat.ofNat (nat_lit 0)`) are keyed the same way
+    because unification sometimes converts the former into the latter.
+
+    For example, the unification problem `Nat.zero =?= ?n + 0` is resolved by setting `?n := 0`,
+    *not* `?n := Nat.zero`-/
+private def normalizeNatConst (e : Expr) : Expr :=
+  if let some n := natConst? e then
+    mkNatLit n
   else
-    none
-where
-  loop (e : Expr) : OptionT Id Nat := do
-    let f := e.getAppFn
-    match f with
-    | .lit (.natVal n) => return n
-    | .const fName .. =>
-      if fName == ``Nat.succ && e.getAppNumArgs == 1 then
-        let r ← loop e.appArg!
-        return r+1
-      else if fName == ``OfNat.ofNat && e.getAppNumArgs == 3 then
-        loop (e.getArg! 1)
-      else if fName == ``Nat.zero && e.getAppNumArgs == 0 then
-        return 0
-      else
-        failure
-    | _ => failure
+    e
 
 private def isNatType (e : Expr) : MetaM Bool :=
   return (← whnf e).isConstOf ``Nat
@@ -310,7 +307,10 @@ where
 
 /-- whnf for the discrimination tree module -/
 def reduceDT (e : Expr) (root : Bool) (config : WhnfCoreConfig) : MetaM Expr :=
-  if root then reduceUntilBadKey e config else reduce e config
+  if root then
+    reduceUntilBadKey e config
+  else
+    normalizeNatConst <$> reduce e config
 
 /- Remark: we use `shouldAddAsStar` only for nested terms, and `root == false` for nested terms -/
 
@@ -350,8 +350,6 @@ private def pushArgs (root : Bool) (todo : Array Expr) (e : Expr) (config : Whnf
       return (.lit v, todo)
     | .const c _ =>
       unless root do
-        if let some v := toNatLit? e then
-          return (.lit v, todo)
         if (← shouldAddAsStar c e) then
           return (.star, todo)
       let nargs := e.getAppNumArgs
@@ -468,10 +466,6 @@ def insertIfSpecific [BEq α] (d : DiscrTree α) (e : Expr) (v : α) (config : W
 
 private def getKeyArgs (e : Expr) (isMatch root : Bool) (config : WhnfCoreConfig) : MetaM (Key × Array Expr) := do
   let e ← reduceDT e root config
-  unless root do
-    -- See pushArgs
-    if let some v := toNatLit? e then
-      return (.lit v, #[])
   match e.getAppFn with
   | .lit v         => return (.lit v, #[])
   | .const c _     =>

tests/lean/run/discrTreeNumeral.lean

@@ -0,0 +1,144 @@
+/-!
+  # Simp lemmas can match numerals and constant offsets
+-/
+
+def MyProp (_ : Nat) : Prop := True
+
+variable (n : Nat)
+
+theorem myProp_zero : MyProp 0 = True := rfl
+theorem myProp_add_one : MyProp (n + 1) = True := rfl
+theorem myProp_ofNat : MyProp (OfNat.ofNat n) = True := rfl
+theorem myProp_ofNat_add_one : MyProp (OfNat.ofNat (n + 1)) = True := rfl
+
+/-- Lemmas about a specific numeral match that numeral. -/
+example : MyProp 0 := by
+  dsimp only [myProp_zero]
+
+/-- Lemmas about all numerals match a specific numeral.
+
+  (Regression test for https://github.com/leanprover/lean4/issues/2867)-/
+example : MyProp 0 := by
+  dsimp only [myProp_ofNat]
+
+/-- Lemmas about any natural number plus a constant offset match their left-hand side. -/
+example : MyProp (n + 1) := by
+  dsimp only [myProp_add_one]
+
+/-- Lemmas about any natural number plus a constant offset match a larger offset. -/
+example : MyProp (n + 2) := by
+  dsimp only [myProp_add_one]
+
+/-- Lemmas about any natural number plus a constant offset match the constant offset. -/
+example : MyProp 1 := by
+  dsimp only [myProp_add_one]
+
+/-- Lemmas about any natural number plus a constant offset match constants exceeding the offset. -/
+example : MyProp 2 := by
+  dsimp only [myProp_add_one]
+
+/-- Lemmas about numerals that express a lower-bound
+  through a constant offset match their left-hand side.-/
+example : MyProp (OfNat.ofNat (n + 1)) := by
+  dsimp only [myProp_ofNat_add_one]
+
+/-- Lemmas about numerals that express a lower-bound
+  through a constant offset match a larger offset.-/
+example : MyProp (OfNat.ofNat (n + 2)) := by
+  dsimp only [myProp_ofNat_add_one]
+
+/-- Lemmas about numerals that express a lower-bound
+  through a constant offset match the constant offset.-/
+example : MyProp 1 := by
+  dsimp only [myProp_ofNat_add_one]
+
+/-- Lemmas about numerals that express a lower-bound
+  through a constant offset match numerals exceeding the constant offset.-/
+example : MyProp 2 := by
+  dsimp only [myProp_ofNat_add_one]
+
+/-- Lemmas about `0` match `Nat.zero` since `Nat.zero = 0` is in the default simp set.-/
+example : MyProp Nat.zero := by
+  dsimp [myProp_zero]
+
+/-- Lemmas about `0` match `nat_lit 0` since `nat_lit 0 = 0` is in the default simp set.-/
+example : MyProp (nat_lit 0) := by
+  dsimp [myProp_zero]
+
+/-- Lemmas about any natural number plus a constant offset match their LHS in terms of `Nat.succ`
+  since `Nat.succ n = n + 1` is in the default simp set. -/
+example : MyProp (Nat.succ n) := by
+  dsimp [myProp_add_one]
+
+/-- If a general lemma about numerals is specialized to a particular one, it still applies.
+
+  This is a tricky case because it results in nested `OfNat.ofNat`.
+  For example, `(myProp_ofNat 0 : MyProp (OfNat.ofNat (OfNat.ofNat (nat_lit 0)))`.
+
+  This currently works because we're passing an expression instead of a declaration name,
+  so `simp` considers it a local hypothesis and bypasses precise discrimination tree matching.
+  -/
+example : MyProp 0 := by
+  dsimp only [myProp_ofNat 0]
+
+/-- If a general lemma about natural numbers plus a constant offset is specialized
+  to a constant base, it still applies.
+
+  This is a tricky case because it results in a non-simp-normal lemma statement.
+  For example, `(myProp_add_one 0 : MyProp (0 + 1))` can be simplified to `MyProp 1`.-/
+example : MyProp 1 := by
+  dsimp only [myProp_add_one 0]
+
+/-- If a general lemma about numerals that expresses a lower-bound
+  through a constant offset is specialized to a constant base, it still applies.-/
+example : MyProp 1 := by
+  dsimp only [myProp_ofNat_add_one 0]
+
+class MyClass (n : Nat) : Prop where
+
+instance myClass_succ [MyClass n] : MyClass (n.succ) := ⟨⟩
+
+section BaseZero
+
+local instance myClass_nat_zero : MyClass (Nat.zero) := ⟨⟩
+
+/-- A base instance for `P Nat.zero` and a step instance for `P n → P (n.succ)` can combine to reach any
+  sequence of `Nat.zero...succ.succ`.
+
+  Because the kernel has special handling for unifying `Nat.succ`,
+  this requires that they can also combine to reach any numeral.
+
+  For example, to reach `P (Nat.zero.succ.succ)`,
+  the first recursive step contains the unification problem `P (Nat.zero.succ.succ) =?= P (?n.succ)`,
+  which the kernel solves by setting `?n := 1`, *not* `?n := Nat.zero.succ`.
+  Thereafter, we need to be able to reach the numeral `1`. -/
+example : MyClass (Nat.zero.succ.succ) :=
+  inferInstance
+
+#check show Nat.zero.succ.succ = Nat.succ _ from rfl -- Nat.zero.succ.succ = Nat.succ 1
+
+example : MyClass 1 :=
+  inferInstance
+
+example : MyClass 0 :=
+  inferInstance
+
+end BaseZero
+
+section BaseOne
+
+local instance myClass_succ_nat_zero : MyClass (Nat.zero.succ) := ⟨⟩
+
+/-- A base instance for `P Nat.zero.succ`
+  and a step instance for `P n → P (n.succ)` can combine
+  to reach any sequence of `Nat.zero.succ...succ.succ`.
+
+  Because the kernel has special handling for unifying `Nat.succ`,
+  this requires that they can also combine to reach any numeral greater than `0`.-/
+example : MyClass (Nat.zero.succ.succ) :=
+  inferInstance
+
+example : MyClass 1 :=
+  inferInstance
+
+end BaseOne