From 7cce64ee70b299c8ed774be7ce666f26b1edf478 Mon Sep 17 00:00:00 2001
From: Scott Morrison <scott.morrison@gmail.com>
Date: Wed, 28 Feb 2024 16:41:29 +1100
Subject: [PATCH] feat: omega doesn't check for defeq atoms (#3525)

```
example (a : Nat) :
    (((a + (2 ^ 64 - 1)) % 2 ^ 64 + 1) * 8 - 1 - (a + (2 ^ 64 - 1)) % 2 ^ 64 * 8 + 1) = 8 := by
  omega
```
used to time out, and now is fast.

(We will probably make separate changes later so the defeq checks would
be fast in any case here.)
---
 src/Lean/Elab/Tactic/Omega/Frontend.lean |  2 +-
 src/Lean/Elab/Tactic/Omega/OmegaM.lean   | 16 +++++++++-------
 tests/lean/run/omega.lean                |  6 ++++--
 3 files changed, 14 insertions(+), 10 deletions(-)

diff --git a/src/Lean/Elab/Tactic/Omega/Frontend.lean b/src/Lean/Elab/Tactic/Omega/Frontend.lean
index 734aed0306..718a515057 100644
--- a/src/Lean/Elab/Tactic/Omega/Frontend.lean
+++ b/src/Lean/Elab/Tactic/Omega/Frontend.lean
@@ -68,7 +68,7 @@ def mkEvalRflProof (e : Expr) (lc : LinearCombo) : OmegaM Expr := do
 `e = (coordinate n).eval atoms`. -/
 def mkCoordinateEvalAtomsEq (e : Expr) (n : Nat) : OmegaM Expr := do
   if n < 10 then
-    let atoms := (← getThe State).atoms
+    let atoms ← atoms
     let tail ← mkListLit (.const ``Int []) atoms[n+1:].toArray.toList
     let lem := .str ``LinearCombo s!"coordinate_eval_{n}"
     mkEqSymm (mkAppN (.const lem []) (atoms[:n+1].toArray.push tail))
diff --git a/src/Lean/Elab/Tactic/Omega/OmegaM.lean b/src/Lean/Elab/Tactic/Omega/OmegaM.lean
index a73d859214..6e58833f13 100644
--- a/src/Lean/Elab/Tactic/Omega/OmegaM.lean
+++ b/src/Lean/Elab/Tactic/Omega/OmegaM.lean
@@ -51,7 +51,7 @@ structure Context where
 /-- The internal state for the `OmegaM` monad, recording previously encountered atoms. -/
 structure State where
   /-- The atoms up-to-defeq encountered so far. -/
-  atoms : Array Expr := #[]
+  atoms : HashMap Expr Nat := {}
 
 /-- An intermediate layer in the `OmegaM` monad. -/
 abbrev OmegaM' := StateRefT State (ReaderT Context MetaM)
@@ -76,10 +76,11 @@ def OmegaM.run (m : OmegaM α) (cfg : OmegaConfig) : MetaM α :=
 def cfg : OmegaM OmegaConfig := do pure (← read).cfg
 
 /-- Retrieve the list of atoms. -/
-def atoms : OmegaM (List Expr) := return (← getThe State).atoms.toList
+def atoms : OmegaM (Array Expr) := do
+  return (← getThe State).atoms.toArray.qsort (·.2 < ·.2) |>.map (·.1)
 
 /-- Return the `Expr` representing the list of atoms. -/
-def atomsList : OmegaM Expr := do mkListLit (.const ``Int []) (← atoms)
+def atomsList : OmegaM Expr := do mkListLit (.const ``Int []) (← atoms).toList
 
 /-- Return the `Expr` representing the list of atoms as a `Coeffs`. -/
 def atomsCoeffs : OmegaM Expr := do
@@ -243,15 +244,16 @@ Return its index, and, if it is new, a collection of interesting facts about the
 -/
 def lookup (e : Expr) : OmegaM (Nat × Option (HashSet Expr)) := do
   let c ← getThe State
-  for h : i in [:c.atoms.size] do
-    if ← isDefEq e c.atoms[i] then
-      return (i, none)
+  match c.atoms.find? e with
+  | some i => return (i, none)
+  | none =>
   trace[omega] "New atom: {e}"
   let facts ← analyzeAtom e
   if ← isTracingEnabledFor `omega then
     unless facts.isEmpty do
       trace[omega] "New facts: {← facts.toList.mapM fun e => inferType e}"
-  let i ← modifyGetThe State fun c => (c.atoms.size, { c with atoms := c.atoms.push e })
+  let i ← modifyGetThe State fun c =>
+    (c.atoms.size, { c with atoms := c.atoms.insert e c.atoms.size })
   return (i, some facts)
 
 end Omega
diff --git a/tests/lean/run/omega.lean b/tests/lean/run/omega.lean
index fa9f855f9c..f258fd6390 100644
--- a/tests/lean/run/omega.lean
+++ b/tests/lean/run/omega.lean
@@ -268,8 +268,6 @@ example (x y : Int) (h : x < y) : x ≠ y := by omega
 
 example (x y : Int) (h : x < y) : ¬ x = y := by omega
 
-example (x : Int) : id x ≥ x := by omega
-
 example (prime : Nat → Prop) (x y z : Int) (h1 : 2 * x + ((-3) * y) < 0) (h2 : (-4) * x + 2*  z < 0)
     (h3 : 12 * y + (-4) * z < 0) (_ : prime 7) : False := by omega
 
@@ -426,6 +424,10 @@ example (x : Int) (e : Nat) (hx : x < (2 : Int)^(e+1)) : x < 2^e * 2 := by omega
 example : 2^7 < 165 := by omega
 example (_ : x % 2^7 < 3) : x % 128 < 5 := by omega
 
+example (a : Nat) :
+    (((a + (2 ^ 64 - 1)) % 2 ^ 64 + 1) * 8 - 1 - (a + (2 ^ 64 - 1)) % 2 ^ 64 * 8 + 1) = 8 := by
+  omega
+
 /-! ### BitVec -/
 open BitVec