feat: omega understands multiplication on the right by a constant (#504)

Std/Tactic/Omega/Coeffs/IntList.lean

@@ -32,6 +32,8 @@ namespace Coeffs
 abbrev toList (xs : Coeffs) : List Int := xs
 /-- Identity, turning `List Int` into `Coeffs`. -/
 abbrev ofList (xs : List Int) : Coeffs := xs
+/-- Are the coefficients all zero? -/
+abbrev isZero (xs : Coeffs) : Prop := ∀ x, x ∈ xs → x = 0
 /-- Shim for `IntList.set`. -/
 abbrev set (xs : Coeffs) (i : Nat) (y : Int) : Coeffs := IntList.set xs i y
 /-- Shim for `IntList.get`. -/

Std/Tactic/Omega/Frontend.lean

@@ -157,17 +157,22 @@ partial def asLinearComboImpl (e : Expr) : OmegaM (LinearCombo × OmegaM Expr ×
         (← mkAppM ``Int.neg_congr #[← prf])
         (← mkEqSymm neg_eval)
     pure (-l, prf', facts)
-  | (``HMul.hMul, #[_, _, _, _, n, e']) =>
-    match intCast? n with
-    | some n' =>
-      let (l, prf, facts) ← asLinearCombo e'
-      let prf' : OmegaM Expr := do
-        let smul_eval := mkApp3 (.const ``LinearCombo.smul_eval []) (toExpr l) n (← atomsCoeffs)
+  | (``HMul.hMul, #[_, _, _, _, x, y]) =>
+    let (xl, xprf, xfacts) ← asLinearCombo x
+    let (yl, yprf, yfacts) ← asLinearCombo y
+    if xl.coeffs.isZero ∨ yl.coeffs.isZero then
+      let prf : OmegaM Expr := do
+        let h ← mkDecideProof (mkApp2 (.const ``Or [])
+          (.app (.const ``Coeffs.isZero []) (toExpr xl.coeffs))
+          (.app (.const ``Coeffs.isZero []) (toExpr yl.coeffs)))
+        let mul_eval :=
+          mkApp4 (.const ``LinearCombo.mul_eval []) (toExpr xl) (toExpr yl) (← atomsCoeffs) h
         mkEqTrans
-          (← mkAppM ``Int.mul_congr_right #[n, ← prf])
-          (← mkEqSymm smul_eval)
-      pure (l.smul n', prf', facts)
-    | none => mkAtomLinearCombo e
+          (← mkAppM ``Int.mul_congr #[← xprf, ← yprf])
+          (← mkEqSymm mul_eval)
+      pure (LinearCombo.mul xl yl, prf, xfacts.merge yfacts)
+    else
+      mkAtomLinearCombo e
   | (``HMod.hMod, #[_, _, _, _, n, k]) => rewrite e (mkApp2 (.const ``Int.emod_def []) n k)
   | (``HDiv.hDiv, #[_, _, _, _, x, z]) =>
     match intCast? z with

Std/Tactic/Omega/Int.lean

@@ -34,8 +34,11 @@ protected alias ⟨_, le_lt_asymm⟩ := Int.not_lt
 theorem add_congr {a b c d : Int} (h₁ : a = b) (h₂ : c = d) : a + c = b + d := by
   subst h₁; subst h₂; rfl
 
-theorem mul_congr_right (a : Int) {c d : Int} (h₂ : c = d) : a * c = a * d := by
-  subst h₂; rfl
+theorem mul_congr {a b c d : Int} (h₁ : a = b) (h₂ : c = d) : a * c = b * d := by
+  subst h₁; subst h₂; rfl
+
+theorem mul_congr_left {a b : Int} (h₁ : a = b)  (c : Int) : a * c = b * c := by
+  subst h₁; rfl
 
 theorem sub_congr {a b c d : Int} (h₁ : a = b) (h₂ : c = d) : a - c = b - d := by
   subst h₁; subst h₂; rfl

Std/Tactic/Omega/LinearCombo.lean

@@ -151,3 +151,33 @@ instance : HMul Int LinearCombo LinearCombo := ⟨fun i lc => lc.smul i⟩
     (i * lc).eval v = i * lc.eval v := by
   rcases lc with ⟨a, coeffs⟩
   simp [eval, Int.mul_add]
+
+theorem smul_eval_comm (lc : LinearCombo) (i : Int) (v : Coeffs) :
+    (i * lc).eval v = lc.eval v * i := by
+  simp [Int.mul_comm]
+
+/--
+Multiplication of two linear combinations.
+This is useful only if at least one of the linear combinations is constant,
+and otherwise should be considered as a junk value.
+-/
+def mul (l₁ l₂ : LinearCombo) : LinearCombo :=
+  l₂.const * l₁ + l₁.const * l₂ - { const := l₁.const * l₂.const }
+
+theorem mul_eval_of_const_left (l₁ l₂ : LinearCombo) (v : Coeffs) (w : l₁.coeffs.isZero) :
+    (mul l₁ l₂).eval v = l₁.eval v * l₂.eval v := by
+  have : Coeffs.dot l₁.coeffs v = 0 := IntList.dot_of_left_zero w
+  simp [mul, eval, this, Coeffs.sub_eq_add_neg, Coeffs.dot_distrib_left, Int.add_mul, Int.mul_add,
+    Int.mul_comm]
+
+theorem mul_eval_of_const_right (l₁ l₂ : LinearCombo) (v : Coeffs) (w : l₂.coeffs.isZero) :
+    (mul l₁ l₂).eval v = l₁.eval v * l₂.eval v := by
+  have : Coeffs.dot l₂.coeffs v = 0 := IntList.dot_of_left_zero w
+  simp [mul, eval, this, Coeffs.sub_eq_add_neg, Coeffs.dot_distrib_left, Int.add_mul, Int.mul_add,
+    Int.mul_comm]
+
+theorem mul_eval (l₁ l₂ : LinearCombo) (v : Coeffs) (w : l₁.coeffs.isZero ∨ l₂.coeffs.isZero) :
+    (mul l₁ l₂).eval v = l₁.eval v * l₂.eval v := by
+  rcases w with w | w
+  · rw [mul_eval_of_const_left _ _ _ w]
+  · rw [mul_eval_of_const_right _ _ _ w]

test/omega/test.lean

@@ -48,6 +48,11 @@ example {x y : Int} (_ : 6 * x + 7 * y = 5) : True := by (fail_if_success omega)
 
 example {x y : Nat} (_ : 6 * x + 7 * y = 5) : False := by omega
 
+example {x y : Nat} (_ : x * 6 + y * 7 = 5) : False := by omega
+example {x y : Nat} (_ : 2 * (3 * x) + y * 7 = 5) : False := by omega
+example {x y : Nat} (_ : 2 * x * 3 + y * 7 = 5) : False := by omega
+example {x y : Nat} (_ : 2 * 3 * x + y * 7 = 5) : False := by omega
+
 example {x : Nat} (_ : x < 0) : False := by omega
 
 example {x y z : Int} (_ : x + y > z) (_ : x < 0) (_ : y < 0) (_ : z > 0) : False := by omega