----------------------------------------------------------------------------
-- |
-- Module      :  Data.SBV.SMT.SMTLib2
-- Copyright   :  (c) Levent Erkok
-- License     :  BSD3
-- Maintainer  :  erkokl@gmail.com
-- Stability   :  experimental
-- Portability :  portable
--
-- Conversion of symbolic programs to SMTLib format, Using v2 of the standard
-----------------------------------------------------------------------------
{-# LANGUAGE PatternGuards #-}

module Data.SBV.SMT.SMTLib2(cvt, addNonEqConstraints) where

import Data.Bits (bit)
import qualified Data.Foldable as F (toList)
import qualified Data.Map      as M
import qualified Data.IntMap   as IM
import Data.List (intercalate, partition)
import Numeric (showHex)

import Data.SBV.BitVectors.Data

addNonEqConstraints :: [(Quantifier, NamedSymVar)] -> [[(String, CW)]] -> SMTLibPgm -> Maybe String
addNonEqConstraints qinps allNonEqConstraints (SMTLibPgm _ (aliasTable, pre, post))
  | null allNonEqConstraints
  = Just $ intercalate "\n" $ pre ++ post
  | null refutedModel
  = Nothing
  | True
  = Just $ intercalate "\n" $ pre
    ++ [ "; --- refuted-models ---" ]
    ++ concatMap nonEqs (map (map intName) nonEqConstraints)
    ++ post
 where refutedModel = concatMap nonEqs (map (map intName) nonEqConstraints)
       intName (s, c)
          | Just sw <- s `lookup` aliasTable = (show sw, c)
          | True                             = (s, c)
       -- with existentials, we only add top-level existentials to the refuted-models list
       nonEqConstraints = filter (not . null) $ map (filter (\(s, _) -> s `elem` topUnivs)) allNonEqConstraints
       topUnivs = [s | (_, (_, s)) <- takeWhile (\p -> fst p == EX) qinps]

nonEqs :: [(String, CW)] -> [String]
nonEqs []     =  []
nonEqs [sc]   =  ["(assert " ++ nonEq sc ++ ")"]
nonEqs (sc:r) =  ["(assert (or " ++ nonEq sc]
              ++ map (("            " ++) . nonEq) r
              ++ ["        ))"]

nonEq :: (String, CW) -> String
nonEq (s, c) = "(not (= " ++ s ++ " " ++ cvtCW c ++ "))"

tbd :: String -> a
tbd e = error $ "SBV.SMTLib2: Not-yet-supported: " ++ e

cvt :: Bool                                        -- ^ has infinite precision values
    -> Bool                                        -- ^ is this a sat problem?
    -> [String]                                    -- ^ extra comments to place on top
    -> [(Quantifier, NamedSymVar)]                 -- ^ inputs
    -> [Either SW (SW, [SW])]                      -- ^ skolemized version inputs
    -> [(SW, CW)]                                  -- ^ constants
    -> [((Int, (Bool, Size), (Bool, Size)), [SW])] -- ^ auto-generated tables
    -> [(Int, ArrayInfo)]                          -- ^ user specified arrays
    -> [(String, SBVType)]                         -- ^ uninterpreted functions/constants
    -> [(String, [String])]                        -- ^ user given axioms
    -> Pgm                                         -- ^ assignments
    -> [SW]                                        -- ^ extra constraints
    -> SW                                          -- ^ output variable
    -> ([String], [String])
cvt hasInf isSat comments _inps skolemInps consts tbls arrs uis axs asgnsSeq cstrs out = (pre, [])
  where -- the logic is an over-approaximation
        logic
          | hasInf = ["; Has unbounded Integers; no logic specified."]   -- combination, let the solver pick
          | True   = ["(set-logic " ++ qs ++ as ++ ufs ++ "BV)"]
          where qs  | null foralls && null axs = "QF_"  -- axioms are likely to contain quantifiers
                    | True                     = ""
                as  | null arrs                = ""
                    | True                     = "A"
                ufs | null uis && null tbls    = ""     -- we represent tables as UFs
                    | True                     = "UF"
        pre  =  ["; Automatically generated by SBV. Do not edit."]
             ++ map ("; " ++) comments
             ++ logic
             ++ [ "(set-option :produce-models true)"
                , "; --- literal constants ---"
                ]
             ++ map declConst consts
             ++ [ "; --- skolem constants ---" ]
             ++ [ "(declare-fun " ++ show s ++ " " ++ swFunType ss s ++ ")" | Right (s, ss) <- skolemInps]
             ++ [ "; --- constant tables ---" ]
             ++ concatMap constTable constTables
             ++ [ "; --- skolemized tables ---" ]
             ++ map (skolemTable (unwords (map swType foralls))) skolemTables
             ++ [ "; --- arrays ---" ]
             ++ concat arrayConstants
             ++ [ "; --- uninterpreted constants ---" ]
             ++ concatMap declUI uis
             ++ [ "; --- user given axioms ---" ]
             ++ map declAx axs
             ++ [ "; --- formula ---" ]
             ++ [if null foralls
                 then "(assert ; no quantifiers"
                 else "(assert (forall (" ++ intercalate "\n                 "
                                             ["(" ++ show s ++ " " ++ swType s ++ ")" | s <- foralls] ++ ")"]
             ++ map (letAlign . mkLet) asgns
             ++ map letAlign (if null delayedEqualities then [] else ("(and " ++ deH) : map (align 5) deTs)
             ++ [ impAlign (letAlign assertOut) ++ replicate noOfCloseParens ')' ]
        noOfCloseParens = length asgns + (if null foralls then 1 else 2) + (if null delayedEqualities then 0 else 1)
        (constTables, skolemTables) = ([(t, d) | (t, Left d) <- allTables], [(t, d) | (t, Right d) <- allTables])
        allTables = [(t, genTableData skolemMap (not (null foralls), forallArgs) (map fst consts) t) | t <- tbls]
        (arrayConstants, allArrayDelayeds) = unzip $ map (declArray (not (null foralls)) (map fst consts) skolemMap) arrs
        delayedEqualities@(~(deH:deTs)) = concatMap snd skolemTables ++ concat allArrayDelayeds
        foralls = [s | Left s <- skolemInps]
        forallArgs = concatMap ((" " ++) . show) foralls
        letAlign s
          | null foralls = "   " ++ s
          | True         = "            " ++ s
        impAlign s
          | null delayedEqualities = s
          | True                   = "     " ++ s
        align n s = replicate n ' ' ++ s
        -- if sat,   we assert cstrs /\ out
        -- if prove, we assert ~(cstrs => out) = cstrs /\ not out
        assertOut
           | null cstrs = o
           | True       = "(and " ++ unwords (map mkConj cstrs ++ [o]) ++ ")"
           where mkConj x = "(= " ++ cvtSW skolemMap x ++ " #b1)"
                 o | isSat =            mkConj out
                   | True  = "(not " ++ mkConj out ++ ")"
        skolemMap = M.fromList [(s, ss) | Right (s, ss) <- skolemInps, not (null ss)]
        tableMap  = IM.fromList $ map mkConstTable constTables ++ map mkSkTable skolemTables
          where mkConstTable (((t, _, _), _), _) = (t, "table" ++ show t)
                mkSkTable    (((t, _, _), _), _) = (t, "table" ++ show t ++ forallArgs)
        asgns = F.toList asgnsSeq
        mkLet (s, e) = "(let ((" ++ show s ++ " " ++ cvtExp skolemMap tableMap e ++ "))"
        declConst (s, c) = "(define-fun " ++ show s ++ " " ++ swFunType [] s ++ " " ++ cvtCW c ++ ")"

declUI :: (String, SBVType) -> [String]
declUI (i, t) = ["(declare-fun uninterpreted_" ++ i ++ " " ++ cvtType t ++ ")"]

-- NB. We perform no check to as to whether the axiom is meaningful in any way.
declAx :: (String, [String]) -> String
declAx (nm, ls) = (";; -- user given axiom: " ++ nm ++ "\n   ") ++ intercalate "\n" ls

constTable :: (((Int, (Bool, Size), (Bool, Size)), [SW]), [String]) -> [String]
constTable (((i, (_, atSz), (_, rtSz)), _elts), is) = decl : map wrap is
  where t       = "table" ++ show i
        decl    = "(declare-fun " ++ t ++ " (" ++ smtType atSz ++ ") " ++ smtType rtSz ++ ")"
        wrap  s = "(assert " ++ s ++ ")"

skolemTable :: String -> (((Int, (Bool, Size), (Bool, Size)), [SW]), [String]) -> String
skolemTable qsIn (((i, (_, atSz), (_, rtSz)), _elts), _) = decl
  where qs   = if null qsIn then "" else qsIn ++ " "
        t    = "table" ++ show i
        decl = "(declare-fun " ++ t ++ " (" ++ qs ++ smtType atSz ++ ") " ++ smtType rtSz ++ ")"

-- Left if all constants, Right if otherwise
genTableData :: SkolemMap -> (Bool, String) -> [SW] -> ((Int, (Bool, Size), (Bool, Size)), [SW]) -> Either [String] [String]
genTableData skolemMap (_quantified, args) consts ((i, (sa, at), (_, _rt)), elts)
  | null post = Left  (map (topLevel . snd) pre)
  | True      = Right (map (nested   . snd) (pre ++ post))
  where ssw = cvtSW skolemMap
        (pre, post) = partition fst (zipWith mkElt elts [(0::Int)..])
        t           = "table" ++ show i
        mkElt x k   = (isReady, (idx, ssw x))
          where idx = cvtCW (mkConstCW (sa, at) k)
                isReady = x `elem` consts
        topLevel (idx, v) = "(= (" ++ t ++ " " ++ idx ++ ") " ++ v ++ ")"
        nested   (idx, v) = "(= (" ++ t ++ args ++ " " ++ idx ++ ") " ++ v ++ ")"

-- TODO: We currently do not support non-constant arrays when quantifiers are present, as
-- we might have to skolemize those. Implement this properly.
-- The difficulty is with the ArrayReset/Mutate/Merge: We have to postpone an init if
-- the components are themselves postponed, so this cannot be implemented as a simple map.
declArray :: Bool -> [SW] -> SkolemMap -> (Int, ArrayInfo) -> ([String], [String])
declArray quantified consts skolemMap (i, (_, ((_, atSz), (_, rtSz)), ctx)) = (adecl : map wrap pre, map snd post)
  where topLevel = not quantified || case ctx of
                                       ArrayFree Nothing -> True
                                       ArrayFree (Just sw) -> sw `elem` consts
                                       ArrayReset _ sw     -> sw `elem` consts
                                       ArrayMutate _ a b   -> all (`elem` consts) [a, b]
                                       ArrayMerge c _ _    -> c `elem` consts
        (pre, post) = partition fst ctxInfo
        nm = "array_" ++ show i
        ssw sw
         | topLevel || sw `elem` consts
         = cvtSW skolemMap sw
         | True
         = tbd "Non-constant array initializer in a quantified context"
        adecl = "(declare-fun " ++ nm ++ "() (Array " ++ smtType atSz ++ " " ++ smtType rtSz ++ "))"
        ctxInfo = case ctx of
                    ArrayFree Nothing   -> []
                    ArrayFree (Just sw) -> declA sw
                    ArrayReset _ sw     -> declA sw
                    ArrayMutate j a b -> [(all (`elem` consts) [a, b], "(= " ++ nm ++ " (store array_" ++ show j ++ " " ++ ssw a ++ " " ++ ssw b ++ "))")]
                    ArrayMerge  t j k -> [(t `elem` consts,            "(= " ++ nm ++ " (ite (= #b1 " ++ ssw t ++ ") array_" ++ show j ++ " array_" ++ show k ++ "))")]
        declA sw = let iv = nm ++ "_freeInitializer"
                   in [ (True,             "(declare-fun " ++ iv ++ "() " ++ smtType atSz ++ ")")
                      , (sw `elem` consts, "(= (select " ++ nm ++ " " ++ iv ++ ") " ++ ssw sw ++ ")")
                      ]
        wrap (False, s) = s
        wrap (True, s)  = "(assert " ++ s ++ ")"

swType :: SW -> String
swType s = smtType (sizeOf s)

swFunType :: [SW] -> SW -> String
swFunType ss s = "(" ++ unwords (map swType ss) ++ ") " ++ swType s

smtType :: Size -> String
smtType (Size Nothing)   = "Int"
smtType (Size (Just sz)) = "(_ BitVec " ++ show sz ++ ")"

cvtType :: SBVType -> String
cvtType (SBVType []) = error "SBV.SMT.SMTLib2.cvtType: internal: received an empty type!"
cvtType (SBVType xs) = "(" ++ unwords (map smtType body) ++ ") " ++ smtType ret
  where szs         = map snd xs
        (body, ret) = (init szs, last szs)

type SkolemMap = M.Map  SW [SW]
type TableMap  = IM.IntMap String

cvtSW :: SkolemMap -> SW -> String
cvtSW skolemMap s
  | Just ss <- s `M.lookup` skolemMap
  = "(" ++ show s ++ concatMap ((" " ++) . show) ss ++ ")"
  | True
  = show s

-- NB. The following works with SMTLib2 since all sizes are multiples of 4 (or just 1, which is specially handled)
hex :: Int -> Integer -> String
hex 1  v = "#b" ++ show v
hex sz v = "#x" ++ pad (sz `div` 4) (showHex v "")
  where pad n s = replicate (n - length s) '0' ++ s

cvtCW :: CW -> String
cvtCW x | isInfPrec x     = if w >= 0 then show w else "(- " ++ show (abs w) ++ ")"
  where w = cwVal x
cvtCW x | not (hasSign x) = hex (intSizeOf x) (cwVal x)
-- signed numbers (with 2's complement representation) is problematic
-- since there's no way to put a bvneg over a positive number to get minBound..
-- Hence, we punt and use binary notation in that particular case
cvtCW x | cwVal x == least = mkMinBound (intSizeOf x)
  where least = negate (2 ^ intSizeOf x)
cvtCW x = negIf (w < 0) $ hex (intSizeOf x) (abs w)
  where w = cwVal x

negIf :: Bool -> String -> String
negIf True  a = "(bvneg " ++ a ++ ")"
negIf False a = a

-- anamoly at the 2's complement min value! Have to use binary notation here
-- as there is no positive value we can provide to make the bvneg work.. (see above)
mkMinBound :: Int -> String
mkMinBound i = "#b1" ++ replicate (i-1) '0'

getTable :: TableMap -> Int -> String
getTable m i
  | Just tn <- i `IM.lookup` m = tn
  | True                       = error $ "SBV.SMTLib2: Cannot locate table " ++ show i

unbounded :: SBVExpr -> a
unbounded expr = error $ "SBV.SMTLib2: Unsupported operation on unbounded integers: " ++ show expr

cvtExp :: SkolemMap -> TableMap -> SBVExpr -> String
cvtExp skolemMap tableMap expr@(SBVApp _ arguments) = sh expr
  where ssw = cvtSW skolemMap
        hasInfPrecArgs = any isInfPrec arguments
        ensureBV       = not hasInfPrecArgs || unbounded expr
        lift2  o _ [x, y] = "(" ++ o ++ " " ++ x ++ " " ++ y ++ ")"
        lift2  o _ sbvs   = error $ "SBV.SMTLib2.sh.lift2: Unexpected arguments: "   ++ show (o, sbvs)
        lift2B oU oS sgn sbvs = "(ite " ++ lift2S oU oS sgn sbvs ++ " #b1 #b0)"
        lift2S oU oS sgn sbvs
          | sgn
          = lift2 oS sgn sbvs
          | True
          = lift2 oU sgn sbvs
        lift2N o sgn sbvs = "(bvnot " ++ lift2 o sgn sbvs ++ ")"
        lift1  o _ [x]    = "(" ++ o ++ " " ++ x ++ ")"
        lift1  o _ sbvs   = error $ "SBV.SMT.SMTLib2.sh.lift1: Unexpected arguments: "   ++ show (o, sbvs)
        sh (SBVApp Ite [a, b, c]) = "(ite (= #b1 " ++ ssw a ++ ") " ++ ssw b ++ " " ++ ssw c ++ ")"
        sh (SBVApp (LkUp (t, (_, atSz), _, l) i e) [])
          | needsCheck = "(ite " ++ cond ++ ssw e ++ " " ++ lkUp ++ ")"
          | True       = lkUp
          where needsCheck = maybe True (\at -> (2::Integer)^at > fromIntegral l) (unSize atSz)
                lkUp = "(" ++ getTable tableMap t ++ " " ++ ssw i ++ ")"
                cond
                 | hasSign i = "(or " ++ le0 ++ " " ++ gtl ++ ") "
                 | True      = gtl ++ " "
                (less, leq) = case atSz of
                                 Size Nothing -> ("<", "<=")
                                 _            -> if hasSign i then ("bvslt", "bvsle") else ("bvult", "bvule")
                mkCnst = cvtCW . mkConstCW (hasSign i, sizeOf i)
                le0  = "(" ++ less ++ " " ++ ssw i ++ " " ++ mkCnst 0 ++ ")"
                gtl  = "(" ++ leq  ++ " " ++ mkCnst l ++ " " ++ ssw i ++ ")"
        sh (SBVApp (ArrEq i j) []) = "(ite (= array_" ++ show i ++ " array_" ++ show j ++") #b1 #b0)"
        sh (SBVApp (ArrRead i) [a]) = "(select array_" ++ show i ++ " " ++ ssw a ++ ")"
        sh (SBVApp (Uninterpreted nm) [])   = "uninterpreted_" ++ nm
        sh (SBVApp (Uninterpreted nm) args) = "(uninterpreted_" ++ nm ++ " " ++ unwords (map ssw args) ++ ")"
        sh (SBVApp (Extract i j) [a]) | ensureBV = "((_ extract " ++ show i ++ " " ++ show j ++ ") " ++ ssw a ++ ")"
        sh (SBVApp (Rol i) [a])
           | not hasInfPrecArgs = rot  ssw "rotate_left"  i a
           | True               = sh (SBVApp (Shl i) [a])     -- Haskell treats rotateL as shiftL for unbounded values
        sh (SBVApp (Ror i) [a])
           | not hasInfPrecArgs = rot  ssw "rotate_right" i a
           | True               = sh (SBVApp (Shr i) [a])     -- Haskell treats rotateR as shiftR for unbounded values
        sh (SBVApp (Shl i) [a])
           | not hasInfPrecArgs = shft ssw "bvshl"  "bvshl"  i a
           | i < 0              = sh (SBVApp (Shr (-i)) [a])  -- flip sign/direction
           | True               = "(* " ++ ssw a ++ " " ++ show (bit i :: Integer) ++ ")"  -- Implement shiftL by multiplication by 2^i
        sh (SBVApp (Shr i) [a])
           | not hasInfPrecArgs = shft ssw "bvlshr" "bvashr" i a
           | i < 0              = sh (SBVApp (Shl (-i)) [a])  -- flip sign/direction
           | True               = "(div " ++ ssw a ++ " " ++ show (bit i :: Integer) ++ ")"  -- Implement shiftR by division by 2^i
        sh (SBVApp op args)
          | Just f <- lookup op smtBVOpTable, ensureBV
          = f (any hasSign args) (map ssw args)
          where -- The first 4 operators below do make sense for Integer's in Haskell, but there's
                -- no obvious counterpart for them in the SMTLib translation.
                -- TODO: provide support for these.
                smtBVOpTable = [ (And,  lift2 "bvand")
                               , (Or,   lift2 "bvor")
                               , (XOr,  lift2 "bvxor")
                               , (Not,  lift1 "bvnot")
                               , (Join, lift2 "concat")
                               ]
        sh inp@(SBVApp op args)
          | hasInfPrecArgs
          = case lookup op smtOpIntTable of
              Just f -> f True (map ssw args)
              _      -> unbounded inp
          | Just f <- lookup op smtOpBVTable
          = f (any hasSign args) (map ssw args)
          | True
          = error $ "SBV.SMT.SMTLib2.cvtExp.sh: impossible happened; can't translate: " ++ show inp
          where smtOpBVTable  = [ (Plus,          lift2   "bvadd")
                                , (Minus,         lift2   "bvsub")
                                , (Times,         lift2   "bvmul")
                                , (Quot,          lift2S  "bvudiv" "bvsdiv")
                                , (Rem,           lift2S  "bvurem" "bvsrem")
                                , (Equal,         lift2   "bvcomp")
                                , (NotEqual,      lift2N  "bvcomp")
                                , (LessThan,      lift2B  "bvult" "bvslt")
                                , (GreaterThan,   lift2B  "bvugt" "bvsgt")
                                , (LessEq,        lift2B  "bvule" "bvsle")
                                , (GreaterEq,     lift2B  "bvuge" "bvsge")
                                ]
                smtOpIntTable = [ (Plus,          lift2   "+")
                                , (Minus,         lift2   "-")
                                , (Times,         lift2   "*")
                                , (Quot,          lift2   "div")
                                , (Rem,           lift2   "mod")
                                , (Equal,         lift2B  "=" "=")
                                , (NotEqual,      lift2B  "distinct" "distinct")
                                , (LessThan,      lift2B  "<"  "<")
                                , (GreaterThan,   lift2B  ">"  ">")
                                , (LessEq,        lift2B  "<=" "<=")
                                , (GreaterEq,     lift2B  ">=" ">=")
                                ]

rot :: (SW -> String) -> String -> Int -> SW -> String
rot ssw o c x = "((_ " ++ o ++ " " ++ show c ++ ") " ++ ssw x ++ ")"

shft :: (SW -> String) -> String -> String -> Int -> SW -> String
shft ssw oW oS c x = "(" ++ o ++ " " ++ ssw x ++ " " ++ cvtCW c' ++ ")"
   where s  = hasSign x
         c' = mkConstCW (s, sizeOf x) c
         o  = if hasSign x then oS else oW