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

module Language.Hasmtlib.Type.Pipe where

import Language.Hasmtlib.Type.SMT
import Language.Hasmtlib.Type.MonadSMT
import Language.Hasmtlib.Internal.Expr
import Language.Hasmtlib.Internal.Render
import Language.Hasmtlib.Type.Solution
import Language.Hasmtlib.Codec
import Language.Hasmtlib.Internal.Parser hiding (var, constant)
import qualified SMTLIB.Backends as B
import Data.List (isPrefixOf)
import Data.IntMap (singleton)
import Data.Coerce
import Data.ByteString.Builder
import Data.ByteString.Lazy hiding (filter, singleton, isPrefixOf)
import Data.Attoparsec.ByteString hiding (Result)
import Control.Monad.State
import Control.Lens hiding (List)

-- | A pipe to the solver.
--   If 'B.Solver' is 'B.Queuing' then all commands that do not expect an answer are sent to the queue.
--   All commands that expect an answer have the queue to be sent to the solver before sending the command itself.
--   If 'B.Solver' is not 'B.Queuing', all commands are sent to the solver immediately.
data Pipe = Pipe
  { Pipe -> Int
_lastPipeVarId :: {-# UNPACK #-} !Int              -- ^ Last Id assigned to a new var
  , Pipe -> Maybe String
_mPipeLogic    :: Maybe String                     -- ^ Logic for the SMT-Solver
  , Pipe -> Solver
_pipe          :: !B.Solver                        -- ^ Active pipe to the backend
  }

$(makeLenses ''Pipe)

instance (MonadState Pipe m, MonadIO m) => MonadSMT Pipe 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)) -> Pipe -> (Int, Pipe)
Lens' Pipe Int
lastPipeVarId ((Int -> (Int, Int)) -> Pipe -> (Int, Pipe)) -> 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
    Pipe
smt <- m Pipe
forall s (m :: * -> *). MonadState s m => m s
get
    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
    IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ Solver -> Builder -> IO ()
B.command_ (Pipe
smtPipe -> Getting Solver Pipe Solver -> Solver
forall s a. s -> Getting a s a -> a
^.Getting Solver Pipe Solver
Lens' Pipe Solver
pipe) (Builder -> IO ()) -> Builder -> IO ()
forall a b. (a -> b) -> a -> b
$ SMTVar t -> Builder
forall (t :: SMTSort). KnownSMTSort t => SMTVar t -> Builder
renderDeclareVar 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
    Pipe
smt <- m Pipe
forall s (m :: * -> *). MonadState s m => m s
get
    Expr 'BoolSort
qExpr <- case Pipe
smtPipe -> Getting (Maybe String) Pipe (Maybe String) -> Maybe String
forall s a. s -> Getting a s a -> a
^.Getting (Maybe String) Pipe (Maybe String)
Lens' Pipe (Maybe String)
mPipeLogic 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
    IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ Solver -> Builder -> IO ()
B.command_ (Pipe
smtPipe -> Getting Solver Pipe Solver -> Solver
forall s a. s -> Getting a s a -> a
^.Getting Solver Pipe Solver
Lens' Pipe Solver
pipe) (Builder -> IO ()) -> Builder -> IO ()
forall a b. (a -> b) -> a -> b
$ Expr 'BoolSort -> Builder
renderAssert Expr 'BoolSort
qExpr
  {-# INLINEABLE assert #-}

  setOption :: SMTOption -> m ()
setOption SMTOption
opt = do
    Pipe
smt <- m Pipe
forall s (m :: * -> *). MonadState s m => m s
get
    IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ Solver -> Builder -> IO ()
B.command_ (Pipe
smtPipe -> Getting Solver Pipe Solver -> Solver
forall s a. s -> Getting a s a -> a
^.Getting Solver Pipe Solver
Lens' Pipe Solver
pipe) (Builder -> IO ()) -> Builder -> IO ()
forall a b. (a -> b) -> a -> b
$ SMTOption -> Builder
forall a. Render a => a -> Builder
render SMTOption
opt

  setLogic :: String -> m ()
setLogic String
l = do
    (Maybe String -> Identity (Maybe String)) -> Pipe -> Identity Pipe
Lens' Pipe (Maybe String)
mPipeLogic ((Maybe String -> Identity (Maybe String))
 -> Pipe -> Identity Pipe)
-> String -> m ()
forall s (m :: * -> *) a b.
MonadState s m =>
ASetter s s a (Maybe b) -> b -> m ()
?= String
l
    Pipe
smt <- m Pipe
forall s (m :: * -> *). MonadState s m => m s
get
    IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ Solver -> Builder -> IO ()
B.command_ (Pipe
smtPipe -> Getting Solver Pipe Solver -> Solver
forall s a. s -> Getting a s a -> a
^.Getting Solver Pipe Solver
Lens' Pipe Solver
pipe) (Builder -> IO ()) -> Builder -> IO ()
forall a b. (a -> b) -> a -> b
$ Builder -> Builder
renderSetLogic (String -> Builder
stringUtf8 String
l)

instance (MonadState Pipe m, MonadIO m) => MonadIncrSMT Pipe m where
  push :: m ()
push = do
    Pipe
smt <- m Pipe
forall s (m :: * -> *). MonadState s m => m s
get
    IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ Solver -> Builder -> IO ()
B.command_ (Pipe
smtPipe -> Getting Solver Pipe Solver -> Solver
forall s a. s -> Getting a s a -> a
^.Getting Solver Pipe Solver
Lens' Pipe Solver
pipe) Builder
"(push 1)"
  {-# INLINE push #-}

  pop :: m ()
pop = do
    Pipe
smt <- m Pipe
forall s (m :: * -> *). MonadState s m => m s
get
    IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ Solver -> Builder -> IO ()
B.command_ (Pipe
smtPipe -> Getting Solver Pipe Solver -> Solver
forall s a. s -> Getting a s a -> a
^.Getting Solver Pipe Solver
Lens' Pipe Solver
pipe) Builder
"(pop 1)"
  {-# INLINE pop #-}

  checkSat :: m Result
checkSat = do
    Pipe
smt <- m Pipe
forall s (m :: * -> *). MonadState s m => m s
get
    ByteString
result <- IO ByteString -> m ByteString
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO ByteString -> m ByteString) -> IO ByteString -> m ByteString
forall a b. (a -> b) -> a -> b
$ Solver -> Builder -> IO ByteString
B.command (Pipe
smtPipe -> Getting Solver Pipe Solver -> Solver
forall s a. s -> Getting a s a -> a
^.Getting Solver Pipe Solver
Lens' Pipe Solver
pipe) Builder
"(check-sat)"
    case Parser Result -> ByteString -> Either String Result
forall a. Parser a -> ByteString -> Either String a
parseOnly Parser Result
resultParser (ByteString -> ByteString
toStrict ByteString
result) of
      Left String
e    -> IO Result -> m Result
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO Result -> m Result) -> IO Result -> m Result
forall a b. (a -> b) -> a -> b
$ do
        ByteString -> IO ()
forall a. Show a => a -> IO ()
print ByteString
result
        String -> IO Result
forall a. HasCallStack => String -> a
error String
e
      Right Result
res -> Result -> m Result
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Result
res

  getModel :: m Solution
getModel = do
    Pipe
smt   <- m Pipe
forall s (m :: * -> *). MonadState s m => m s
get
    ByteString
model <- IO ByteString -> m ByteString
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO ByteString -> m ByteString) -> IO ByteString -> m ByteString
forall a b. (a -> b) -> a -> b
$ Solver -> Builder -> IO ByteString
B.command (Pipe
smtPipe -> Getting Solver Pipe Solver -> Solver
forall s a. s -> Getting a s a -> a
^.Getting Solver Pipe Solver
Lens' Pipe Solver
pipe) Builder
"(get-model)"
    case Parser Solution -> ByteString -> Either String Solution
forall a. Parser a -> ByteString -> Either String a
parseOnly Parser Solution
anyModelParser (ByteString -> ByteString
toStrict ByteString
model) of
      Left String
e    -> IO Solution -> m Solution
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO Solution -> m Solution) -> IO Solution -> m Solution
forall a b. (a -> b) -> a -> b
$ do
        ByteString -> IO ()
forall a. Show a => a -> IO ()
print ByteString
model
        String -> IO Solution
forall a. HasCallStack => String -> a
error String
e
      Right Solution
sol -> Solution -> m Solution
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Solution
sol

  getValue :: forall t. KnownSMTSort t => Expr t -> m (Maybe (Decoded (Expr t)))
  getValue :: forall (t :: SMTSort).
KnownSMTSort t =>
Expr t -> m (Maybe (Decoded (Expr t)))
getValue v :: Expr t
v@(Var SMTVar t
x) = do
    Pipe
smt   <- m Pipe
forall s (m :: * -> *). MonadState s m => m s
get
    ByteString
model <- IO ByteString -> m ByteString
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO ByteString -> m ByteString) -> IO ByteString -> m ByteString
forall a b. (a -> b) -> a -> b
$ Solver -> Builder -> IO ByteString
B.command (Pipe
smtPipe -> Getting Solver Pipe Solver -> Solver
forall s a. s -> Getting a s a -> a
^.Getting Solver Pipe Solver
Lens' Pipe Solver
pipe) (Builder -> IO ByteString) -> Builder -> IO ByteString
forall a b. (a -> b) -> a -> b
$ Builder -> Builder -> Builder
forall a. Render a => Builder -> a -> Builder
renderUnary Builder
"get-value" (Builder -> Builder) -> Builder -> Builder
forall a b. (a -> b) -> a -> b
$ Builder
"(" Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> SMTVar t -> Builder
forall a. Render a => a -> Builder
render SMTVar t
x Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> Builder
")"
    case Parser (SMTVarSol t) -> ByteString -> Either String (SMTVarSol t)
forall a. Parser a -> ByteString -> Either String a
parseOnly (forall (t :: SMTSort).
KnownSMTSort t =>
SMTVar t -> Parser (SMTVarSol t)
getValueParser @t SMTVar t
x) (ByteString -> ByteString
toStrict ByteString
model) of
      Left String
e    -> IO (Maybe (HaskellType t)) -> m (Maybe (HaskellType t))
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Maybe (HaskellType t)) -> m (Maybe (HaskellType t)))
-> IO (Maybe (HaskellType t)) -> m (Maybe (HaskellType t))
forall a b. (a -> b) -> a -> b
$ do
        ByteString -> IO ()
forall a. Show a => a -> IO ()
print ByteString
model
        String -> IO (Maybe (HaskellType t))
forall a. HasCallStack => String -> a
error String
e
      Right SMTVarSol t
sol -> Maybe (HaskellType t) -> m (Maybe (HaskellType t))
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe (HaskellType t) -> m (Maybe (HaskellType t)))
-> Maybe (HaskellType t) -> m (Maybe (HaskellType t))
forall a b. (a -> b) -> a -> b
$ Solution -> Expr t -> Maybe (Decoded (Expr t))
forall a. Codec a => Solution -> a -> Maybe (Decoded a)
decode (Int -> SomeKnownSMTSort SMTVarSol -> Solution
forall a. Int -> a -> IntMap a
singleton (SMTVarSol t
solSMTVarSol t -> Getting Int (SMTVarSol t) Int -> Int
forall s a. s -> Getting a s a -> a
^.(SMTVar t -> Const Int (SMTVar t))
-> SMTVarSol t -> Const Int (SMTVarSol t)
forall (t :: SMTSort) (f :: * -> *).
Functor f =>
(SMTVar t -> f (SMTVar t)) -> SMTVarSol t -> f (SMTVarSol t)
solVar((SMTVar t -> Const Int (SMTVar t))
 -> SMTVarSol t -> Const Int (SMTVarSol t))
-> ((Int -> Const Int Int) -> SMTVar t -> Const Int (SMTVar t))
-> Getting Int (SMTVarSol t) Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
.(Int -> Const Int Int) -> SMTVar t -> Const Int (SMTVar t)
forall (t1 :: SMTSort) (t2 :: SMTSort) (p :: * -> * -> *)
       (f :: * -> *).
(Profunctor p, Functor f) =>
p Int (f Int) -> p (SMTVar t1) (f (SMTVar t2))
varId) (SMTVarSol t -> SomeKnownSMTSort SMTVarSol
forall (t :: SMTSort) (f :: SMTSort -> *).
KnownSMTSort t =>
f t -> SomeKnownSMTSort f
SomeKnownSMTSort SMTVarSol t
sol)) Expr t
v
  getValue Expr t
expr = do
    Solution
model <- m Solution
forall s (m :: * -> *). MonadIncrSMT s m => m Solution
getModel
    Maybe (HaskellType t) -> m (Maybe (HaskellType t))
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe (HaskellType t) -> m (Maybe (HaskellType t)))
-> Maybe (HaskellType t) -> m (Maybe (HaskellType t))
forall a b. (a -> b) -> a -> b
$ Solution -> Expr t -> Maybe (Decoded (Expr t))
forall a. Codec a => Solution -> a -> Maybe (Decoded a)
decode Solution
model Expr t
expr
  {-# INLINEABLE getValue #-}