Implement nat reduction in normalization

Lean/Egg/Core/Normalize.lean

@@ -4,15 +4,19 @@ open Lean Meta Core
 
 namespace Egg
 
--- Performs ζ-, β- and η-reduction, converts `Expr.proj`s to `Expr.app`s and removes `Expr.mdata`s.
+-- Performs ζ-, converts `Expr.proj`s to `Expr.app`s, and removes `Expr.mdata`s.
+-- Depending on the `cfg`, also performs β- and η-reduction and reduction of internalized natural
+-- number operations.
 partial def normalize (e : Expr) (cfg : Config.Normalization) : MetaM Expr :=
   go e
 where
   go : Expr → MetaM Expr
     | .mdata _ e    => go e
     | .app fn arg   => do
-      let app := .app (← go fn) (← go arg)
-      if cfg.betaReduceRws then betaReduce app else return app
+      let mut app := .app (← go fn) (← go arg)
+      if cfg.betaReduceRws then app ← betaReduce app
+      if cfg.natReduceRws  then app ← natReduce app
+      return app
     | .lam n ty b i => do
       withLocalDecl n i (← go ty) fun fvar => do
         let body ← mkLambdaFVars #[fvar] (← go <| b.instantiate1 fvar)
@@ -27,3 +31,8 @@ where
   expandProj (ty : Name) (ctor : Nat) (b : Expr) : MetaM Expr := do
     let some field := (getStructureFields (← getEnv) ty)[ctor]? | throwError "'Egg.normalize' failed to reduce proj"
     mkProjection b field
+
+  natReduce (e : Expr) : MetaM Expr := do
+    if let some n ← evalNat e
+    then return mkNatLit n
+    else return e

Lean/Egg/Tests/NatLit.lean

@@ -67,6 +67,25 @@ example : 12345 % 67890 = 12345 := by
 example : 12345 % 0 = 12345 := by
   egg
 
--- TODO: We currently only β- and η-reduce rewrites. We should also perform the nat-lit reductions.
+set_option egg.natReduceRws false in
 example (h : ∀ f : Nat → Nat, f (1 + 1) = x) : id 2 = x := by
-  rw [←h] -- egg [h]
+  fail_if_success egg [h]
+  apply h
+
+example (h : ∀ f : Nat → Nat, f (1 + 1) = x) : id 2 = x := by
+  egg [h]
+
+example (h : ∀ f : Nat → Nat, f (3 - 2) = x) : id 1 = x := by
+  egg [h]
+
+example (h : ∀ f : Nat → Nat, f (2 * 3) = x) : id 6 = x := by
+  egg [h]
+
+example (h : ∀ f : Nat → Nat, f (4 / 2) = x) : id 2 = x := by
+  egg [h]
+
+example (h : ∀ f : Nat → Nat, f (5 % 3) = x) : id 2 = x := by
+  egg [h]
+
+example (h : ∀ f : Nat → Nat, f (2 ^ 3) = x) : id 8 = x := by
+  egg [h]