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

module Language.Hasmtlib.Type.SMT where

import Language.Hasmtlib.Internal.Expr
import Language.Hasmtlib.Internal.Render
import Language.Hasmtlib.Type.MonadSMT
import Language.Hasmtlib.Type.SMTSort
import Language.Hasmtlib.Type.Option
import Data.List (isPrefixOf)
import Data.Default
import Data.Coerce
import Data.Sequence hiding ((|>), filter)
import Data.Data (toConstr, showConstr)
import Data.ByteString.Builder
import Control.Monad.State
import Control.Lens hiding (List)

-- | The state of the SMT-problem.
data SMT = SMT
  { SMT -> Int
_lastVarId :: {-# UNPACK #-} !Int                     -- ^ Last Id assigned to a new var
  , SMT -> Seq (SomeKnownSMTSort SMTVar)
_vars     :: !(Seq (SomeKnownSMTSort SMTVar))         -- ^ All constructed variables
  , SMT -> Seq (Expr 'BoolSort)
_formulas :: !(Seq (Expr BoolSort))                   -- ^ All asserted formulas
  , SMT -> Maybe String
_mlogic   :: Maybe String                             -- ^ Logic for the SMT-Solver
  , SMT -> [SMTOption]
_options  :: [SMTOption]                              -- ^ All manually configured SMT-Solver-Options
  }
$(makeLenses ''SMT)

instance Default SMT where
  def :: SMT
def = Int
-> Seq (SomeKnownSMTSort SMTVar)
-> Seq (Expr 'BoolSort)
-> Maybe String
-> [SMTOption]
-> SMT
SMT Int
0 Seq (SomeKnownSMTSort SMTVar)
forall a. Monoid a => a
mempty Seq (Expr 'BoolSort)
forall a. Monoid a => a
mempty Maybe String
forall a. Monoid a => a
mempty [Bool -> SMTOption
ProduceModels Bool
True]

instance MonadState SMT m => MonadSMT SMT 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
$ (Int -> (Int, Int)) -> SMT -> (Int, SMT)
Lens' SMT Int
lastVarId ((Int -> (Int, Int)) -> SMT -> (Int, SMT)) -> 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
    (Seq (SomeKnownSMTSort SMTVar)
 -> Identity (Seq (SomeKnownSMTSort SMTVar)))
-> SMT -> Identity SMT
Lens' SMT (Seq (SomeKnownSMTSort SMTVar))
vars ((Seq (SomeKnownSMTSort SMTVar)
  -> Identity (Seq (SomeKnownSMTSort SMTVar)))
 -> SMT -> Identity SMT)
-> (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 (cs :: [SMTSort -> Constraint]) (f :: SMTSort -> *)
       (t :: SMTSort).
AllC cs t =>
f t -> SomeSMTSort cs f
SomeSMTSort 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). SMTVar t -> Expr t
Var SMTVar t
newVar
  {-# INLINEABLE var' #-}

  assert :: Expr 'BoolSort -> m ()
assert Expr 'BoolSort
expr = do
    SMT
smt <- m SMT
forall s (m :: * -> *). MonadState s m => m s
get
    Expr 'BoolSort
qExpr <- case SMT
smtSMT -> Getting (Maybe String) SMT (Maybe String) -> Maybe String
forall s a. s -> Getting a s a -> a
^.Getting (Maybe String) SMT (Maybe String)
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
expr
      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
expr else Expr 'BoolSort -> m (Expr 'BoolSort)
forall s (m :: * -> *) (t :: SMTSort).
MonadSMT s m =>
Expr t -> m (Expr t)
quantify Expr 'BoolSort
expr
    (SMT -> SMT) -> m ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((SMT -> SMT) -> m ()) -> (SMT -> SMT) -> m ()
forall a b. (a -> b) -> a -> b
$ \SMT
s -> SMT
s SMT -> (SMT -> SMT) -> SMT
forall a b. a -> (a -> b) -> b
& (Seq (Expr 'BoolSort) -> Identity (Seq (Expr 'BoolSort)))
-> SMT -> Identity SMT
Lens' SMT (Seq (Expr 'BoolSort))
formulas ((Seq (Expr 'BoolSort) -> Identity (Seq (Expr 'BoolSort)))
 -> SMT -> Identity SMT)
-> (Seq (Expr 'BoolSort) -> Seq (Expr 'BoolSort)) -> SMT -> SMT
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 = ([SMTOption] -> Identity [SMTOption]) -> SMT -> Identity SMT
Lens' SMT [SMTOption]
options (([SMTOption] -> Identity [SMTOption]) -> SMT -> Identity SMT)
-> ([SMTOption] -> [SMTOption]) -> m ()
forall s (m :: * -> *) a b.
MonadState s m =>
ASetter s s a b -> (a -> b) -> m ()
%= ((SMTOption
opt:) ([SMTOption] -> [SMTOption])
-> ([SMTOption] -> [SMTOption]) -> [SMTOption] -> [SMTOption]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (SMTOption -> Bool) -> [SMTOption] -> [SMTOption]
forall a. (a -> Bool) -> [a] -> [a]
filter (Bool -> Bool
not (Bool -> Bool) -> (SMTOption -> Bool) -> SMTOption -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SMTOption -> SMTOption -> Bool
eqCon SMTOption
opt))
    where
      eqCon :: SMTOption -> SMTOption -> Bool
      eqCon :: SMTOption -> SMTOption -> Bool
eqCon SMTOption
l SMTOption
r = Constr -> String
showConstr (SMTOption -> Constr
forall a. Data a => a -> Constr
toConstr SMTOption
l) String -> String -> Bool
forall a. Eq a => a -> a -> Bool
== Constr -> String
showConstr (SMTOption -> Constr
forall a. Data a => a -> Constr
toConstr SMTOption
r)

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

instance RenderSeq SMT where
  renderSeq :: SMT -> Seq Builder
renderSeq SMT
smt =
       [Builder] -> Seq Builder
forall a. [a] -> Seq a
fromList (SMTOption -> Builder
forall a. Render a => a -> Builder
render (SMTOption -> Builder) -> [SMTOption] -> [Builder]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SMT
smtSMT -> Getting [SMTOption] SMT [SMTOption] -> [SMTOption]
forall s a. s -> Getting a s a -> a
^.Getting [SMTOption] SMT [SMTOption]
Lens' SMT [SMTOption]
options)
       Seq Builder -> Seq Builder -> Seq Builder
forall a. Seq a -> Seq a -> Seq a
>< Seq Builder
-> (String -> Seq Builder) -> Maybe String -> Seq Builder
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Seq Builder
forall a. Monoid a => a
mempty (Builder -> Seq Builder
forall a. a -> Seq a
singleton (Builder -> Seq Builder)
-> (String -> Builder) -> String -> Seq Builder
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Builder -> Builder
renderSetLogic (Builder -> Builder) -> (String -> Builder) -> String -> Builder
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> Builder
stringUtf8) (SMT
smtSMT -> Getting (Maybe String) SMT (Maybe String) -> Maybe String
forall s a. s -> Getting a s a -> a
^.Getting (Maybe String) SMT (Maybe String)
Lens' SMT (Maybe String)
mlogic)
       Seq Builder -> Seq Builder -> Seq Builder
forall a. Seq a -> Seq a -> Seq a
>< Seq (SomeKnownSMTSort SMTVar) -> Seq Builder
renderVars (SMT
smtSMT
-> Getting
     (Seq (SomeKnownSMTSort SMTVar)) SMT (Seq (SomeKnownSMTSort SMTVar))
-> Seq (SomeKnownSMTSort SMTVar)
forall s a. s -> Getting a s a -> a
^.Getting
  (Seq (SomeKnownSMTSort SMTVar)) SMT (Seq (SomeKnownSMTSort SMTVar))
Lens' SMT (Seq (SomeKnownSMTSort SMTVar))
vars)
       Seq Builder -> Seq Builder -> Seq Builder
forall a. Seq a -> Seq a -> Seq a
>< (Expr 'BoolSort -> Builder) -> Seq (Expr 'BoolSort) -> Seq Builder
forall a b. (a -> b) -> Seq a -> Seq b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Expr 'BoolSort -> Builder
renderAssert (SMT
smtSMT
-> Getting (Seq (Expr 'BoolSort)) SMT (Seq (Expr 'BoolSort))
-> Seq (Expr 'BoolSort)
forall s a. s -> Getting a s a -> a
^.Getting (Seq (Expr 'BoolSort)) SMT (Seq (Expr 'BoolSort))
Lens' SMT (Seq (Expr 'BoolSort))
formulas)

renderSetLogic :: Builder -> Builder
renderSetLogic :: Builder -> Builder
renderSetLogic = Builder -> Builder -> Builder
forall a. Render a => Builder -> a -> Builder
renderUnary Builder
"set-logic"

renderDeclareVar :: forall t. KnownSMTSort t => SMTVar t -> Builder
renderDeclareVar :: forall (t :: SMTSort). KnownSMTSort t => SMTVar t -> Builder
renderDeclareVar SMTVar t
v = Builder -> SMTVar t -> Builder -> SSMTSort t -> Builder
forall a b c.
(Render a, Render b, Render c) =>
Builder -> a -> b -> c -> Builder
renderTernary Builder
"declare-fun" SMTVar t
v (Builder
"()" :: Builder) (forall (t :: SMTSort). KnownSMTSort t => SSMTSort t
sortSing @t)
{-# INLINEABLE renderDeclareVar #-}

renderAssert :: Expr BoolSort -> Builder
renderAssert :: Expr 'BoolSort -> Builder
renderAssert = Builder -> Expr 'BoolSort -> Builder
forall a. Render a => Builder -> a -> Builder
renderUnary Builder
"assert"
{-# INLINEABLE renderAssert #-}

renderVars :: Seq (SomeKnownSMTSort SMTVar) -> Seq Builder
renderVars :: Seq (SomeKnownSMTSort SMTVar) -> Seq Builder
renderVars = (SomeKnownSMTSort SMTVar -> Builder)
-> Seq (SomeKnownSMTSort SMTVar) -> Seq Builder
forall a b. (a -> b) -> Seq a -> Seq b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\(SomeSMTSort SMTVar t
v) -> SMTVar t -> Builder
forall (t :: SMTSort). KnownSMTSort t => SMTVar t -> Builder
renderDeclareVar SMTVar t
v)
{-# INLINEABLE renderVars #-}