{-# LANGUAGE OverloadedStrings #-}

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

import Data.ByteString.Builder (Builder)
import qualified Data.ByteString.Lazy.Char8 as LBS
import Data.Char (isSpace)
import Data.IORef (IORef, atomicModifyIORef, newIORef)
import Data.List (intersperse)
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.
    Backend -> Builder -> IO ByteString
send :: Builder -> IO LBS.ByteString
  }

type Queue = IORef Builder

-- | 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.
putQueue :: Queue -> Builder -> IO ()
putQueue :: Queue -> Builder -> IO ()
putQueue Queue
q Builder
cmd = Queue -> (Builder -> (Builder, ())) -> IO ()
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef Queue
q ((Builder -> (Builder, ())) -> IO ())
-> (Builder -> (Builder, ())) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Builder
cmds ->
  (Builder
cmds 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.
flushQueue :: Queue -> IO Builder
flushQueue :: Queue -> IO Builder
flushQueue Queue
q = Queue -> (Builder -> (Builder, Builder)) -> IO Builder
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef Queue
q ((Builder -> (Builder, Builder)) -> IO Builder)
-> (Builder -> (Builder, Builder)) -> IO Builder
forall a b. (a -> b) -> a -> b
$ \Builder
cmds ->
  (Builder
forall a. Monoid a => a
mempty, Builder
cmds)

-- | A solver is essentially a wrapper around a solver backend. It also comes with
-- a function for logging the solver's activity, and an optional queue of commands
-- to send to the backend.
--
-- A solver can either be in eager mode or lazy mode. In eager mode, the queue of
-- commands isn't used and the commands are sent to the backend immediately. In
-- lazy mode, commands whose output are not strictly necessary for the rest of the
-- computation (typically the ones whose output should just be "success") and that
-- are sent through 'command_' are not sent to the backend immediately, but
-- rather written on the solver's queue. When a command whose output is actually
-- necessary needs to be sent, the queue is flushed and sent as a batch to the
-- backend.
--
-- Lazy mode should be faster as there usually is a non-negligible constant
-- overhead in sending a command to the backend. But since the commands are sent by
-- batches, a command sent to the solver will only produce an error when the queue
-- is flushed, i.e. when a command with interesting output is sent. You thus
-- probably want to stick with eager mode when debugging. Moreover, when commands
-- are sent by batches, only the last command in the batch may produce an output
-- for parsing to work properly. Hence the ":print-success" option is disabled in
-- lazy mode, and this should not be overriden manually.
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 lazy mode. This should
-- not be overriden manually.
initSolver ::
  Backend ->
  -- | whether to enable lazy mode (see 'Solver' for the meaning of this flag)
  Bool ->
  IO Solver
initSolver :: Backend -> Bool -> IO Solver
initSolver Backend
solverBackend Bool
lazy = do
  Maybe Queue
solverQueue <-
    if Bool
lazy
      then 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 (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
      else Maybe Queue -> IO (Maybe Queue)
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
  if Bool
lazy
    then () -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
    else -- 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 -> Builder -> IO ()
setOption Solver
solver Builder
"print-success" Builder
"true"
  Solver -> Builder -> Builder -> IO ()
setOption Solver
solver Builder
"produce-models" Builder
"true"
  Solver -> IO Solver
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.
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 (m :: * -> *) a. Monad m => a -> m a
return (Builder -> IO Builder) -> Builder -> IO Builder
forall a b. (a -> b) -> a -> b
$ 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
flushQueue Queue
q

-- | A command with no interesting result.
-- In eager mode, the result is checked for correctness.
-- In lazy 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.
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 (m :: * -> *) a. Monad m => a -> m a
return ()
        else
          [Char] -> IO ()
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 ()
putQueue 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

setOption :: Solver -> Builder -> Builder -> IO ()
setOption :: Solver -> Builder -> Builder -> IO ()
setOption Solver
solver Builder
name Builder
value = Solver -> Builder -> IO ()
command_ Solver
solver (Builder -> IO ()) -> Builder -> IO ()
forall a b. (a -> b) -> a -> b
$ [Builder] -> Builder
list [Builder
"set-option", Builder
":" Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> Builder
name, Builder
value]

list :: [Builder] -> Builder
list :: [Builder] -> Builder
list [Builder]
bs = Builder
"(" Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> [Builder] -> Builder
forall a. Monoid a => [a] -> a
mconcat (Builder -> [Builder] -> [Builder]
forall a. a -> [a] -> [a]
intersperse Builder
" " [Builder]
bs) Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> Builder
")"