module Z3.Tagged.Eval (EvalAst, eval, evalT, mapEval,
                       evalBool, evalInt, evalReal, evalBv, evalFunc,
                       runWithModel) where

import Control.Monad.Trans.Maybe
import Control.Monad.Trans.Reader
import Z3.Tagged (Z3, Model, AST, FuncDecl, FuncModel)
import qualified Z3.Tagged as T

-------------------------------------------------
-- ** Helpers

-- | Type of an evaluation function for '(AST s)'.
--
-- Evaluation may fail (i.e. return 'Nothing') for a few
-- reasons, see 'modelEval'.
type EvalAst s a = AST s -> Eval s a
type Eval s = ReaderT (Model s) (MaybeT (Z3 s))

-- | An alias for 'modelEval' with model completion enabled.
eval :: EvalAst s (AST s)
eval :: EvalAst s (AST s)
eval = (Model s -> AST s -> ReaderT (Z3Env s) (ST s) (Maybe (AST s)))
-> EvalAst s (AST s)
forall r a (m :: * -> *) b.
(r -> a -> m (Maybe b)) -> a -> ReaderT r (MaybeT m) b
liftEval Model s -> AST s -> ReaderT (Z3Env s) (ST s) (Maybe (AST s))
forall s. EvalAst s (AST s)
T.eval

-- | Evaluate an (AST s) node of sort /bool/ in the given model.
--
-- See 'modelEval' and 'getBool'.
evalBool :: EvalAst s Bool
evalBool :: EvalAst s Bool
evalBool = (Model s -> AST s -> ReaderT (Z3Env s) (ST s) (Maybe Bool))
-> EvalAst s Bool
forall r a (m :: * -> *) b.
(r -> a -> m (Maybe b)) -> a -> ReaderT r (MaybeT m) b
liftEval Model s -> AST s -> ReaderT (Z3Env s) (ST s) (Maybe Bool)
forall s. EvalAst s Bool
T.evalBool

-- | Evaluate an (AST s) node of sort /int/ in the given model.
--
-- See 'modelEval' and 'getInt'.
evalInt :: EvalAst s Integer
evalInt :: EvalAst s Integer
evalInt = (Model s -> AST s -> ReaderT (Z3Env s) (ST s) (Maybe Integer))
-> EvalAst s Integer
forall r a (m :: * -> *) b.
(r -> a -> m (Maybe b)) -> a -> ReaderT r (MaybeT m) b
liftEval Model s -> AST s -> ReaderT (Z3Env s) (ST s) (Maybe Integer)
forall s. EvalAst s Integer
T.evalInt

-- | Evaluate an (AST s) node of sort /real/ in the given model.
--
-- See 'modelEval' and 'getReal'.
evalReal :: EvalAst s Rational
evalReal :: EvalAst s Rational
evalReal = (Model s -> AST s -> ReaderT (Z3Env s) (ST s) (Maybe Rational))
-> EvalAst s Rational
forall r a (m :: * -> *) b.
(r -> a -> m (Maybe b)) -> a -> ReaderT r (MaybeT m) b
liftEval Model s -> AST s -> ReaderT (Z3Env s) (ST s) (Maybe Rational)
forall s. EvalAst s Rational
T.evalReal

-- | Evaluate an (AST s) node of sort /bit-vector/ in the given model.
--
-- The flag /signed/ decides whether the bit-vector value is
-- interpreted as a signed or unsigned integer.
--
-- See 'modelEval' and 'getBv'.
evalBv :: Bool -- ^ signed?
       -> EvalAst s Integer
evalBv :: Bool -> EvalAst s Integer
evalBv = (Model s -> AST s -> ReaderT (Z3Env s) (ST s) (Maybe Integer))
-> EvalAst s Integer
forall r a (m :: * -> *) b.
(r -> a -> m (Maybe b)) -> a -> ReaderT r (MaybeT m) b
liftEval ((Model s -> AST s -> ReaderT (Z3Env s) (ST s) (Maybe Integer))
 -> EvalAst s Integer)
-> (Bool
    -> Model s -> AST s -> ReaderT (Z3Env s) (ST s) (Maybe Integer))
-> Bool
-> EvalAst s Integer
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Bool
-> Model s -> AST s -> ReaderT (Z3Env s) (ST s) (Maybe Integer)
forall s. Bool -> EvalAst s Integer
T.evalBv

-- | Evaluate a collection of (AST s) nodes in the given model.
evalT :: (Traversable t) => t (AST s) -> Eval s (t (AST s))
evalT :: t (AST s) -> Eval s (t (AST s))
evalT = (Model s
 -> t (AST s) -> ReaderT (Z3Env s) (ST s) (Maybe (t (AST s))))
-> t (AST s) -> Eval s (t (AST s))
forall r a (m :: * -> *) b.
(r -> a -> m (Maybe b)) -> a -> ReaderT r (MaybeT m) b
liftEval Model s
-> t (AST s) -> ReaderT (Z3Env s) (ST s) (Maybe (t (AST s)))
forall (t :: * -> *) s.
Traversable t =>
Model s -> t (AST s) -> Z3 s (Maybe (t (AST s)))
T.evalT

-- | Run a evaluation function on a 'Traversable' data structure of '(AST s)'s
-- (e.g. @[(AST s)]@, @Vector (AST s)@, @Maybe (AST s)@, etc).
--
-- This a generic version of 'evalT' which can be used in combination with
-- other helpers. For instance, @mapEval evalInt@ can be used to obtain
-- the 'Integer' interpretation of a list of '(AST s)' of sort /int/.
mapEval :: (Traversable t) => EvalAst s a -> t (AST s) -> Eval s (t a)
mapEval :: EvalAst s a -> t (AST s) -> Eval s (t a)
mapEval f :: EvalAst s a
f = EvalAst s a -> t (AST s) -> Eval s (t a)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse EvalAst s a
f

-- | Get function as a list of argument/value pairs.
evalFunc :: FuncDecl s -> Eval s FuncModel
evalFunc :: FuncDecl s -> Eval s FuncModel
evalFunc = (Model s
 -> FuncDecl s -> ReaderT (Z3Env s) (ST s) (Maybe FuncModel))
-> FuncDecl s -> Eval s FuncModel
forall r a (m :: * -> *) b.
(r -> a -> m (Maybe b)) -> a -> ReaderT r (MaybeT m) b
liftEval Model s -> FuncDecl s -> ReaderT (Z3Env s) (ST s) (Maybe FuncModel)
forall s. Model s -> FuncDecl s -> Z3 s (Maybe FuncModel)
T.evalFunc

liftEval :: (r -> a -> m (Maybe b)) -> a -> ReaderT r (MaybeT m) b
liftEval :: (r -> a -> m (Maybe b)) -> a -> ReaderT r (MaybeT m) b
liftEval e :: r -> a -> m (Maybe b)
e a :: a
a = (r -> MaybeT m b) -> ReaderT r (MaybeT m) b
forall r (m :: * -> *) a. (r -> m a) -> ReaderT r m a
ReaderT ((r -> MaybeT m b) -> ReaderT r (MaybeT m) b)
-> (r -> MaybeT m b) -> ReaderT r (MaybeT m) b
forall a b. (a -> b) -> a -> b
$ m (Maybe b) -> MaybeT m b
forall (m :: * -> *) a. m (Maybe a) -> MaybeT m a
MaybeT (m (Maybe b) -> MaybeT m b)
-> (r -> m (Maybe b)) -> r -> MaybeT m b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (r -> a -> m (Maybe b)) -> a -> r -> m (Maybe b)
forall a b c. (a -> b -> c) -> b -> a -> c
flip r -> a -> m (Maybe b)
e a
a

runWithModel :: Eval s a -> Z3 s (Maybe a)
runWithModel :: Eval s a -> Z3 s (Maybe a)
runWithModel (ReaderT f :: Model s -> MaybeT (Z3 s) a
f) = MaybeT (Z3 s) a -> Z3 s (Maybe a)
forall (m :: * -> *) a. MaybeT m a -> m (Maybe a)
runMaybeT (MaybeT (Z3 s) a -> Z3 s (Maybe a))
-> MaybeT (Z3 s) a -> Z3 s (Maybe a)
forall a b. (a -> b) -> a -> b
$ ReaderT (Z3Env s) (ST s) (Maybe (Model s))
-> MaybeT (Z3 s) (Model s)
forall (m :: * -> *) a. m (Maybe a) -> MaybeT m a
MaybeT ((Result, Maybe (Model s)) -> Maybe (Model s)
forall a b. (a, b) -> b
snd ((Result, Maybe (Model s)) -> Maybe (Model s))
-> ReaderT (Z3Env s) (ST s) (Result, Maybe (Model s))
-> ReaderT (Z3Env s) (ST s) (Maybe (Model s))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ReaderT (Z3Env s) (ST s) (Result, Maybe (Model s))
forall s. Z3 s (Result, Maybe (Model s))
T.getModel) MaybeT (Z3 s) (Model s)
-> (Model s -> MaybeT (Z3 s) a) -> MaybeT (Z3 s) a
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Model s -> MaybeT (Z3 s) a
f