{-# OPTIONS_GHC -Wall #-}
-----------------------------------------------------------------------------
-- |
-- Module      :  ToySolver.Arith.VirtualSubstitution
-- Copyright   :  (c) Masahiro Sakai 2014
-- License     :  BSD-style
--
-- Maintainer  :  masahiro.sakai@gmail.com
-- Stability   :  provisional
-- Portability :  portable
--
-- Naive implementation of virtual substitution
--
-- Reference:
--
-- * V. Weispfenning. The complexity of linear problems in fields.
--   Journal of Symbolic Computation, 5(1-2): 3-27, Feb.-Apr. 1988.
--
-- * Hirokazu Anai, Shinji Hara. Parametric Robust Control by Quantifier Elimination.
--   J.JSSAC, Vol. 10, No. 1, pp. 41-51, 2003.
--
-----------------------------------------------------------------------------
module ToySolver.Arith.VirtualSubstitution
  ( QFFormula
  , Model
  , Eval (..)

  -- * Projection
  , project
  , projectN
  , projectCases
  , projectCasesN

  -- * Constraint solving
  , solve
  , solveQFFormula
  ) where

import Control.Exception
import Control.Monad
import qualified Data.Foldable as Foldable
import qualified Data.IntMap as IM
import qualified Data.IntSet as IS
import Data.Set (Set)
import qualified Data.Set as Set
import Data.Maybe
import Data.VectorSpace hiding (project)

import ToySolver.Data.OrdRel
import ToySolver.Data.Boolean
import ToySolver.Data.BoolExpr (BoolExpr (..))
import qualified ToySolver.Data.BoolExpr as BoolExpr
import qualified ToySolver.Data.LA as LA
import ToySolver.Data.IntVar

-- | Quantifier-free formula of LRA
type QFFormula = BoolExpr (LA.Atom Rational)

{-| @'project' x φ@ returns @(ψ, lift)@ such that:

* @⊢_LRA ∀y1, …, yn. (∃x. φ) ↔ ψ@ where @{y1, …, yn} = FV(φ) \\ {x}@, and

* if @M ⊧_LRA ψ@ then @lift M ⊧ φ@.
-}
project :: Var -> QFFormula -> (QFFormula, Model Rational -> Model Rational)
project :: Var -> QFFormula -> (QFFormula, Model Rational -> Model Rational)
project Var
x QFFormula
formula = (QFFormula
formula', Model Rational -> Model Rational
mt)
  where
    xs :: [(QFFormula, Model Rational -> Model Rational)]
xs = Var -> QFFormula -> [(QFFormula, Model Rational -> Model Rational)]
projectCases Var
x QFFormula
formula
    formula' :: QFFormula
formula' = forall {a}. [BoolExpr a] -> BoolExpr a
orB' [QFFormula
phi | (QFFormula
phi,Model Rational -> Model Rational
_) <- [(QFFormula, Model Rational -> Model Rational)]
xs]
    mt :: Model Rational -> Model Rational
mt Model Rational
m = forall a. [a] -> a
head forall a b. (a -> b) -> a -> b
$ do
      (QFFormula
phi, Model Rational -> Model Rational
mt') <- [(QFFormula, Model Rational -> Model Rational)]
xs
      forall (f :: * -> *). Alternative f => Bool -> f ()
guard forall a b. (a -> b) -> a -> b
$ forall m e v. Eval m e v => m -> e -> v
eval Model Rational
m QFFormula
phi
      forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ Model Rational -> Model Rational
mt' Model Rational
m
    orB' :: [BoolExpr a] -> BoolExpr a
orB' = forall a. MonotoneBoolean a => [a] -> a
orB forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap forall {a}. BoolExpr a -> [BoolExpr a]
f
      where
        f :: BoolExpr a -> [BoolExpr a]
f (Or [BoolExpr a]
xs) = forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap BoolExpr a -> [BoolExpr a]
f [BoolExpr a]
xs
        f BoolExpr a
x = [BoolExpr a
x]

{-| @'projectN' {x1,…,xm} φ@ returns @(ψ, lift)@ such that:

* @⊢_LRA ∀y1, …, yn. (∃x1, …, xm. φ) ↔ ψ@ where @{y1, …, yn} = FV(φ) \\ {x1,…,xm}@, and

* if @M ⊧_LRA ψ@ then @lift M ⊧_LRA φ@.
-}
projectN :: VarSet -> QFFormula -> (QFFormula, Model Rational -> Model Rational)
projectN :: VarSet
-> QFFormula -> (QFFormula, Model Rational -> Model Rational)
projectN VarSet
vs2 = [Var] -> QFFormula -> (QFFormula, Model Rational -> Model Rational)
f (VarSet -> [Var]
IS.toList VarSet
vs2)
  where
    f :: [Var] -> QFFormula -> (QFFormula, Model Rational -> Model Rational)
    f :: [Var] -> QFFormula -> (QFFormula, Model Rational -> Model Rational)
f [] QFFormula
formula     = (QFFormula
formula, forall a. a -> a
id)
    f (Var
v:[Var]
vs) QFFormula
formula = (QFFormula
formula3, Model Rational -> Model Rational
mt1 forall b c a. (b -> c) -> (a -> b) -> a -> c
. Model Rational -> Model Rational
mt2)
      where
        (QFFormula
formula2, Model Rational -> Model Rational
mt1) = Var -> QFFormula -> (QFFormula, Model Rational -> Model Rational)
project Var
v QFFormula
formula
        (QFFormula
formula3, Model Rational -> Model Rational
mt2) = [Var] -> QFFormula -> (QFFormula, Model Rational -> Model Rational)
f [Var]
vs QFFormula
formula2

{-| @'projectCases' x φ@ returns @[(ψ_1, lift_1), …, (ψ_m, lift_m)]@ such that:

* @⊢_LRA ∀y1, …, yn. (∃x. φ) ↔ (ψ_1 ∨ … ∨ φ_m)@ where @{y1, …, yn} = FV(φ) \\ {x}@, and

* if @M ⊧_LRA ψ_i@ then @lift_i M ⊧_LRA φ@.
-}
projectCases :: Var -> QFFormula -> [(QFFormula, Model Rational -> Model Rational)]
projectCases :: Var -> QFFormula -> [(QFFormula, Model Rational -> Model Rational)]
projectCases Var
v QFFormula
phi = [(QFFormula
psi, \Model Rational
m -> forall a. Var -> a -> IntMap a -> IntMap a
IM.insert Var
v (forall m e v. Eval m e v => m -> e -> v
LA.eval Model Rational
m Expr Rational
t) Model Rational
m) | (QFFormula
psi, Expr Rational
t) <- Var -> QFFormula -> [(QFFormula, Expr Rational)]
projectCases' Var
v QFFormula
phi]

{-| @'projectCases' x φ@ returns @[(ψ_1, lift_1), …, (ψ_m, lift_m)]@ such that:

* @⊢_LRA ∀y1, …, yn. (∃x. φ) ↔ (ψ_1 ∨ … ∨ φ_m)@ where @{y1, …, yn} = FV(φ) \\ {x}@, and

* if @M ⊧_LRA ψ_i@ then @lift_i M ⊧_LRA φ@.

Note that

> (∃x. φ(x)) ⇔ ∨_{t∈S} φ(t)
>   where
>     Ψ = {a_i x - b_i ρ_i 0 | i ∈ I, ρ_i ∈ {=, ≠, ≤, <}} the set of atomic subformulas in φ(x)
>     S = { b_i / a_i, b_i / a_i + 1, b_i / a_i - 1 | i∈I } ∪ {1/2 (b_i / a_i + b_j / a_j) | i,j∈I, i≠j}
-}
projectCases' :: Var -> QFFormula -> [(QFFormula, LA.Expr Rational)]
projectCases' :: Var -> QFFormula -> [(QFFormula, Expr Rational)]
projectCases' Var
v QFFormula
phi
  | QFFormula
phi' forall a. Eq a => a -> a -> Bool
== forall a. MonotoneBoolean a => a
false = []
  | forall a. Set a -> Bool
Set.null Set (Expr Rational)
xs = [(QFFormula
phi', forall r. (Num r, Eq r) => r -> Expr r
LA.constant Rational
0)]
  | Bool
otherwise   = [(QFFormula
phi'', Expr Rational
t) | Expr Rational
t <- forall a. Set a -> [a]
Set.toList Set (Expr Rational)
s, let phi'' :: QFFormula
phi'' = Var -> Expr Rational -> QFFormula -> QFFormula
applySubst1 Var
v Expr Rational
t QFFormula
phi', QFFormula
phi'' forall a. Eq a => a -> a -> Bool
/= forall a. MonotoneBoolean a => a
false]
  where
    phi' :: QFFormula
phi' = QFFormula -> QFFormula
simplify QFFormula
phi
    xs :: Set (Expr Rational)
xs = Var -> QFFormula -> Set (Expr Rational)
collect Var
v QFFormula
phi'
    s :: Set (Expr Rational)
s = forall (f :: * -> *) a. (Foldable f, Ord a) => f (Set a) -> Set a
Set.unions
        [ Set (Expr Rational)
xs
        , forall a. Ord a => [a] -> Set a
Set.fromList [Expr Rational
e forall v. AdditiveGroup v => v -> v -> v
^+^ forall r. (Num r, Eq r) => r -> Expr r
LA.constant Rational
1 | Expr Rational
e <- forall a. Set a -> [a]
Set.toList Set (Expr Rational)
xs]
        , forall a. Ord a => [a] -> Set a
Set.fromList [Expr Rational
e forall v. AdditiveGroup v => v -> v -> v
^-^ forall r. (Num r, Eq r) => r -> Expr r
LA.constant Rational
1 | Expr Rational
e <- forall a. Set a -> [a]
Set.toList Set (Expr Rational)
xs]
        , forall a. Ord a => [a] -> Set a
Set.fromList [(Expr Rational
e1 forall v. AdditiveGroup v => v -> v -> v
^+^ Expr Rational
e2) forall v s.
(VectorSpace v, s ~ Scalar v, Fractional s) =>
v -> s -> v
^/ Rational
2 | (Expr Rational
e1,Expr Rational
e2) <- forall a. [a] -> [(a, a)]
pairs (forall a. Set a -> [a]
Set.toList Set (Expr Rational)
xs)]
        ]

{-| @'projectCasesN' {x1,…,xm} φ@ returns @[(ψ_1, lift_1), …, (ψ_n, lift_n)]@ such that:

* @⊢_LRA ∀y1, …, yp. (∃x. φ) ↔ (ψ_1 ∨ … ∨ φ_n)@ where @{y1, …, yp} = FV(φ) \\ {x1,…,xm}@, and

* if @M ⊧_LRA ψ_i@ then @lift_i M ⊧_LRA φ@.
-}
projectCasesN :: VarSet -> QFFormula -> [(QFFormula, Model Rational -> Model Rational)]
projectCasesN :: VarSet
-> QFFormula -> [(QFFormula, Model Rational -> Model Rational)]
projectCasesN VarSet
vs = [Var]
-> QFFormula -> [(QFFormula, Model Rational -> Model Rational)]
f (VarSet -> [Var]
IS.toList VarSet
vs)
  where
    f :: [Var]
-> QFFormula -> [(QFFormula, Model Rational -> Model Rational)]
f [] QFFormula
phi = forall (m :: * -> *) a. Monad m => a -> m a
return (QFFormula
phi, forall a. a -> a
id)
    f (Var
v:[Var]
vs) QFFormula
phi = do
      (QFFormula
phi2, Model Rational -> Model Rational
mt1) <- Var -> QFFormula -> [(QFFormula, Model Rational -> Model Rational)]
projectCases Var
v QFFormula
phi
      (QFFormula
phi3, Model Rational -> Model Rational
mt2) <- [Var]
-> QFFormula -> [(QFFormula, Model Rational -> Model Rational)]
f [Var]
vs QFFormula
phi2
      forall (m :: * -> *) a. Monad m => a -> m a
return (QFFormula
phi3, Model Rational -> Model Rational
mt1 forall b c a. (b -> c) -> (a -> b) -> a -> c
. Model Rational -> Model Rational
mt2)

simplify :: QFFormula -> QFFormula
simplify :: QFFormula -> QFFormula
simplify = forall a. BoolExpr a -> BoolExpr a
BoolExpr.simplify forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall b atom. Boolean b => (atom -> b) -> BoolExpr atom -> b
BoolExpr.fold OrdRel (Expr Rational) -> QFFormula
simplifyLit

simplifyLit :: LA.Atom Rational -> QFFormula
simplifyLit :: OrdRel (Expr Rational) -> QFFormula
simplifyLit (OrdRel Expr Rational
lhs RelOp
op Expr Rational
rhs) =
  case forall r. Num r => Expr r -> Maybe r
LA.asConst Expr Rational
e of
    Just Rational
c -> if forall a. Ord a => RelOp -> a -> a -> Bool
evalOp RelOp
op Rational
c Rational
0 then forall a. MonotoneBoolean a => a
true else forall a. MonotoneBoolean a => a
false
    Maybe Rational
Nothing -> forall a. a -> BoolExpr a
Atom (forall e. e -> RelOp -> e -> OrdRel e
OrdRel Expr Rational
e RelOp
op (forall r. (Num r, Eq r) => r -> Expr r
LA.constant Rational
0))
  where
    e :: Expr Rational
e = Expr Rational
lhs forall v. AdditiveGroup v => v -> v -> v
^-^ Expr Rational
rhs

collect :: Var -> QFFormula -> Set (LA.Expr Rational)
collect :: Var -> QFFormula -> Set (Expr Rational)
collect Var
v = forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
Foldable.foldMap forall {s}. (Eq s, Fractional s) => OrdRel (Expr s) -> Set (Expr s)
f
  where
    f :: OrdRel (Expr s) -> Set (Expr s)
f (OrdRel Expr s
lhs RelOp
_ Expr s
rhs) = forall a. (?callStack::CallStack) => Bool -> a -> a
assert (Expr s
rhs forall a. Eq a => a -> a -> Bool
== forall r. (Num r, Eq r) => r -> Expr r
LA.constant s
0) forall a b. (a -> b) -> a -> b
$
      case forall r. Num r => Var -> Expr r -> Maybe (r, Expr r)
LA.extractMaybe Var
v Expr s
lhs of
        Maybe (s, Expr s)
Nothing -> forall a. Set a
Set.empty
        Just (s
a,Expr s
b) -> forall a. a -> Set a
Set.singleton (forall v. AdditiveGroup v => v -> v
negateV (Expr s
b forall v s.
(VectorSpace v, s ~ Scalar v, Fractional s) =>
v -> s -> v
^/ s
a))

applySubst1 :: Var -> LA.Expr Rational -> QFFormula -> QFFormula
applySubst1 :: Var -> Expr Rational -> QFFormula -> QFFormula
applySubst1 Var
v Expr Rational
t = forall b atom. Boolean b => (atom -> b) -> BoolExpr atom -> b
BoolExpr.fold OrdRel (Expr Rational) -> QFFormula
f
  where
    f :: OrdRel (Expr Rational) -> QFFormula
f OrdRel (Expr Rational)
rel = forall a. a -> BoolExpr a
Atom (forall r. (Num r, Eq r) => Var -> Expr r -> Atom r -> Atom r
LA.applySubst1Atom Var
v Expr Rational
t OrdRel (Expr Rational)
rel)

pairs :: [a] -> [(a,a)]
pairs :: forall a. [a] -> [(a, a)]
pairs [] = []
pairs (a
x:[a]
xs) = [(a
x,a
x2) | a
x2 <- [a]
xs] forall a. [a] -> [a] -> [a]
++ forall a. [a] -> [(a, a)]
pairs [a]
xs

-- | @'solveQFFormula' {x1,…,xn} φ@ returns @Just M@ such that @M ⊧_LRA φ@ when
-- such @M@ exists, returns @Nothing@ otherwise.
--
-- @FV(φ)@ must be a subset of @{x1,…,xn}@.
--
solveQFFormula :: VarSet -> QFFormula -> Maybe (Model Rational)
solveQFFormula :: VarSet -> QFFormula -> Maybe (Model Rational)
solveQFFormula VarSet
vs QFFormula
formula = forall a. [a] -> Maybe a
listToMaybe forall a b. (a -> b) -> a -> b
$ do
  (QFFormula
formula2, Model Rational -> Model Rational
mt) <- VarSet
-> QFFormula -> [(QFFormula, Model Rational -> Model Rational)]
projectCasesN VarSet
vs QFFormula
formula
  let m :: Model Rational
      m :: Model Rational
m = forall a. IntMap a
IM.empty
  forall (f :: * -> *). Alternative f => Bool -> f ()
guard forall a b. (a -> b) -> a -> b
$ forall m e v. Eval m e v => m -> e -> v
eval Model Rational
m QFFormula
formula2
  forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ Model Rational -> Model Rational
mt Model Rational
m

-- | @'solve' {x1,…,xn} φ@ returns @Just M@ such that @M ⊧_LRA φ@ when
-- such @M@ exists, returns @Nothing@ otherwise.
--
-- @FV(φ)@ must be a subset of @{x1,…,xn}@.
--
solve :: VarSet -> [LA.Atom Rational] -> Maybe (Model Rational)
solve :: VarSet -> [OrdRel (Expr Rational)] -> Maybe (Model Rational)
solve VarSet
vs [OrdRel (Expr Rational)]
cs = VarSet -> QFFormula -> Maybe (Model Rational)
solveQFFormula VarSet
vs (forall a. MonotoneBoolean a => [a] -> a
andB [forall a. a -> BoolExpr a
Atom OrdRel (Expr Rational)
c | OrdRel (Expr Rational)
c <- [OrdRel (Expr Rational)]
cs])