-- SPDX-FileCopyrightText: 2025 Sören Tempel <soeren+git@soeren-tempel.net>
--
-- SPDX-License-Identifier: GPL-3.0-only

module Language.QBE.Backend.Store
  ( Store (cValues),
    Assign,
    empty,
    sexprs,
    finalize,
    setModel,
    getConcolic,
  )
where

import Data.Map qualified as Map
import Language.QBE.Backend.Model qualified as Model
import Language.QBE.Simulator.Concolic.Expression qualified as CE
import Language.QBE.Simulator.Default.Expression qualified as DE
import Language.QBE.Simulator.Expression qualified as E
import Language.QBE.Simulator.Symbolic.Expression qualified as SE
import Language.QBE.Types qualified as QBE
import SimpleBV qualified as SMT
import System.Random (StdGen, genWord64R)

-- | Concrete variable assignment.
type Assign = Map.Map String DE.RegVal

-- A variable store mapping variable names to concrete values.
data Store
  = Store
  { Store -> Assign
cValues :: Assign,
    Store -> Map String BitVector
sValues :: Map.Map String SE.BitVector,
    Store -> Map String BitVector
defined :: Map.Map String SE.BitVector,
    Store -> StdGen
randGen :: StdGen
  }

-- | Create a new (empty) store.
empty :: StdGen -> Store
empty :: StdGen -> Store
empty = Assign
-> Map String BitVector -> Map String BitVector -> StdGen -> Store
Store Assign
forall k a. Map k a
Map.empty Map String BitVector
forall k a. Map k a
Map.empty Map String BitVector
forall k a. Map k a
Map.empty

-- | Obtain symbolic values as a list of "SimpleBV" expressions.
sexprs :: Store -> [SMT.SExpr]
sexprs :: Store -> [SExpr]
sexprs = (BitVector -> SExpr) -> [BitVector] -> [SExpr]
forall a b. (a -> b) -> [a] -> [b]
map BitVector -> SExpr
SE.toSExpr ([BitVector] -> [SExpr])
-> (Store -> [BitVector]) -> Store -> [SExpr]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map String BitVector -> [BitVector]
forall k a. Map k a -> [a]
Map.elems (Map String BitVector -> [BitVector])
-> (Store -> Map String BitVector) -> Store -> [BitVector]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Store -> Map String BitVector
sValues

-- | Finalize all pending symbolic variable declarations.
finalize :: SMT.Solver -> Store -> IO Store
finalize :: Solver -> Store -> IO Store
finalize Solver
solver store :: Store
store@(Store {sValues :: Store -> Map String BitVector
sValues = Map String BitVector
m, defined :: Store -> Map String BitVector
defined = Map String BitVector
defs}) = do
  let new :: Map String BitVector
new = Map String BitVector
m Map String BitVector
-> Map String BitVector -> Map String BitVector
forall k a b. Ord k => Map k a -> Map k b -> Map k a
`Map.difference` Map String BitVector
defs
  ((String, BitVector) -> IO SExpr) -> [(String, BitVector)] -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ ((String -> BitVector -> IO SExpr)
-> (String, BitVector) -> IO SExpr
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry String -> BitVector -> IO SExpr
declareSymbolic) ([(String, BitVector)] -> IO ()) -> [(String, BitVector)] -> IO ()
forall a b. (a -> b) -> a -> b
$ Map String BitVector -> [(String, BitVector)]
forall k a. Map k a -> [(k, a)]
Map.toList Map String BitVector
new

  Store -> IO Store
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
    Store
store
      { defined = Map.union defs new,
        sValues = Map.empty
      }
  where
    declareSymbolic :: String -> BitVector -> IO SExpr
declareSymbolic String
n BitVector
v =
      Solver -> String -> Int -> IO SExpr
SMT.declareBV Solver
solver String
n (Int -> IO SExpr) -> Int -> IO SExpr
forall a b. (a -> b) -> a -> b
$ BitVector -> Int
SE.bitSize BitVector
v

-- | Create a variable store from a 'Model.Model'.
setModel :: Store -> Model.Model -> Store
setModel :: Store -> Model -> Store
setModel Store
store Model
model =
  Store
store {cValues = Map.fromList $ Model.toList model}

-- | Lookup the variable name in the store, if it doesn't exist return
-- an unconstrained 'CE.Concolic' value with a random concrete part.
getConcolic :: Store -> String -> QBE.ExtType -> (Store, CE.Concolic DE.RegVal)
getConcolic :: Store -> String -> ExtType -> (Store, Concolic RegVal)
getConcolic store :: Store
store@Store {randGen :: Store -> StdGen
randGen = StdGen
rand} String
name ExtType
ty =
  ( Store
store
      { sValues = newSymVars,
        cValues = newConVars,
        randGen = nextRand
      },
    RegVal -> Maybe BitVector -> Concolic RegVal
forall v. v -> Maybe BitVector -> Concolic v
CE.Concolic RegVal
concrete (BitVector -> Maybe BitVector
forall a. a -> Maybe a
Just BitVector
symbolic)
  )
  where
    (BitVector
symbolic, Map String BitVector
newSymVars) =
      let bv :: BitVector
bv = String -> ExtType -> BitVector
SE.symbolic String
name ExtType
ty
       in (BitVector
bv, String -> BitVector -> Map String BitVector -> Map String BitVector
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert String
name BitVector
bv (Map String BitVector -> Map String BitVector)
-> Map String BitVector -> Map String BitVector
forall a b. (a -> b) -> a -> b
$ Store -> Map String BitVector
sValues Store
store)

    (RegVal
concrete, Assign
newConVars, StdGen
nextRand) =
      let cm :: Assign
cm = Store -> Assign
cValues Store
store
       in case String -> Assign -> Maybe RegVal
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup String
name Assign
cm of
            Just RegVal
cv -> (RegVal
cv, Assign
cm, StdGen
rand)
            Maybe RegVal
Nothing ->
              let maxValue :: Word64
maxValue = (Word64
2 Word64 -> Int -> Word64
forall a b. (Num a, Integral b) => a -> b -> a
^ ExtType -> Int
QBE.extTypeBitSize ExtType
ty) Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
- Word64
1
                  (Word64
rv, StdGen
nr) = Word64 -> StdGen -> (Word64, StdGen)
forall g. RandomGen g => Word64 -> g -> (Word64, g)
genWord64R Word64
maxValue StdGen
rand
                  conValue :: RegVal
conValue = ExtType -> Word64 -> RegVal
forall v. ValueRepr v => ExtType -> Word64 -> v
E.fromLit ExtType
ty Word64
rv
               in (RegVal
conValue, String -> RegVal -> Assign -> Assign
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert String
name RegVal
conValue Assign
cm, StdGen
nr)