From 5fb302db33cd84d487cf74a4bb1e0dee125143da Mon Sep 17 00:00:00 2001
From: Scott Morrison <scott.morrison@gmail.com>
Date: Fri, 15 Dec 2023 11:00:11 +1100
Subject: [PATCH] Expr.nat

---
 Std/Lean/Expr.lean | 38 ++++++++++++++++++++++++++++++++++++++
 1 file changed, 38 insertions(+)

diff --git a/Std/Lean/Expr.lean b/Std/Lean/Expr.lean
index b6768f190b..a2296fe1e0 100644
--- a/Std/Lean/Expr.lean
+++ b/Std/Lean/Expr.lean
@@ -135,3 +135,41 @@ def constName (e : Expr) : Name :=
 /-- Return the function (name) and arguments of an application. -/
 def getAppFnArgs (e : Expr) : Name × Array Expr :=
   withApp e λ e a => (e.constName, a)
+
+
+/-- Turns an expression that is a natural number literal into a natural number. -/
+def natLit! : Expr → Nat
+  | lit (Literal.natVal v) => v
+  | _                      => panic! "nat literal expected"
+
+/-- Turns an expression that is a constructor of `Int` applied to a natural number literal
+into an integer. -/
+def intLit! (e : Expr) : Int :=
+  if e.isAppOfArity ``Int.ofNat 1 then
+    e.appArg!.natLit!
+  else if e.isAppOfArity ``Int.negOfNat 1 then
+    .negOfNat e.appArg!.natLit!
+  else
+    panic! "not a raw integer literal"
+
+/--
+Checks if an expression is a "natural number in normal form",
+i.e. of the form `OfNat n`, where `n` matches `.lit (.natVal lit)` for some `lit`.
+and if so returns `lit`.
+-/
+-- Note that an `Expr.lit (.natVal n)` is not considered in normal form!
+def nat? (e : Expr) : Option Nat := do
+  guard <| e.isAppOfArity ``OfNat.ofNat 3
+  let lit (.natVal n) := e.appFn!.appArg! | none
+  n
+
+/--
+Checks if an expression is an "integer in normal form",
+i.e. either a natural number in normal form, or the negation of one,
+and if so returns the integer.
+-/
+def int? (e : Expr) : Option Int :=
+  if e.isAppOfArity ``Neg.neg 3 then
+    (- ·) <$> e.appArg!.nat?
+  else
+    e.nat?