{-  CAO Compiler
    Copyright (C) 2014 Cryptography and Information Security Group, HASLab - INESC TEC and Universidade do Minho

    This program is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program.  If not, see <http://www.gnu.org/licenses/>. -}

{-# LANGUAGE BangPatterns               #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FlexibleContexts #-}

{-
Module      :  $Header$
Description :  Compiler global state.
Copyright   :  (C) 2014 Cryptography and Information Security Group, HASLab - INESC TEC and Universidade do Minho
License     :  GPL

Maintainer  :  Paulo Silva <paufil@di.uminho.pt>
Stability   :  experimental
Portability :  non-portable

Compiler global state.
-}

module Language.CAO.Common.State
    ( CaoState

    , getLastVar

    , tmpFromPool

    , fileName

    , withTcST
    , keepScope
    , keepState
    , keepGScope
    , withSrcLoc
    , setSrcLoc
    , getSrcLoc
    , getHeap
    , putHeap
    , updateHeap
    , addHypothesis
    , getHypothesis
    , getsTcS
    , getMode
    , setMode
    , setYices
    , getYices
    , newUniq
    , nextTyVarId

    , reTypVar -- Simplify
    , lookupReTypVar -- Simplify
    , withSimplifyST
    , resetSimplifyST

    , allConsts -- Target
    , getConst -- Target
    , withTargetST
    , resetTargetST

    , putFieldProj -- PreC
    , putFunType -- PreC
    , getFunType -- PreC
    , putRetStruct -- PreC
    , getRetStruct -- PreC
    , setRefVar -- PreC
    , getRefVar -- PreC
    , lookupFieldProj -- PreC
    , withPreCST
    , resetPreCST

    , getAllocVars -- C
    , getAllAllocVars -- C
    , storeAllocVar -- C
    , withCST 
    , resetCST
    , allocScope

    , initialState -- Monad, Main/CAO
    , lastVar -- SSA, Monad, MonadState CaoState m => 
    , getFileName -- Monad, Parser, C
    , setFileName -- Monad, Interpreter, MonadState CaoState m => , Main/CAO
    , getTmpVars -- Monad, Target, C, PreC
    , storeTmpVar -- Indist, Target, C, PreC
 ) where

import Control.Monad.State.Strict

import qualified Data.List as List
import Data.IntMap (IntMap)
import qualified Data.IntMap as IMap
import Data.Map(Map)
import qualified Data.Map as Map
import Data.Maybe (fromJust)

import Language.CAO.Common.Literal
import Language.CAO.Common.Outputable
import Language.CAO.Common.SrcLoc
import Language.CAO.Common.Representation
import Language.CAO.Common.Utils
import Language.CAO.Common.Var

import Language.CAO.Index
import Language.CAO.Index.Eval

import Language.CAO.Type

import Language.CAO.Typechecker.Heap

import Main.Flags (RunMode(..))

data CaoState = CaoState 
    { fileName   :: !String
    , lastVar    :: {-# UNPACK #-} !Int -- Last generated variable identifier
    , mode       :: !RunMode
    , yicesSMT   :: Maybe FilePath

    , tmpVars    :: [Var]
    , tcst       :: (Maybe TcS)
    , simplifyST :: (Maybe SimplifyST)
    , targetST   :: (Maybe TargetST)
    , precST     :: (Maybe PreCST)
    , cst        :: (Maybe CST)
    }

initialState :: CaoState
initialState = CaoState 
    { fileName   = "<Unknown file name>"
    , lastVar    = iNIT_VAR_ID
    , mode       = CAO
    , yicesSMT   = Nothing
    , tmpVars    = []
    , tcst       = Nothing
    , simplifyST = Nothing
    , targetST   = Nothing
    , precST     = Nothing
    , cst        = Nothing
 }

setFileName :: MonadState CaoState m => String -> m ()
setFileName fn = modify $! \s -> s { fileName = fn }

getFileName :: MonadState CaoState m => m String
getFileName = gets fileName

getLastVar :: MonadState CaoState m => m Int
getLastVar = do
    s <- get
    let i = lastVar s
    put $! s { lastVar = i + 1 }
    return i

setMode :: MonadState CaoState m => RunMode -> m ()
setMode m = modify $! \ st -> st { mode = m }

getMode :: MonadState CaoState m => m RunMode
getMode = gets mode

setYices :: MonadState CaoState m => Maybe FilePath -> m ()
setYices y = modify $  \st -> st { yicesSMT = y }

getYices :: MonadState CaoState m => m (Maybe FilePath)
getYices = gets yicesSMT

--------------------------------------------------------------------------------

storeTmpVar :: MonadState CaoState m => Var -> m ()
storeTmpVar lvar = modify (\ s -> s { tmpVars = lvar : tmpVars s } )

tmpFromPool :: MonadState CaoState m => Type Var -> m (Maybe Var)
tmpFromPool typ = do
    s <- get
    let pool = tmpVars s
        (v, pool') = findAndDelete ((== typ) . varType) pool
    put $ s { tmpVars = pool' }
    return v

getTmpVars :: MonadState CaoState m => m [Var]
getTmpVars = gets tmpVars

--------------------------------------------------------------------------------
data TcS
    = TcS { curLoc    :: !SrcLoc
          , heap      :: {-# UNPACK #-} !Heap
          }

emptyTcS :: TcS
emptyTcS 
    = TcS { curLoc    = unkSrcLoc
          , heap      = emptyHeap
          }

withTcST :: MonadState CaoState m => m a -> m a
withTcST m = do
    modify ( \ s -> s { tcst = Just emptyTcS } )
    r <- m
    modify ( \ s -> s { tcst = Nothing } )
    return r

getsTcS :: MonadState CaoState m => (TcS -> a) -> m a
getsTcS f = liftM (f . fromJust) (gets tcst)

modifyTcS :: MonadState CaoState m => (TcS -> TcS) -> m ()
modifyTcS f = modify (\ st -> st { tcst = fmap f (tcst st) } )

withSrcLoc :: MonadState CaoState m => SrcLoc -> m a -> m a
withSrcLoc loc m = do
    st0 <- get
    let tc0 = fromJust $ tcst st0
        oldLoc = curLoc tc0
    put $! st0 { tcst = Just $ tc0 { curLoc = loc } }
    r <- m
    st1 <- get
    let tc1 = fromJust $ tcst st1
    put $! st1 { tcst = Just $ tc1 { curLoc = oldLoc } }
    return r

{-# INLINE getSrcLoc #-}
getSrcLoc :: MonadState CaoState m => m SrcLoc
getSrcLoc = getsTcS curLoc

setSrcLoc :: MonadState CaoState m => SrcLoc -> m ()
setSrcLoc loc = modifyTcS (\ st -> st { curLoc = loc })

{-# INLINE nextTyVarId #-}
nextTyVarId :: MonadState CaoState m => m TyVarId
nextTyVarId = getLastVar

{-# INLINE newUniq #-}
newUniq :: MonadState CaoState m => m Int
newUniq = getLastVar

keepScope :: MonadState CaoState m => m a -> m a
keepScope m = do
    h   <- getHeap
    r   <- m
    putHeap h
    return r

keepState :: MonadState CaoState m => m a -> m a
keepState m = do
    st <- get
    r  <- m
    put $! st
    return r

keepGScope :: MonadState CaoState m => m a -> m a
keepGScope m = do
    h   <- getHeap
    r   <- m
    updateHeap ( replaceGlobalHeap h )
    return r

--------------------------------------------------------------------------------
-- Heap
--------------------------------------------------------------------------------

{-# INLINE getHeap #-}
getHeap :: MonadState CaoState m => m Heap
getHeap = getsTcS heap

putHeap :: MonadState CaoState m => Heap -> m ()
putHeap h = modifyTcS (\ st -> st { heap = h })

updateHeap :: MonadState CaoState m => (Heap -> Heap) -> m ()
updateHeap f = modifyTcS ( \ st -> st { heap = f (heap st) } )

--------------------------------------------------------------------------------
-- XXX: What if introducing false hypothesis?
-- XXX: consider removing evalCond. Add it as pre-condition?
addHypothesis :: MonadState CaoState m => [ICond Var] -> m ()
addHypothesis i = modifyTcS (\ st -> st { heap = addHyp (heap st) (map evalCond i) } )

getHypothesis :: MonadState CaoState m => m [ICond Var]
getHypothesis = liftM getHyp getHeap

--------------------------------------------------------------------------------
-- Simplify

newtype SimplifyST = SimplifyST { reTypVars :: IntMap Var }

emptySimplifyST :: SimplifyST
emptySimplifyST = SimplifyST IMap.empty

withSimplifyST :: MonadState CaoState m => m a -> m a
withSimplifyST m = do
    modify ( \ s -> s { simplifyST = Just emptySimplifyST } )
    r <- m
    modify ( \ s -> s { simplifyST = Nothing } )
    return r

resetSimplifyST :: MonadState CaoState m => m ()
resetSimplifyST = modify (\ s -> s { simplifyST = Just emptySimplifyST } )

reTypVar :: MonadState CaoState m => Var -> m ()
reTypVar v = modify (\ s -> s { simplifyST = fmap aux (simplifyST s) } )
    where 
    aux sst = sst { reTypVars = IMap.insert (varId v) v (reTypVars sst) }

lookupReTypVar :: MonadState CaoState m => Var -> m (Maybe Var)
lookupReTypVar v = liftM (IMap.lookup (varId v) . reTypVars . fromJust) (gets simplifyST)

--------------------------------------------------------------------------------
-- Target

newtype TargetST = TargetST
    -- Constants to allow for global declaration and reuse of constants
    { consts :: Map (Literal Var) [Var] }

emptyTargetST :: TargetST
emptyTargetST = TargetST Map.empty

withTargetST :: MonadState CaoState m => m a -> m a
withTargetST m = do
    modify ( \ s -> s { targetST = Just emptyTargetST } )
    r <- m
    modify ( \ s -> s { targetST = Nothing } )
    return r

resetTargetST :: MonadState CaoState m => m ()
resetTargetST  = modify ( \ s -> s { tmpVars  = [] } )

getsConst :: MonadState CaoState m => m (Map (Literal Var) [Var])
getsConst = liftM (consts . fromJust) $ gets targetST 

allConsts :: MonadState CaoState m => m [(Var, Literal Var)]
allConsts = liftM (concat . map (uncurry zip . mapSnd repeat) . map swap . Map.assocs) $ getsConst

-- Kind of hash function for literals
-- TODO: mods can be dependent types. This means that some unification instead of
-- equality may be needed.
getConst :: MonadState CaoState m => Type Var -> Literal Var -> m Var
getConst ty lit = do
    cm <- getsConst
    maybe (aux cm) (litExists cm) (Map.lookup lit cm) 
    where
    litExists cm = maybe (aux cm) return . List.find ((ty ==) . varType)

    aux cm = do
         u <- getLastVar
         let litV = mkGConst (mkVarName $ "const_" ++ showTy ty ++ toString lit) u ty c
         modify (\s -> s { targetST = Just $ TargetST $ Map.insertWith (++) lit [litV] cm } )
         return litV

    c = case lit of
            ILit n -> Just $ IInt n
            _      -> Nothing

-- TODO: missing cases
showTy :: PP a => Type a -> String
showTy Int = "int_"
showTy RInt = "rint_"
showTy Bool = "bool_"
showTy (Bits U (IInt n)) = "ubits" ++ show n ++ "_"
showTy (Bits S (IInt n)) = "sbits" ++ show n ++ "_"
showTy _ = ""

--------------------------------------------------------------------------------
-- PreC

data PreCST = PreCST
    { fieldProj :: [(Var, Integer)]
    , refVar    :: (Maybe Var)
    , funType   :: (Map.Map Var (Type Var))
    , retStruct :: (Maybe Var)
    }

emptyPreCST :: PreCST
emptyPreCST = PreCST [] Nothing Map.empty Nothing

withPreCST :: MonadState CaoState m => m a -> m a
withPreCST m = do
    modify ( \ s -> s { precST = Just emptyPreCST } )
    r <- m
    modify ( \ s -> s { precST = Nothing } )
    return r

resetPreCST :: MonadState CaoState m => m ()
resetPreCST = do
    modify ( \ s -> s { tmpVars = [] } )
    modifyPreCST ( \ p -> p { refVar = Nothing, retStruct = Nothing } )

modifyPreCST :: MonadState CaoState m => (PreCST -> PreCST) -> m ()
modifyPreCST f = modify ( \ s -> s { precST = fmap f ( precST s ) } )

putFieldProj :: MonadState CaoState m => (Var, Integer) -> m ()
putFieldProj (v,i) = modifyPreCST (\ s -> s { fieldProj = (v',i) : fieldProj s } )
    where 
    SField _ ty = varType v
    v'          = setType ty v

{-# INLINE getsPreCST #-}
getsPreCST :: MonadState CaoState m => (PreCST -> a) -> m a
getsPreCST f = liftM (f . fromJust) $ gets precST

lookupFieldProj :: MonadState CaoState m => Var -> m (Maybe Integer)
lookupFieldProj v = liftM (lookup v) $ getsPreCST fieldProj

getRefVar :: MonadState CaoState m => m (Maybe Var)
getRefVar = getsPreCST refVar

setRefVar :: MonadState CaoState m => Var -> m ()
setRefVar str = modifyPreCST (\ s -> s { refVar = Just str })

putFunType :: MonadState CaoState m => Var -> Type Var -> m ()
putFunType f tf = modifyPreCST $ \ s -> s { funType = Map.insert f tf (funType s) }

getFunType :: MonadState CaoState m => Var -> m (Maybe (Type Var))
getFunType v = liftM (Map.lookup v) $ getsPreCST funType

putRetStruct :: MonadState CaoState m => Var -> m ()
putRetStruct ts = modifyPreCST $ \ s -> s { retStruct = Just ts }

getRetStruct :: MonadState CaoState m => m (Maybe Var)
getRetStruct = getsPreCST retStruct


--------------------------------------------------------------------------------
-- C

newtype CST = CST { allocVars    :: [[Var]] }

emptyCST :: CST
emptyCST = CST [[]]

withCST :: MonadState CaoState m => m a -> m a
withCST m = do
    modify ( \ s -> s { cst = Just emptyCST } )
    r <- m
    modify ( \ s -> s { cst = Nothing } )
    return r

resetCST :: MonadState CaoState m => m ()
resetCST = modify ( \ s -> s 
    { tmpVars = []
    , cst     = Just emptyCST
    } )

modifyCST :: MonadState CaoState m => (CST -> CST) -> m ()
modifyCST f = modify ( \ s -> s { cst = fmap f ( cst s ) } )

storeAllocVar :: MonadState CaoState m => Var -> m ()
storeAllocVar v = modifyCST (\ c -> c { allocVars = aux (allocVars c) })
    where
    aux (h : t) = (v : h) : t
    aux _ = error ""

getAllocVars :: MonadState CaoState m => m [Var]
getAllocVars = liftM (head . allocVars . fromJust) (gets cst)

getAllAllocVars :: MonadState CaoState m => m [Var]
getAllAllocVars = liftM (concat . allocVars . fromJust) (gets cst)

allocScope :: MonadState CaoState m => m a -> m a
allocScope m = do
    modifyCST (\ c -> c { allocVars = [] : allocVars c } )
    r <- m
    modifyCST (\ c -> c { allocVars = tail (allocVars c) } )
    return r