------------------------------------------------------------------------------
--- This module provides an abstract representation of a variant
--- of the SMT-LIB language appropriate for checking Curry programs.
--- In particular, polymorphic function declarations are supported.
---
--- It might be later replaced by an extended version of the
--- SMT-LIB language specified in the package `smtlib`.
---
--- @author  Michael Hanus
--- @version March 2019
------------------------------------------------------------------------------

module ESMT where

import List ( (\\), intercalate, isPrefixOf, union )

import Data.FiniteMap
import Text.Pretty

------------------------------------------------------------------------------
--- An SMT-LIB script consists of a list of commands.
data SMTLib = SMTLib [Command]
  deriving (Eq, Show)

type SVar  = Int  -- variables 
type Ident = String -- identifiers (e.g., functions)


--- Sorts

data Sort = SComb Ident [Sort]
  deriving (Eq, Show)

-- Shows a sort as a string.
showSort :: Sort -> String
showSort (SComb s ss) = s ++ intercalate "_" (map showSort ss)

--- Does the sort represent a type parameter (`TVar i`)?
isTypeParameter :: Sort -> Bool
isTypeParameter (SComb s ss) = null ss && "TVar" `isPrefixOf` s && length s > 4

----------------------------------------------------------------------------
--- Terms

--- A literal is a natural number, float, or string.
data TLiteral = TInt    Int
              | TFloat  Float
              | TString String
 deriving (Eq, Show)

--- An identifier possibly with a sort attached.
data QIdent = Id Ident
            | As Ident Sort
 deriving (Eq, Show)

--- The identifier of a possibly sorted identifier.
qidName :: QIdent -> Ident
qidName (Id n  ) = n
qidName (As n _) = n

--- Sorted variables used in quantifiers.
data SortedVar = SV SVar Sort
 deriving (Eq, Show)

--data Pattern = PComb SVar [SVar]
--  deriving (Eq, Show)

--- Terms occurring in formulas
data Term = TConst TLiteral
          | TSVar  SVar
          | TComb  QIdent [Term]
          | Let    [(SVar, Term)] Term
          | Forall [SortedVar] Term
          | Exists [SortedVar] Term
          --| Match  Term [(Pattern, Term)]
          --| Annot  Term [Attribute]
 deriving (Eq, Show)

-- Smart constructors:

--- Combined term with string identifier.
tComb :: Ident -> [Term] -> Term
tComb f ts = TComb (Id f) ts

--- Conjunction
tConj :: [Term] -> Term
tConj = tComb "and"

--- Disjunction
tDisj :: [Term] -> Term
tDisj = tComb "or"

--- Negation
tNot :: Term -> Term
tNot t = tComb "not" [t]

--- A Boolean true.
tTrue :: Term
tTrue = tComb "true" []

--- A Boolean false.
tFalse :: Term
tFalse = tComb "false" []

--- Equality between two terms.
tEqu :: Term -> Term -> Term
tEqu t1 t2 = tComb "=" [t1, t2]

--- Equality between a variable and a term.
tEquVar :: SVar -> Term -> Term
tEquVar v t = tEqu (TSVar v) t

--- A constant with a sort declaration.
sortedConst :: Ident -> Sort -> Term
sortedConst c s = TComb (As c s) []

----------------------------------------------------------------------------
--- Datatype declaration consisting of type variables and constructor decls.
data DTDecl = DT [Ident] [DTCons] -- polymorphic type
 deriving (Eq, Show)

data DTCons = DCons Ident [(Ident,Sort)]
 deriving (Eq, Show)

--- The signature of a function declaration.
data FunSig = FunSig Ident [Sort] Sort
  deriving (Eq, Show)

--- The signature and arguments of a function declaration.
data FunDec = FunDec Ident [SortedVar] Sort
  deriving (Eq, Show)

--- Commands in SMT script.
--- The command `DefineSigsRec` is new. It is similar to `DefineFunsRec`,
--- but the associated term is the axiomatization of the function
--- definition. Thus, it contains all quantifiers and the equation
--- `(= (f x1...xn) rhs-term)` so that it can also be used to specify,
--- by exploiting disjunctions, the meaning of non-deterministic functions.
--- Moreover, it supports parametric polymoprhism by providing a list
--- of type paramters which can be used in the type signature and term.
--- Such polymorphic declarations are replaced by type-instantiated
--- monomorphic definitions before printing an SMT script.
data Command = Assert              Term
             | CheckSat
             | Comment             String
             | DeclareVar          SortedVar
             | DeclareDatatypes    [(Ident, Int, DTDecl)] -- name/arity/decls
             | DeclareFun          Ident [Sort] Sort
             | DeclareSort         Ident Int
             | DefineFunsRec       [(FunDec, Term)]
             | DefineSigsRec       [([Ident], FunSig, Term)]
             | EmptyLine
  deriving (Eq, Show)

-- Smart constructors:

--- Assert a simplified formula.
sAssert :: Term -> Command
sAssert = Assert . simpTerm

--- Examples:
{-
listType =
  DeclareDatatypes 
    [("List",1,
      DT ["T"]
         [ DCons "nil" [],
           DCons "cons" [("car", SComb "T" []),
                         ("cdr", SComb "List" [SComb "T" []])]])]

maybeType =
  DeclareDatatypes 
    [("Maybe",1,
      DT ["T"]
         [ DCons "Nothing" [],
           DCons "Just" [("just", SComb "T" [])]])]

pairType =
  DeclareDatatypes 
    [("Pair",2,
      DT ["X", "Y"]
         [ DCons "nil" [],
           DCons "mk-pair" [("first",  SComb "X" []),
                            ("second", SComb "Y" [])]])]
-}
























---------------------------------------------------------------------------
-- All possibly sorted identifiers occurring in a SMT term.
allQIdsOfTerm :: Term -> [QIdent]
allQIdsOfTerm (TConst _)     = []
allQIdsOfTerm (TSVar _)      = []
allQIdsOfTerm (TComb f args) = foldr union [f] (map allQIdsOfTerm args)
allQIdsOfTerm (Forall _ arg) = allQIdsOfTerm arg
allQIdsOfTerm (Exists _ arg) = allQIdsOfTerm arg
allQIdsOfTerm (Let bs e)     =
  foldr union [] (map allQIdsOfTerm (e : map snd bs))

-- TODO: should be extended to all commands but currently sufficient
allQIdsOfAsserts :: [Command] -> [QIdent]
allQIdsOfAsserts = foldr union [] . map allQIdsOfAssert
 where allQIdsOfAssert cmd = case cmd of Assert t -> allQIdsOfTerm t
                                         _        -> []

---------------------------------------------------------------------------
--- All type parameters occurring in a sort.
typeParamsOfSort :: Sort -> [Ident]
typeParamsOfSort s@(SComb sn ss) =
  if isTypeParameter s then [sn]
                       else foldr union [] (map typeParamsOfSort ss)

--- All type parameters contained in a term.
typeParamsOfTerm :: Term -> [Ident]
typeParamsOfTerm (TConst _) = []
typeParamsOfTerm (TSVar  _) = []
typeParamsOfTerm (TComb f args) = foldr union (typeParamsOfQId f)
                                        (map typeParamsOfTerm args)
typeParamsOfTerm (Forall svs arg) =
  foldr union (typeParamsOfTerm arg) (map typeParamsOfSV svs)
typeParamsOfTerm (Exists svs arg) =
  foldr union (typeParamsOfTerm arg) (map typeParamsOfSV svs)
typeParamsOfTerm (Let bs e)     =
  foldr union [] (map typeParamsOfTerm (e : map snd bs))

typeParamsOfQId :: QIdent -> [Ident]
typeParamsOfQId (Id _  ) = []
typeParamsOfQId (As _ s) = typeParamsOfSort s

typeParamsOfSV :: SortedVar -> [Ident]
typeParamsOfSV (SV _ s) = typeParamsOfSort s

typeParamsOfFunSig :: FunSig -> [Ident]
typeParamsOfFunSig (FunSig _ ss s) =
  foldr union [] (map typeParamsOfSort (ss++[s]))

--- A type paramter substitution.
type TPSubst = FM Ident Sort

--- The empty substitution
emptyTPSubst :: TPSubst
emptyTPSubst = emptyFM (<)

----------------------------------------------------------------------------
--- Compute sort matching, i.e., if `matchSort t1 t2 = s`, then `t2 = s(t1)`.
matchSort :: Sort -> Sort -> Maybe TPSubst
matchSort s1@(SComb sn1 ss1) s2@(SComb sn2 ss2)
 | isTypeParameter s1
 = Just $ if s1 == s2 then emptyTPSubst
                      else addToFM emptyTPSubst (head (typeParamsOfSort s1)) s2
 | otherwise
 = if sn1 == sn2 then matchSorts ss1 ss2 else Nothing

matchSorts :: [Sort] -> [Sort] -> Maybe TPSubst
matchSorts []       []       = Just emptyTPSubst
matchSorts []       (_:_)    = Nothing
matchSorts (_:_)    []       = Nothing
matchSorts (t1:ts1) (t2:ts2) = do
  s <- matchSort t1 t2
  t <- matchSorts (map (substSort s) ts1)(map (substSort s) ts2)
  return (plusFM s t)

--- Applies a sort substitution to a term.
substSort :: TPSubst -> Sort -> Sort
substSort sub (SComb sn ss) =
  maybe (SComb sn (map (substSort sub) ss)) id (lookupFM sub sn)

--- All type parameters contained in a term.
substTerm :: TPSubst -> Term -> Term
substTerm sub term = case term of
  TConst _ -> term
  TSVar  _ -> term
  TComb f args -> TComb (substQId sub f) (map (substTerm sub) args)
  Forall svs arg -> Forall (map (substSV sub) svs) (substTerm sub arg)
  Exists svs arg -> Forall (map (substSV sub) svs) (substTerm sub arg)
  Let bs e -> Let (map (\ (v,s) -> (v, substTerm sub s)) bs) (substTerm sub e)

substQId :: TPSubst -> QIdent -> QIdent
substQId _ qid@(Id _) = qid
substQId sub (As n s) = As n (substSort sub s)

substSV :: TPSubst -> SortedVar -> SortedVar
substSV sub (SV v s) = SV v (substSort sub s)

substFunSig :: TPSubst -> FunSig -> FunSig
substFunSig sub (FunSig fn ss s) =
  FunSig fn (map (substSort sub) ss) (substSort sub s)

substDefSig :: TPSubst -> ([Ident], FunSig, Term) -> ([Ident], FunSig, Term)
substDefSig tsub (ps, fsig, term) =
  (ps \\ keysFM tsub, substFunSig tsub fsig, substTerm tsub term)

--------------------------------------------------------------------------
-- Rename identifiers.

rnmTerm :: (Ident -> Ident) -> Term -> Term
rnmTerm rnm term = case term of
  TConst _ -> term
  TSVar  _ -> term
  TComb f args -> TComb (rnmQId rnm f) (map (rnmTerm rnm) args)
  Forall svs arg -> Forall svs (rnmTerm rnm arg)
  Exists svs arg -> Forall svs (rnmTerm rnm arg)
  Let bs e -> Let (map (\ (v,s) -> (v, rnmTerm rnm s)) bs) (rnmTerm rnm e)

rnmQId :: (Ident -> Ident) -> QIdent -> QIdent
rnmQId rnm (Id n)   = Id (rnm n)
rnmQId rnm (As n s) = As (rnm n) s

rnmFunSig :: (Ident -> Ident) -> FunSig -> FunSig
rnmFunSig rnm (FunSig fn ss s) = FunSig (rnm fn) ss s

rnmDefSig :: (Ident -> Ident) -> ([Ident],FunSig,Term) -> ([Ident],FunSig,Term)
rnmDefSig rnm (ps, fsig, term) =
  (ps, rnmFunSig rnm fsig, rnmTerm rnm term)

--------------------------------------------------------------------------
-- A simplifier for terms:
simpTerm :: Term -> Term
simpTerm (TConst l) = TConst l
simpTerm (TSVar  v) = TSVar v
simpTerm (Let bs t) = if null bs then t'
                                 else Let bs' t'
 where bs' = map (\ (v,tm) -> (v, simpTerm tm)) bs
       t'  = simpTerm t
simpTerm (Forall vs t) = if null vs then t' else Forall vs t'
 where t' = simpTerm t
simpTerm (Exists vs t) = if null vs then t' else Exists vs t'
 where t' = simpTerm t
simpTerm (TComb f ts)
 | f == Id "apply" && not (null ts')
 = case head ts' of TComb s' ts0 -> TComb s' (ts0 ++ tail ts')
                    _            -> fts
 | f == Id "not"
 = case ts' of [TComb s' [ts0]] -> if s' == f then ts0 else fts
               _                -> fts
 | f == Id "and"
 = case filter (/= tTrue) ts' of
          []  -> tTrue
          cjs -> if tFalse `elem` cjs
                   then tFalse
                   else TComb f (concatMap joinSame cjs)
 | f == Id "or"
 = case filter (/= tFalse) ts' of
          []  -> tFalse
          djs -> if tTrue `elem` djs
                   then tTrue
                   else TComb f (concatMap joinSame djs)
 | otherwise = fts
 where
  ts' = map simpTerm ts
  fts = TComb f ts'

  joinSame arg = case arg of TComb f' args | f==f' -> args
                             _                     -> [arg]

--------------------------------------------------------------------------
-- Remove As-identifiers if they are functions (for better readability):
reduceAsInTerm :: Term -> Term
reduceAsInTerm (TConst l) = TConst l
reduceAsInTerm (TSVar  v) = TSVar v
reduceAsInTerm (Let bs t) = Let (map (\ (v,tm) -> (v, reduceAsInTerm tm)) bs)
                                (reduceAsInTerm t)
reduceAsInTerm (Forall vs t) = Forall vs (reduceAsInTerm t)
reduceAsInTerm (Exists vs t) = Exists vs (reduceAsInTerm t)
reduceAsInTerm (TComb f ts) = TComb (simpAs f) (map reduceAsInTerm ts)
 where
  simpAs qid = case qid of As n (SComb s _) | s == "Func" -> Id n
                           _ -> qid

--------------------------------------------------------------------------
-- Remove parametric polymorphism (provided by `DefineSigsRec`)
-- in an SMT script.
unpoly :: [Command] -> [Command]
unpoly cmds =
   let allsigs = map sigNameSort (allSigs cmds)
       allqids = filter ((`elem` (map fst allsigs)) . qidName)
                         (allQIdsOfAsserts cmds)
       newcmds = map (addInstsCmd allqids) cmds
   in map (unpolyCmd allsigs) newcmds
 where
  addInstsCmd qids cmd = case cmd of
    DefineSigsRec fts -> DefineSigsRec $ concatMap (addSigInstances qids) fts
    _ -> cmd

  unpolyCmd sigs cmd = case cmd of
    DefineSigsRec fts -> DefineSigsRec $ map unpolySig fts
    Assert term -> Assert (rnmQIdWithTInstTerm sigs term)
    _ -> cmd

  unpolySig (ps, sig, term) =
    let sub = addListToFM emptyTPSubst (map (\p -> (p,SComb "TVar" [])) ps)
    in ([], substFunSig sub sig, substTerm sub term)

-- Add to a given (polymorphic) define-sig element all its type instances
-- required by qualified identifiers occurring in the first argument.
addSigInstances :: [QIdent] -> ([Ident], FunSig, Term)
                -> [([Ident], FunSig, Term)]
addSigInstances qids fts@(ps, (FunSig fn ss rs), _) =
  fts : if null ps then [] -- nothing to be added for monomorphic types
                   else concatMap addSigInst qids
 where
  addSigInst (Id _) = []
  addSigInst (As n s) =
    if n == fn
      then maybe [] (\tsub -> [(rnmDefSig (toTInstName fn ps tsub)
                                          (substDefSig tsub fts))])
                 (matchSort (sigTypeAsSort ss rs) s)
      else []

-- Rename a sorted name w.r.t. its type instance of the polymorphic function.
rnmQIdWithTInst :: [(Ident, ([Ident],Sort))] -> QIdent -> QIdent
rnmQIdWithTInst _ (Id n) = Id n
rnmQIdWithTInst sigs qid@(As n s) =
  maybe qid
        (\ (ps,psort) -> maybe qid
                               (\tsub -> As (addTInstName ps tsub n) s)
                               (matchSort psort s))
        (lookup n sigs)

rnmQIdWithTInstTerm :: [(Ident, ([Ident],Sort))] -> Term -> Term
rnmQIdWithTInstTerm sigs term = case term of
  TConst _ -> term
  TSVar  _ -> term
  TComb f args -> TComb (rnmQIdWithTInst sigs f)
                        (map (rnmQIdWithTInstTerm sigs) args)
  Forall svs arg -> Forall svs (rnmQIdWithTInstTerm sigs arg)
  Exists svs arg -> Forall svs (rnmQIdWithTInstTerm sigs arg)
  Let bs e -> Let (map (\ (v,s) -> (v, rnmQIdWithTInstTerm sigs s)) bs)
                  (rnmQIdWithTInstTerm sigs e)

-- Renaming operation which changes the name of a given (polymorphic)
-- function w.r.t. a list of type parameters and a substitution.
toTInstName :: Ident -> [Ident] -> TPSubst -> Ident -> Ident
toTInstName fn ps tsub n | fn == n   = addTInstName ps tsub n
                         | otherwise = n

-- Add a sort index to a name of a (polymorphic) function w.r.t.
-- a list of type parameters and a substitution.
addTInstName :: [Ident] -> TPSubst -> Ident -> Ident
addTInstName ps tsub n =
  n ++ concatMap (\p -> maybe p (('_':) . showSort) (lookupFM tsub p)) ps

allSigs :: [Command] -> [([Ident], FunSig, Term)]
allSigs = concatMap sigOfCmd
 where sigOfCmd cmd = case cmd of DefineSigsRec fts -> fts
                                  _                 -> []

-- The name of a signature.
nameOfSig :: ([Ident], FunSig, Term) -> Ident
nameOfSig (_, FunSig n _ _, _) = n

-- The name and sort of a signature.
sigNameSort :: ([Ident], FunSig, Term) -> (Ident, ([Ident],Sort))
sigNameSort (ps, FunSig n ss s, _) = (n, (ps, sigTypeAsSort ss s))

sigTypeAsSort :: [Sort] -> Sort -> Sort
sigTypeAsSort [] s = s
sigTypeAsSort (t:ts) s = SComb "Func" [t, sigTypeAsSort ts s]

--------------------------------------------------------------------------
-- Pretty printing:

--- Show an SMT-LIB script with a newline
showSMT :: [Command] -> String
showSMT cmds = pPrint (pretty (SMTLib (unpoly cmds))) ++ "\n"

instance Pretty SMTLib where
  pretty (SMTLib cmds) = vsep (map pretty cmds)

instance Pretty Sort where
  pretty (SComb i ss) = parensIf (not $ null ss) $
                          text i <+> (hsep (map pretty ss))

instance Pretty TLiteral where
  pretty (TInt    n) = int n
  pretty (TFloat  f) = float f
  pretty (TString s) = text s

instance Pretty QIdent where
  pretty (Id i  ) = text i
  pretty (As i s) = parent [text "as", text i, pretty s]

instance Pretty SortedVar where
  pretty (SV v s) = parent [prettyVar v, pretty s]


instance Pretty Term where
  pretty (TConst c) = pretty c
  pretty (TSVar  v) = prettyVar v
  pretty (TComb  qi ts) = parensIf (not $ null ts) $
                           pretty qi <+> (hsep (map pretty ts))
  pretty (Let     bs t) = parent [text "let", parent (map ppBind bs), pretty t]
    where ppBind (v, t') = parent [prettyVar v, pretty t']
  pretty (Forall svs t) = parent [ text "forall"
                                 , parent (map pretty svs)
                                 , pretty t
                                 ]
  pretty (Exists svs t) = parent [ text "exists"
                                 , parent (map pretty svs)
                                 , pretty t
                                 ]

instance Pretty DTDecl where
  pretty (DT tys cs) = if null tys
                         then parent (map pretty cs)
                         else parent [ text "par"
                                     , parent (map text tys)
                                     , parent (map pretty cs)
                                     ]

instance Pretty DTCons where
  pretty (DCons sym sels) = parent [text sym, (hsep (map prettySel sels))]
   where
    prettySel (n,s) = parent [text n, pretty s]

instance Pretty FunSig where
  pretty (FunSig fn ss s) = parent (ppCmd (DeclareFun fn ss s))

instance Pretty FunDec where
  pretty (FunDec fn svs s) = parent [ text fn, parent (map pretty svs)
                                    , pretty s
                                    ]

instance Pretty Command where
  pretty cmd = case cmd of
    Comment cmt -> semi <+> text cmt
    EmptyLine   -> text ""
    DefineSigsRec fts -> vsep $ map (pretty . (\ (_,t,_) -> t)) fts ++
                                map ppSigBody fts
    _           -> parent $ ppCmd cmd

ppSigBody :: ([Ident],FunSig,Term) -> Doc
ppSigBody (_, FunSig fn _ _, term) = vsep $ map pretty
  [ EmptyLine, Comment $ "Axiomatization of function '" ++ fn ++ "'"
  , sAssert term ]

--- Pretty printing of SMT-LIB commands.
ppCmd :: Command -> [Doc]
ppCmd cmd = case cmd of
  Assert   t -> [text "assert", pretty (reduceAsInTerm t)] -- for nicer printing
  CheckSat   -> [text "check-sat"]
  DeclareVar (SV v s)  -> [text "declare-const", prettyVar v, pretty s]
  DeclareDatatypes sds -> -- use old SMT syntax:
    if length sds /= 1
      then error "Datatype declaration with more than one type!"
      else let (tc, _, DT tvs cs) = head sds
           in [ text "declare-datatypes"
              , parent (map text tvs)
              , parent [parent (text tc : map pretty cs)]
              ]
  --DeclareDatatypes sds -> let (ss, ars, ds) = unzip3 sds in
  --                        [ text "declare-datatypes"
  --                        , parent (map (\ (s,a) -> parent [text s, int a])
  --                                      (zip ss ars))
  --                        , parent (map pretty ds)
  --                        ]
  DeclareFun fn ss s -> [ text "declare-fun"
                        , text fn
                        , parent (map pretty ss)
                        , pretty s
                        ]
  DefineFunsRec fts -> let (fs, ts) = unzip fts in
                       [ text "define-funs-rec"
                       , parent (map pretty fs)
                       , parent (map pretty ts)
                       ]
  DeclareSort sym n -> [text "declare-sort", text sym, int n]
  _          -> error $ "ppCmd: command '" ++ show cmd ++ "' not reachable!"

prettyVar :: SVar -> Doc
prettyVar v = text ('x' : show v)

--- Pretty print the given documents separated with spaces and parenthesized
parent :: [Doc] -> Doc
parent = encloseSep lparen rparen space