refactor: reduce omega's dependency on fvar IDs (#9723)

datokrat · web-flow · commit f15d531acbc1 · 2025-08-11T07:17:24.000Z
This PR replaces some `HashSet Expr`-typed collections of facts in
`omega`'s implementation with plain lists. This change makes some
`omega` calls faster, some slower, but the advantage is that `omega`'s
performance is more independent the state of the name generator that
produces fvar IDs.

I've created this PR for discussion and am happy to hear opinions on
whether this should be merged or not. A good reason *not* to merge is
that it causes regressions in some places and `grind` is expected to
supersede `omega` either way. A good reason to merge is that `omega` is
used all over the place and its flaky performance increases the noise in
future benchmarks.
diff --git a/src/Lean/Elab/Tactic/Omega/Frontend.lean b/src/Lean/Elab/Tactic/Omega/Frontend.lean
@@ -86,7 +86,7 @@ def mkCoordinateEvalAtomsEq (e : Expr) (n : Nat) : OmegaM Expr := do
     mkEqTrans eq (← mkEqSymm (mkApp2 (.const ``LinearCombo.coordinate_eval []) n atoms))
 
 /-- Construct the linear combination (and its associated proof and new facts) for an atom. -/
-def mkAtomLinearCombo (e : Expr) : OmegaM (LinearCombo × OmegaM Expr × Std.HashSet Expr) := do
+def mkAtomLinearCombo (e : Expr) : OmegaM (LinearCombo × OmegaM Expr × List Expr) := do
   let (n, facts) ← lookup e
   return ⟨LinearCombo.coordinate n, mkCoordinateEvalAtomsEq e n, facts.getD ∅⟩
 
@@ -100,7 +100,7 @@ Gives a small (10%) speedup in testing.
 I tried using a pointer based cache,
 but there was never enough subexpression sharing to make it effective.
 -/
-partial def asLinearCombo (e : Expr) : OmegaM (LinearCombo × OmegaM Expr × Std.HashSet Expr) := do
+partial def asLinearCombo (e : Expr) : OmegaM (LinearCombo × OmegaM Expr × List Expr) := do
   let cache ← get
   match cache.get? e with
   | some (lc, prf) =>
@@ -126,7 +126,7 @@ We also transform the expression as we descend into it:
 * pushing coercions: `↑(x + y)`, `↑(x * y)`, `↑(x / k)`, `↑(x % k)`, `↑k`
 * unfolding `emod`: `x % k` → `x - x / k`
 -/
-partial def asLinearComboImpl (e : Expr) : OmegaM (LinearCombo × OmegaM Expr × Std.HashSet Expr) := do
+partial def asLinearComboImpl (e : Expr) : OmegaM (LinearCombo × OmegaM Expr × List Expr) := do
   trace[omega] "processing {e}"
   match groundInt? e with
   | some i =>
@@ -148,7 +148,7 @@ partial def asLinearComboImpl (e : Expr) : OmegaM (LinearCombo × OmegaM Expr ×
       mkEqTrans
         (← mkAppM ``Int.add_congr #[← prf₁, ← prf₂])
         (← mkEqSymm add_eval)
-    pure (l₁ + l₂, prf, facts₁.union facts₂)
+    pure (l₁ + l₂, prf, facts₁ ++ facts₂)
   | (``HSub.hSub, #[_, _, _, _, e₁, e₂]) => do
     let (l₁, prf₁, facts₁) ← asLinearCombo e₁
     let (l₂, prf₂, facts₂) ← asLinearCombo e₂
@@ -158,7 +158,7 @@ partial def asLinearComboImpl (e : Expr) : OmegaM (LinearCombo × OmegaM Expr ×
       mkEqTrans
         (← mkAppM ``Int.sub_congr #[← prf₁, ← prf₂])
         (← mkEqSymm sub_eval)
-    pure (l₁ - l₂, prf, facts₁.union facts₂)
+    pure (l₁ - l₂, prf, facts₁ ++ facts₂)
   | (``Neg.neg, #[_, _, e']) => do
     let (l, prf, facts) ← asLinearCombo e'
     let prf' : OmegaM Expr := do
@@ -184,7 +184,7 @@ partial def asLinearComboImpl (e : Expr) : OmegaM (LinearCombo × OmegaM Expr ×
           mkEqTrans
             (← mkAppM ``Int.mul_congr #[← xprf, ← yprf])
             (← mkEqSymm mul_eval)
-        pure (some (LinearCombo.mul xl yl, prf, xfacts.union yfacts), true)
+        pure (some (LinearCombo.mul xl yl, prf, xfacts ++ yfacts), true)
       else
         pure (none, false)
     match r? with
@@ -242,15 +242,15 @@ where
   Apply a rewrite rule to an expression, and interpret the result as a `LinearCombo`.
   (We're not rewriting any subexpressions here, just the top level, for efficiency.)
   -/
-  rewrite (lhs rw : Expr) : OmegaM (LinearCombo × OmegaM Expr × Std.HashSet Expr) := do
+  rewrite (lhs rw : Expr) : OmegaM (LinearCombo × OmegaM Expr × List Expr) := do
     trace[omega] "rewriting {lhs} via {rw} : {← inferType rw}"
     match (← inferType rw).eq? with
     | some (_, _lhs', rhs) =>
       let (lc, prf, facts) ← asLinearCombo rhs
       let prf' : OmegaM Expr := do mkEqTrans rw (← prf)
       pure (lc, prf', facts)
     | none => panic! "Invalid rewrite rule in 'asLinearCombo'"
-  handleNatCast (e i n : Expr) : OmegaM (LinearCombo × OmegaM Expr × Std.HashSet Expr) := do
+  handleNatCast (e i n : Expr) : OmegaM (LinearCombo × OmegaM Expr × List Expr) := do
     match n with
     | .fvar h =>
       if let some v ← h.getValue? then
@@ -297,7 +297,7 @@ where
     | (``Fin.val, #[n, x]) =>
       handleFinVal e i n x
     | _ => mkAtomLinearCombo e
-  handleFinVal (e i n x : Expr) : OmegaM (LinearCombo × OmegaM Expr × Std.HashSet Expr) := do
+  handleFinVal (e i n x : Expr) : OmegaM (LinearCombo × OmegaM Expr × List Expr) := do
     match x with
     | .fvar h =>
       if let some v ← h.getValue? then
@@ -343,12 +343,11 @@ We solve equalities as they are discovered, as this often results in an earlier
 -/
 def addIntEquality (p : MetaProblem) (h x : Expr) : OmegaM MetaProblem := do
   let (lc, prf, facts) ← asLinearCombo x
-  let newFacts : Std.HashSet Expr := facts.fold (init := ∅) fun s e =>
-    if p.processedFacts.contains e then s else s.insert e
+  let newFacts : List Expr := facts.filter (p.processedFacts.contains · = false)
   trace[omega] "Adding proof of {lc} = 0"
   pure <|
   { p with
-    facts := newFacts.toList ++ p.facts
+    facts := newFacts ++ p.facts
     problem := ← (p.problem.addEquality lc.const lc.coeffs
       (some do mkEqTrans (← mkEqSymm (← prf)) h)) |>.solveEqualities }
 
@@ -359,12 +358,11 @@ We solve equalities as they are discovered, as this often results in an earlier
 -/
 def addIntInequality (p : MetaProblem) (h y : Expr) : OmegaM MetaProblem := do
   let (lc, prf, facts) ← asLinearCombo y
-  let newFacts : Std.HashSet Expr := facts.fold (init := ∅) fun s e =>
-    if p.processedFacts.contains e then s else s.insert e
+  let newFacts : List Expr := facts.filter (p.processedFacts.contains · = false)
   trace[omega] "Adding proof of {lc} ≥ 0"
   pure <|
   { p with
-    facts := newFacts.toList ++ p.facts
+    facts := newFacts ++ p.facts
     problem := ← (p.problem.addInequality lc.const lc.coeffs
       (some do mkAppM ``le_of_le_of_eq #[h, (← prf)])) |>.solveEqualities }
 
diff --git a/src/Lean/Elab/Tactic/Omega/OmegaM.lean b/src/Lean/Elab/Tactic/Omega/OmegaM.lean
@@ -168,11 +168,11 @@ def mkEqReflWithExpectedType (a b : Expr) : MetaM Expr := do
 Analyzes a newly recorded atom,
 returning a collection of interesting facts about it that should be added to the context.
 -/
-def analyzeAtom (e : Expr) : OmegaM (Std.HashSet Expr) := do
+def analyzeAtom (e : Expr) : OmegaM (List Expr) := do
   match e.getAppFnArgs with
   | (``Nat.cast, #[.const ``Int [], _, e']) =>
     -- Casts of natural numbers are non-negative.
-    let mut r := (∅ : Std.HashSet Expr).insert (Expr.app (.const ``Int.ofNat_nonneg []) e')
+    let mut r := [Expr.app (.const ``Int.ofNat_nonneg []) e']
     match (← cfg).splitNatSub, e'.getAppFnArgs with
       | true, (``HSub.hSub, #[_, _, _, _, a, b]) =>
         -- `((a - b : Nat) : Int)` gives a dichotomy
@@ -194,9 +194,8 @@ def analyzeAtom (e : Expr) : OmegaM (Std.HashSet Expr) := do
       let ne_zero := mkApp3 (.const ``Ne [1]) (.const ``Int []) k (toExpr (0 : Int))
       let pos := mkApp4 (.const ``LT.lt [0]) (.const ``Int []) (.const ``Int.instLTInt [])
         (toExpr (0 : Int)) k
-      pure <| (∅ : Std.HashSet Expr).insert
-        (mkApp3 (.const ``Int.mul_ediv_self_le []) x k (← mkDecideProof ne_zero)) |>.insert
-          (mkApp3 (.const ``Int.lt_mul_ediv_self_add []) x k (← mkDecideProof pos))
+      pure [mkApp3 (.const ``Int.mul_ediv_self_le []) x k (← mkDecideProof ne_zero),
+            mkApp3 (.const ``Int.lt_mul_ediv_self_add []) x k (← mkDecideProof pos)]
   | (``HMod.hMod, #[_, _, _, _, x, k]) =>
     match k.getAppFnArgs with
     | (``HPow.hPow, #[_, _, _, _, b, exp]) => match natCast? b with
@@ -206,10 +205,9 @@ def analyzeAtom (e : Expr) : OmegaM (Std.HashSet Expr) := do
         let b_pos := mkApp4 (.const ``LT.lt [0]) (.const ``Int []) (.const ``Int.instLTInt [])
           (toExpr (0 : Int)) b
         let pow_pos := mkApp3 (.const ``Lean.Omega.Int.pos_pow_of_pos []) b exp (← mkDecideProof b_pos)
-        pure <| (∅ : Std.HashSet Expr).insert
-          (mkApp3 (.const ``Int.emod_nonneg []) x k
-              (mkApp3 (.const ``Int.ne_of_gt []) k (toExpr (0 : Int)) pow_pos)) |>.insert
-            (mkApp3 (.const ``Int.emod_lt_of_pos []) x k pow_pos)
+        pure [mkApp3 (.const ``Int.emod_nonneg []) x k
+                (mkApp3 (.const ``Int.ne_of_gt []) k (toExpr (0 : Int)) pow_pos),
+              mkApp3 (.const ``Int.emod_lt_of_pos []) x k pow_pos]
     | (``Nat.cast, #[.const ``Int [], _, k']) =>
       match k'.getAppFnArgs with
       | (``HPow.hPow, #[_, _, _, _, b, exp]) => match natCast? b with
@@ -220,28 +218,25 @@ def analyzeAtom (e : Expr) : OmegaM (Std.HashSet Expr) := do
             (toExpr (0 : Nat)) b
           let pow_pos := mkApp3 (.const ``Nat.pos_pow_of_pos []) b exp (← mkDecideProof b_pos)
           let cast_pos := mkApp2 (.const ``Int.ofNat_pos_of_pos []) k' pow_pos
-          pure <| (∅ : Std.HashSet Expr).insert
-            (mkApp3 (.const ``Int.emod_nonneg []) x k
-                (mkApp3 (.const ``Int.ne_of_gt []) k (toExpr (0 : Int)) cast_pos)) |>.insert
-              (mkApp3 (.const ``Int.emod_lt_of_pos []) x k cast_pos)
+          pure [mkApp3 (.const ``Int.emod_nonneg []) x k
+                  (mkApp3 (.const ``Int.ne_of_gt []) k (toExpr (0 : Int)) cast_pos),
+                mkApp3 (.const ``Int.emod_lt_of_pos []) x k cast_pos]
       | _ =>  match x.getAppFnArgs with
         | (``Nat.cast, #[.const ``Int [], _, x']) =>
           -- Since we push coercions inside `%`, we need to record here that
           -- `(x : Int) % (y : Int)` is non-negative.
-          pure <| (∅ : Std.HashSet Expr).insert (mkApp2 (.const ``Int.emod_ofNat_nonneg []) x' k)
+          pure [mkApp2 (.const ``Int.emod_ofNat_nonneg []) x' k]
         | _ => pure ∅
     | _ => pure ∅
   | (``Min.min, #[_, _, x, y]) =>
-    pure <| (∅ : Std.HashSet Expr).insert (mkApp2 (.const ``Int.min_le_left []) x y) |>.insert
-      (mkApp2 (.const ``Int.min_le_right []) x y)
+    pure [mkApp2 (.const ``Int.min_le_left []) x y, mkApp2 (.const ``Int.min_le_right []) x y]
   | (``Max.max, #[_, _, x, y]) =>
-    pure <| (∅ : Std.HashSet Expr).insert (mkApp2 (.const ``Int.le_max_left []) x y) |>.insert
-      (mkApp2 (.const ``Int.le_max_right []) x y)
+    pure [mkApp2 (.const ``Int.le_max_left []) x y, mkApp2 (.const ``Int.le_max_right []) x y]
   | (``ite, #[α, i, dec, t, e]) =>
       if α == (.const ``Int []) then
-        pure <| (∅ : Std.HashSet Expr).insert <| mkApp5 (.const ``ite_disjunction [0]) α i dec t e
+        pure [mkApp5 (.const ``ite_disjunction [0]) α i dec t e]
       else
-        pure {}
+        pure []
   | _ => pure ∅
 
 /--
@@ -254,7 +249,7 @@ Return its index, and, if it is new, a collection of interesting facts about the
 * for each new atom of the form `((a - b : Nat) : Int)`, the fact:
   `b ≤ a ∧ ((a - b : Nat) : Int) = a - b ∨ a < b ∧ ((a - b : Nat) : Int) = 0`
 -/
-def lookup (e : Expr) : OmegaM (Nat × Option (Std.HashSet Expr)) := do
+def lookup (e : Expr) : OmegaM (Nat × Option (List Expr)) := do
   let c ← getThe State
   let e ← canon e
   match c.atoms[e]? with
@@ -264,7 +259,7 @@ def lookup (e : Expr) : OmegaM (Nat × Option (Std.HashSet Expr)) := do
   let facts ← analyzeAtom e
   if ← isTracingEnabledFor `omega then
     unless facts.isEmpty do
-      trace[omega] "New facts: {← facts.toList.mapM fun e => inferType e}"
+      trace[omega] "New facts: {← facts.mapM fun e => inferType e}"
   let i ← modifyGetThe State fun c =>
     (c.atoms.size, { c with atoms := c.atoms.insert e c.atoms.size })
   return (i, some facts)
