{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE LambdaCase #-}

{- |
This module provides a concrete implementation for 'MonadOMT' with it's state 'OMT'.
-}
module Language.Hasmtlib.Type.OMT
(
  -- * SoftFormula
  -- ** Type
  SoftFormula(..)

  -- ** Lens
, formula, mWeight, mGroupId

  -- * Optimization targets
, Minimize(..), Maximize(..)

  -- * OMT
  -- ** Type
, OMT(..)

  -- ** Lens
, smt, targetMinimize, targetMaximize, softFormulas
)
where

import Language.Hasmtlib.Internal.Sharing
import Language.Hasmtlib.Type.MonadSMT
import Language.Hasmtlib.Type.SMTSort
import Language.Hasmtlib.Type.Expr
import Language.Hasmtlib.Type.SMT
import Data.Some.Constraint
import Data.List (isPrefixOf)
import Data.Default
import Data.Coerce
import Data.Sequence hiding ((|>), filter)
import Control.Monad.State
import Control.Lens hiding (List)

-- | An assertion of a booolean expression in OMT that may be weighted.
data SoftFormula = SoftFormula
  { SoftFormula -> Expr 'BoolSort
_formula  :: Expr BoolSort    -- ^ The underlying soft formula
  , SoftFormula -> Maybe Double
_mWeight  :: Maybe Double     -- ^ Weight of the soft formula
  , SoftFormula -> Maybe String
_mGroupId :: Maybe String     -- ^ Group-Id of the soft formula
  }
$(makeLenses ''SoftFormula)

-- | A newtype for numerical expressions that are target of a minimization.
newtype Minimize t = Minimize { forall (t :: SMTSort). Minimize t -> Expr t
_targetMin :: Expr t }

-- | A newtype for numerical expressions that are target of a maximization.
newtype Maximize t = Maximize { forall (t :: SMTSort). Maximize t -> Expr t
_targetMax :: Expr t }

-- | The state of the OMT-problem.
data OMT = OMT
  { OMT -> SMT
_smt            :: !SMT                                 -- ^ The underlying 'SMT'-Problem
  , OMT -> Seq (SomeKnownSMTSort Minimize)
_targetMinimize :: !(Seq (SomeKnownSMTSort Minimize))   -- ^ All expressions to minimize
  , OMT -> Seq (SomeKnownSMTSort Maximize)
_targetMaximize :: !(Seq (SomeKnownSMTSort Maximize))   -- ^ All expressions to maximize
  , OMT -> Seq SoftFormula
_softFormulas   :: !(Seq SoftFormula)                   -- ^ All soft assertions of boolean expressions
  }
$(makeLenses ''OMT)

instance Default OMT where
  def :: OMT
def = SMT
-> Seq (SomeKnownSMTSort Minimize)
-> Seq (SomeKnownSMTSort Maximize)
-> Seq SoftFormula
-> OMT
OMT SMT
forall a. Default a => a
def Seq (SomeKnownSMTSort Minimize)
forall a. Monoid a => a
mempty Seq (SomeKnownSMTSort Maximize)
forall a. Monoid a => a
mempty Seq SoftFormula
forall a. Monoid a => a
mempty

instance Sharing OMT where
  type SharingMonad OMT = Monad
  stableMap :: Lens' OMT (HashMap (StableName ()) (SomeKnownSMTSort Expr))
stableMap = (SMT -> f SMT) -> OMT -> f OMT
Lens' OMT SMT
smt((SMT -> f SMT) -> OMT -> f OMT)
-> ((HashMap (StableName ()) (SomeKnownSMTSort Expr)
     -> f (HashMap (StableName ()) (SomeKnownSMTSort Expr)))
    -> SMT -> f SMT)
-> (HashMap (StableName ()) (SomeKnownSMTSort Expr)
    -> f (HashMap (StableName ()) (SomeKnownSMTSort Expr)))
-> OMT
-> f OMT
forall b c a. (b -> c) -> (a -> b) -> a -> c
.(HashMap (StableName ()) (SomeKnownSMTSort Expr)
 -> f (HashMap (StableName ()) (SomeKnownSMTSort Expr)))
-> SMT -> f SMT
Lens' SMT (HashMap (StableName ()) (SomeKnownSMTSort Expr))
Language.Hasmtlib.Type.SMT.stableMap
  assertSharedNode :: forall (m :: * -> *).
(MonadState OMT m, SharingMonad OMT m) =>
StableName () -> Expr 'BoolSort -> m ()
assertSharedNode StableName ()
_ Expr 'BoolSort
expr = ASetter OMT OMT (Seq (Expr 'BoolSort)) (Seq (Expr 'BoolSort))
-> (Seq (Expr 'BoolSort) -> Seq (Expr 'BoolSort)) -> m ()
forall s (m :: * -> *) a b.
MonadState s m =>
ASetter s s a b -> (a -> b) -> m ()
modifying ((SMT -> Identity SMT) -> OMT -> Identity OMT
Lens' OMT SMT
smt((SMT -> Identity SMT) -> OMT -> Identity OMT)
-> ((Seq (Expr 'BoolSort) -> Identity (Seq (Expr 'BoolSort)))
    -> SMT -> Identity SMT)
-> ASetter OMT OMT (Seq (Expr 'BoolSort)) (Seq (Expr 'BoolSort))
forall b c a. (b -> c) -> (a -> b) -> a -> c
.(Seq (Expr 'BoolSort) -> Identity (Seq (Expr 'BoolSort)))
-> SMT -> Identity SMT
Lens' SMT (Seq (Expr 'BoolSort))
formulas) (Seq (Expr 'BoolSort) -> Expr 'BoolSort -> Seq (Expr 'BoolSort)
forall s a. Snoc s s a a => s -> a -> s
|> Expr 'BoolSort
expr)
  setSharingMode :: forall (m :: * -> *). MonadState OMT m => SharingMode -> m ()
setSharingMode SharingMode
sm = (SMT -> Identity SMT) -> OMT -> Identity OMT
Lens' OMT SMT
smt((SMT -> Identity SMT) -> OMT -> Identity OMT)
-> ((SharingMode -> Identity SharingMode) -> SMT -> Identity SMT)
-> (SharingMode -> Identity SharingMode)
-> OMT
-> Identity OMT
forall b c a. (b -> c) -> (a -> b) -> a -> c
.(SharingMode -> Identity SharingMode) -> SMT -> Identity SMT
Lens' SMT SharingMode
sharingMode ((SharingMode -> Identity SharingMode) -> OMT -> Identity OMT)
-> SharingMode -> m ()
forall s (m :: * -> *) a b.
MonadState s m =>
ASetter s s a b -> b -> m ()
.= SharingMode
sm

instance MonadState OMT m => MonadSMT OMT m where
  smtvar' :: forall (t :: SMTSort). KnownSMTSort t => Proxy t -> m (SMTVar t)
smtvar' Proxy t
_ = (Int -> SMTVar t) -> m Int -> m (SMTVar t)
forall a b. (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Int -> SMTVar t
forall a b. Coercible a b => a -> b
coerce (m Int -> m (SMTVar t)) -> m Int -> m (SMTVar t)
forall a b. (a -> b) -> a -> b
$ ((SMT -> (Int, SMT)) -> OMT -> (Int, OMT)
Lens' OMT SMT
smt((SMT -> (Int, SMT)) -> OMT -> (Int, OMT))
-> ((Int -> (Int, Int)) -> SMT -> (Int, SMT))
-> (Int -> (Int, Int))
-> OMT
-> (Int, OMT)
forall b c a. (b -> c) -> (a -> b) -> a -> c
.(Int -> (Int, Int)) -> SMT -> (Int, SMT)
Lens' SMT Int
lastVarId) ((Int -> (Int, Int)) -> OMT -> (Int, OMT)) -> Int -> m Int
forall s (m :: * -> *) a.
(MonadState s m, Num a) =>
LensLike' ((,) a) s a -> a -> m a
<+= Int
1
  {-# INLINE smtvar' #-}

  var' :: forall (t :: SMTSort). KnownSMTSort t => Proxy t -> m (Expr t)
var' Proxy t
p = do
    SMTVar t
newVar <- Proxy t -> m (SMTVar t)
forall s (m :: * -> *) (t :: SMTSort).
(MonadSMT s m, KnownSMTSort t) =>
Proxy t -> m (SMTVar t)
forall (t :: SMTSort). KnownSMTSort t => Proxy t -> m (SMTVar t)
smtvar' Proxy t
p
    (SMT -> Identity SMT) -> OMT -> Identity OMT
Lens' OMT SMT
smt((SMT -> Identity SMT) -> OMT -> Identity OMT)
-> ((Seq (SomeKnownSMTSort SMTVar)
     -> Identity (Seq (SomeKnownSMTSort SMTVar)))
    -> SMT -> Identity SMT)
-> (Seq (SomeKnownSMTSort SMTVar)
    -> Identity (Seq (SomeKnownSMTSort SMTVar)))
-> OMT
-> Identity OMT
forall b c a. (b -> c) -> (a -> b) -> a -> c
.(Seq (SomeKnownSMTSort SMTVar)
 -> Identity (Seq (SomeKnownSMTSort SMTVar)))
-> SMT -> Identity SMT
Lens' SMT (Seq (SomeKnownSMTSort SMTVar))
vars ((Seq (SomeKnownSMTSort SMTVar)
  -> Identity (Seq (SomeKnownSMTSort SMTVar)))
 -> OMT -> Identity OMT)
-> (Seq (SomeKnownSMTSort SMTVar) -> Seq (SomeKnownSMTSort SMTVar))
-> m ()
forall s (m :: * -> *) a b.
MonadState s m =>
ASetter s s a b -> (a -> b) -> m ()
%= (Seq (SomeKnownSMTSort SMTVar)
-> SomeKnownSMTSort SMTVar -> Seq (SomeKnownSMTSort SMTVar)
forall s a. Snoc s s a a => s -> a -> s
|> SMTVar t -> SomeKnownSMTSort SMTVar
forall k (csf :: [(k -> *) -> Constraint])
       (csa :: [k -> Constraint]) (f :: k -> *) (a :: k).
(AllC csf f, AllC csa a) =>
f a -> Somes1 csf csa
Some1 SMTVar t
newVar)
    Expr t -> m (Expr t)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Expr t -> m (Expr t)) -> Expr t -> m (Expr t)
forall a b. (a -> b) -> a -> b
$ SMTVar t -> Expr t
forall (t :: SMTSort). KnownSMTSort t => SMTVar t -> Expr t
Var SMTVar t
newVar
  {-# INLINE var' #-}

  assert :: Expr 'BoolSort -> m ()
assert Expr 'BoolSort
expr = do
    OMT
omt <- m OMT
forall s (m :: * -> *). MonadState s m => m s
get
    Expr 'BoolSort
sExpr <- SharingMode -> Expr 'BoolSort -> m (Expr 'BoolSort)
forall (t :: SMTSort) s (m :: * -> *).
(KnownSMTSort t, MonadSMT s m, Sharing s, SharingMonad s m) =>
SharingMode -> Expr t -> m (Expr t)
runSharing (OMT
omtOMT -> Getting SharingMode OMT SharingMode -> SharingMode
forall s a. s -> Getting a s a -> a
^.(SMT -> Const SharingMode SMT) -> OMT -> Const SharingMode OMT
Lens' OMT SMT
smt((SMT -> Const SharingMode SMT) -> OMT -> Const SharingMode OMT)
-> ((SharingMode -> Const SharingMode SharingMode)
    -> SMT -> Const SharingMode SMT)
-> Getting SharingMode OMT SharingMode
forall b c a. (b -> c) -> (a -> b) -> a -> c
.(SharingMode -> Const SharingMode SharingMode)
-> SMT -> Const SharingMode SMT
Lens' SMT SharingMode
sharingMode) Expr 'BoolSort
expr
    Expr 'BoolSort
qExpr <- case OMT
omtOMT -> Getting (Maybe String) OMT (Maybe String) -> Maybe String
forall s a. s -> Getting a s a -> a
^.(SMT -> Const (Maybe String) SMT)
-> OMT -> Const (Maybe String) OMT
Lens' OMT SMT
smt((SMT -> Const (Maybe String) SMT)
 -> OMT -> Const (Maybe String) OMT)
-> ((Maybe String -> Const (Maybe String) (Maybe String))
    -> SMT -> Const (Maybe String) SMT)
-> Getting (Maybe String) OMT (Maybe String)
forall b c a. (b -> c) -> (a -> b) -> a -> c
.(Maybe String -> Const (Maybe String) (Maybe String))
-> SMT -> Const (Maybe String) SMT
Lens' SMT (Maybe String)
mlogic of
      Maybe String
Nothing    -> Expr 'BoolSort -> m (Expr 'BoolSort)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Expr 'BoolSort
sExpr
      Just String
logic -> if String
"QF" String -> String -> Bool
forall a. Eq a => [a] -> [a] -> Bool
`isPrefixOf` String
logic then Expr 'BoolSort -> m (Expr 'BoolSort)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Expr 'BoolSort
sExpr else Expr 'BoolSort -> m (Expr 'BoolSort)
forall s (m :: * -> *) (t :: SMTSort).
(MonadSMT s m, KnownSMTSort t) =>
Expr t -> m (Expr t)
quantify Expr 'BoolSort
sExpr
    (OMT -> OMT) -> m ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((OMT -> OMT) -> m ()) -> (OMT -> OMT) -> m ()
forall a b. (a -> b) -> a -> b
$ \OMT
s -> OMT
s OMT -> (OMT -> OMT) -> OMT
forall a b. a -> (a -> b) -> b
& ((SMT -> Identity SMT) -> OMT -> Identity OMT
Lens' OMT SMT
smt((SMT -> Identity SMT) -> OMT -> Identity OMT)
-> ((Seq (Expr 'BoolSort) -> Identity (Seq (Expr 'BoolSort)))
    -> SMT -> Identity SMT)
-> ASetter OMT OMT (Seq (Expr 'BoolSort)) (Seq (Expr 'BoolSort))
forall b c a. (b -> c) -> (a -> b) -> a -> c
.(Seq (Expr 'BoolSort) -> Identity (Seq (Expr 'BoolSort)))
-> SMT -> Identity SMT
Lens' SMT (Seq (Expr 'BoolSort))
formulas) ASetter OMT OMT (Seq (Expr 'BoolSort)) (Seq (Expr 'BoolSort))
-> (Seq (Expr 'BoolSort) -> Seq (Expr 'BoolSort)) -> OMT -> OMT
forall s t a b. ASetter s t a b -> (a -> b) -> s -> t
%~ (Seq (Expr 'BoolSort) -> Expr 'BoolSort -> Seq (Expr 'BoolSort)
forall s a. Snoc s s a a => s -> a -> s
|> Expr 'BoolSort
qExpr)
  {-# INLINE assert #-}

  setOption :: SMTOption -> m ()
setOption SMTOption
opt = (SMT -> Identity SMT) -> OMT -> Identity OMT
Lens' OMT SMT
smt((SMT -> Identity SMT) -> OMT -> Identity OMT)
-> (([SMTOption] -> Identity [SMTOption]) -> SMT -> Identity SMT)
-> ([SMTOption] -> Identity [SMTOption])
-> OMT
-> Identity OMT
forall b c a. (b -> c) -> (a -> b) -> a -> c
.([SMTOption] -> Identity [SMTOption]) -> SMT -> Identity SMT
Lens' SMT [SMTOption]
options (([SMTOption] -> Identity [SMTOption]) -> OMT -> Identity OMT)
-> [SMTOption] -> m ()
forall s (m :: * -> *) a.
(MonadState s m, Semigroup a) =>
ASetter' s a -> a -> m ()
<>= SMTOption -> [SMTOption]
forall a. a -> [a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure SMTOption
opt

  setLogic :: String -> m ()
setLogic String
l = (SMT -> Identity SMT) -> OMT -> Identity OMT
Lens' OMT SMT
smt((SMT -> Identity SMT) -> OMT -> Identity OMT)
-> ((Maybe String -> Identity (Maybe String))
    -> SMT -> Identity SMT)
-> (Maybe String -> Identity (Maybe String))
-> OMT
-> Identity OMT
forall b c a. (b -> c) -> (a -> b) -> a -> c
.(Maybe String -> Identity (Maybe String)) -> SMT -> Identity SMT
Lens' SMT (Maybe String)
mlogic ((Maybe String -> Identity (Maybe String)) -> OMT -> Identity OMT)
-> String -> m ()
forall s (m :: * -> *) a b.
MonadState s m =>
ASetter s s a (Maybe b) -> b -> m ()
?= String
l

instance MonadSMT OMT m => MonadOMT OMT m where
  minimize :: forall (t :: SMTSort).
(KnownSMTSort t, Num (Expr t)) =>
Expr t -> m ()
minimize Expr t
expr = do
    SharingMode
sm <- Getting SharingMode OMT SharingMode -> m SharingMode
forall s (m :: * -> *) a. MonadState s m => Getting a s a -> m a
use ((SMT -> Const SharingMode SMT) -> OMT -> Const SharingMode OMT
Lens' OMT SMT
smt((SMT -> Const SharingMode SMT) -> OMT -> Const SharingMode OMT)
-> ((SharingMode -> Const SharingMode SharingMode)
    -> SMT -> Const SharingMode SMT)
-> Getting SharingMode OMT SharingMode
forall b c a. (b -> c) -> (a -> b) -> a -> c
.(SharingMode -> Const SharingMode SharingMode)
-> SMT -> Const SharingMode SMT
Lens' SMT SharingMode
sharingMode)
    Expr t
sExpr <- SharingMode -> Expr t -> m (Expr t)
forall (t :: SMTSort) s (m :: * -> *).
(KnownSMTSort t, MonadSMT s m, Sharing s, SharingMonad s m) =>
SharingMode -> Expr t -> m (Expr t)
runSharing SharingMode
sm Expr t
expr
    ASetter
  OMT
  OMT
  (Seq (SomeKnownSMTSort Minimize))
  (Seq (SomeKnownSMTSort Minimize))
-> (Seq (SomeKnownSMTSort Minimize)
    -> Seq (SomeKnownSMTSort Minimize))
-> m ()
forall s (m :: * -> *) a b.
MonadState s m =>
ASetter s s a b -> (a -> b) -> m ()
modifying ASetter
  OMT
  OMT
  (Seq (SomeKnownSMTSort Minimize))
  (Seq (SomeKnownSMTSort Minimize))
Lens' OMT (Seq (SomeKnownSMTSort Minimize))
targetMinimize (Seq (SomeKnownSMTSort Minimize)
-> SomeKnownSMTSort Minimize -> Seq (SomeKnownSMTSort Minimize)
forall s a. Snoc s s a a => s -> a -> s
|> Minimize t -> SomeKnownSMTSort Minimize
forall k (csf :: [(k -> *) -> Constraint])
       (csa :: [k -> Constraint]) (f :: k -> *) (a :: k).
(AllC csf f, AllC csa a) =>
f a -> Somes1 csf csa
Some1 (Expr t -> Minimize t
forall (t :: SMTSort). Expr t -> Minimize t
Minimize Expr t
sExpr))
  maximize :: forall (t :: SMTSort).
(KnownSMTSort t, Num (Expr t)) =>
Expr t -> m ()
maximize Expr t
expr = do
    SharingMode
sm <- Getting SharingMode OMT SharingMode -> m SharingMode
forall s (m :: * -> *) a. MonadState s m => Getting a s a -> m a
use ((SMT -> Const SharingMode SMT) -> OMT -> Const SharingMode OMT
Lens' OMT SMT
smt((SMT -> Const SharingMode SMT) -> OMT -> Const SharingMode OMT)
-> ((SharingMode -> Const SharingMode SharingMode)
    -> SMT -> Const SharingMode SMT)
-> Getting SharingMode OMT SharingMode
forall b c a. (b -> c) -> (a -> b) -> a -> c
.(SharingMode -> Const SharingMode SharingMode)
-> SMT -> Const SharingMode SMT
Lens' SMT SharingMode
sharingMode)
    Expr t
sExpr <- SharingMode -> Expr t -> m (Expr t)
forall (t :: SMTSort) s (m :: * -> *).
(KnownSMTSort t, MonadSMT s m, Sharing s, SharingMonad s m) =>
SharingMode -> Expr t -> m (Expr t)
runSharing SharingMode
sm Expr t
expr
    ASetter
  OMT
  OMT
  (Seq (SomeKnownSMTSort Maximize))
  (Seq (SomeKnownSMTSort Maximize))
-> (Seq (SomeKnownSMTSort Maximize)
    -> Seq (SomeKnownSMTSort Maximize))
-> m ()
forall s (m :: * -> *) a b.
MonadState s m =>
ASetter s s a b -> (a -> b) -> m ()
modifying ASetter
  OMT
  OMT
  (Seq (SomeKnownSMTSort Maximize))
  (Seq (SomeKnownSMTSort Maximize))
Lens' OMT (Seq (SomeKnownSMTSort Maximize))
targetMaximize (Seq (SomeKnownSMTSort Maximize)
-> SomeKnownSMTSort Maximize -> Seq (SomeKnownSMTSort Maximize)
forall s a. Snoc s s a a => s -> a -> s
|> Maximize t -> SomeKnownSMTSort Maximize
forall k (csf :: [(k -> *) -> Constraint])
       (csa :: [k -> Constraint]) (f :: k -> *) (a :: k).
(AllC csf f, AllC csa a) =>
f a -> Somes1 csf csa
Some1 (Expr t -> Maximize t
forall (t :: SMTSort). Expr t -> Maximize t
Maximize Expr t
sExpr))
  assertSoft :: Expr 'BoolSort -> Maybe Double -> Maybe String -> m ()
assertSoft Expr 'BoolSort
expr Maybe Double
w Maybe String
gid = do
    SharingMode
sm <- Getting SharingMode OMT SharingMode -> m SharingMode
forall s (m :: * -> *) a. MonadState s m => Getting a s a -> m a
use ((SMT -> Const SharingMode SMT) -> OMT -> Const SharingMode OMT
Lens' OMT SMT
smt((SMT -> Const SharingMode SMT) -> OMT -> Const SharingMode OMT)
-> ((SharingMode -> Const SharingMode SharingMode)
    -> SMT -> Const SharingMode SMT)
-> Getting SharingMode OMT SharingMode
forall b c a. (b -> c) -> (a -> b) -> a -> c
.(SharingMode -> Const SharingMode SharingMode)
-> SMT -> Const SharingMode SMT
Lens' SMT SharingMode
sharingMode)
    Expr 'BoolSort
sExpr <- SharingMode -> Expr 'BoolSort -> m (Expr 'BoolSort)
forall (t :: SMTSort) s (m :: * -> *).
(KnownSMTSort t, MonadSMT s m, Sharing s, SharingMonad s m) =>
SharingMode -> Expr t -> m (Expr t)
runSharing SharingMode
sm Expr 'BoolSort
expr
    ASetter OMT OMT (Seq SoftFormula) (Seq SoftFormula)
-> (Seq SoftFormula -> Seq SoftFormula) -> m ()
forall s (m :: * -> *) a b.
MonadState s m =>
ASetter s s a b -> (a -> b) -> m ()
modifying ASetter OMT OMT (Seq SoftFormula) (Seq SoftFormula)
Lens' OMT (Seq SoftFormula)
softFormulas (Seq SoftFormula -> SoftFormula -> Seq SoftFormula
forall s a. Snoc s s a a => s -> a -> s
|> Expr 'BoolSort -> Maybe Double -> Maybe String -> SoftFormula
SoftFormula Expr 'BoolSort
sExpr Maybe Double
w Maybe String
gid)