Solvers.hs

{- |
    Module: EVM.Solvers
    Description: Solver orchestration
-}
module EVM.Solvers where

import Prelude hiding (LT, GT)

import GHC.Natural
import GHC.IO.Handle (Handle, hFlush, hSetBuffering, BufferMode(..))
import Control.Concurrent.Chan (Chan, newChan, writeChan, readChan)
import Control.Concurrent (forkIO, killThread)
import Control.Concurrent.STM (writeTChan, newTChan, TChan, tryReadTChan, atomically)
import Control.Monad
import Control.Monad.State.Strict
import Control.Monad.IO.Unlift
import Data.Map (Map)
import Data.Map qualified as Map
import Data.Set (Set, isSubsetOf, fromList, toList)
import Data.Maybe (fromMaybe, isJust, fromJust)
import Data.Either (isLeft)
import Data.Text qualified as TStrict
import Data.Text.Lazy (Text)
import Data.Text.Lazy qualified as T
import Data.Text.Lazy.IO qualified as T
import Data.Text.Lazy.Builder
import qualified Data.Text.Lazy.Builder.Int (decimal)
import System.Process (createProcess, cleanupProcess, proc, ProcessHandle, std_in, std_out, std_err, StdStream(..), createPipe)
import System.FilePath ((</>))
import EVM.Effects
import Data.Bits ((.&.))
import Numeric (showHex)
import EVM.Expr (simplifyProps)

import EVM.Keccak qualified as Keccak (concreteKeccaks)
import EVM.SMT
import EVM.Types

-- | Supported solvers
data Solver
  = Z3
  | CVC5
  | Bitwuzla
  | EmptySolver
  | Custom Text

instance Show Solver where
  show Z3 = "z3"
  show CVC5 = "cvc5"
  show Bitwuzla = "bitwuzla"
  show EmptySolver = "empty-smt-solver"
  show (Custom s) = T.unpack s


-- | A running solver instance
data SolverInstance = SolverInstance
  { solvertype :: Solver
  , stdin      :: Handle
  , stdout     :: Handle
  , process    :: ProcessHandle
  }

-- | A channel representing a group of solvers
newtype SolverGroup = SolverGroup (Chan Task)

data MultiSol = MultiSol
  { maxSols :: Int
  , numBytes :: Int
  , var :: String
  }

-- | A script to be executed, a list of models to be extracted in the case of a sat result, and a channel where the result should be written
data Task = TaskSingle SingleData | TaskMulti MultiData

newtype CacheEntry = CacheEntry [Prop]
  deriving (Show, Eq)

data MultiData = MultiData
  { smt2 :: SMT2
  , multiSol :: MultiSol
  , resultChan :: Chan (Maybe [W256])
  }

data SingleData = SingleData
  { smt2 :: SMT2
  , props :: Maybe [Prop]
  , resultChan :: Chan SMTResult
  }

-- returns True if a is a superset of any of the sets in bs
supersetAny :: Set Prop -> [Set Prop] -> Bool
supersetAny a bs = any (`isSubsetOf` a) bs

checkMulti :: SolverGroup -> Err SMT2 -> MultiSol -> IO (Maybe [W256])
checkMulti (SolverGroup taskq) smt2 multiSol = do
  if isLeft smt2 then pure Nothing
  else do
    -- prepare result channel
    resChan <- newChan
    -- send task to solver group
    writeChan taskq (TaskMulti (MultiData (getNonError smt2) multiSol resChan))
    -- collect result
    readChan resChan

checkSatWithProps :: App m => SolverGroup -> [Prop] -> m (SMTResult)
checkSatWithProps sg props = do
  conf <- readConfig
  let psSimp = if conf.simp then simplifyProps props else props
  if psSimp == [PBool False] then pure Qed
  else do
    let concreteKeccaks = fmap (\(buf,val) -> PEq (Lit val) (Keccak buf)) (toList $ Keccak.concreteKeccaks props)
    let smt2 = assertProps conf (if conf.simp then psSimp <> concreteKeccaks else psSimp)
    if isLeft smt2 then pure $ Error $ getError smt2
    else liftIO $ checkSat sg (Just props) smt2

-- When props is Nothing, the cache will not be filled or used
checkSat :: SolverGroup -> Maybe [Prop] -> Err SMT2 -> IO SMTResult
checkSat (SolverGroup taskq) props smt2 = do
  if isLeft smt2 then pure $ Error $ getError smt2
  else do
    -- prepare result channel
    resChan <- newChan
    -- send task to solver group
    writeChan taskq (TaskSingle (SingleData (getNonError smt2) props resChan))
    -- collect result
    readChan resChan

writeSMT2File :: SMT2 -> FilePath -> String -> IO ()
writeSMT2File smt2 path postfix = do
    let content = formatSMT2 smt2 <> "\n\n(check-sat)"
        fullPath = path </> "query-" <> postfix <> ".smt2"
    T.writeFile fullPath content

withSolvers :: App m => Solver -> Natural -> Natural -> Maybe Natural -> (SolverGroup -> m a) -> m a
withSolvers solver count threads timeout cont = do
    -- spawn solvers
    instances <- mapM (const $ liftIO $ spawnSolver solver threads timeout) [1..count]
    -- spawn orchestration thread
    taskq <- liftIO newChan
    cacheq <- liftIO . atomically $ newTChan
    availableInstances <- liftIO newChan
    liftIO $ forM_ instances (writeChan availableInstances)
    orchestrate' <- toIO $ orchestrate taskq cacheq availableInstances [] 0
    orchestrateId <- liftIO $ forkIO orchestrate'

    -- run continuation with task queue
    res <- cont (SolverGroup taskq)

    -- cleanup and return results
    liftIO $ mapM_ (stopSolver) instances
    liftIO $ killThread orchestrateId
    pure res
  where
    orchestrate :: App m => Chan Task -> TChan CacheEntry -> Chan SolverInstance -> [Set Prop] -> Int -> m b
    orchestrate taskq cacheq avail knownUnsat fileCounter = do
      conf <- readConfig
      mx <- liftIO . atomically $ tryReadTChan cacheq
      case mx of
        Just (CacheEntry props)  -> do
          let knownUnsat' = (fromList props):knownUnsat
          when conf.debug $ liftIO $ putStrLn "   adding UNSAT cache"
          orchestrate taskq cacheq avail knownUnsat' fileCounter
        Nothing -> do
          task <- liftIO $ readChan taskq
          case task of
            TaskSingle (SingleData _ props r) | isJust props && supersetAny (fromList (fromJust props)) knownUnsat -> do
              liftIO $ writeChan r Qed
              when conf.debug $ liftIO $ putStrLn "   Qed found via cache!"
              orchestrate taskq cacheq avail knownUnsat fileCounter
            _ -> do
              inst <- liftIO $ readChan avail
              runTask' <- case task of
                TaskSingle (SingleData smt2 props r) -> toIO $ getOneSol smt2 props r cacheq inst avail fileCounter
                TaskMulti (MultiData smt2 multiSol r) -> toIO $ getMultiSol smt2 multiSol r inst avail fileCounter
              _ <- liftIO $ forkIO runTask'
              orchestrate taskq cacheq avail knownUnsat (fileCounter + 1)

getMultiSol :: forall m. (MonadIO m, ReadConfig m) => SMT2 -> MultiSol -> (Chan (Maybe [W256])) -> SolverInstance -> Chan SolverInstance -> Int -> m ()
getMultiSol smt2@(SMT2 cmds cexvars _) multiSol r inst availableInstances fileCounter = do
  conf <- readConfig
  when conf.dumpQueries $ liftIO $ writeSMT2File smt2 "." (show fileCounter)
  -- reset solver and send all lines of provided script
  out <- liftIO $ do
    resetRes <- sendScript inst $ SMTScript [SMTCommand "(reset)"]
    case resetRes of
      e@(Left _) -> pure e
      _ -> sendScript inst cmds
  case out of
    Left err -> liftIO $ do
      when conf.debug $ putStrLn $ "Unable to write SMT to solver: " <> (T.unpack err)
      writeChan r Nothing
    Right _ -> do
      sat <- liftIO $ sendCommand inst $ SMTCommand "(check-sat)"
      when conf.dumpQueries $ liftIO $ writeSMT2File smt2 "." (show fileCounter <> "-origquery")
      subRun [] smt2 sat
  -- put the instance back in the list of available instances
  liftIO $ writeChan availableInstances inst
  where
    maskFromBytesCount k
      | k <= 32 = (2 ^ (8 * k) - 1)
      | otherwise = internalError "Byte length exceeds 256-bit capacity"
    subRun :: (MonadIO m, ReadConfig m) => [W256] -> SMT2 -> Text -> m ()
    subRun vals fullSmt sat = do
      conf <- readConfig
      case sat of
        "unsat" -> liftIO $ do
          when conf.debug $ putStrLn $ "No more solutions to query, returning: " <> show vals
          liftIO $ writeChan r (Just vals)
        "timeout" -> liftIO $ do
           when conf.debug $ putStrLn "Timeout inside SMT solver."
           writeChan r Nothing
        "unknown" -> liftIO $ do
           when conf.debug $ putStrLn "Unknown result by SMT solver."
           dumpUnsolved fullSmt fileCounter conf.dumpUnsolved
           writeChan r Nothing
        "sat" -> do
          if length vals >= multiSol.maxSols then liftIO $ do
            when conf.debug $ putStrLn "Too many solutions to symbolic query."
            writeChan r Nothing
          else do
            cex <- liftIO $ getModel inst cexvars
            case Map.lookup (Var (TStrict.pack multiSol.var)) cex.vars of
              Just v -> do
                let hexMask = maskFromBytesCount multiSol.numBytes
                    maskedVal = v .&. hexMask
                    toSMT n = Data.Text.Lazy.Builder.Int.decimal n
                    maskedVar = "(bvand " <> fromString multiSol.var <> " (_ bv" <> toSMT hexMask <> " 256))"
                    restrict = SMTCommand $ "(assert (not (= " <> maskedVar <> " (_ bv" <> toSMT maskedVal <> " 256))))"
                    newSmt = fullSmt <> SMT2 (SMTScript [restrict]) mempty mempty
                when conf.debug $ liftIO $ putStrLn $ "Got one solution to symbolic query, val: 0x" <> (showHex maskedVal "") <>
                  " now have " <> show (length vals + 1) <> " solution(s), max is: " <> show multiSol.maxSols
                when conf.dumpQueries $ liftIO $ writeSMT2File newSmt "." (show fileCounter <> "-sol" <> show (length vals))
                out <- liftIO $ sendCommand inst restrict
                case out of
                  "success" -> do
                    out2 <- liftIO $ sendCommand inst  (SMTCommand "(check-sat)")
                    subRun (maskedVal:vals) newSmt out2
                  err -> liftIO $ do
                    when conf.debug $ putStrLn $ "Unable to write SMT to solver: " <> (T.unpack err)
                    writeChan r Nothing
              Nothing -> internalError $ "variable " <>  multiSol.var <> " not part of model (i.e. cex) ... that's not possible"
        err -> liftIO $ do
          when conf.debug $ putStrLn $ "Unable to write SMT to solver: " <> (T.unpack err)
          writeChan r Nothing

getOneSol :: (MonadIO m, ReadConfig m) => SMT2 -> Maybe [Prop] -> Chan SMTResult -> TChan CacheEntry -> SolverInstance -> Chan SolverInstance -> Int -> m ()
getOneSol smt2@(SMT2 cmds cexvars _) props r cacheq inst availableInstances fileCounter = do
  conf <- readConfig
  liftIO $ do
    when (conf.dumpQueries) $ writeSMT2File smt2 "." (show fileCounter)
    -- reset solver and send all lines of provided script
    out <- do
      resetRes <- sendScript inst $ SMTScript [SMTCommand "(reset)"]
      case resetRes of
        e@(Left _) -> pure e
        _ -> sendScript inst cmds
    case out of
      -- if we got an error then return it
      Left e -> writeChan r (Error $ "Error while writing SMT to solver: " <> T.unpack e)
      -- otherwise call (check-sat), parse the result, and send it down the result channel
      Right () -> do
        sat <- sendCommand inst $ SMTCommand "(check-sat)"
        res <- do
            case sat of
              "unsat" -> do
                when (isJust props) $ liftIO . atomically $ writeTChan cacheq (CacheEntry (fromJust props))
                pure Qed
              "timeout" -> pure $ Unknown "Result timeout by SMT solver"
              "unknown" -> do
                dumpUnsolved smt2 fileCounter conf.dumpUnsolved
                pure $ Unknown "Result unknown by SMT solver"
              "sat" -> Cex <$> getModel inst cexvars
              _ -> pure . Error $ "Unable to parse SMT solver output: " <> T.unpack sat
        writeChan r res

    -- put the instance back in the list of available instances
    writeChan availableInstances inst

dumpUnsolved :: SMT2 -> Int -> Maybe FilePath -> IO ()
dumpUnsolved fullSmt fileCounter dump = do
   case dump of
     Just path -> writeSMT2File fullSmt path $ "unsolved-" <> show fileCounter
     Nothing -> pure ()

getModel :: SolverInstance -> CexVars -> IO SMTCex
getModel inst cexvars = do
  -- get an initial version of the model from the solver
  initialModel <- getRaw
  -- check the sizes of buffer models and shrink if needed
  if bufsUsable initialModel
  then pure initialModel
  else do
    -- get concrete values for each buffers max read index
    hints <- capHints <$> queryMaxReads (getValue inst) cexvars.buffers
    snd <$> runStateT (shrinkModel hints) initialModel
  where
    getRaw :: IO SMTCex
    getRaw = do
      vars <- getVars parseVar (getValue inst) (fmap T.toStrict cexvars.calldata)
      addrs <- getAddrs parseEAddr (getValue inst) (fmap T.toStrict cexvars.addrs)
      buffers <- getBufs (getValue inst) (Map.keys cexvars.buffers)
      storage <- getStore (getValue inst) cexvars.storeReads
      blockctx <- getVars parseBlockCtx (getValue inst) (fmap T.toStrict cexvars.blockContext)
      txctx <- getVars parseTxCtx (getValue inst) (fmap T.toStrict cexvars.txContext)
      pure $ SMTCex vars addrs buffers storage blockctx txctx

    -- sometimes the solver might give us back a model for the max read index
    -- that is too high to be a useful cex (e.g. in the case of reads from a
    -- symbolic index), so we cap the max value of the starting point to be 1024
    capHints :: Map Text W256 -> Map Text W256
    capHints = fmap (min 1024)

    -- shrink all the buffers in a model
    shrinkModel :: Map Text W256 -> StateT SMTCex IO ()
    shrinkModel hints = do
      m <- get
      -- iterate over all the buffers in the model, and shrink each one in turn if needed
      forM_ (Map.keys m.buffers) $ \case
        AbstractBuf b -> do
          let name = T.fromStrict b
              hint = fromMaybe
                       (internalError $ "Could not find hint for buffer: " <> T.unpack name)
                       (Map.lookup name hints)
          shrinkBuf name hint
        _ -> internalError "Received model from solver for non AbstractBuf"

    -- starting with some guess at the max useful size for a buffer, cap
    -- it's size to that value, and ask the solver to check satisfiability. If
    -- it's still sat with the new constraint, leave that constraint on the
    -- stack and return a new model, if it's unsat, double the size of the hint
    -- and try again.
    shrinkBuf :: Text -> W256 -> StateT SMTCex IO ()
    shrinkBuf buf hint = do
      let encBound = "(_ bv" <> Data.Text.Lazy.Builder.Int.decimal hint <> " 256)"
      answer <- liftIO $ do
        checkCommand inst $ SMTCommand "(push 1)"
        checkCommand inst $ SMTCommand ("(assert (bvule " <> (fromLazyText buf) <> "_length " <> encBound <> "))")
        sendCommand inst $ SMTCommand "(check-sat)"
      case answer of
        "sat" -> do
          model <- liftIO getRaw
          put model
        "unsat" -> do
          liftIO $ checkCommand inst $ SMTCommand "(pop 1)"
          let nextHint = if hint == 0 then 1 else hint * 2
          if nextHint < hint || nextHint > 1_073_741_824
            then pure () -- overflow or over 1GB
            else shrinkBuf buf nextHint
        _ -> do -- unexpected answer -> clean up and do not change the model
          liftIO $ checkCommand inst $ SMTCommand "(pop 1)"
          pure ()


    -- we set a pretty arbitrary upper limit (of 32) to decide if we need to do some shrinking
    bufsUsable :: SMTCex -> Bool
    bufsUsable model = any (go . snd) (Map.toList model.buffers)
      where
        go (Flat _) = True
        go (Comp c) = case c of
          (Base _ sz) -> sz <= 32
          -- TODO: do I need to check the write idx here?
          (Write _ idx next) -> idx <= 32 && go (Comp next)

mkTimeout :: Maybe Natural -> Text
mkTimeout t = T.pack $ show $ (1000 *)$ case t of
  Nothing -> 300 :: Natural
  Just t' -> t'

-- | Arguments used when spawning a solver instance
solverArgs :: Solver -> Natural -> Maybe Natural -> [Text]
solverArgs solver threads timeout = case solver of
  Bitwuzla ->
    [ "--lang=smt2"
    , "--produce-models"
    , "--time-limit-per=" <> mkTimeout timeout
    , "--bv-solver=preprop"
    , "--bv-output-format=16"
    ]
  Z3 ->
    [ "-st"
    , "smt.threads=" <> (T.pack $ show threads)
    , "-in" ]
  CVC5 ->
    [ "--lang=smt"
    , "--produce-models"
    , "--print-success"
    , "--interactive"
    , "--incremental"
    , "--tlimit-per=" <> mkTimeout timeout
    , "--arrays-exp"
    ]
  EmptySolver -> []
  Custom _ -> []

-- | Spawns a solver instance, and sets the various global config options that we use for our queries
spawnSolver :: Solver -> Natural -> Maybe (Natural) -> IO SolverInstance
spawnSolver solver threads timeout = do
  (readout, writeout) <- createPipe
  let cmd
        = (proc (show solver) (fmap T.unpack $ solverArgs solver threads timeout ))
            { std_in = CreatePipe
            , std_out = UseHandle writeout
            , std_err = UseHandle writeout
            }
  (Just stdin, Nothing, Nothing, process) <- createProcess cmd
  hSetBuffering stdin (BlockBuffering (Just 1000000))
  let solverInstance = SolverInstance solver stdin readout process

  case solver of
    CVC5 -> pure solverInstance
    Bitwuzla -> do
      _ <- sendCommand solverInstance $ SMTCommand "(set-option :print-success true)"
      pure solverInstance
    EmptySolver -> pure solverInstance
    Z3 -> do
      _ <- sendCommand solverInstance $ SMTCommand "(set-option :print-success true)"
      _ <- sendCommand solverInstance $ SMTCommand ("(set-option :timeout " <> (fromLazyText $ mkTimeout timeout) <> ")")
      pure solverInstance
    Custom _ -> pure solverInstance

-- | Cleanly shutdown a running solver instance
stopSolver :: SolverInstance -> IO ()
stopSolver (SolverInstance _ stdin stdout process) = cleanupProcess (Just stdin, Just stdout, Nothing, process)

-- | Sends a list of commands to the solver. Returns the first error, if there was one.
sendScript :: SolverInstance -> SMTScript -> IO (Either Text ())
sendScript solver (SMTScript entries) = do
  go entries
  where
    go [] = pure $ Right ()
    go (SMTComment _ : rest) = go rest
    go (c@(SMTCommand command):cs) = do
      out <- sendCommand solver c
      case out of
        "success" -> go cs
        e -> pure $ Left $ "Solver returned an error:\n" <> e <> "\nwhile sending the following command: " <> toLazyText command

checkCommand :: SolverInstance -> SMTEntry -> IO ()
checkCommand inst cmd = do
  res <- sendCommand inst cmd
  case res of
    "success" -> pure ()
    _ -> internalError $ "Unexpected solver output: " <> T.unpack res


-- | Strips trailing \r, if present
stripCarriageReturn :: Text -> Text
stripCarriageReturn t = fromMaybe t $ T.stripSuffix "\r" t

-- | Sends a string to the solver and appends a newline, returns the first available line from the output buffer
sendCommand :: SolverInstance -> SMTEntry -> IO Text
sendCommand _ (SMTComment _) = internalError "Attempting to send a comment as a command to SMT solver"
sendCommand (SolverInstance _ stdin stdout _) (SMTCommand cmd) = do
  T.hPutStrLn stdin $ toLazyText cmd
  hFlush stdin
  stripCarriageReturn <$> (T.hGetLine stdout)

-- | Returns a string representation of the model for the requested variable
getValue :: SolverInstance -> Text -> IO Text
getValue (SolverInstance _ stdin stdout _) var = do
  T.hPutStrLn stdin (T.append (T.append "(get-value (" var) "))")
  hFlush stdin
  fmap (T.unlines . reverse) (readSExpr stdout)

-- | Reads lines from h until we have a balanced sexpr
readSExpr :: Handle -> IO [Text]
readSExpr h = go 0 0 []
  where
    go 0 0 _ = do
      line <- T.hGetLine h
      let cleanLine = stripCarriageReturn line
          ls = T.length $ T.filter (== '(') cleanLine
          rs = T.length $ T.filter (== ')') cleanLine
      if ls == rs
         then pure [cleanLine]
         else go ls rs [cleanLine]
    go ls rs prev = do
      line <- T.hGetLine h
      let cleanLine = stripCarriageReturn line
          ls' = T.length $ T.filter (== '(') cleanLine
          rs' = T.length $ T.filter (== ')') cleanLine
      if (ls + ls') == (rs + rs')
         then pure $ cleanLine : prev
         else go (ls + ls') (rs + rs') (cleanLine : prev)