203 lines
7.3 KiB
Text
203 lines
7.3 KiB
Text
/-
|
||
Copyright (c) 2022 Microsoft Corporation. All rights reserved.
|
||
Released under Apache 2.0 license as described in the file LICENSE.
|
||
Authors: Leonardo de Moura
|
||
-/
|
||
import Lean.Meta.Check
|
||
import Lean.Meta.Offset
|
||
|
||
namespace Lean.Meta.Linear.Nat
|
||
|
||
deriving instance Repr for Nat.Linear.Expr
|
||
|
||
abbrev LinearExpr := Nat.Linear.Expr
|
||
abbrev LinearCnstr := Nat.Linear.ExprCnstr
|
||
abbrev PolyExpr := Nat.Linear.Poly
|
||
|
||
def LinearExpr.toExpr (e : LinearExpr) : Expr :=
|
||
open Nat.Linear.Expr in
|
||
match e with
|
||
| num v => mkApp (mkConst ``num) (mkNatLit v)
|
||
| var i => mkApp (mkConst ``var) (mkNatLit i)
|
||
| add a b => mkApp2 (mkConst ``add) (toExpr a) (toExpr b)
|
||
| mulL k a => mkApp2 (mkConst ``mulL) (mkNatLit k) (toExpr a)
|
||
| mulR a k => mkApp2 (mkConst ``mulL) (toExpr a) (mkNatLit k)
|
||
|
||
instance : ToExpr LinearExpr where
|
||
toExpr a := a.toExpr
|
||
toTypeExpr := mkConst ``Nat.Linear.Expr
|
||
|
||
protected def LinearCnstr.toExpr (c : LinearCnstr) : Expr :=
|
||
mkApp3 (mkConst ``Nat.Linear.ExprCnstr.mk) (toExpr c.eq) (LinearExpr.toExpr c.lhs) (LinearExpr.toExpr c.rhs)
|
||
|
||
instance : ToExpr LinearCnstr where
|
||
toExpr a := a.toExpr
|
||
toTypeExpr := mkConst ``Nat.Linear.ExprCnstr
|
||
|
||
open Nat.Linear.Expr in
|
||
def LinearExpr.toArith (ctx : Array Expr) (e : LinearExpr) : MetaM Expr := do
|
||
match e with
|
||
| num v => return mkNatLit v
|
||
| var i => return ctx[i]
|
||
| add a b => mkAdd (← toArith ctx a) (← toArith ctx b)
|
||
| mulL k a => mkMul (mkNatLit k) (← toArith ctx a)
|
||
| mulR a k => mkMul (← toArith ctx a) (mkNatLit k)
|
||
|
||
def LinearCnstr.toArith (ctx : Array Expr) (c : LinearCnstr) : MetaM Expr := do
|
||
if c.eq then
|
||
mkEq (← LinearExpr.toArith ctx c.lhs) (← LinearExpr.toArith ctx c.rhs)
|
||
else
|
||
return mkApp4 (mkConst ``LE.le [levelZero]) (mkConst ``Nat) (mkConst ``instLENat) (← LinearExpr.toArith ctx c.lhs) (← LinearExpr.toArith ctx c.rhs)
|
||
|
||
namespace ToLinear
|
||
|
||
structure State where
|
||
varMap : ExprMap Nat := {}
|
||
vars : Array Expr := #[]
|
||
|
||
abbrev M := StateRefT State MetaM
|
||
|
||
open Nat.Linear.Expr
|
||
|
||
def addAsVar (e : Expr) : M LinearExpr := do
|
||
if let some x := (← get).varMap.find? e then
|
||
return var x
|
||
else
|
||
let x := (← get).vars.size
|
||
modify fun s => { varMap := s.varMap.insert e x, vars := s.vars.push e }
|
||
return var x
|
||
|
||
partial def toLinearExpr (e : Expr) : M LinearExpr := do
|
||
match e with
|
||
| Expr.lit (Literal.natVal n) _ => return num n
|
||
| Expr.mdata _ e _ => toLinearExpr e
|
||
| Expr.const ``Nat.zero .. => return num 0
|
||
| Expr.app .. => visit e
|
||
| Expr.mvar .. => visit e
|
||
| _ => addAsVar e
|
||
where
|
||
visit (e : Expr) : M LinearExpr := do
|
||
let f := e.getAppFn
|
||
match f with
|
||
| Expr.mvar .. =>
|
||
let eNew ← instantiateMVars e
|
||
if eNew != e then
|
||
toLinearExpr eNew
|
||
else
|
||
addAsVar e
|
||
| Expr.const declName .. =>
|
||
let numArgs := e.getAppNumArgs
|
||
if declName == ``Nat.succ && numArgs == 1 then
|
||
return inc (← toLinearExpr e.appArg!)
|
||
else if declName == ``Nat.add && numArgs == 2 then
|
||
return add (← toLinearExpr (e.getArg! 0)) (← toLinearExpr (e.getArg! 1))
|
||
else if declName == ``Nat.mul && numArgs == 2 then
|
||
match (← evalNat (e.getArg! 0) |>.run) with
|
||
| some k => return mulL k (← toLinearExpr (e.getArg! 1))
|
||
| none => match (← evalNat (e.getArg! 1) |>.run) with
|
||
| some k => return mulR (← toLinearExpr (e.getArg! 0)) k
|
||
| none => addAsVar e
|
||
else if isNatProjInst declName numArgs then
|
||
if let some e ← unfoldProjInst? e then
|
||
toLinearExpr e
|
||
else
|
||
addAsVar e
|
||
else
|
||
addAsVar e
|
||
| _ => addAsVar e
|
||
|
||
partial def toLinearCnstr? (e : Expr) : M (Option LinearCnstr) := do
|
||
let f := e.getAppFn
|
||
match f with
|
||
| Expr.mvar .. =>
|
||
let eNew ← instantiateMVars e
|
||
if eNew != e then
|
||
toLinearCnstr? eNew
|
||
else
|
||
return none
|
||
| Expr.const declName .. =>
|
||
let numArgs := e.getAppNumArgs
|
||
if declName == ``Eq && numArgs == 3 then
|
||
return some { eq := true, lhs := (← toLinearExpr (e.getArg! 1)), rhs := (← toLinearExpr (e.getArg! 2)) }
|
||
else if declName == ``Nat.le && numArgs == 2 then
|
||
return some { eq := false, lhs := (← toLinearExpr (e.getArg! 0)), rhs := (← toLinearExpr (e.getArg! 1)) }
|
||
else if declName == ``Nat.lt && numArgs == 2 then
|
||
return some { eq := false, lhs := (← toLinearExpr (e.getArg! 0)).inc, rhs := (← toLinearExpr (e.getArg! 1)) }
|
||
else if numArgs == 4 && (declName == ``GE.ge || declName == ``GT.gt) then
|
||
if let some e ← unfoldDefinition? e then
|
||
toLinearCnstr? e
|
||
else
|
||
return none
|
||
else if numArgs == 4 && (declName == ``LE.le || declName == ``LT.lt) then
|
||
if (← isDefEq (e.getArg! 0) (mkConst ``Nat)) then
|
||
if let some e ← unfoldProjInst? e then
|
||
toLinearCnstr? e
|
||
else
|
||
return none
|
||
else
|
||
return none
|
||
else
|
||
return none
|
||
| _ => return none
|
||
|
||
def run (x : M α) : MetaM (α × Array Expr) := do
|
||
let (a, s) ← x.run {}
|
||
return (a, s.vars)
|
||
|
||
end ToLinear
|
||
|
||
open ToLinear (toLinearCnstr? toLinearExpr)
|
||
|
||
def toContextExpr (ctx : Array Expr) : MetaM Expr := do
|
||
mkListLit (mkConst ``Nat) ctx.toList
|
||
|
||
def reflTrue : Expr :=
|
||
mkApp2 (mkConst ``Eq.refl [levelOne]) (mkConst ``Bool) (mkConst ``Bool.true)
|
||
|
||
def simpCnstrPos? (e : Expr) : MetaM (Option (Expr × Expr)) := do
|
||
let (some c, ctx) ← ToLinear.run (ToLinear.toLinearCnstr? e) | return none
|
||
let c₁ := c.toPoly
|
||
let c₂ := c₁.norm
|
||
if c₂.isUnsat then
|
||
let p := mkApp3 (mkConst ``Nat.Linear.ExprCnstr.eq_false_of_isUnsat) (← toContextExpr ctx) (toExpr c) reflTrue
|
||
return some (mkConst ``False, p)
|
||
else if c₂.isValid then
|
||
let p := mkApp3 (mkConst ``Nat.Linear.ExprCnstr.eq_true_of_isValid) (← toContextExpr ctx) (toExpr c) reflTrue
|
||
return some (mkConst ``True, p)
|
||
else if c₂.hasFewerMonomials c₁ then
|
||
let c₂ : LinearCnstr := c₂.toExpr
|
||
let p := mkApp4 (mkConst ``Nat.Linear.ExprCnstr.eq_of_toNormPoly_eq) (← toContextExpr ctx) (toExpr c) (toExpr c₂) reflTrue
|
||
let r ← c₂.toArith ctx
|
||
return some (r, p)
|
||
else
|
||
return none
|
||
|
||
def simpCnstr? (e : Expr) : MetaM (Option (Expr × Expr)) := do
|
||
if let some arg := e.not? then
|
||
let mut eNew? := none
|
||
let mut thmName := Name.anonymous
|
||
if arg.isAppOfArity ``LE.le 4 then
|
||
eNew? := some (← mkLE (← mkAdd (arg.getArg! 3) (mkNatLit 1)) (arg.getArg! 2))
|
||
thmName := ``Nat.not_le_eq
|
||
else if arg.isAppOfArity ``GE.ge 4 then
|
||
eNew? := some (← mkLE (← mkAdd (arg.getArg! 2) (mkNatLit 1)) (arg.getArg! 3))
|
||
thmName := ``Nat.not_ge_eq
|
||
else if arg.isAppOfArity ``LT.lt 4 then
|
||
eNew? := some (← mkLE (arg.getArg! 3) (arg.getArg! 2))
|
||
thmName := ``Nat.not_lt_eq
|
||
else if arg.isAppOfArity ``GT.gt 4 then
|
||
eNew? := some (← mkLE (arg.getArg! 2) (arg.getArg! 3))
|
||
thmName := ``Nat.not_gt_eq
|
||
if let some eNew := eNew? then
|
||
if let some (eNew', h₂) ← simpCnstrPos? eNew then
|
||
let h₁ := mkApp2 (mkConst thmName) (arg.getArg! 2) (arg.getArg! 3)
|
||
let h := mkApp6 (mkConst ``Eq.trans [levelOne]) (mkSort levelZero) e eNew eNew' h₁ h₂
|
||
return some (eNew', h)
|
||
else
|
||
return none
|
||
else
|
||
return none
|
||
else
|
||
simpCnstrPos? e
|
||
|
||
end Lean.Meta.Linear.Nat
|