-- | Various properties and transformations of Boogie program elements
module Boogie.Util ( 
  -- * Types
  TypeBinding,
  typeSubst,
  renameTypeVars,
  fromTVNames,
  isFreeIn,
  isTypeVar,
  unifier,
  freshTVName,
  -- * Expressions
  freeVarsTwoState,
  freeVars,
  freeOldVars,
  VarBinding,
  exprSubst,
  paramSubst,
  mapRefs,
  refSelections,
  applications,
  removeBoundClashes,
  -- * Specs
  preconditions,
  postconditions,
  modifies,
  assumePreconditions,
  assumePostconditions,
  -- * Functions and procedures
  PSig (..),
  psigParams,
  psigArgTypes,
  psigRetTypes,
  psigModifies,
  psigRequires,
  psigEnsures,
  psigType,
  PDef (..),
  -- * Code generation
  num, eneg, enot,
  (|+|), (|-|), (|*|), (|/|), (|%|), (|=|), (|!=|), (|<|), (|<=|), (|>|), (|>=|), (|&|), (|||), (|=>|), (|<=>|),
  conjunction,
  disjunction,
  guardWith,
  assume,
  -- * Special names
  nullaryType,
  noneType,
  anyType,
  tupleType,
  functionConst,
  functionExpr,
  functionFromConst,
  removeClashesWith,
  -- * Misc
  fromRight,
  maybe2,
  deleteAll,
  restrictDomain,
  removeDomain,
  partitionDomain,
  mapItwType,
  allM,
  anyM,
  findM,
  changeState,
  withLocalState
) where

import Boogie.AST
import Boogie.Position
import Boogie.Tokens
import Pretty
import Boogie.PrettyAST
import Data.Maybe
import Data.List
import Data.Map (Map, (!))
import qualified Data.Map as M
import Data.Set (Set)
import qualified Data.Set as S
import Control.Applicative
import Control.Monad.State
import Control.Lens

{- Types -}

-- | Mapping from type variables to types
type TypeBinding = Map Id Type

-- | 'typeSubst' @binding t@ :
-- Substitute all free type variables in @t@ according to binding;
-- all variables in the domain of @bindings@ are considered free if not explicitly bound
typeSubst :: TypeBinding -> Type -> Type
typeSubst _ BoolType = BoolType
typeSubst _ IntType = IntType
typeSubst binding (IdType id []) = case M.lookup id binding of
  Just t -> t
  Nothing -> IdType id []
typeSubst binding (IdType id args) = IdType id (map (typeSubst binding) args)
typeSubst binding (MapType bv domains range) = MapType bv (map (typeSubst removeBound) domains) (typeSubst removeBound range)
  where removeBound = deleteAll bv binding
  
-- | 'renameTypeVars' @tv newTV binding@ : @binding@ with each occurrence of one of @tv@ replaced with corresponding @newTV@ 
-- (in both domain and range)
renameTypeVars :: [Id] -> [Id] -> TypeBinding -> TypeBinding
renameTypeVars tv newTV binding = let
    tvMap = M.fromList $ zip tv newTV 
    replace tv = M.findWithDefault tv tv tvMap
    tvToType = fromTVNames tv newTV
  in M.map (typeSubst tvToType) (M.mapKeys replace binding)
    
-- | @x@ `isFreeIn` @t@ : does @x@ occur free in @t@?
isFreeIn :: Id -> Type -> Bool
x `isFreeIn` (IdType y []) = x == y
x `isFreeIn` (IdType y args) = any (x `isFreeIn`) args
x `isFreeIn` (MapType bv domains range) = x `notElem` bv && any (x `isFreeIn`) (range:domains)
_ `isFreeIn` _ = False  
  
-- | 'fromTVNames' @tvs tvs'@ : type binding that replaces type variables @tvs@ with type variables @tvs'@
fromTVNames :: [Id] -> [Id] -> TypeBinding
fromTVNames tvs tvs' = M.fromList (zip tvs (map nullaryType tvs'))

-- | 'freshTVName @n@ : Fresh type variable with a unique identifier n
freshTVName n = nonIdChar : show n

-- | 'isTypeVar' @contextTypeVars v@ : Is @v@ either one of  @contextTypeVars@ or a freash type variable generated by 'freshTVName'?
isTypeVar :: [Id] -> Id -> Bool
isTypeVar contextTypeVars v = head v == nonIdChar || v `elem` contextTypeVars
    
-- | 'unifier' @fv xs ys@ : most general unifier of @xs@ and @ys@ with shared free type variables of the context @fv@
unifier :: [Id] -> [Type] -> [Type] -> Maybe TypeBinding
unifier _ [] [] = Just M.empty
unifier fv (IntType:xs) (IntType:ys) = unifier fv xs ys
unifier fv (BoolType:xs) (BoolType:ys) = unifier fv xs ys
unifier fv ((IdType id1 args1):xs) ((IdType id2 args2):ys) | id1 == id2 = unifier fv (args1 ++ xs) (args2 ++ ys)
unifier fv ((IdType id []):xs) (y:ys) | isTypeVar fv id = 
  if id `isFreeIn` y then Nothing 
  else M.insert id y <$> unifier fv (update xs) (update ys)
    where update = map (typeSubst (M.singleton id y))
unifier fv (x:xs) ((IdType id []):ys) | isTypeVar fv id = 
  if id `isFreeIn` x then Nothing 
  else M.insert id x <$> unifier fv (update xs) (update ys)
    where update = map (typeSubst (M.singleton id x))
unifier fv ((MapType bv1 domains1 range1):xs) ((MapType bv2 domains2 range2):ys) =
  case forallUnifier fv bv1 (range1:domains1) bv2 (range2:domains2) of
    Nothing -> Nothing
    Just u -> M.union u <$> (unifier fv (update u xs) (update u ys))
  where
    update u = map (typeSubst u)
unifier _ _ _ = Nothing
    
-- | 'forallUnifier' @fv bv1 xs bv2 ys@ :   
-- Most general unifier of @xs@ and @ys@,
-- where @bv1@ are universally quantified type variables in @xs@ and @bv2@ are universally quantified type variables in @ys@,
-- and @fv@ are free type variables of the enclosing context
forallUnifier :: [Id] -> [Id] -> [Type] -> [Id] -> [Type] -> Maybe TypeBinding
forallUnifier fv bv1 xs bv2 ys = if length bv1 /= length bv2 || length xs /= length ys 
  then Nothing
  else case unifier (fv ++ bv1) xs (map withFreshBV ys) of
    Nothing -> Nothing
    Just u -> let (boundU, freeU) = M.partitionWithKey (\k _ -> k `elem` bv1) u
      in if all isFreshBV (M.elems boundU) && not (any hasFreshBV (M.elems freeU))
        then Just freeU
        else Nothing
  where
    freshBV = bv2 `removeClashesWith` bv1
    withFreshBV = typeSubst (fromTVNames bv2 freshBV)
    -- does a type correspond to one of the renamed bound variables?
    isFreshBV (IdType id []) = id `elem` freshBV
    isFreshBV _ = False
    -- does type t contain any fresh bound variables of m2?
    hasFreshBV t = any (`isFreeIn` t) freshBV
  
{- Expressions -}

-- | Free variables in an expression, referred to in current state and old state
freeVarsTwoState :: Expression -> ([Id], [Id])
freeVarsTwoState e = case node e of
  Literal _ -> ([], [])
  Logical _ _ -> ([], [])
  Var x -> ([x], [])
  Application name args -> freeVarsTwoState $ attachPos (position e) $ MapSelection (functionExpr name (position e)) args
  MapSelection m args ->  over both (nub . concat) (unzip (map freeVarsTwoState (m : args)))
  MapUpdate m args val ->  over both (nub . concat) (unzip (map freeVarsTwoState (val : m : args)))
  Old e -> let (state, old) = freeVarsTwoState e in ([], state ++ old)
  IfExpr cond e1 e2 -> over both (nub . concat) (unzip [freeVarsTwoState cond, freeVarsTwoState e1, freeVarsTwoState e2])
  Coercion e _ -> freeVarsTwoState e
  UnaryExpression _ e -> freeVarsTwoState e
  BinaryExpression _ e1 e2 -> over both (nub . concat) (unzip [freeVarsTwoState e1, freeVarsTwoState e2])
  Quantified _ _ _ boundVars e -> let (state, old) = freeVarsTwoState e in (state \\ map fst boundVars, old)

-- | Free variables in an expression, in current state
freeVars = fst . freeVarsTwoState
-- | Free variables in an expression, in old state
freeOldVars = snd . freeVarsTwoState

-- | Mapping from variables to expressions
type VarBinding = Map Id Expression

-- | 'exprSubst' @binding e@ : substitute all free variables in @e@ according to @binding@;
-- all variables in the domain of @bindings@ are considered free if not explicitly bound
exprSubst :: VarBinding -> Expression -> Expression
exprSubst binding (Pos pos e) = attachPos pos $ exprSubst' binding e

exprSubst' binding (Var id) = case M.lookup id binding of
  Nothing -> Var id
  Just e -> (node e)
exprSubst' binding (Application id args) = Application id (map (exprSubst binding) args)
exprSubst' binding (MapSelection m args) = MapSelection (exprSubst binding m) (map (exprSubst binding) args)
exprSubst' binding (MapUpdate m args val) = MapUpdate (exprSubst binding m) (map (exprSubst binding) args) (exprSubst binding val)
exprSubst' binding (Old e) = Old (exprSubst binding e)
exprSubst' binding (IfExpr cond e1 e2) = IfExpr (exprSubst binding cond) (exprSubst binding e1) (exprSubst binding e2)
exprSubst' binding (Coercion e t) = Coercion (exprSubst binding e) t
exprSubst' binding (UnaryExpression op e) = UnaryExpression op (exprSubst binding e)
exprSubst' binding (BinaryExpression op e1 e2) = BinaryExpression op (exprSubst binding e1) (exprSubst binding e2)
exprSubst' binding (Quantified t qop tv boundVars e) = Quantified t qop tv boundVars (exprSubst binding' e)
  where binding' = deleteAll (map fst boundVars) binding
exprSubst' _ e = e

-- | 'paramBinding' @sig def@ :
-- Binding of parameter names from procedure signature @sig@ to their equivalents from procedure definition @def@
paramBinding :: PSig -> PDef -> VarBinding
paramBinding sig def = M.fromList $ zip (sigIns ++ sigOuts) (defIns ++ defOuts)
  where
    sigIns = map itwId $ psigArgs sig
    sigOuts = map itwId $ psigRets sig
    defIns = map (gen . Var) $ pdefIns def
    defOuts = map (gen . Var) $ pdefOuts def
  
-- | 'paramSubst' @sig def@ :
-- Substitute parameter names from @sig@ in an expression with their equivalents from @def@
paramSubst :: PSig -> PDef -> Expression -> Expression  
paramSubst sig def = if not (pdefParamsRenamed def) 
  then id 
  else exprSubst (paramBinding sig def)

-- | 'mapRefs' @expr@ : all map references that occur in @expr@
mapRefs :: Expression -> [Ref]
mapRefs expr = case node expr of
  Literal (Reference _ r) -> [r]
  Literal _ -> []
  Logical _ _ -> []
  Var x -> []
  Application name args -> nub . concat $ map mapRefs args
  MapSelection m args -> nub . concat $ map mapRefs (m : args)
  MapUpdate m args val ->  nub . concat $ map mapRefs (val : m : args)
  Old e -> internalError $ text "mapRefs should only be applied in single-state context"
  IfExpr cond e1 e2 -> nub . concat $ [mapRefs cond, mapRefs e1, mapRefs e2]
  Coercion e _ -> mapRefs e
  UnaryExpression _ e -> mapRefs e
  BinaryExpression _ e1 e2 -> nub . concat $ [mapRefs e1, mapRefs e2]
  Quantified _ _ _ _ e -> mapRefs e  
  
-- | 'refSelections' @expr@ : all map reference selections that occur in @expr@, where the map is a free variable
refSelections :: Expression -> [(Ref, [Expression])]
refSelections expr = case node expr of
  Literal _ -> []
  Logical _ _ -> []
  Var x -> []
  Application name args -> nub . concat $ map refSelections args
  MapSelection m args -> case node m of 
   Literal (Reference _ r) -> (r, args) : (nub . concat $ map refSelections args)
   _ -> nub . concat $ map refSelections (m : args)
  MapUpdate m args val ->  nub . concat $ map refSelections (val : m : args)  
  Old e -> internalError $ text "refSelections should only be applied in single-state context"
  IfExpr cond e1 e2 -> nub . concat $ [refSelections cond, refSelections e1, refSelections e2]
  Coercion e _ -> refSelections e
  UnaryExpression _ e -> refSelections e
  BinaryExpression _ e1 e2 -> nub . concat $ [refSelections e1, refSelections e2]
  Quantified _ _ _ boundVars e -> refSelections e
  
-- | 'applications' @expr@ : all function applications that occur in @expr@
applications :: Expression -> [(Id, [Expression])]
applications expr = case node expr of
  Literal _-> []
  Var x-> []
  Application name args-> (name, args) : (nub . concat $ map applications args)
  MapSelection m args-> nub . concat $ map applications (m : args)
  MapUpdate m args val->  nub . concat $ map applications (val : m : args)
  Old e-> internalError $ text "applications should only be applied in single-state context"
  IfExpr cond e1 e2-> nub . concat $ [applications cond, applications e1, applications e2]
  Coercion e _-> applications e
  UnaryExpression _ e-> applications e
  BinaryExpression _ e1 e2-> nub . concat $ [applications e1, applications e2]
  Quantified _ _ _ _ e-> applications e

-- | 'removeBoundClashes' @names expr@: replace any bound variables in @expr@ that clash with @names@
removeBoundClashes :: [Id] -> Expression -> Expression
removeBoundClashes names (Pos p expr) = attachPos p $ case expr of
  Literal _ -> expr
  Var x -> expr
  Application name args -> Application name (map go args)
  MapSelection m args -> MapSelection (go m) (map go args)
  MapUpdate m args val ->  MapUpdate (go m) (map go args) (go val)
  Old e -> Old (go e)
  IfExpr cond e1 e2 -> IfExpr (go cond) (go e1) (go e2)
  Coercion e t -> Coercion (go e) t
  UnaryExpression op e -> UnaryExpression op (go e)
  BinaryExpression op e1 e2 -> BinaryExpression op (go e1) (go e2)
  Quantified t op tv bv e -> let
      vars = map fst bv
      types = map snd bv
      vars' = removeClashesWith vars names
      e' = exprSubst (M.fromList $ zip vars (map (gen . Var) vars')) e
    in Quantified t op tv (zip vars' types) (removeBoundClashes (names ++ vars') e')
  where
    go = removeBoundClashes names
  
{- Specs -}

-- | 'preconditions' @specs@ : all precondition clauses in @specs@  
preconditions :: [Contract] -> [SpecClause]
preconditions specs = catMaybes (map extractPre specs)
  where 
    extractPre (Requires f e) = Just (SpecClause Precondition f e)
    extractPre _ = Nothing

-- | 'postconditions' @specs@ : all postcondition clauses in @specs@     
postconditions :: [Contract] -> [SpecClause]
postconditions specs = catMaybes (map extractPost specs)
  where 
    extractPost (Ensures f e) = Just (SpecClause Postcondition f e)
    extractPost _ = Nothing
   
-- | 'modifies' @specs@ : all modifies clauses in @specs@   
modifies :: [Contract] -> [Id]
modifies specs = (nub . concat . catMaybes) (map extractMod specs)
  where
    extractMod (Modifies _ ids) = Just ids
    extractMod _ = Nothing
  
-- | Make all preconditions in contracts free  
assumePreconditions :: PSig -> PSig
assumePreconditions sig = sig { psigContracts = map assumePrecondition (psigContracts sig) }
  where
    assumePrecondition (Requires _ e) = Requires True e
    assumePrecondition c = c
    
-- | Make all postconditions in contracts free  
assumePostconditions :: PSig -> PSig
assumePostconditions sig = sig { psigContracts = map assumePostcondition (psigContracts sig) }
  where
    assumePostcondition (Ensures _ e) = Ensures True e
    assumePostcondition c = c

{- Procedures -}
     
-- | Procedure signature 
data PSig = PSig {
    psigName :: Id,               -- ^ Procedure name
    psigTypeVars :: [Id],         -- ^ Type variables
    psigArgs :: [IdTypeWhere],    -- ^ In-parameters
    psigRets :: [IdTypeWhere],    -- ^ Out-parameters
    psigContracts :: [Contract]   -- ^ Contracts
  }
            
instance Show PSig where
  show = show . psigName
  
instance Eq PSig where
  s1 == s2 = psigName s1 == psigName s2  
  
-- | All parameters of a procedure signature 
psigParams sig = psigArgs sig ++ psigRets sig
-- | Types of in-parameters of a procedure signature
psigArgTypes = (map itwType) . psigArgs
-- | Types of out-parameters of a procedure signature
psigRetTypes = (map itwType) . psigRets
-- | Procedure signature as a map type
psigType sig = MapType (psigTypeVars sig) (psigArgTypes sig) (tupleType $ psigRetTypes sig) 
-- | Modifies clauses of a procedure signature
psigModifies = modifies . psigContracts
-- | Preconditions of a procedure signature
psigRequires = preconditions . psigContracts
-- | Postconditions of a procedure signature
psigEnsures = postconditions . psigContracts
  
-- | Procedure definition;
-- a single procedure might have multiple definitions (one per body)
data PDef = PDef { 
    pdefIns :: [Id],                  -- ^ In-parameter names (in the same order as 'psigArgs' in the corresponding signature)
    pdefOuts :: [Id],                 -- ^ Out-parameter names (in the same order as 'psigRets' in the corresponding signature)
    pdefParamsRenamed :: Bool,        -- ^ Are any parameter names in this definition different for the procedure signature? (used for optimizing parameter renaming, True is a safe default)
    pdefBody :: BasicBody,            -- ^ Body
    pdefPos :: SourcePos              -- ^ Location of the (first line of the) procedure definition in the source
  }

{- Code generation -}

num i = gen $ Literal (IntValue i)
eneg e = inheritPos (UnaryExpression Neg) e
enot e = inheritPos (UnaryExpression Not) e
e1 |+|    e2 = inheritPos2 (BinaryExpression Plus) e1 e2
e1 |-|    e2 = inheritPos2 (BinaryExpression Minus) e1 e2
e1 |*|    e2 = inheritPos2 (BinaryExpression Times) e1 e2
e1 |/|    e2 = inheritPos2 (BinaryExpression Div) e1 e2
e1 |%|    e2 = inheritPos2 (BinaryExpression Mod) e1 e2
e1 |=|    e2 = inheritPos2 (BinaryExpression Eq) e1 e2
e1 |!=|   e2 = inheritPos2 (BinaryExpression Neq) e1 e2
e1 |<|    e2 = inheritPos2 (BinaryExpression Ls) e1 e2
e1 |<=|   e2 = inheritPos2 (BinaryExpression Leq) e1 e2
e1 |>|    e2 = inheritPos2 (BinaryExpression Gt) e1 e2
e1 |>=|   e2 = inheritPos2 (BinaryExpression Geq) e1 e2
e1 |&|    e2 = if node e1 == tt 
                  then e2
                  else if node e2 == tt
                    then e1
                    else inheritPos2 (BinaryExpression And) e1 e2
e1 |||    e2 = if node e1 == ff 
                  then e2
                  else if node e2 == ff
                    then e1
                    else inheritPos2 (BinaryExpression Or) e1 e2
e1 |=>|   e2 = if node e1 == tt 
                  then e2
                  else inheritPos2 (BinaryExpression Implies) e1 e2
e1 |<=>|  e2 = inheritPos2 (BinaryExpression Equiv) e1 e2
assume e = attachPos (position e) (Predicate [] (SpecClause Inline True e))

conjunction es = foldl (|&|) (gen tt) es
disjunction es = foldl (|||) (gen ff) es
guardWith gs e = conjunction gs |=>| e

{- Special names -}

-- | 'nullaryType' @id@ : type denoted by @id@ without arguments
nullaryType id = IdType id []

-- | Dummy type used during type checking to denote error
noneType = nullaryType ("NONE" ++ [nonIdChar])

-- | Dummy type used when the type does not matter
anyType = nullaryType ("ANY" ++ [nonIdChar])

-- | Dummy type used to represent procedure returns as a single type
tupleType = IdType ("TUPLE" ++ [nonIdChar])

functionFrefix = "function "

-- | 'functionConst' @name@ : name of a map constant that corresponds function @name@
-- (must be distinct from all global names)
functionConst name = functionFrefix ++ name

functionExpr name pos = attachPos pos . Var $ functionConst name

-- | 'functionFromConst' @name@ : reverse of 'functionConst', returns Nothing if @name@ does not have the right form
functionFromConst name = let (prefix, suffix) = splitAt (length functionFrefix) name
  in if prefix == functionFrefix
      then Just suffix
      else Nothing
      
-- | 'removeClashesWith' @names names'@ :
-- A version of @names@ that is disjoint from @names'@
-- (if @names@ does not have duplicates, then result also does not have duplicates)
removeClashesWith :: [Id] -> [Id] -> [Id]
removeClashesWith names names' = map changeName names
  where
    -- new name for tv that does not coincide with any names'
    changeName tv = if tv `elem` names' then tv ++ replicate (level + 1) nonIdChar else tv
    -- maximum number of nonIdChar characters at the end of any names or names'; 
    -- by appending (level + 1) nonIdChar charactes to tv we make is different from all names' and unchanged names
    level = maximum [fromJust (findIndex (\c -> c /= nonIdChar) (reverse id)) | id <- names ++ names']      
  
{- Misc -}

-- | Extract the element out of a 'Right' and throw an error if its argument is 'Left'
fromRight :: Either a b -> b
fromRight (Right x) = x

-- | Two-argument equivalent of 'maybe'
maybe2 :: c -> (a -> b -> c) -> (Maybe a) -> (Maybe b) -> c  
maybe2 def f mx my = maybe def (maybe (const def) f mx) my  

-- | 'deleteAll' @keys m@ : map @m@ with @keys@ removed from its domain
deleteAll :: Ord k => [k] -> Map k a -> Map k a
deleteAll keys m = foldr M.delete m keys

-- | 'restrictDomain' @keys m@ : map @m@ restricted on the set of keys @keys@
restrictDomain :: Ord k => Set k -> Map k a -> Map k a
restrictDomain keys m = fst $ partitionDomain keys m

-- | 'removeDomain' @keys m@ : map @m@ with the set of keys @keys@ removed from its domain
removeDomain :: Ord k => Set k -> Map k a -> Map k a
removeDomain keys m = snd $ partitionDomain keys m

-- | 'partitionDomain' @keys m@ : map @m@ partitioned into two maps, restricted to @keys@ and the rest
partitionDomain :: Ord k => Set k -> Map k a -> (Map k a, Map k a)
partitionDomain keys m = M.partitionWithKey (\k _ -> k `S.member` keys) m

mapItwType f (IdTypeWhere i t w) = IdTypeWhere i (f t) w

-- | Monadic version of 'any'
anyM :: (Functor m, Monad m) => (a -> m Bool) -> [a] -> m Bool
anyM pred xs = isJust <$> findM pred xs
  
-- | Monadic version of 'all'
allM :: (Functor m, Monad m) => (a -> m Bool) -> [a] -> m Bool
allM pred xs = isNothing <$> findM (\x -> not <$> pred x) xs
  
-- | Monadic version of 'find' (finds the first element in a list for which a computation evaluates to True) 
findM :: (Functor m, Monad m) => (a -> m Bool) -> [a] -> m (Maybe a)
findM _ [] = return Nothing
findM pred (x : xs) = do
  res <- pred x
  if res then return (Just x) else findM pred xs  

-- | Execute a computation with state of type @t@ inside a computation with state of type @s@
changeState :: Monad m => (s -> t) -> (t -> s -> s) -> StateT t m a -> StateT s m a
changeState getter modifier e = do
  st <- gets getter
  (res, st') <- lift $ runStateT e st
  modify $ modifier st'
  return res  

-- | 'withLocalState' @localState e@ :
-- Execute @e@ in current state modified by @localState@, and then restore current state
withLocalState :: Monad m => (s -> t) -> StateT t m a -> StateT s m a
withLocalState localState e = changeState localState (flip const) e

