-- SPDX-FileCopyrightText: 2024 University of Bremen
-- SPDX-FileCopyrightText: 2025 Sören Tempel <soeren+git@soeren-tempel.net>
--
-- SPDX-License-Identifier: MIT AND GPL-3.0-only

module Language.QBE.Backend.Tracer
  ( Branch (Branch),
    newBranch,
    fromBranch,
    ExecTrace,
    newExecTrace,
    toSExprs,
    appendBranch,
    appendCons,
    solveTrace,
  )
where

import Control.Exception (throwIO)
import Control.Monad (when)
import Language.QBE.Backend (SolverError (UnknownResult), prefixLength)
import Language.QBE.Backend.Model qualified as Model
import Language.QBE.Simulator.Symbolic.Expression qualified as SE
import SimpleBV qualified as SMT

-- Represents a branch condition in the executed code
data Branch
  = Branch
      Bool -- Whether negation of the branch was attempted
      SE.BitVector -- The symbolic branch condition
  deriving (Int -> Branch -> ShowS
[Branch] -> ShowS
Branch -> String
(Int -> Branch -> ShowS)
-> (Branch -> String) -> ([Branch] -> ShowS) -> Show Branch
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Branch -> ShowS
showsPrec :: Int -> Branch -> ShowS
$cshow :: Branch -> String
show :: Branch -> String
$cshowList :: [Branch] -> ShowS
showList :: [Branch] -> ShowS
Show, Branch -> Branch -> Bool
(Branch -> Branch -> Bool)
-> (Branch -> Branch -> Bool) -> Eq Branch
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Branch -> Branch -> Bool
== :: Branch -> Branch -> Bool
$c/= :: Branch -> Branch -> Bool
/= :: Branch -> Branch -> Bool
Eq)

-- Create a new branch condition.
newBranch :: SE.BitVector -> Branch
newBranch :: BitVector -> Branch
newBranch = Bool -> BitVector -> Branch
Branch Bool
False

-- Create a new branch from an existing branch, thereby updating its metadata.
-- It is assumed that the condition, encoded in the branches, is equal.
fromBranch :: Branch -> Branch -> Branch
fromBranch :: Branch -> Branch -> Branch
fromBranch (Branch Bool
wasNeg' BitVector
_) (Branch Bool
wasNeg BitVector
ast) =
  Bool -> BitVector -> Branch
Branch (Bool
wasNeg Bool -> Bool -> Bool
|| Bool
wasNeg') BitVector
ast

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

-- Represents a single execution through a program, tracking for each
-- symbolic branch condition if it was 'True' or 'False'.
type ExecTrace = [(Bool, Branch)]

-- Create a new empty execution tree.
newExecTrace :: ExecTrace
newExecTrace :: ExecTrace
newExecTrace = []

-- Return all branch conditions of an 'ExecTrace'.
toSExprs :: ExecTrace -> [SMT.SExpr]
toSExprs :: ExecTrace -> [SExpr]
toSExprs = ((Bool, Branch) -> SExpr) -> ExecTrace -> [SExpr]
forall a b. (a -> b) -> [a] -> [b]
map (\(Bool
_, Branch Bool
_ BitVector
bv) -> BitVector -> SExpr
SE.toSExpr BitVector
bv)

-- Append a branch to the execution trace, denoting via a 'Bool'
-- if the branch was taken or if it was not taken.
appendBranch :: ExecTrace -> Bool -> Branch -> ExecTrace
appendBranch :: ExecTrace -> Bool -> Branch -> ExecTrace
appendBranch ExecTrace
trace Bool
wasTrue Branch
branch = ExecTrace
trace ExecTrace -> ExecTrace -> ExecTrace
forall a. [a] -> [a] -> [a]
++ [(Bool
wasTrue, Branch
branch)]

-- Append a constraint to the execution tree. This constraint must
-- be true and, contrary to appendBranch, negation will not be
-- attempted for it.
appendCons :: ExecTrace -> SE.BitVector -> ExecTrace
appendCons :: ExecTrace -> BitVector -> ExecTrace
appendCons ExecTrace
trace BitVector
cons = ExecTrace
trace ExecTrace -> ExecTrace -> ExecTrace
forall a. [a] -> [a] -> [a]
++ [(Bool
True, Bool -> BitVector -> Branch
Branch Bool
True BitVector
cons)]

-- For a given execution trace, return an assignment (represented
-- as a 'Model.Model') which statisfies all symbolic branch conditions.
-- If such an assignment does not exist, then 'Nothing' is returned.
--
-- Throws a 'SolverError' on an unknown solver result (e.g., on timeout).
solveTrace :: SMT.Solver -> [SMT.SExpr] -> ExecTrace -> ExecTrace -> IO (Maybe Model.Model)
solveTrace :: Solver -> [SExpr] -> ExecTrace -> ExecTrace -> IO (Maybe Model)
solveTrace Solver
solver [SExpr]
inputVars ExecTrace
oldTrace ExecTrace
newTrace = do
  -- Determine the common prefix of the current trace and the old trace
  -- drop constraints beyond this common prefix from the current solver
  -- context. Thereby, keeping the common prefix and making use of
  -- incremental solving capabilities.
  let prefix :: Int
prefix = ExecTrace -> ExecTrace -> Int
forall a. Eq a => [a] -> [a] -> Int
prefixLength ExecTrace
newTrace ExecTrace
oldTrace
  let toDrop :: Int
toDrop = ExecTrace -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ExecTrace
oldTrace Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
prefix

  -- Micro optimization: When we don't have anything to drop, then
  -- don't call .popMany thereby avoiding communication with the solver.
  Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
toDrop Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
0) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$
    Solver -> Integer -> IO ()
SMT.popMany Solver
solver (Int -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
toDrop)

  -- Only enforce new constraints, i.e. those beyond the common prefix.
  ExecTrace -> IO ()
assertTrace (Int -> ExecTrace -> ExecTrace
forall a. Int -> [a] -> [a]
drop Int
prefix ExecTrace
newTrace)

  Result
isSat <- Solver -> IO Result
SMT.check Solver
solver
  case Result
isSat of
    Result
SMT.Sat -> Model -> Maybe Model
forall a. a -> Maybe a
Just (Model -> Maybe Model) -> IO Model -> IO (Maybe Model)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Solver -> [SExpr] -> IO Model
Model.getModel Solver
solver [SExpr]
inputVars
    Result
SMT.Unsat -> Maybe Model -> IO (Maybe Model)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe Model
forall a. Maybe a
Nothing
    Result
SMT.Unknown -> SolverError -> IO (Maybe Model)
forall e a. Exception e => e -> IO a
throwIO SolverError
UnknownResult
  where
    -- Add all conditions enforced by the given 'ExecTrace' to the solver.
    -- Returns a list of all asserted conditions.
    assertTrace :: ExecTrace -> IO ()
    assertTrace :: ExecTrace -> IO ()
assertTrace [] = () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
    assertTrace ExecTrace
t = do
      let conds :: [SExpr]
conds = ((Bool, Branch) -> SExpr) -> ExecTrace -> [SExpr]
forall a b. (a -> b) -> [a] -> [b]
map (\(Bool
b, Branch Bool
_ BitVector
c) -> Bool -> BitVector -> SExpr
SE.toCond Bool
b BitVector
c) ExecTrace
t
      (SExpr -> IO ()) -> [SExpr] -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (\SExpr
c -> Solver -> IO ()
SMT.push Solver
solver IO () -> IO () -> IO ()
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Solver -> SExpr -> IO ()
SMT.assert Solver
solver SExpr
c) [SExpr]
conds