smtlib-tools

  1module Main where
  2
  3import Data.Maybe (isNothing)
  4import Control.Monad.IO.Class (liftIO)
  5import Control.Monad (forM_, void, when, unless)
  6import Control.Monad.State.Strict (StateT, gets, modify, runStateT)
  7import Data.Map.Strict qualified as Map
  8import Data.Maybe (fromMaybe)
  9import SimpleSMT qualified as SMT
 10import System.IO (stdin, hGetContents)
 11import Options.Applicative qualified as OPT
 12
 13data Opts = Opts
 14  { optElimAnds :: Bool,
 15    optLogic :: Maybe String,
 16    optTrackAssert :: Bool }
 17
 18optsParser :: OPT.Parser Opts
 19optsParser =
 20  Opts
 21    <$> OPT.switch
 22      ( OPT.long "eliminate-ands"
 23          <> OPT.short 'e'
 24          <> OPT.help "Eliminate outer and expressions in check-sat-assuming"
 25      )
 26    <*> OPT.optional
 27      ( OPT.strOption $
 28          OPT.long "logic"
 29            <> OPT.short 'l'
 30            <> OPT.metavar "LOGIC"
 31            <> OPT.help "Overwrite logic"
 32      )
 33    <*> OPT.switch
 34      ( OPT.long "track-assertions"
 35          <> OPT.short 'a'
 36          <> OPT.help "Track push/pop assertion stacks"
 37      )
 38
 39------------------------------------------------------------------------
 40
 41data UnwindEnv
 42  = UnwindEnv
 43  { trackAsserts :: Bool,
 44    elimAnds :: Bool,
 45    setLogic :: Maybe String,
 46    logicWritten :: Bool,
 47    declaredVars :: Map.Map String SMT.SExpr,
 48    assertStack :: [[SMT.SExpr]],
 49    exprs :: [SMT.SExpr]
 50  }
 51  deriving (Show, Eq)
 52
 53mkUnwindEnv :: Bool -> Bool -> Maybe String -> UnwindEnv
 54mkUnwindEnv trackAsserts andElim setLogic =
 55  UnwindEnv trackAsserts andElim setLogic False Map.empty [] []
 56
 57------------------------------------------------------------------------
 58
 59buildMap :: [SMT.SExpr] -> Map.Map String SMT.SExpr
 60buildMap exprs =
 61  Map.fromList $
 62    foldl (\acc x -> toVarPair x : acc) [] exprs
 63
 64toVarPair :: SMT.SExpr -> (String, SMT.SExpr)
 65toVarPair (SMT.List [SMT.Atom varName, varValue]) = (varName, varValue)
 66toVarPair _ = error "invalid let binding"
 67
 68substVars :: Map.Map String SMT.SExpr -> SMT.SExpr -> SMT.SExpr
 69substVars varMap atom@(SMT.Atom name) =
 70  fromMaybe atom $ Map.lookup name varMap
 71substVars varMap (SMT.List lst) =
 72  SMT.List $ map (substVars varMap) lst
 73
 74inlineLet' :: Map.Map String SMT.SExpr -> SMT.SExpr -> SMT.SExpr
 75inlineLet' varMap (SMT.List [SMT.Atom "let", SMT.List letBind, letBody]) =
 76  let varMap' = Map.union (buildMap letBind) varMap
 77    in substVars varMap' (inlineLet' varMap' letBody)
 78  where
 79inlineLet' varMap atom = substVars varMap atom
 80
 81inlineLet :: SMT.SExpr -> SMT.SExpr
 82inlineLet = inlineLet' Map.empty
 83
 84------------------------------------------------------------------------
 85
 86newAssertLevel :: StateT UnwindEnv IO ()
 87newAssertLevel =
 88  modify (\s -> s {assertStack = [] : assertStack s})
 89
 90popAssertLevel :: StateT UnwindEnv IO ()
 91popAssertLevel = modify go
 92  where
 93    go s@UnwindEnv {assertStack = []} = s
 94    go s@UnwindEnv {assertStack = _ : xs} = s {assertStack = xs}
 95
 96addAssertion :: SMT.SExpr -> StateT UnwindEnv IO ()
 97addAssertion assertions' = do
 98  stk <- gets assertStack
 99  let newStk = case stk of
100        (x : xs) -> (x ++ assertions) : xs
101        [] -> [assertions]
102  modify (\s -> s {assertStack = newStk})
103  where
104    assertions = [inlineLet assertions']
105
106addExpr :: SMT.SExpr -> StateT UnwindEnv IO ()
107addExpr expr =
108  modify (\s -> s {exprs = exprs s ++ [expr]})
109
110-- | Retrieve a list of conjunction from and expressions for check-sat-assuming.
111subsumeAnds :: SMT.SExpr -> [SMT.SExpr]
112subsumeAnds (SMT.List [SMT.Atom "and", lhs@(SMT.List (SMT.Atom "and" : _)), rhs])
113  = subsumeAnds lhs ++ subsumeAnds rhs
114subsumeAnds (SMT.List [SMT.Atom "and", lhs, rhs@(SMT.List (SMT.Atom "and" : _))])
115  = subsumeAnds lhs ++ subsumeAnds rhs
116subsumeAnds (SMT.List [SMT.Atom "and", lhs, rhs])
117  = subsumeAnds lhs ++ subsumeAnds rhs
118subsumeAnds expr = [expr] -- Stop recursion on first non-and expression.
119
120getAsserts :: StateT UnwindEnv IO [SMT.SExpr]
121getAsserts = do
122  andOpt <- gets elimAnds
123  consLs <- gets (concat . reverse . assertStack)
124  if andOpt
125    then pure [SMT.List (inlineAnds consLs)]
126    else pure [SMT.List consLs]
127  where
128    inlineAnds :: [SMT.SExpr] -> [SMT.SExpr]
129    inlineAnds = concat . map subsumeAnds
130
131completeQuery :: StateT UnwindEnv IO ()
132completeQuery = do
133  query <- gets exprs
134  liftIO $ putStr (serialize query)
135
136  track <- gets trackAsserts
137  if track
138    then modify ( \s -> s { exprs = [] })
139    else modify ( \s -> s { exprs = [], assertStack = [] })
140  where
141    serialize :: [SMT.SExpr] -> String
142    serialize = unlines . map (`SMT.showsSExpr` "")
143
144addVariable :: String -> SMT.SExpr -> StateT UnwindEnv IO ()
145addVariable name expr = do
146  vars <- gets declaredVars
147  when (isNothing $ Map.lookup name vars) $ do
148    liftIO $ putStrLn (SMT.showsSExpr expr "")
149    modify (\s -> s {declaredVars = Map.insert name expr vars})
150
151transExpr :: SMT.SExpr -> StateT UnwindEnv IO ()
152transExpr (SMT.List [SMT.Atom "push", SMT.Atom arg]) = do
153  let num = (read arg :: Integer)
154  forM_ [1 .. num] (const newAssertLevel)
155transExpr (SMT.List [SMT.Atom "pop", SMT.Atom arg]) = do
156  let num = (read arg :: Integer)
157  forM_ [1 .. num] (const popAssertLevel)
158transExpr (SMT.List [SMT.Atom "assert", xs]) =
159  addAssertion xs
160transExpr (SMT.List [SMT.Atom "check-sat"]) = do
161  asserts <- getAsserts
162  addExpr $ SMT.List (SMT.Atom "check-sat-assuming" : asserts)
163  completeQuery
164transExpr expr@(SMT.List ((SMT.Atom "set-logic") : _)) = do
165  wasWritten <- gets logicWritten
166  unless (wasWritten) $ do
167    mayLogic <- gets setLogic
168    case mayLogic of
169      Just l ->
170        liftIO $ putStrLn ((SMT.showsSExpr $ SMT.List [SMT.Atom "set-logic", SMT.Atom l]) "")
171      Nothing ->
172        liftIO $ putStrLn ((SMT.showsSExpr expr) "")
173    modify (\s -> s {logicWritten = True})
174transExpr expr@(SMT.List ((SMT.Atom "declare-fun") : (SMT.Atom name) : _)) = addVariable name expr
175transExpr (SMT.List ((SMT.Atom "set-option" : _))) = pure ()
176transExpr (SMT.List ((SMT.Atom "set-info") : _)) = pure ()
177transExpr (SMT.List ((SMT.Atom "get-value") : _)) = pure ()
178transExpr (SMT.List ((SMT.Atom "exit") : _)) = modify (\s -> s { assertStack = [] })
179transExpr expr = addExpr expr
180
181transform :: [SMT.SExpr] -> StateT UnwindEnv IO ()
182transform sexprs = forM_ sexprs transExpr
183
184------------------------------------------------------------------------
185
186readSExprs :: String -> [SMT.SExpr]
187readSExprs str = go (SMT.readSExpr str)
188  where
189    go :: Maybe (SMT.SExpr, String) -> [SMT.SExpr]
190    go Nothing = []
191    go (Just (acc, rest)) = acc : go (SMT.readSExpr rest)
192
193getQueries :: Opts -> [SMT.SExpr] -> IO ()
194getQueries opts exprs = void $ runTransform exprs
195  where
196    runTransform e =
197      runStateT
198        (transform e)
199        $ mkUnwindEnv (optTrackAssert opts) (optElimAnds opts) (optLogic opts)
200
201main :: IO ()
202main = do
203  args  <- OPT.execParser cmd
204  exprs <- readSExprs <$> hGetContents stdin
205
206  getQueries args exprs
207  where
208    cmd :: OPT.ParserInfo Opts
209    cmd =
210      OPT.info
211        (optsParser OPT.<**> OPT.helper)
212        ( OPT.fullDesc
213            <> OPT.progDesc "normalize smt-lib inputs"
214        )