quebex

  1-- SPDX-FileCopyrightText: 2025 Sören Tempel <soeren+git@soeren-tempel.net>
  2--
  3-- SPDX-License-Identifier: MIT AND GPL-3.0-only
  4
  5module SMTUnwind (unwind) where
  6
  7import Control.Monad (forM_)
  8import Control.Monad.State (State, gets, modify, runState)
  9import Data.Functor ((<&>))
 10import SimpleSMT qualified as SMT
 11
 12data UnwindEnv
 13  = UnwindEnv
 14  { assertStack :: [[SMT.SExpr]],
 15    exprs :: [SMT.SExpr],
 16    queries :: [[SMT.SExpr]]
 17  }
 18  deriving (Show, Eq)
 19
 20mkUnwindEnv :: UnwindEnv
 21mkUnwindEnv = UnwindEnv [] [] []
 22
 23------------------------------------------------------------------------
 24
 25newAssertLevel :: State UnwindEnv ()
 26newAssertLevel =
 27  modify (\s -> s {assertStack = [] : assertStack s})
 28
 29popAssertLevel :: State UnwindEnv ()
 30popAssertLevel = modify go
 31  where
 32    go s@UnwindEnv {assertStack = []} = s
 33    go s@UnwindEnv {assertStack = _ : xs} = s {assertStack = xs}
 34
 35addAssertion :: [SMT.SExpr] -> State UnwindEnv ()
 36addAssertion assertions = do
 37  stk <- gets assertStack
 38  let newStk = case stk of
 39        (x : xs) -> (x ++ assertions) : xs
 40        [] -> [assertions]
 41  modify (\s -> s {assertStack = newStk})
 42
 43addExpr :: SMT.SExpr -> State UnwindEnv ()
 44addExpr expr =
 45  modify (\s -> s {exprs = exprs s ++ [expr]})
 46
 47getAsserts :: State UnwindEnv [SMT.SExpr]
 48getAsserts = gets (concat . assertStack)
 49
 50completeQuery :: State UnwindEnv ()
 51completeQuery =
 52  modify
 53    ( \s ->
 54        s
 55          { queries = queries s ++ [exprs s],
 56            exprs = []
 57          }
 58    )
 59
 60transExpr :: SMT.SExpr -> State UnwindEnv ()
 61transExpr (SMT.List [SMT.Atom "push", SMT.Atom arg]) = do
 62  let num = (read arg :: Integer)
 63  forM_ [1 .. num] (const newAssertLevel)
 64transExpr (SMT.List [SMT.Atom "pop", SMT.Atom arg]) = do
 65  let num = (read arg :: Integer)
 66  forM_ [1 .. num] (const popAssertLevel)
 67transExpr (SMT.List ((SMT.Atom "assert") : xs)) =
 68  addAssertion xs
 69transExpr (SMT.List [SMT.Atom "check-sat"]) = do
 70  asserts <- getAsserts
 71  addExpr $ SMT.List [SMT.Atom "check-sat-assuming", SMT.List asserts]
 72  completeQuery
 73transExpr (SMT.List ((SMT.Atom "get-value") : _)) = pure ()
 74transExpr expr = addExpr expr
 75
 76transform :: [SMT.SExpr] -> State UnwindEnv ()
 77transform sexprs = forM_ sexprs transExpr
 78
 79------------------------------------------------------------------------
 80
 81readSExprs :: String -> [SMT.SExpr]
 82readSExprs str = go (SMT.readSExpr str)
 83  where
 84    go :: Maybe (SMT.SExpr, String) -> [SMT.SExpr]
 85    go Nothing = []
 86    go (Just (acc, rest)) = acc : go (SMT.readSExpr rest)
 87
 88getQueries :: [SMT.SExpr] -> [[SMT.SExpr]]
 89getQueries exprs = queries (snd $ runTransform exprs)
 90  where
 91    runTransform e = runState (transform e) mkUnwindEnv
 92
 93unwind :: FilePath -> IO String
 94unwind origFp = do
 95  exprs <- readFile origFp <&> readSExprs
 96  let queries = getQueries exprs
 97  pure $ serialize (concat queries)
 98  where
 99    serialize :: [SMT.SExpr] -> String
100    serialize = unlines . map (`SMT.showsSExpr` "")