Optimize Rat's polynorm and fix bugs. (#138)

* Optimize Rat's polynorm and fix bugs.

* Fix lean-cvc5 version.

* Compile Smt.Rat.

* Fix lean-cvc5 version.

* Fix Rat bug.

.github/workflows/ci.yml

@@ -25,6 +25,7 @@ jobs:
         lake update
         lake build
         lake build +Smt:dynlib
+        lake build +Smt.Rat:dynlib
         lake build +Smt.Real:dynlib
     - name: Test lean-smt
       run: lake run test

Smt/Reconstruct.lean

@@ -136,7 +136,7 @@ def addThm (type : Expr) (val : Expr) : ReconstructM Expr := do
   let name := Name.num `s (← incCount)
   let mv ← Meta.mkFreshExprMVar type .natural name
   mv.mvarId!.assign val
-  trace[smt.reconstruct.proof] "have {name} : {type} := {mv}"
+  trace[smt.reconstruct.proof] "have {name} : {type} := {val}"
   return mv
 
 def addTac (type : Expr) (tac : MVarId → MetaM Unit) : ReconstructM Expr := do
@@ -202,6 +202,7 @@ def solve (query : String) (timeout : Option Nat) : MetaM (Except Error cvc5.Pro
     Solver.setOption "dag-thresh" "0"
     Solver.setOption "simplification" "none"
     Solver.setOption "enum-inst" "true"
+    Solver.setOption "cegqi-midpoint" "true"
     Solver.setOption "produce-models" "true"
     Solver.setOption "produce-proofs" "true"
     Solver.setOption "proof-elim-subtypes" "true"

Smt/Reconstruct/Int/Polynorm.lean

@@ -256,8 +256,8 @@ inductive Expr where
   | var (v : Nat)
   | neg (a : Expr)
   | add (a b : Expr)
-  | mul (a b : Expr)
   | sub (a b : Expr)
+  | mul (a b : Expr)
 deriving Inhabited, Repr
 
 namespace Expr
@@ -372,7 +372,7 @@ def nativePolyNorm (mv : MVarId) : MetaM Unit := do
   let ctx : Q(PolyNorm.Context) ← pure q((«$is».getD · 0))
   let ⟨cache, el, _⟩ ← toQPolyNormExpr ctx es l {}
   let ⟨_, er, _⟩ ← toQPolyNormExpr ctx es r cache
-  let hp  ← nativeDecide q(«$el».toPolynomial = «$er».toPolynomial)
+  let hp ← nativeDecide q(«$el».toPolynomial = «$er».toPolynomial)
   let he := q(@PolyNorm.Expr.denote_eq_from_toPolynomial_eq $ctx $el $er $hp)
   mv.assign he
 where

Smt/Reconstruct/Prop.lean

@@ -246,9 +246,9 @@ def reconstructResolution (c₁ c₂ : Array cvc5.Term) (pol l : cvc5.Term) (hps
     let hij : Q($ps[$i] = ¬$qs[$j]) :=
       if pol.getBooleanValue then .app q(Prop.eq_not_not) q($ps[$i])
       else .app q(@Eq.refl Prop) q($ps[$i])
-    addThm (← rightAssocOp q(Or) (getResolutionResult c₁ c₂ pol l)) q(Prop.orN_resolution $hps $hqs $hi $hj $hij)
+    return q(Prop.orN_resolution $hps $hqs $hi $hj $hij)
   else
-    addThm (← rightAssocOp q(Or) (c₁ ++ c₂)) q(@Prop.orN_append_left $ps $qs $hps)
+    return q(@Prop.orN_append_left $ps $qs $hps)
 where
   rightAssocOp (op : Expr) (ts : Array cvc5.Term) : ReconstructM Expr := do
     if ts.isEmpty then
@@ -285,12 +285,12 @@ def reconstructChainResolution (cs as : Array cvc5.Term) (ps : Array Expr) : Rec
     let c₂ := clausify pf.getChildren[1]!.getResult l
     let hps ← reconstructProof pf.getChildren[0]!
     let hqs ← reconstructProof pf.getChildren[1]!
-    reconstructResolution c₁ c₂ p l hps hqs
+    addThm (← reconstructTerm pf.getResult) (← reconstructResolution c₁ c₂ p l hps hqs)
   | .CHAIN_RESOLUTION =>
     let cs := pf.getChildren.map (·.getResult)
     let as := pf.getArguments
     let ps ← pf.getChildren.mapM reconstructProof
-    reconstructChainResolution cs as ps
+    addThm (← reconstructTerm pf.getResult) (← reconstructChainResolution cs as ps)
   | .FACTORING =>
     let p : Q(Prop) ← reconstructTerm pf.getChildren[0]!.getResult
     let q : Q(Prop) ← reconstructTerm pf.getResult
@@ -345,11 +345,9 @@ def reconstructChainResolution (cs as : Array cvc5.Term) (ps : Array Expr) : Rec
     let f := fun pf ⟨q, hq⟩ => do
       let p : Q(Prop) ← reconstructTerm pf.getResult
       let hp : Q($p) ← reconstructProof pf
-      let hq ← addThm q($p ∧ $q) q(And.intro $hp $hq)
-      let q := q($p ∧ $q)
-      return ⟨q, hq⟩
-    let ⟨_, hq⟩ ← cpfs.pop.foldrM f (⟨q, hq⟩ : Σ q : Q(Prop), Q($q))
-    return hq
+      return ⟨q($p ∧ $q), q(And.intro $hp $hq)⟩
+    let ⟨q, hq⟩ ← cpfs.pop.foldrM f (⟨q, hq⟩ : Σ q : Q(Prop), Q($q))
+    addThm q hq
   | .NOT_OR_ELIM =>
     let f t ps := do
       let p : Q(Prop) ← reconstructTerm t

Smt/Reconstruct/Rat.lean

@@ -422,7 +422,6 @@ where
       | .GT =>
         addThm q(($x₁ > $x₂) = ($y₁ > $y₂)) q(Rat.gt_of_sub_eq_neg $hcx $hcy $h)
       | _   => return none
-    return none
   | .ARITH_MULT_SIGN =>
     if pf.getResult[1]![0]!.getSort.isInteger then return none
     reconstructMulSign pf

Smt/Reconstruct/Rat/Polynorm.lean

@@ -18,14 +18,32 @@ theorem Rat.neg_mul_eq_neg_mul (a b : Rat) : -(a * b) = -a * b := by
 
 namespace Smt.Reconstruct.Rat.PolyNorm
 
-abbrev Var := Nat
+structure Var where
+  type : Bool
+  val  : Nat
+deriving DecidableEq, Repr
+
+instance : LE Var where
+  le v₁ v₂ := v₁.type < v₂.type ∨ (v₁.type = v₂.type ∧ v₁.val ≤ v₂.val)
+
+instance : LT Var where
+  lt v₁ v₂ := v₁.type < v₂.type ∨ (v₁.type = v₂.type ∧ v₁.val < v₂.val)
+
+instance (v₁ v₂ : Var) : Decidable (v₁ ≤ v₂) :=
+  if h : v₁.type < v₂.type ∨ (v₁.type = v₂.type ∧ v₁.val ≤ v₂.val) then isTrue h else isFalse h
+
+instance (v₁ v₂ : Var) : Decidable (v₁ < v₂) :=
+  if h : v₁.type < v₂.type ∨ (v₁.type = v₂.type ∧ v₁.val < v₂.val) then isTrue h else isFalse h
 
 def Context := Var → Rat
 
+def IntContext := Nat → Int
+def RatContext := Nat → Rat
+
 structure Monomial where
   coeff : Rat
   vars : List Var
-deriving Inhabited, Repr
+deriving Inhabited, Repr, DecidableEq
 
 namespace Monomial
 
@@ -158,37 +176,87 @@ theorem denote_divConst {p : Polynomial} : (p.divConst c).denote ctx = p.denote
 
 end Polynomial
 
-inductive Expr where
-  | val (v : Rat)
+inductive IntExpr where
+  | val (v : Int)
   | var (v : Nat)
-  | neg (a : Expr)
-  | add (a b : Expr)
-  | sub (a b : Expr)
-  | mul (a b : Expr)
-  | divConst (a : Expr) (c : Rat)
+  | neg (a : IntExpr)
+  | add (a b : IntExpr)
+  | sub (a b : IntExpr)
+  | mul (a b : IntExpr)
 deriving Inhabited, Repr
 
-namespace Expr
+namespace IntExpr
 
-def toPolynomial : Expr → Polynomial
+def toPolynomial : IntExpr → Polynomial
   | val v => if v = 0 then [] else [{ coeff := v, vars := [] }]
-  | var v => [{ coeff := 1, vars := [v] }]
+  | var v => [{ coeff := 1, vars := [⟨false, v⟩] }]
   | neg a => a.toPolynomial.neg
   | add a b => Polynomial.add a.toPolynomial b.toPolynomial
   | sub a b => Polynomial.sub a.toPolynomial b.toPolynomial
   | mul a b => Polynomial.mul a.toPolynomial b.toPolynomial
-  | divConst a c => Polynomial.divConst a.toPolynomial c
 
-def denote (ctx : Context) : Expr → Rat
+def denote (ctx : IntContext) : IntExpr → Int
   | val v => v
   | var v => ctx v
   | neg a => -a.denote ctx
   | add a b => a.denote ctx + b.denote ctx
   | sub a b => a.denote ctx - b.denote ctx
   | mul a b => a.denote ctx * b.denote ctx
-  | divConst a c => a.denote ctx / c
 
-theorem denote_toPolynomial {e : Expr} : e.denote ctx = e.toPolynomial.denote ctx := by
+theorem denote_toPolynomial {rctx : RatContext} {e : IntExpr} : e.denote ictx = e.toPolynomial.denote (fun ⟨b, n⟩ => if b then rctx n else ictx n) := by
+  induction e with
+  | val v =>
+    simp only [denote, toPolynomial]
+    split <;> rename_i hv
+    · rewrite [hv]; rfl
+    · simp [Polynomial.denote, Monomial.denote]
+  | var v =>
+    simp [denote, toPolynomial, Polynomial.denote, Monomial.denote]
+  | neg a ih =>
+    simp only [denote, toPolynomial, Polynomial.denote_neg, Rat.intCast_neg, ih]
+  | add a b ih₁ ih₂ =>
+    simp only [denote, toPolynomial, Polynomial.denote_add, Rat.intCast_add, ih₁, ih₂]
+  | sub a b ih₁ ih₂ =>
+    simp only [denote, toPolynomial, Polynomial.denote_sub, Rat.intCast_sub, ih₁, ih₂]
+  | mul a b ih₁ ih₂ =>
+    simp only [denote, toPolynomial, Polynomial.denote_mul, Rat.intCast_mul, ih₁, ih₂]
+
+end IntExpr
+
+inductive RatExpr where
+  | val (v : Rat)
+  | var (v : Nat)
+  | neg (a : RatExpr)
+  | add (a b : RatExpr)
+  | sub (a b : RatExpr)
+  | mul (a b : RatExpr)
+  | divConst (a : RatExpr) (c : Rat)
+  | cast (a : IntExpr)
+deriving Inhabited, Repr
+
+namespace RatExpr
+
+def toPolynomial : RatExpr → Polynomial
+  | val v => if v = 0 then [] else [{ coeff := v, vars := [] }]
+  | var v => [{ coeff := 1, vars := [⟨true, v⟩] }]
+  | neg a => a.toPolynomial.neg
+  | add a b => Polynomial.add a.toPolynomial b.toPolynomial
+  | sub a b => Polynomial.sub a.toPolynomial b.toPolynomial
+  | mul a b => Polynomial.mul a.toPolynomial b.toPolynomial
+  | divConst a c => Polynomial.divConst a.toPolynomial c
+  | cast a => a.toPolynomial
+
+def denote (ictx : IntContext) (rctx : RatContext)  : RatExpr → Rat
+  | val v => v
+  | var v => rctx v
+  | neg a => -a.denote ictx rctx
+  | add a b => a.denote ictx rctx + b.denote ictx rctx
+  | sub a b => a.denote ictx rctx - b.denote ictx rctx
+  | mul a b => a.denote ictx rctx * b.denote ictx rctx
+  | divConst a c => a.denote ictx rctx / c
+  | cast a => a.denote ictx
+
+theorem denote_toPolynomial {e : RatExpr} : e.denote ictx rctx = e.toPolynomial.denote (fun ⟨b, n⟩ => if b then rctx n else ictx n) := by
   induction e with
   | val v =>
     simp only [denote, toPolynomial]
@@ -207,24 +275,35 @@ theorem denote_toPolynomial {e : Expr} : e.denote ctx = e.toPolynomial.denote ct
     simp only [denote, toPolynomial, Polynomial.denote_mul, ih₁, ih₂]
   | divConst a c ih =>
     simp only [denote, toPolynomial, Polynomial.denote_divConst, ih]
+  | cast a =>
+    simpa only [denote] using IntExpr.denote_toPolynomial
 
-theorem denote_eq_from_toPolynomial_eq {e₁ e₂ : Expr} (h : e₁.toPolynomial = e₂.toPolynomial) : e₁.denote ctx = e₂.denote ctx := by
+theorem denote_eq_from_toPolynomial_eq {e₁ e₂ : RatExpr} (h : e₁.toPolynomial = e₂.toPolynomial) : e₁.denote ictx rctx = e₂.denote ictx rctx := by
   rw [denote_toPolynomial, denote_toPolynomial, h]
 
-end PolyNorm.Expr
+end PolyNorm.RatExpr
 
 open Lean hiding Rat
 open Qq
 
-abbrev PolyM := StateT (Array Q(Rat)) MetaM
+abbrev PolyM := StateT (Array Q(Int) × Array Q(Rat)) MetaM
+
+def getRatIndex (e : Q(Rat)) : PolyM Nat := do
+  let ⟨is, rs⟩ ← get
+  if let some i := rs.findIdx? (· == e) then
+    return i
+  else
+    let size := rs.size
+    set (is, rs.push e)
+    return size
 
-def getIndex (e : Q(Rat)) : PolyM Nat := do
-  let es ← get
-  if let some i := es.findIdx? (· == e) then
+def getIntIndex (e : Q(Rat)) : PolyM Nat := do
+  let ⟨is, rs⟩ ← get
+  if let some i := is.findIdx? (· == e) then
     return i
   else
-    let size := es.size
-    set (es.push e)
+    let size := is.size
+    set (is.push e, rs)
     return size
 
 partial def toRatConst (e : Q(Rat)) : PolyM Rat := do
@@ -237,40 +316,64 @@ partial def toRatConst (e : Q(Rat)) : PolyM Rat := do
   | ~q($x / $y) => pure ((← toRatConst x) / (← toRatConst y))
   | e => throwError "[poly_norm] expected a rational number, got {e}"
 
-partial def toPolyNormExpr (e : Q(Rat)) : PolyM PolyNorm.Expr := do
+partial def toQPolyNormIntExpr (e : Q(Int)) : PolyM Q(PolyNorm.IntExpr) := do
   match e with
-  | ~q(OfNat.ofNat $x) => pure (.val x.rawNatLit?.get!)
-  | ~q(-$x) => pure (.neg (← toPolyNormExpr x))
-  | ~q($x + $y) => pure (.add (← toPolyNormExpr x) (← toPolyNormExpr y))
-  | ~q($x - $y) => pure (.sub (← toPolyNormExpr x) (← toPolyNormExpr y))
-  | ~q($x * $y) => pure (.mul (← toPolyNormExpr x) (← toPolyNormExpr y))
-  | ~q($x / $y) => pure (.divConst (← toPolyNormExpr x) (← toRatConst y))
-  | e => let v : Nat ← getIndex e; pure (.var v)
-
-partial def toQPolyNormExpr (e : Q(Rat)) : PolyM Q(PolyNorm.Expr) := do
+  | ~q(OfNat.ofNat $n) => pure q(.val (@OfNat.ofNat Int $n _))
+  | ~q(-$x) => pure q(.neg $(← toQPolyNormIntExpr x))
+  | ~q($x + $y) => pure q(.add $(← toQPolyNormIntExpr x) $(← toQPolyNormIntExpr y))
+  | ~q($x - $y) => pure q(.sub $(← toQPolyNormIntExpr x) $(← toQPolyNormIntExpr y))
+  | ~q($x * $y) => pure q(.mul $(← toQPolyNormIntExpr x) $(← toQPolyNormIntExpr y))
+  | e => let v : Nat ← getIntIndex e; pure q(.var $v)
+
+partial def toQPolyNormRatExpr (e : Q(Rat)) : PolyM Q(PolyNorm.RatExpr) := do
   match e with
   | ~q(OfNat.ofNat $n) => pure q(.val (@OfNat.ofNat Rat $n _))
-  | ~q(-$x) => pure q(.neg $(← toQPolyNormExpr x))
-  | ~q($x + $y) => pure q(.add $(← toQPolyNormExpr x) $(← toQPolyNormExpr y))
-  | ~q($x - $y) => pure q(.sub $(← toQPolyNormExpr x) $(← toQPolyNormExpr y))
-  | ~q($x * $y) => pure q(.mul $(← toQPolyNormExpr x) $(← toQPolyNormExpr y))
-  | ~q($x / $y) => pure q(.divConst $(← toQPolyNormExpr x) $(PolyNorm.Monomial.toExpr.toExprCoeff (← toRatConst y)))
-  | e => let v : Nat ← getIndex e; pure q(.var $v)
+  | ~q(-$x) => pure q(.neg $(← toQPolyNormRatExpr x))
+  | ~q($x + $y) => pure q(.add $(← toQPolyNormRatExpr x) $(← toQPolyNormRatExpr y))
+  | ~q($x - $y) => pure q(.sub $(← toQPolyNormRatExpr x) $(← toQPolyNormRatExpr y))
+  | ~q($x * $y) => pure q(.mul $(← toQPolyNormRatExpr x) $(← toQPolyNormRatExpr y))
+  | ~q($x / $y) => pure q(.divConst $(← toQPolyNormRatExpr x) $(PolyNorm.Monomial.toExpr.toExprCoeff (← toRatConst y)))
+  | ~q(($x : Int)) => pure q(.cast $(← toQPolyNormIntExpr x))
+  | e => let v : Nat ← getRatIndex e; pure q(.var $v)
 
 def polyNorm (mv : MVarId) : MetaM Unit := do
   let some (_, l, r) := (← mv.getType).eq?
     | throwError "[poly_norm] expected an equality, got {← mv.getType}"
-  let (l, es) ← (toQPolyNormExpr l).run #[]
-  let (r, es) ← (toQPolyNormExpr r).run es
-  let is : Q(Array Rat) := es.foldl (fun acc e => q(«$acc».push $e)) q(#[])
-  let ctx : Q(PolyNorm.Context) := q(fun v => if h : v < «$is».size then $is[v] else 0)
+  let (l, (is, rs)) ← (toQPolyNormRatExpr l).run (#[], #[])
+  let (r, (is, rs)) ← (toQPolyNormRatExpr r).run (is, rs)
+  let is : Q(Array Int) ← pure (is.foldl (fun acc e => q(«$acc».push $e)) q(#[]))
+  let rs : Q(Array Rat) ← pure (rs.foldl (fun acc e => q(«$acc».push $e)) q(#[]))
+  let ictx : Q(PolyNorm.IntContext) := q((«$is».getD · 0))
+  let rctx : Q(PolyNorm.RatContext) := q((«$rs».getD · 0))
   let h : Q(«$l».toPolynomial = «$r».toPolynomial) := .app q(@Eq.refl.{1} PolyNorm.Polynomial) q(«$l».toPolynomial)
-  mv.assign q(@PolyNorm.Expr.denote_eq_from_toPolynomial_eq $ctx $l $r $h)
+  mv.assign q(@PolyNorm.RatExpr.denote_eq_from_toPolynomial_eq $ictx $rctx $l $r $h)
+
+def nativePolyNorm (mv : MVarId) : MetaM Unit := do
+  let some (_, l, r) := (← mv.getType).eq?
+    | throwError "[poly_norm] expected an equality, got {← mv.getType}"
+  let (l, (is, rs)) ← (toQPolyNormRatExpr l).run (#[], #[])
+  let (r, (is, rs)) ← (toQPolyNormRatExpr r).run (is, rs)
+  let is : Q(Array Int) ← pure (is.foldl (fun acc e => q(«$acc».push $e)) q(#[]))
+  let rs : Q(Array Rat) ← pure (rs.foldl (fun acc e => q(«$acc».push $e)) q(#[]))
+  let ictx : Q(PolyNorm.IntContext) := q((«$is».getD · 0))
+  let rctx : Q(PolyNorm.RatContext) := q((«$rs».getD · 0))
+  let h ← nativeDecide q(«$l».toPolynomial = «$r».toPolynomial)
+  mv.assign q(@PolyNorm.RatExpr.denote_eq_from_toPolynomial_eq $ictx $rctx $l $r $h)
 where
-  logPolynomial (e : Q(PolyNorm.Expr)) (es : Array Q(Rat)) := do
-  let p ← unsafe Meta.evalExpr PolyNorm.Expr q(PolyNorm.Expr) e
-  let ppCtx := fun v => if h : v < es.size then es[v] else q(0)
-  logInfo m!"poly := {PolyNorm.Polynomial.toExpr p.toPolynomial ppCtx}"
+  nativeDecide (p : Q(Prop)) : MetaM Q($p) := do
+    let hp : Q(Decidable $p) ← Meta.synthInstance q(Decidable $p)
+    let auxDeclName ← mkNativeAuxDecl `_nativePolynorm q(Bool) q(decide $p)
+    let b : Q(Bool) := .const auxDeclName []
+    return .app q(@of_decide_eq_true $p $hp) (.app q(Lean.ofReduceBool $b true) q(Eq.refl true))
+  mkNativeAuxDecl (baseName : Name) (type value : Expr) : MetaM Name := do
+    let auxName ← Lean.mkAuxName baseName 1
+    let decl := Declaration.defnDecl {
+      name := auxName, levelParams := [], type, value
+      hints := .abbrev
+      safety := .safe
+    }
+    addAndCompile decl
+    pure auxName
 
 namespace Tactic
 
@@ -283,7 +386,24 @@ open Lean.Elab Tactic in
     Rat.polyNorm mv
     replaceMainGoal []
 
+syntax (name := nativePolyNorm) "native_poly_norm" : tactic
+
+open Lean.Elab Tactic in
+@[tactic nativePolyNorm] def evalNativePolyNorm : Tactic := fun _ =>
+  withMainContext do
+    let mv ← Tactic.getMainGoal
+    Rat.nativePolyNorm mv
+    replaceMainGoal []
+
 end Smt.Reconstruct.Rat.Tactic
 
+open Smt.Reconstruct.Rat.PolyNorm.RatExpr
+
 example (x y z : Rat) : 1 * (x + y) * z / 4  = 1/(2 * 2) * (z * y + x * z) := by
   poly_norm
+
+example (x y : Int) (z : Rat) : 1 * (x + y) * z / 4  = 1/(2 * 2) * (z * y + x * z) := by
+  poly_norm
+
+example (x y : Int) (z : Rat) : 1 * (x + y) * z / 4  = 1/(2 * 2) * (z * y + x * z) := by
+  native_poly_norm