-- SPDX-FileCopyrightText: 2025-2026 Sören Tempel <soeren+git@soeren-tempel.net>
--
-- SPDX-License-Identifier: GPL-3.0-only
module Language.QBE.Simulator.Explorer
  ( defSolver,
    logSolver,
    PathResult (..),
    Engine (expLastPath),
    newEngine,
    explorePath,
    exploreFunc,
  )
where

import Control.Applicative (empty, (<|>))
import Control.Monad.Catch (try)
import Control.Monad.IO.Class (liftIO)
import Control.Monad.State.Strict (StateT, evalStateT, get, lift, modify, put)
import Data.Map qualified as Map
import Language.QBE.Backend.DFS (PathSel, findUnexplored, newPathSel, trackTrace)
import Language.QBE.Backend.Model (Model)
import Language.QBE.Backend.Store qualified as ST
import Language.QBE.Backend.Tracer qualified as T
import Language.QBE.Simulator (execFunc)
import Language.QBE.Simulator.Concolic.State
  ( Env (envStore),
    ErrorPath (pathError, pathInput),
    ErrorState (errStore, errTracer),
    SimState (..),
    makeConcolic,
    runPath,
  )
import Language.QBE.Simulator.Error (EvalError)
import Language.QBE.Types qualified as QBE
import SimpleBV qualified as SMT
import System.Directory (findExecutable)
import System.IO (Handle)

logic :: String
logic :: String
logic = String
"QF_BV"

findSolver :: IO (String, [String])
findSolver :: IO (String, [String])
findSolver =
  String -> [String] -> IO (String, [String])
solver String
"bitwuzla" []
    IO (String, [String])
-> IO (String, [String]) -> IO (String, [String])
forall a. IO a -> IO a -> IO a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> String -> [String] -> IO (String, [String])
solver String
"z3" [String
"-smt2", String
"-in"]
    IO (String, [String])
-> IO (String, [String]) -> IO (String, [String])
forall a. IO a -> IO a -> IO a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> String -> [String] -> IO (String, [String])
solver String
"cvc5" [String
"--incremental"]
    IO (String, [String])
-> IO (String, [String]) -> IO (String, [String])
forall a. IO a -> IO a -> IO a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> String -> IO (String, [String])
forall a. String -> IO a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"no suitable sover found in PATH"
  where
    solver :: String -> [String] -> IO (String, [String])
    solver :: String -> [String] -> IO (String, [String])
solver String
exec [String]
args = do
      Maybe String
r <- String -> IO (Maybe String)
findExecutable String
exec
      IO (String, [String])
-> (String -> IO (String, [String]))
-> Maybe String
-> IO (String, [String])
forall b a. b -> (a -> b) -> Maybe a -> b
maybe IO (String, [String])
forall a. IO a
forall (f :: * -> *) a. Alternative f => f a
empty (\String
_ -> (String, [String]) -> IO (String, [String])
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (String
exec, [String]
args)) Maybe String
r

defSolver :: IO SMT.Solver
defSolver :: IO Solver
defSolver = do
  -- l <- SMT.newLogger 0
  (String
solver, [String]
args) <- IO (String, [String])
findSolver
  Solver
s <- String -> [String] -> Maybe Logger -> IO Solver
SMT.newSolver String
solver [String]
args Maybe Logger
forall a. Maybe a
Nothing
  Solver -> String -> IO ()
SMT.setLogic Solver
s String
logic
  Solver -> IO Solver
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Solver
s

logSolver :: Handle -> IO SMT.Solver
logSolver :: Handle -> IO Solver
logSolver Handle
handle = do
  Logger
l <- Handle -> Int -> IO Logger
SMT.newLoggerWithHandle Handle
handle Int
0
  (String
solver, [String]
args) <- IO (String, [String])
findSolver
  Solver
s <-
    Config -> IO Solver
SMT.newSolverWithConfig
      (String -> [String] -> Config
SMT.defaultConfig String
solver [String]
args)
        { SMT.solverLogger = SMT.smtSolverLogger l
        }
  Solver -> String -> IO ()
SMT.setLogic Solver
s String
logic
  Solver -> IO Solver
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Solver
s

------------------------------------------------------------------------

data PathResult
  = PathResult
  { PathResult -> Maybe EvalError
pathErr :: Maybe EvalError,
    PathResult -> ExecTrace
pathTrace :: T.ExecTrace,
    PathResult -> Assign
pathVars :: ST.Assign
  }
  deriving (Int -> PathResult -> ShowS
[PathResult] -> ShowS
PathResult -> String
(Int -> PathResult -> ShowS)
-> (PathResult -> String)
-> ([PathResult] -> ShowS)
-> Show PathResult
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> PathResult -> ShowS
showsPrec :: Int -> PathResult -> ShowS
$cshow :: PathResult -> String
show :: PathResult -> String
$cshowList :: [PathResult] -> ShowS
showList :: [PathResult] -> ShowS
Show, PathResult -> PathResult -> Bool
(PathResult -> PathResult -> Bool)
-> (PathResult -> PathResult -> Bool) -> Eq PathResult
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: PathResult -> PathResult -> Bool
== :: PathResult -> PathResult -> Bool
$c/= :: PathResult -> PathResult -> Bool
/= :: PathResult -> PathResult -> Bool
Eq)

initPath :: PathResult
initPath :: PathResult
initPath = Maybe EvalError -> ExecTrace -> Assign -> PathResult
PathResult Maybe EvalError
forall a. Maybe a
Nothing [] Assign
forall k a. Map k a
Map.empty

data Engine
  = Engine
  { Engine -> Solver
expSolver :: SMT.Solver,
    Engine -> PathSel
expPathSel :: PathSel,
    Engine -> Env
expEnv :: Env,
    Engine -> PathResult
expLastPath :: PathResult
  }

newEngine :: Env -> SMT.Solver -> Engine
newEngine :: Env -> Solver -> Engine
newEngine Env
env Solver
solver =
  Engine
    { expSolver :: Solver
expSolver = Solver
solver,
      expPathSel :: PathSel
expPathSel = PathSel
newPathSel,
      expEnv :: Env
expEnv = Env
env,
      expLastPath :: PathResult
expLastPath = PathResult
initPath
    }

findNext :: [SMT.SExpr] -> T.ExecTrace -> StateT Engine IO (Maybe Model)
findNext :: [SExpr] -> ExecTrace -> StateT Engine IO (Maybe Model)
findNext [SExpr]
symVars ExecTrace
eTrace = do
  Engine
engine <- StateT Engine IO Engine
forall s (m :: * -> *). MonadState s m => m s
get

  let pathSel :: PathSel
pathSel = PathSel -> ExecTrace -> PathSel
trackTrace (Engine -> PathSel
expPathSel Engine
engine) ExecTrace
eTrace
  (Maybe Model
model, PathSel
nextPathSel) <-
    IO (Maybe Model, PathSel)
-> StateT Engine IO (Maybe Model, PathSel)
forall a. IO a -> StateT Engine IO a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Maybe Model, PathSel)
 -> StateT Engine IO (Maybe Model, PathSel))
-> IO (Maybe Model, PathSel)
-> StateT Engine IO (Maybe Model, PathSel)
forall a b. (a -> b) -> a -> b
$ Solver -> [SExpr] -> PathSel -> IO (Maybe Model, PathSel)
findUnexplored (Engine -> Solver
expSolver Engine
engine) [SExpr]
symVars PathSel
pathSel

  Engine -> StateT Engine IO ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put (Engine -> StateT Engine IO ()) -> Engine -> StateT Engine IO ()
forall a b. (a -> b) -> a -> b
$ Engine
engine {expPathSel = nextPathSel}
  Maybe Model -> StateT Engine IO (Maybe Model)
forall a. a -> StateT Engine IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe Model
model

-- TODO: Consider modelling changes of the PathSel (via findNext) and
-- changes of the Store (via ST.finalize and ST.setModel) as a StateT.
explorePath :: SimState a -> StateT Engine IO Bool
explorePath :: forall a. SimState a -> StateT Engine IO Bool
explorePath SimState a
simState = do
  engine :: Engine
engine@(Engine {expEnv :: Engine -> Env
expEnv = Env
env}) <- StateT Engine IO Engine
forall s (m :: * -> *). MonadState s m => m s
get
  Either ErrorPath (ExecTrace, Store)
maybePath <- StateT Engine IO (ExecTrace, Store)
-> StateT Engine IO (Either ErrorPath (ExecTrace, Store))
forall (m :: * -> *) e a.
(HasCallStack, MonadCatch m, Exception e) =>
m a -> m (Either e a)
try (StateT Engine IO (ExecTrace, Store)
 -> StateT Engine IO (Either ErrorPath (ExecTrace, Store)))
-> StateT Engine IO (ExecTrace, Store)
-> StateT Engine IO (Either ErrorPath (ExecTrace, Store))
forall a b. (a -> b) -> a -> b
$ Env -> StateT Engine IO (ExecTrace, Store)
forall {t :: (* -> *) -> * -> *}.
MonadTrans t =>
Env -> t IO (ExecTrace, Store)
run Env
env
  let (Maybe EvalError
mayErr, ExecTrace
eTrace, Store
nStore) =
        case Either ErrorPath (ExecTrace, Store)
maybePath of
          Left (ErrorPath
err :: ErrorPath) ->
            let st :: ErrorState
st = ErrorPath -> ErrorState
pathInput ErrorPath
err
             in (EvalError -> Maybe EvalError
forall a. a -> Maybe a
Just (EvalError -> Maybe EvalError) -> EvalError -> Maybe EvalError
forall a b. (a -> b) -> a -> b
$ ErrorPath -> EvalError
pathError ErrorPath
err, ErrorState -> ExecTrace
errTracer ErrorState
st, ErrorState -> Store
errStore ErrorState
st)
          Right (ExecTrace
t, Store
s) -> (Maybe EvalError
forall a. Maybe a
Nothing, ExecTrace
t, Store
s)

  -- Before finalizing the store, we can extract the variables we encountered
  -- during this concrete execution, as well as the concrete values used for
  -- these variables during the execution.
  let inputVars :: [SExpr]
inputVars = Store -> [SExpr]
ST.sexprs Store
nStore
      varAssign :: Assign
varAssign = Store -> Assign
ST.cValues Store
nStore
  Engine -> StateT Engine IO ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put (Engine -> StateT Engine IO ()) -> Engine -> StateT Engine IO ()
forall a b. (a -> b) -> a -> b
$ Engine
engine {expLastPath = PathResult mayErr eTrace varAssign}

  -- Finalize the store (declare new symbolic vars in solver) and then,
  -- based on the new solver state, solve constraints to find a new input.
  Store
store <- IO Store -> StateT Engine IO Store
forall a. IO a -> StateT Engine IO a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO Store -> StateT Engine IO Store)
-> IO Store -> StateT Engine IO Store
forall a b. (a -> b) -> a -> b
$ Solver -> Store -> IO Store
ST.finalize (Engine -> Solver
expSolver Engine
engine) Store
nStore
  Maybe Model
model <- [SExpr] -> ExecTrace -> StateT Engine IO (Maybe Model)
findNext [SExpr]
inputVars ExecTrace
eTrace
  case Maybe Model
model of
    Maybe Model
Nothing -> Bool -> StateT Engine IO Bool
forall a. a -> StateT Engine IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
False
    Just Model
newModel -> do
      let nEnv :: Env
nEnv = Env
env {envStore = ST.setModel store newModel}
       in (Engine -> Engine) -> StateT Engine IO ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify (\Engine
e -> Engine
e {expEnv = nEnv})
      Bool -> StateT Engine IO Bool
forall a. a -> StateT Engine IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
True
  where
    run :: Env -> t IO (ExecTrace, Store)
run Env
env = IO (ExecTrace, Store) -> t IO (ExecTrace, Store)
forall (m :: * -> *) a. Monad m => m a -> t m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (IO (ExecTrace, Store) -> t IO (ExecTrace, Store))
-> IO (ExecTrace, Store) -> t IO (ExecTrace, Store)
forall a b. (a -> b) -> a -> b
$ StateT Env IO (ExecTrace, Store) -> Env -> IO (ExecTrace, Store)
forall (m :: * -> *) s a. Monad m => StateT s m a -> s -> m a
evalStateT (SimState (ExecTrace, Store) -> StateT Env IO (ExecTrace, Store)
forall a. SimState a -> StateT Env IO a
unSimState (SimState (ExecTrace, Store) -> StateT Env IO (ExecTrace, Store))
-> SimState (ExecTrace, Store) -> StateT Env IO (ExecTrace, Store)
forall a b. (a -> b) -> a -> b
$ SimState a -> SimState (ExecTrace, Store)
forall a. SimState a -> SimState (ExecTrace, Store)
runPath SimState a
simState) Env
env

------------------------------------------------------------------------

exploreFunc ::
  Engine ->
  QBE.FuncDef ->
  [(String, QBE.ExtType)] ->
  IO [PathResult]
exploreFunc :: Engine -> FuncDef -> [(String, ExtType)] -> IO [PathResult]
exploreFunc Engine
engine FuncDef
entry [(String, ExtType)]
params = do
  let funcState :: SimState (Maybe (Concolic RegVal))
funcState = ((String, ExtType) -> SimState (Concolic RegVal))
-> [(String, ExtType)] -> SimState [Concolic RegVal]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM ((String -> ExtType -> SimState (Concolic RegVal))
-> (String, ExtType) -> SimState (Concolic RegVal)
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry String -> ExtType -> SimState (Concolic RegVal)
makeConcolic) [(String, ExtType)]
params SimState [Concolic RegVal]
-> ([Concolic RegVal] -> SimState (Maybe (Concolic RegVal)))
-> SimState (Maybe (Concolic RegVal))
forall a b. SimState a -> (a -> SimState b) -> SimState b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= FuncDef -> [Concolic RegVal] -> SimState (Maybe (Concolic RegVal))
forall (m :: * -> *) v.
Simulator m v =>
FuncDef -> [v] -> m (Maybe v)
execFunc FuncDef
entry
  StateT Engine IO [PathResult] -> Engine -> IO [PathResult]
forall (m :: * -> *) s a. Monad m => StateT s m a -> s -> m a
evalStateT (SimState (Maybe (Concolic RegVal)) -> StateT Engine IO [PathResult]
forall {a}. SimState a -> StateT Engine IO [PathResult]
exploreFunc' SimState (Maybe (Concolic RegVal))
funcState) Engine
engine
  where
    exploreFunc' :: SimState a -> StateT Engine IO [PathResult]
exploreFunc' SimState a
st = do
      Bool
morePaths <- SimState a -> StateT Engine IO Bool
forall a. SimState a -> StateT Engine IO Bool
explorePath SimState a
st
      Engine
curEngine <- StateT Engine IO Engine
forall s (m :: * -> *). MonadState s m => m s
get

      let ret :: PathResult
ret = Engine -> PathResult
expLastPath Engine
curEngine
       in if Bool
morePaths
            then (PathResult
ret :) ([PathResult] -> [PathResult])
-> StateT Engine IO [PathResult] -> StateT Engine IO [PathResult]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SimState a -> StateT Engine IO [PathResult]
exploreFunc' SimState a
st
            else [PathResult] -> StateT Engine IO [PathResult]
forall a. a -> StateT Engine IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [PathResult
ret]