fix: Nat literal bug at DiscrTree.lean

src/Lean/Meta/DiscrTree.lean

@@ -207,13 +207,16 @@ private def isOffset (fName : Name) (e : Expr) : MetaM Bool := do
   TODO: add hook for users adding their own functions for controlling `shouldAddAsStar`
   Different `DiscrTree` users may populate this set using, for example, attributes.
 
-  Remark: we currently tag `Nat.zero` and "offset" terms to avoid having to add special
+  Remark: we currently tag numeral and "offset" terms to avoid having to add special
   support for `Expr.lit` and offset terms.
   Example, suppose the discrimination tree contains the entry
   `Nat.succ ?m |-> v`, and we are trying to retrieve the matches for `Expr.lit (Literal.natVal 1) _`.
-  In this scenario, we want to retrieve `Nat.succ ?m |-> v` -/
+  In this scenario, we want to retrieve `Nat.succ ?m |-> v`
+
+  TODO: add better support for Nat literals. Using `star` is a cheap trick to avoid different ways of representing them.
+-/
 private def shouldAddAsStar (fName : Name) (e : Expr) : MetaM Bool := do
-  if fName == ``Nat.zero then
+  if isNumeral e then
     return true
   else
     isOffset fName e
@@ -313,8 +316,13 @@ private def pushArgs (root : Bool) (todo : Array Expr) (e : Expr) : MetaM (Key s
       let todo ← pushArgsAux info.paramInfo (nargs-1) e todo
       return (k, todo)
     match fn with
-    | .lit v         => return (.lit v, todo)
-    | .const c _     =>
+    | .lit v     =>
+      unless root do
+        if fn.isNatLit then
+          -- See comment at `shouldAddAsStar`
+          return (.star, todo)
+      return (.lit v, todo)
+    | .const c _ =>
       unless root do
         if (← shouldAddAsStar c e) then
           return (.star, todo)

tests/lean/run/bugNatLitDiscrTree.lean

@@ -0,0 +1,25 @@
+def f' (n : Nat) : Option { r : Nat // r ≤ n } :=
+  match n with
+  | 0   => some ⟨0, Nat.le_refl _⟩
+  | n+1 => match f' n with
+    | some ⟨m, h₁⟩ =>
+      have : m < n+1 := Nat.lt_of_le_of_lt h₁ (Nat.lt_succ_self _)
+      match f' m with
+      | some ⟨r, h₂⟩ => some ⟨r, Nat.le_trans h₂ (Nat.le_trans h₁ (Nat.le_succ _))⟩
+      | none => none
+    | none => none
+
+theorem f'_ne_none (n : Nat) : f' n ≠ none := by
+  match n with
+  | 0   => simp (config := { decide := false }) [f']; done
+  | n+1 =>
+    simp [f']
+    have ih₁ := f'_ne_none n
+    split
+    next m h₁ he =>
+      have : m < n+1 := Nat.lt_of_le_of_lt h₁ (Nat.lt_succ_self _)
+      have ih₂ := f'_ne_none m
+      split
+      next => simp
+      next h => contradiction
+    next => contradiction