{-# 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 as IMap (singleton)
import Data.Dependent.Map as DMap
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 
            (SSMTSort t -> IntValueMap t -> Solution
forall {k1} (k2 :: k1 -> *) (v :: k1) (f :: k1 -> *).
k2 v -> f v -> DMap k2 f
DMap.singleton 
              (forall (t :: SMTSort). KnownSMTSort t => SSMTSort t
sortSing @t) 
              (IntMap (Value t) -> IntValueMap t
forall (t :: SMTSort). IntMap (Value t) -> IntValueMap t
IntValueMap (IntMap (Value t) -> IntValueMap t)
-> IntMap (Value t) -> IntValueMap t
forall a b. (a -> b) -> a -> b
$ Int -> Value t -> IntMap (Value t)
forall a. Int -> a -> IntMap a
IMap.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
solSMTVarSol t -> Getting (Value t) (SMTVarSol t) (Value t) -> Value t
forall s a. s -> Getting a s a -> a
^.Getting (Value t) (SMTVarSol t) (Value t)
forall (t :: SMTSort) (f :: * -> *).
Functor f =>
(Value t -> f (Value t)) -> SMTVarSol t -> f (SMTVarSol t)
solVal))) 
            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 #-}