{-# LANGUAGE OverloadedStrings #-}

module SMTLIB.Backends
  ( Backend (..),
    QueuingFlag (..),
    Solver,
    initSolver,
    command,
    command_,
    flushQueue,
  )
where

import Control.Monad ((<=<))
import Data.ByteString.Builder (Builder)
import qualified Data.ByteString.Lazy.Char8 as LBS
import Data.Char (isSpace)
import Data.IORef (IORef, modifyIORef, newIORef, readIORef, writeIORef)
import Prelude hiding (log)

-- | The type of solver backends. SMTLib2 commands are sent to a backend which
-- processes them and outputs the solver's response.
data Backend = Backend
  { -- | Send a command to the backend.
    -- While the implementation depends on the backend, this function is usually
    -- *not* thread-safe.
    Backend -> Builder -> IO ByteString
send :: Builder -> IO LBS.ByteString,
    -- | Send a command that doesn't produce any response to the backend.
    -- The backend may implement this by not reading the output and leaving it
    -- for a later read, or reading the output and discarding it immediately.
    -- Hence this method should only be used when the command does not produce
    -- any response to be outputted.
    -- Again, this function may not be thread-safe.
    Backend -> Builder -> IO ()
send_ :: Builder -> IO ()
  }

type Queue = IORef Builder

-- | Whether to enable queuing for a 'Solver'.
data QueuingFlag
  = -- | In 'NoQueuing' mode, the 'Solver' has no queue and commands are sent to
    -- the backend immediately.
    --
    -- In this mode, the option @:print-success@ is enabled by 'initSolver' to
    -- monitor that commands are being accepted by the SMT solver.
    NoQueuing
  | -- | In 'Queuing' mode, commands whose output is not strictly necessary for
    -- the rest of the computation (typically the ones whose output should just
    -- be @success@) are not sent to the backend immediately, but rather written
    -- on the solver's queue.
    --
    -- It is the responsibility of the caller to identify these commands and
    -- sent them through the 'command_' call.
    --
    -- When a command is sent whose output is actually necessary, the queue is
    -- flushed and sent as a batch to the backend.
    --
    -- 'Queuing' mode should be faster as there usually is a non-negligible
    -- constant overhead in sending a command to the backend. When commands are
    -- sent in batches, a command sent to the solver will only produce an error
    -- when it is later sent to the backend. Therefore, you probably want to
    -- stick with 'NoQueuing' mode when debugging.
    --
    -- For parsing to work properly, at most one of the commands in the batch
    -- can produce an output. Hence the @:print-success@ option should not be
    -- enabled in 'Queuing' mode.
    Queuing

-- | Push a command on the solver's queue of commands to evaluate.
-- The command must not produce any output when evaluated, unless it is the last
-- command added before the queue is flushed.
-- For a fixed queue, this function is *not* thread-safe.
put :: Queue -> Builder -> IO ()
put :: Queue -> Builder -> IO ()
put Queue
q Builder
cmd = Queue -> (Builder -> Builder) -> IO ()
forall a. IORef a -> (a -> a) -> IO ()
modifyIORef Queue
q (Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> Builder
cmd)

-- | Empty the queue of commands to evaluate and return its content as a bytestring
-- builder.
-- For a fixed queue, this function is *not* thread-safe.
flush :: Queue -> IO Builder
flush :: Queue -> IO Builder
flush Queue
q = do
  Builder
cmds <- Queue -> IO Builder
forall a. IORef a -> IO a
readIORef Queue
q
  Queue -> Builder -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef Queue
q Builder
forall a. Monoid a => a
mempty
  Builder -> IO Builder
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Builder
cmds

-- | A solver is essentially a wrapper around a solver backend. It also comes an
-- optional queue of commands to send to the backend.
data Solver = Solver
  { -- | The backend processing the commands.
    Solver -> Backend
backend :: Backend,
    -- | An optional queue to write commands that are to be sent to the solver lazily.
    Solver -> Maybe Queue
queue :: Maybe Queue
  }

-- | Create a new solver and initialize it with some options so that it behaves
-- correctly for our use.
-- In particular, the "print-success" option is disabled in 'Queuing' mode. This
-- should not be overriden manually.
initSolver ::
  -- | whether to enable 'Queuing' mode
  QueuingFlag ->
  -- | the solver backend
  Backend ->
  IO Solver
initSolver :: QueuingFlag -> Backend -> IO Solver
initSolver QueuingFlag
queuing Backend
solverBackend = do
  Maybe Queue
solverQueue <- case QueuingFlag
queuing of
    QueuingFlag
Queuing -> do
      Queue
ref <- Builder -> IO Queue
forall a. a -> IO (IORef a)
newIORef Builder
forall a. Monoid a => a
mempty
      Maybe Queue -> IO (Maybe Queue)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe Queue -> IO (Maybe Queue))
-> Maybe Queue -> IO (Maybe Queue)
forall a b. (a -> b) -> a -> b
$ Queue -> Maybe Queue
forall a. a -> Maybe a
Just Queue
ref
    QueuingFlag
NoQueuing -> Maybe Queue -> IO (Maybe Queue)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe Queue
forall a. Maybe a
Nothing
  let solver :: Solver
solver = Backend -> Maybe Queue -> Solver
Solver Backend
solverBackend Maybe Queue
solverQueue
  case QueuingFlag
queuing of
    QueuingFlag
Queuing -> () -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
    QueuingFlag
NoQueuing ->
      -- this should not be enabled when the queue is used, as it messes with parsing
      -- the outputs of commands that are actually interesting
      --
      -- TODO checking for correctness and enabling laziness can be made compatible
      -- but it would require the solver backends to return several outputs at once
      -- alternatively, we may consider that the user wanting both features should
      -- implement their own backend that deals with this
      Solver -> Builder -> IO ()
command_ Solver
solver Builder
"(set-option :print-success true)"
  Solver -> IO Solver
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Solver
solver

-- | Have the solver evaluate a SMT-LIB command.
-- This forces the queued commands to be evaluated as well, but their results are
-- *not* checked for correctness.
--
-- Concurrent calls to different solvers are thread-safe, but not concurrent
-- calls on the same solver or the same backend.
--
-- Only one command must be given per invocation, or the multiple commands must
-- together produce the output of one command only.
command :: Solver -> Builder -> IO LBS.ByteString
command :: Solver -> Builder -> IO ByteString
command Solver
solver Builder
cmd = do
  Backend -> Builder -> IO ByteString
send (Solver -> Backend
backend Solver
solver)
    (Builder -> IO ByteString) -> IO Builder -> IO ByteString
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< case Solver -> Maybe Queue
queue Solver
solver of
      Maybe Queue
Nothing -> Builder -> IO Builder
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Builder
cmd
      Just Queue
q -> (Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> Builder
cmd) (Builder -> Builder) -> IO Builder -> IO Builder
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Queue -> IO Builder
flush Queue
q

-- | A command with no interesting result.
-- In 'NoQueuing' mode, the result is checked for correctness. In 'Queuing'
-- mode, (unless the queue is flushed and evaluated right after) the command
-- must not produce any output when evaluated, and its output is thus in
-- particular not checked for correctness.
--
-- Concurrent calls to different solvers are thread-safe, but not concurrent
-- calls on the same solver or the same backend.
--
-- Only one command must be given per invocation, or the multiple commands must
-- together produce the output of one command only.
command_ :: Solver -> Builder -> IO ()
command_ :: Solver -> Builder -> IO ()
command_ Solver
solver Builder
cmd =
  case Solver -> Maybe Queue
queue Solver
solver of
    Maybe Queue
Nothing -> do
      ByteString
res <- Backend -> Builder -> IO ByteString
send (Solver -> Backend
backend Solver
solver) Builder
cmd
      if ByteString -> ByteString
trim ByteString
res ByteString -> ByteString -> Bool
forall a. Eq a => a -> a -> Bool
== ByteString
"success"
        then () -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
        else
          [Char] -> IO ()
forall a. [Char] -> IO a
forall (m :: * -> *) a. MonadFail m => [Char] -> m a
fail ([Char] -> IO ()) -> [Char] -> IO ()
forall a b. (a -> b) -> a -> b
$
            [[Char]] -> [Char]
unlines
              [ [Char]
"Unexpected result from the SMT solver:",
                [Char]
"  Expected: success",
                [Char]
"  Got: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ ByteString -> [Char]
forall a. Show a => a -> [Char]
show ByteString
res
              ]
    Just Queue
q -> Queue -> Builder -> IO ()
put Queue
q Builder
cmd
  where
    trim :: ByteString -> ByteString
trim = (Char -> Bool) -> ByteString -> ByteString
LBS.dropWhile Char -> Bool
isSpace (ByteString -> ByteString)
-> (ByteString -> ByteString) -> ByteString -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> ByteString
LBS.reverse (ByteString -> ByteString)
-> (ByteString -> ByteString) -> ByteString -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Char -> Bool) -> ByteString -> ByteString
LBS.dropWhile Char -> Bool
isSpace (ByteString -> ByteString)
-> (ByteString -> ByteString) -> ByteString -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> ByteString
LBS.reverse

-- | Force the content of the queue to be sent to the solver.
-- Only useful in queuing mode, does nothing in non-queuing mode.
flushQueue :: Solver -> IO ()
flushQueue :: Solver -> IO ()
flushQueue Solver
solver = IO () -> (Queue -> IO ()) -> Maybe Queue -> IO ()
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (() -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ()) (Backend -> Builder -> IO ()
send_ (Solver -> Backend
backend Solver
solver) (Builder -> IO ()) -> (Queue -> IO Builder) -> Queue -> IO ()
forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< Queue -> IO Builder
flush) (Maybe Queue -> IO ()) -> Maybe Queue -> IO ()
forall a b. (a -> b) -> a -> b
$ Solver -> Maybe Queue
queue Solver
solver