-- |
-- Type classes and built-in implementation for primitive Haskell types
-- 

module Ntha.Z3.Class (
    -- ** Types whose values are encodable to Z3 internal AST
    Z3Encoded(..),
    -- ** Types representable as Z3 Sort
    -- XXX: Unsound now
    -- XXX: Too flexible, can be used to encode Type ADT
    Z3Sorted(..),
    -- ** Type proxy helper, used with Z3Sorted
    Z3Sort(..),
    -- ** Types with reserved value for Z3 encoding use
    -- XXX: Magic value for built-in types
    Z3Reserved(..),
    -- ** Monad which can be instantiated into a concrete context
    SMT(..)
) where

import Ntha.Z3.Logic
import Z3.Monad

import Control.Monad.Except

import qualified Data.Map as M
import qualified Data.Set as S

data Z3Sort a = Z3Sort

class Z3Encoded a where
    encode :: SMT m e => a -> m e AST

-- | XXX: Unsound
class Z3Sorted a where
    -- | Map a value to Sort, the value should be a type-level thing
    sort :: SMT m e => a -> m e Sort
    sort _ = sortPhantom (Z3Sort :: Z3Sort a)

    -- | Map a Haskell type to Sort
    sortPhantom :: SMT m e => Z3Sort a -> m e Sort
    sortPhantom _ = smtError "sort error"

class Z3Encoded a => Z3Reserved a where
    def :: a

class (MonadError String (m e), MonadZ3 (m e)) => SMT m e where
    -- | Globally unique id
    genFreshId :: m e Int

    -- | Given data type declarations, extra field, and the SMT monad, return the fallible result in IO monad
    runSMT :: Z3Sorted ty => [(String, [(String, [(String, ty)])])] -> e -> m e a -> IO (Either String a)

    -- | Binding a variable String name to two things: an de Brujin idx as Z3 AST generated by mkBound and binder's Sort
    bindQualified :: String -> AST -> Sort -> m e ()

    -- | Get the above AST
    -- FIXME: The context management need extra -- we need to make sure that old binding wouldn't be destoryed
    -- XXX: We shouldn't expose a Map here. A fallible query interface is better
    getQualifierCtx :: m e (M.Map String (AST, Sort))

    -- | Get the preprocessed datatype context, a map from ADT's type name to its Z3 Sort
    -- XXX: We shouldn't expose a Map here. A fallible query interface is better
    getDataTypeCtx :: m e (M.Map String Sort)

    -- | Get extra
    getExtra :: m e e

    -- | Set extra
    modifyExtra :: (e -> e) -> m e ()

    -- | User don't have to import throwError
    smtError :: String -> m e a
    smtError = throwError

instance Z3Reserved Int where
    def = -1 -- XXX: Magic number

instance Z3Sorted Int where
    sortPhantom _ = mkIntSort

instance Z3Encoded Int where
    encode i = mkIntSort >>= mkInt i

instance Z3Reserved Double where
    def = -1.0 -- XXX: Magic number

instance Z3Sorted Double where
    sortPhantom _ = mkRealSort

instance Z3Encoded Double where
    encode = mkRealNum

instance Z3Reserved Bool where
    def = False -- XXX: Magic number

instance Z3Sorted Bool where
    sortPhantom _ = mkBoolSort

instance Z3Encoded Bool where
    encode = mkBool

-- The basic idea:
-- For each (k, v), assert in Z3 that if we select k from array we will get
-- the same value v
-- HACK: to set a default value for rest fields (or else we always get the last asserted value
--       as default, which is certainly not complying to finite map's definition), thus the
--       user should guarantee that he/she will never never think this value as a vaid one,
--       if not, he/she might get "a valid value mapped to a invalid key" semantics
instance (Z3Sorted k, Z3Encoded k, Z3Sorted v, Z3Reserved v) => Z3Encoded (M.Map k v) where
    encode m = do
        fid <- genFreshId
        arrSort <- sort m
        arr <- mkFreshConst ("map" ++ "_" ++ show fid) arrSort
        mapM_ (\(k, v) -> do
            kast <- encode k
            vast <- encode v
            sel <- mkSelect arr kast
            mkEq sel vast >>= assert) (M.toList m)
        arrValueDef <- mkArrayDefault arr
        vdef <- encode (def :: v)
        mkEq arrValueDef vdef >>= assert
        return arr

instance (Z3Sorted k, Z3Sorted v) => Z3Sorted (M.Map k v) where
    sortPhantom _ = do
        sk <- sortPhantom  (Z3Sort :: Z3Sort k)
        sv <- sortPhantom  (Z3Sort :: Z3Sort v)
        mkArraySort sk sv

-- Basic idea:
-- Set v =def= Map v {0, 1}
-- Thank god, this is much more sound
instance (Z3Sorted v, Z3Encoded v) => Z3Encoded (S.Set v) where
    encode s = do
        setSort <- sort s
        fid <- genFreshId
        arr <- mkFreshConst ("set" ++ "_" ++ show fid) setSort
        mapM_ (\e -> do
            ast <- encode e
            sel <- mkSelect arr ast
            one <- (mkIntSort >>= mkInt 1)
            mkEq sel one >>= assert) (S.toList s)
        arrValueDef <- mkArrayDefault arr
        zero <- (mkIntSort >>= mkInt 0)
        mkEq zero arrValueDef >>= assert
        return arr

instance Z3Sorted v => Z3Sorted (S.Set v) where
    sortPhantom _ = do
        sortElem <- sortPhantom (Z3Sort :: Z3Sort v)
        intSort <- mkIntSort
        mkArraySort sortElem intSort

instance (Z3Sorted t, Z3Sorted ty, Z3Encoded a) => Z3Encoded (Pred t ty a) where
    encode PTrue = mkTrue
    encode PFalse = mkFalse
    encode (PConj p1 p2) = do
        a1 <- encode p1
        a2 <- encode p2
        mkAnd [a1, a2]

    encode (PDisj p1 p2) = do
        a1 <- encode p1
        a2 <- encode p2
        mkOr [a1, a2]

    encode (PXor p1 p2) = do
        a1 <- encode p1
        a2 <- encode p2
        mkXor a1 a2

    encode (PNeg p) = encode p >>= mkNot

    encode (PForAll x ty p) = do
        sym <- mkStringSymbol x
        xsort <- sort ty
        -- "0" is de brujin idx for current binder
        -- it is passed to Z3 which returns an intenal (idx :: AST)
        -- This (idx :: AST) will be used to replace the variable
        -- in the abstraction body when encountered, thus it is stored
        -- in context by bindQualified we provide
        -- XXX: we should save and restore qualifier context here
        idx <- mkBound 0 xsort
        local $ do
            bindQualified x idx xsort
            body <- encode p
            -- The first [] is [Pattern], which is not really useful here
            mkForall [] [sym] [xsort] body

    encode (PExists x ty p) = do
        sym <- mkStringSymbol x
        xsort <- sort ty
        idx <- mkBound 0 xsort
        local $ do
            bindQualified x idx xsort
            a <- encode p
            mkExists [] [sym] [xsort] a

    -- HACK
    encode (PExists2 x y ty p) = do
        sym1 <- mkStringSymbol x
        sym2 <- mkStringSymbol y
        xsort <- sort ty
        idx1 <- mkBound 0 xsort
        idx2 <- mkBound 1 xsort
        local $ do
            bindQualified x idx1 xsort
            bindQualified y idx2 xsort
            a <- encode p
            mkExists [] [sym1, sym2] [xsort, xsort] a

    encode (PImpli p1 p2) = do
        a1 <- encode p1
        a2 <- encode p2
        mkImplies a1 a2

    encode (PIff p1 p2) = do
        a1 <- encode p1
        a2 <- encode p2
        mkIff a1 a2

    encode (PAssert a) = encode a