fix: omega regression (#4989)

This is a better fix to the problem reported at
https://leanprover.zulipchat.com/#narrow/stream/287929-mathlib4/topic/nat.20fighting,
which itself had a problem as reported at
https://leanprover.zulipchat.com/#narrow/stream/270676-lean4/topic/omega.20regression/near/456539091.

(cherry picked from commit 215b4a6)

src/Init/Omega/Int.lean

@@ -83,6 +83,10 @@ theorem neg_congr {a b : Int} (h₁ : a = b) : -a = -b := by
 theorem lt_of_gt {x y : Int} (h : x > y) : y < x := gt_iff_lt.mp h
 theorem le_of_ge {x y : Int} (h : x ≥ y) : y ≤ x := ge_iff_le.mp h
 
+theorem ofNat_mul_nonneg {a b : Nat} : 0 ≤ (a : Int) * b := by
+  rw [← Int.ofNat_mul]
+  exact Int.ofNat_zero_le (a * b)
+
 theorem ofNat_sub_eq_zero {b a : Nat} (h : ¬ b ≤ a) : ((a - b : Nat) : Int) = 0 :=
   Int.ofNat_eq_zero.mpr (Nat.sub_eq_zero_of_le (Nat.le_of_lt (Nat.not_le.mp h)))
 

src/Lean/Elab/Tactic/Omega/Frontend.lean

@@ -101,9 +101,9 @@ partial def asLinearCombo (e : Expr) : OmegaM (LinearCombo × OmegaM Expr × Has
     trace[omega] "Found in cache: {e}"
     return (lc, prf, ∅)
   | none =>
-    let r ← asLinearComboImpl e
-    modifyThe Cache fun cache => (cache.insert e (r.1, r.2.1.run' cache))
-    pure r
+    let (lc, proof, r) ← asLinearComboImpl e
+    modifyThe Cache fun cache => (cache.insert e (lc, proof.run' cache))
+    pure (lc, proof, r)
 
 /--
 Translates an expression into a `LinearCombo`.
@@ -255,16 +255,9 @@ where
     | (``Nat.succ, #[n]) => rewrite e (.app (.const ``Int.ofNat_succ []) n)
     | (``HAdd.hAdd, #[_, _, _, _, a, b]) => rewrite e (mkApp2 (.const ``Int.ofNat_add []) a b)
     | (``HMul.hMul, #[_, _, _, _, a, b]) =>
-      -- Don't push the cast into a multiplication unless it produces a non-trivial linear combination.
-      let r? ← commitWhen do
-        let (lc, prf, r) ← rewrite e (mkApp2 (.const ``Int.ofNat_mul []) a b)
-        if lc.isAtom then
-          pure (none, false)
-        else
-          pure (some (lc, prf, r), true)
-      match r? with
-      | some r => pure r
-      | none => mkAtomLinearCombo e
+      let (lc, prf, r) ← rewrite e (mkApp2 (.const ``Int.ofNat_mul []) a b)
+      -- Add the fact that the multiplication is non-negative.
+      pure (lc, prf, r.insert (mkApp2 (.const ``Int.ofNat_mul_nonneg []) a b))
     | (``HDiv.hDiv, #[_, _, _, _, a, b]) => rewrite e (mkApp2 (.const ``Int.ofNat_ediv []) a b)
     | (``OfNat.ofNat, #[_, n, _]) => rewrite e (.app (.const ``Int.natCast_ofNat []) n)
     | (``HMod.hMod, #[_, _, _, _, a, b]) => rewrite e (mkApp2 (.const ``Int.ofNat_emod []) a b)

tests/lean/run/omega.lean

@@ -398,6 +398,9 @@ example (n : Nat) : 2 * (n * n) = n * n + n * n := by omega
 -- example (n : Nat) : 2 * n * n = n * n + n * n := by omega
 -- example (n : Nat) : n * 2 * n = n * n + n * n := by omega
 
+-- From https://leanprover.zulipchat.com/#narrow/stream/270676-lean4/topic/omega.20regression/near/456539091
+example (a : Nat) : a * 1 = a := by omega
+
 /-! ### Fin -/
 
 -- Test `<`