Solver.hs 3.17 KB
Newer Older
1 2
{-# LANGUAGE TypeSynonymInstances, FlexibleInstances #-}

3
module Solver
4
    (checkSat,checkSatInt,checkSatBool,MModelS,MModelI,MModelB,
Philipp Meyer's avatar
Philipp Meyer committed
5 6
     MModel(..),mVal,mValues,mElemsWith,mElemSum,mElem,mNotElem,
     cElem,cNotElem,
7
     Z3Type(..),mkOr',mkAnd',mkAdd',mkSub',mkMul')
8 9
where

10 11
import Z3.Monad
import Control.Monad
12
import qualified Data.Map as M
13

14
newtype MModel a = MModel { getMap :: M.Map String a }
15

16
instance Show a => Show (MModel a) where
17
        show = show . M.toList . getMap
18

19 20
type MModelS = MModel AST
type MModelI = MModel Integer
21
type MModelB = MModel (Maybe Bool)
22

23
class Z3Type a where
24
        mkVal :: a -> Z3 AST
25
        getVal :: AST -> Z3 a
26 27

instance Z3Type Integer where
28 29 30 31 32 33 34 35 36
        mkVal = mkInt
        getVal = getInt

instance Z3Type (Maybe Bool) where
        mkVal x = case x of
                      Nothing -> error "can not make undefined constant"
                      Just True -> mkTrue
                      Just False -> mkFalse
        getVal  = getBool
37 38

mVal :: MModel a -> String -> a
39
mVal m x = M.findWithDefault (error ("key not found: " ++ x)) x (getMap m)
40

41
mValues :: MModel a -> [a]
42 43
mValues m = M.elems $ getMap m

Philipp Meyer's avatar
Philipp Meyer committed
44 45 46 47 48 49
mElem :: MModelS -> String -> Z3 AST
mElem m x = (mVal m x `mkGt`) =<< mkInt (0::Integer)

mNotElem :: MModelS -> String -> Z3 AST
mNotElem m x = mkEq (mVal m x) =<< mkInt (0::Integer)

50
mElemsWith :: (a -> Bool) -> MModel a -> [String]
51 52
mElemsWith f m = M.keys $ M.filter f $ getMap m

Philipp Meyer's avatar
Philipp Meyer committed
53 54 55 56 57 58 59 60
mElemSum :: MModelS -> [String] -> Z3 AST
mElemSum m xs = mkAdd' $ map (mVal m) xs

cElem :: MModelI -> String -> Bool
cElem m x = mVal m x > 0

cNotElem :: MModelI -> String -> Bool
cNotElem m x = mVal m x == 0
61

62 63 64 65 66 67 68 69
mkOr' :: [AST] -> Z3 AST
mkOr' [] = mkFalse
mkOr' xs = mkOr xs

mkAnd' :: [AST] -> Z3 AST
mkAnd' [] = mkTrue
mkAnd' xs = mkAnd xs

70 71 72 73 74 75 76 77 78 79 80 81
mkAdd' :: [AST] -> Z3 AST
mkAdd' [] = mkInt (0::Integer)
mkAdd' xs = mkAdd xs

mkSub' :: [AST] -> Z3 AST
mkSub' [] = mkInt (0::Integer)
mkSub' xs = mkSub xs

mkMul' :: [AST] -> Z3 AST
mkMul' [] = mkInt (1::Integer)
mkMul' xs = mkMul xs

82 83 84 85 86 87 88 89 90 91 92 93 94 95
checkSat :: Z3Type a => Z3 Sort -> ([String], MModel AST -> Z3 ()) ->
        Z3 (Maybe (MModel a))
checkSat mkSort (vars, constraint) = do
        sort <- mkSort
        syms <- mapM (mkStringSymbol >=> flip mkConst sort) vars
        let smodel = MModel $ M.fromList $ vars `zip` syms
        constraint smodel
        result <- getModel
        case result of
            (Unsat, Nothing) -> return Nothing
            (Sat, Just m) -> do
                ms <- evalT m syms
                case ms of
                    Just xs -> do
96
                        vals <- mapM getVal xs
97 98 99 100 101 102
                        let cmodel = MModel $ M.fromList $ vars `zip` vals
                        return $ Just cmodel
                    Nothing -> error "Prover returned incomplete model"
            (Undef, _) -> error "Prover returned unknown"
            (Unsat, Just _) -> error "Prover returned unsat but a model"
            (Sat, Nothing) -> error "Prover returned sat but no model"
103

104
checkSatInt :: ([String], MModel AST -> Z3 ()) -> IO (Maybe (MModel Integer))
105 106 107 108
checkSatInt = evalZ3 . checkSat mkIntSort

checkSatBool :: ([String], MModel AST -> Z3 ()) -> IO (Maybe (MModel (Maybe Bool)))
checkSatBool = evalZ3 . checkSat mkBoolSort