feat: omega understands Prod.Lex (#511)

Co-authored-by: Mario Carneiro <di.gama@gmail.com>

Std/Data/Prod/Lex.lean

@@ -3,8 +3,20 @@ Copyright (c) 2022 Jannis Limperg. All rights reserved.
 Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Jannis Limperg
 -/
+import Std.Tactic.LeftRight
+import Std.Tactic.RCases
+import Std.Logic
 
-namespace Prod.Lex
+namespace Prod
+
+theorem lex_def (r : α → α → Prop) (s : β → β → Prop) {p q : α × β} :
+    Prod.Lex r s p q ↔ r p.1 q.1 ∨ p.1 = q.1 ∧ s p.2 q.2 :=
+  ⟨fun h => by cases h <;> simp [*], fun h =>
+    match p, q, h with
+    | (a, b), (c, d), Or.inl h => Lex.left _ _ h
+    | (a, b), (c, d), Or.inr ⟨e, h⟩ => by subst e; exact Lex.right _ h⟩
+
+namespace Lex
 
 instance [αeqDec : DecidableEq α] {r : α → α → Prop} [rDec : DecidableRel r]
     {s : β → β → Prop} [sDec : DecidableRel s] : DecidableRel (Prod.Lex r s)

Std/Tactic/Omega/Frontend.lean

@@ -190,6 +190,22 @@ partial def asLinearComboImpl (e : Expr) : OmegaM (LinearCombo × OmegaM Expr ×
     | (``HMod.hMod, #[_, _, _, _, a, b]) => rewrite e (mkApp2 (.const ``Int.ofNat_emod []) a b)
     | (``HSub.hSub, #[_, _, _, _, mkAppN (.const ``HSub.hSub _) #[_, _, _, _, a, b], c]) =>
       rewrite e (mkApp3 (.const ``Int.ofNat_sub_sub []) a b c)
+    | (``Prod.fst, #[_, β, p]) => match p with
+      | .app (.app (.app (.app (.const ``Prod.mk [0, v]) _) _) x) y =>
+        rewrite e (mkApp3 (.const ``Int.ofNat_fst_mk [v]) β x y)
+      | _ => mkAtomLinearCombo e
+    | (``Prod.snd, #[α, _, p]) => match p with
+      | .app (.app (.app (.app (.const ``Prod.mk [u, 0]) _) _) x) y =>
+        rewrite e (mkApp3 (.const ``Int.ofNat_snd_mk [u]) α x y)
+      | _ => mkAtomLinearCombo e
+    | _ => mkAtomLinearCombo e
+  | (``Prod.fst, #[α, β, p]) => match p with
+    | .app (.app (.app (.app (.const ``Prod.mk [u, v]) _) _) x) y =>
+      rewrite e (mkApp4 (.const ``Prod.fst_mk [u, v]) α x β y)
+    | _ => mkAtomLinearCombo e
+  | (``Prod.snd, #[α, β, p]) => match p with
+    | .app (.app (.app (.app (.const ``Prod.mk [u, v]) _) _) x) y =>
+      rewrite e (mkApp4 (.const ``Prod.snd_mk [u, v]) α x β y)
     | _ => mkAtomLinearCombo e
   | _ => mkAtomLinearCombo e
 where
@@ -207,7 +223,6 @@ where
     | none => panic! "Invalid rewrite rule in 'asLinearCombo'"
 
 end
-
 namespace MetaProblem
 
 /-- The trivial `MetaProblem`, with no facts to processs and a trivial `Problem`. -/
@@ -250,6 +265,7 @@ def addIntInequality (p : MetaProblem) (h y : Expr) : OmegaM MetaProblem := do
 
 /-- Given a fact `h` with type `¬ P`, return a more useful fact obtained by pushing the negation. -/
 def pushNot (h P : Expr) : MetaM (Option Expr) := do
+  let P ← whnfR P
   match P with
   | .forallE _ t b _ =>
     if (← isProp t) && (← isProp b) then
@@ -258,32 +274,35 @@ def pushNot (h P : Expr) : MetaM (Option Expr) := do
     else
       return none
   | .app _ _ =>
-    return match P.getAppFnArgs with
-    | (``LT.lt, #[.const ``Int [], _, x, y]) => some (mkApp3 (.const ``Int.le_of_not_lt []) x y h)
-    | (``LE.le, #[.const ``Int [], _, x, y]) => some (mkApp3 (.const ``Int.lt_of_not_le []) x y h)
-    | (``LT.lt, #[.const ``Nat [], _, x, y]) => some (mkApp3 (.const ``Nat.le_of_not_lt []) x y h)
-    | (``LE.le, #[.const ``Nat [], _, x, y]) => some (mkApp3 (.const ``Nat.lt_of_not_le []) x y h)
-    | (``GT.gt, #[.const ``Int [], _, x, y]) => some (mkApp3 (.const ``Int.le_of_not_lt []) y x h)
-    | (``GE.ge, #[.const ``Int [], _, x, y]) => some (mkApp3 (.const ``Int.lt_of_not_le []) y x h)
-    | (``GT.gt, #[.const ``Nat [], _, x, y]) => some (mkApp3 (.const ``Nat.le_of_not_lt []) y x h)
-    | (``GE.ge, #[.const ``Nat [], _, x, y]) => some (mkApp3 (.const ``Nat.lt_of_not_le []) y x h)
-    | (``Eq, #[.const ``Nat [], x, y]) => some (mkApp3 (.const ``Nat.lt_or_gt_of_ne []) x y h)
-    | (``Eq, #[.const ``Int [], x, y]) => some (mkApp3 (.const ``Int.lt_or_gt_of_ne []) x y h)
+    match P.getAppFnArgs with
+    | (``LT.lt, #[.const ``Int [], _, x, y]) =>
+      return some (mkApp3 (.const ``Int.le_of_not_lt []) x y h)
+    | (``LE.le, #[.const ``Int [], _, x, y]) =>
+      return some (mkApp3 (.const ``Int.lt_of_not_le []) x y h)
+    | (``LT.lt, #[.const ``Nat [], _, x, y]) =>
+      return some (mkApp3 (.const ``Nat.le_of_not_lt []) x y h)
+    | (``LE.le, #[.const ``Nat [], _, x, y]) =>
+      return some (mkApp3 (.const ``Nat.lt_of_not_le []) x y h)
+    | (``Eq, #[.const ``Nat [], x, y]) =>
+      return some (mkApp3 (.const ``Nat.lt_or_gt_of_ne []) x y h)
+    | (``Eq, #[.const ``Int [], x, y]) =>
+      return some (mkApp3 (.const ``Int.lt_or_gt_of_ne []) x y h)
+    | (``Prod.Lex, _) => return some (← mkAppM ``Prod.of_not_lex #[h])
     | (``Dvd.dvd, #[.const ``Nat [], _, k, x]) =>
-      some (mkApp3 (.const ``Nat.emod_pos_of_not_dvd []) k x h)
+      return some (mkApp3 (.const ``Nat.emod_pos_of_not_dvd []) k x h)
     | (``Dvd.dvd, #[.const ``Int [], _, k, x]) =>
       -- This introduces a disjunction that could be avoided by checking `k ≠ 0`.
-      some (mkApp3 (.const ``Int.emod_pos_of_not_dvd []) k x h)
-    | (``Or, #[P₁, P₂]) => some (mkApp3 (.const ``and_not_not_of_not_or []) P₁ P₂ h)
+      return some (mkApp3 (.const ``Int.emod_pos_of_not_dvd []) k x h)
+    | (``Or, #[P₁, P₂]) => return some (mkApp3 (.const ``and_not_not_of_not_or []) P₁ P₂ h)
     | (``And, #[P₁, P₂]) =>
-      some (mkApp5 (.const ``Decidable.or_not_not_of_not_and []) P₁ P₂
+      return some (mkApp5 (.const ``Decidable.or_not_not_of_not_and []) P₁ P₂
         (.app (.const ``Classical.propDecidable []) P₁)
         (.app (.const ``Classical.propDecidable []) P₂) h)
     | (``Iff, #[P₁, P₂]) =>
-      some (mkApp5 (.const ``Decidable.and_not_or_not_and_of_not_iff []) P₁ P₂
+      return some (mkApp5 (.const ``Decidable.and_not_or_not_and_of_not_iff []) P₁ P₂
         (.app (.const ``Classical.propDecidable []) P₁)
         (.app (.const ``Classical.propDecidable []) P₂) h)
-    | _ => none
+    | _ => return none
   | _ => return none
 
 /--
@@ -326,6 +345,7 @@ partial def addFact (p : MetaProblem) (h : Expr) : OmegaM (MetaProblem × Nat) :
       p.addFact (mkApp3 (.const ``Int.ofNat_le_of_le []) x y h)
     | (``Ne, #[.const ``Int [], x, y]) =>
       p.addFact (mkApp3 (.const ``Int.lt_or_gt_of_ne []) x y h)
+    | (``Prod.Lex, _) => p.addFact (← mkAppM ``Prod.of_lex #[h])
     | (``Dvd.dvd, #[.const ``Nat [], _, k, x]) =>
       p.addFact (mkApp3 (.const ``Nat.mod_eq_zero_of_dvd []) k x h)
     | (``Dvd.dvd, #[.const ``Int [], _, k, x]) =>

Std/Tactic/Omega/Int.lean

@@ -5,6 +5,7 @@ Authors: Scott Morrison
 -/
 import Std.Classes.Order
 import Std.Data.Int.Lemmas
+import Std.Data.Prod.Lex
 import Std.Tactic.LeftRight
 
 /-!
@@ -94,6 +95,10 @@ theorem add_nonnneg_iff_neg_le (a b : Int) : 0 ≤ a + b ↔ -b ≤ a := by
 theorem add_nonnneg_iff_neg_le' (a b : Int) : 0 ≤ a + b ↔ -a ≤ b := by
   rw [Int.add_comm, add_nonnneg_iff_neg_le]
 
+theorem ofNat_fst_mk {β} {x : Nat} {y : β} : (Prod.mk x y).fst = (x : Int) := rfl
+theorem ofNat_snd_mk {α} {x : α} {y : Nat} : (Prod.mk x y).snd = (y : Int) := rfl
+
+
 end Int
 
 namespace Nat
@@ -102,3 +107,19 @@ theorem lt_of_gt {x y : Nat} (h : x > y) : y < x := gt_iff_lt.mp h
 theorem le_of_ge {x y : Nat} (h : x ≥ y) : y ≤ x := ge_iff_le.mp h
 
 end Nat
+
+namespace Prod
+
+theorem of_lex (w : Prod.Lex r s p q) : r p.fst q.fst ∨ p.fst = q.fst ∧ s p.snd q.snd :=
+  (Prod.lex_def r s).mp w
+
+theorem of_not_lex {α} {r : α → α → Prop} [DecidableEq α] {β} {s : β → β → Prop}
+    {p q : α × β} (w : ¬ Prod.Lex r s p q) :
+    ¬ r p.fst q.fst ∧ (p.fst ≠ q.fst ∨ ¬ s p.snd q.snd) := by
+  rw [Prod.lex_def, not_or, Decidable.not_and] at w
+  exact w
+
+theorem fst_mk : (Prod.mk x y).fst = x := rfl
+theorem snd_mk : (Prod.mk x y).snd = y := rfl
+
+end Prod

test/omega/test.lean

@@ -96,7 +96,6 @@ example {x : Nat} : (x + 4) / 2 ≤ x + 2 := by omega
 example {x : Int} {m : Nat} (_ : 0 < m) (_ : ¬x % ↑m < (↑m + 1) / 2) : -↑m / 2 ≤ x % ↑m - ↑m := by
   omega
 
-
 example (h : (7 : Int) = 0) : False := by omega
 
 example (h : (7 : Int) ≤ 0) : False := by omega
@@ -295,6 +294,29 @@ example (p n p' n' : Nat) (h : p + n' = p' + n) : n + p' = n' + p := by
 
 example (a b c : Int) (h1 : 32 / a < b) (h2 : b < c) : 32 / a < c := by omega
 
+-- Test that we consume expression metadata when necessary.
+example : 0 = 0 := by
+  have : 0 = 0 := by omega
+  omega -- This used to fail.
+
+/-! ### `Prod.Lex` -/
+
+-- This example comes from the termination proof
+-- for `permutationsAux.rec` in `Mathlib.Data.List.Defs`.
+example {x y : Nat} : Prod.Lex (· < ·) (· < ·) (x, x) (Nat.succ y + x, Nat.succ y) := by omega
+
+-- We test the termination proof in-situ:
+def List.permutationsAux.rec' {C : List α → List α → Sort v} (H0 : ∀ is, C [] is)
+    (H1 : ∀ t ts is, C ts (t :: is) → C is [] → C (t :: ts) is) : ∀ l₁ l₂, C l₁ l₂
+  | [], is => H0 is
+  | t :: ts, is =>
+      H1 t ts is (permutationsAux.rec' H0 H1 ts (t :: is)) (permutationsAux.rec' H0 H1 is [])
+  termination_by _ ts is => (length ts + length is, length ts)
+  decreasing_by simp_wf; omega
+
+example {x y w z : Nat} (h : Prod.Lex (· < ·) (· < ·) (x + 1, y + 1) (w, z)) :
+    Prod.Lex (· < ·) (· < ·) (x, y) (w, z) := by omega
+
 -- Verify that we can handle `iff` statements in hypotheses:
 example (a b : Int) (h : a < 0 ↔ b < 0) (w : b > 3) : a ≥ 0 := by omega