smtlib-tools

 1module Main where
 2
 3import Conduit
 4import Data.Functor (($>))
 5import Data.List (uncons)
 6import Data.Text (Text, unpack)
 7import SimpleSMT qualified as SMT
 8import System.Environment (getArgs)
 9import System.Exit (exitFailure)
10import System.IO (hPutStrLn, stderr)
11
12z3Solver :: IO SMT.Solver
13z3Solver = do
14  -- l <- SMT.newLogger 0
15  s <- SMT.newSolver "z3" ["-smt2", "-in"] Nothing
16  return s
17
18simplify :: SMT.Solver -> SMT.SExpr -> IO SMT.SExpr
19simplify s e = SMT.command s $ SMT.List [SMT.Atom "simplify", e]
20
21------------------------------------------------------------------------
22
23fwdCmd :: SMT.Solver -> SMT.SExpr -> IO SMT.SExpr
24fwdCmd solver expr = SMT.command solver expr $> expr
25
26transSExpr :: String -> SMT.Solver -> SMT.SExpr -> IO SMT.SExpr
27transSExpr _ s e@(SMT.List ((SMT.Atom "set-logic") : _)) = fwdCmd s e
28transSExpr _ s e@(SMT.List ((SMT.Atom "declare-fun") : _)) = fwdCmd s e
29transSExpr a s e@(SMT.List xs) =
30  let recur = SMT.List <$> mapM (transSExpr a s) xs
31   in case uncons xs of
32        Just (SMT.Atom name, _) ->
33          if name == a
34            then simplify s e
35            else recur
36        _ -> recur
37transSExpr _ _ e = pure e
38
39yieldSExpr :: ConduitM Text SMT.SExpr IO ()
40yieldSExpr = loop ""
41  where
42    loop rest = await >>= maybe (return ()) (go . (rest ++) . unpack)
43    go x =
44      case SMT.readSExpr x of
45        Just (expr, rest) -> yield expr >> go rest
46        Nothing -> loop x
47
48------------------------------------------------------------------------
49
50transformStdin :: String -> SMT.Solver -> IO ()
51transformStdin atomName solver =
52  runConduit $
53    stdinC
54      .| decodeUtf8C
55      .| yieldSExpr
56      .| mapMC (transSExpr atomName solver)
57      .| mapM_C printSExpr
58  where
59    printSExpr :: SMT.SExpr -> IO ()
60    printSExpr e = putStrLn $ SMT.showsSExpr e ""
61
62main :: IO ()
63main = do
64  args <- getArgs
65  case args of
66    [exprName] -> z3Solver >>= transformStdin exprName
67    _ -> do
68      hPutStrLn stderr "Expected single name of expression to simplify"
69      exitFailure