{-  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/>. -}

{-
Module      :  $Header$
Description :  Tidy CAO variable names.
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

-}

module Language.CAO.Syntax.Tidy
    ( tidyCaoAST
    , showCaoAST
    , showCaoASTDebug
    ) where


import Control.Monad.State

import Data.Map ( Map )
import qualified Data.Map as Map
import qualified Data.Set as Set
import qualified Data.Traversable as T

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

import Language.CAO.Index

import Language.CAO.Syntax
import Language.CAO.Syntax.Utils

import Language.CAO.Type


type TidyM = State TidySt

data TidySt
    = TidySt { symbolCount   :: Map String Int
             , seen          :: Map Var Var
             , globalSymbols :: Map String Int
             }

emptyState :: TidySt
emptyState 
    = TidySt { symbolCount   = Map.empty
             , seen          = Map.empty
             , globalSymbols = Map.empty
             }

resetLocals :: TidyM ()
resetLocals = modify $ \ s -> s { symbolCount = globalSymbols s }

showCaoAST :: Prog Var -> String
showCaoAST = showPpr . tidyCaoAST

showCaoASTDebug :: Prog Var -> String
showCaoASTDebug = showPprDebug . tidyCaoAST

tidyCaoAST :: Prog Var -> Prog Var
tidyCaoAST p@(Prog defs ip) = uncurry Prog (evalState tidyWorker initState)
    where 
    tidyWorker = do
        defs' <- mapM tidyLDef defs
        ip'   <- mapMaybeM tidyFunDef ip
        return (defs', ip')
    initState = emptyState { globalSymbols = Map.fromList glbs }
    glbs      = zip (map getSymbol $ Set.toList $ globals p) (repeat 1)

tidyLDef :: LDef Var -> TidyM (LDef Var)
tidyLDef = tidyLoc tidyDef

tidyDef :: Def Var -> TidyM (Def Var)
tidyDef (VarDef   vd) = resetLocals >> liftM VarDef   (tidyVarDecl   vd)
tidyDef (FunDef   vd) = resetLocals >> liftM FunDef   (tidyFunDef    vd)
tidyDef (TyDef    vd) = resetLocals >> liftM TyDef    (tidyTyDef     vd)
tidyDef (ConstDef cd) = resetLocals >> liftM ConstDef (tidyConstDecl cd)

tidyVarDecl :: VarDecl Var -> TidyM (VarDecl Var)
tidyVarDecl (VarD v td me)
    = liftM3 VarD (tidyLVar v) (tidyTyDecl td) (T.mapM tidyTLExpr me)
tidyVarDecl (MultiD v td)
    = liftM2 MultiD (mapM tidyLVar v) (tidyTyDecl td)
tidyVarDecl (ContD v td es)
    = liftM3 ContD (tidyLVar v) (tidyTyDecl td) (mapM tidyTLExpr es)

tidyConstDecl :: ConstDecl Var -> TidyM (ConstDecl Var)
tidyConstDecl (ConstD v td c) 
    = liftM3 ConstD (tidyLVar v) (tidyTyDecl td) (return c)
tidyConstDecl (MultiConstD v td c) 
    = liftM3 MultiConstD (mapM tidyLVar v) (tidyTyDecl td) (return c)

tidyFunDef :: Fun Var -> TidyM (Fun Var)
tidyFunDef (Fun v args tds lstmts)
    = liftM4 Fun (     tidyLVar   v     )
                 (mapM tidyArg    args  )
                 (mapM tidyTyDecl tds   )
                 (mapM tidyLStmt  lstmts)

tidyTyDef :: TyDef Var -> TidyM (TyDef Var)
tidyTyDef (TySynDef v td)   = liftM2 TySynDef   (tidyLVar v) (tidyTyDecl td)
tidyTyDef (StructDecl v fs) = liftM2 StructDecl (tidyLVar v) (mapM tidyFld fs)

tidyLTyDecl :: LTyDecl Var -> TidyM (LTyDecl Var)
tidyLTyDecl = tidyLoc tidyTyDecl

tidyTyDecl :: TyDecl Var -> TidyM (TyDecl Var)
tidyTyDecl (BitsD s e)      = liftM (BitsD s) $ tidyLExpr e
tidyTyDecl (ModD md)        = liftM ModD      $ tidyMod md
tidyTyDecl (VectorD e td)   = liftM2 VectorD  (tidyLExpr e)
                                              (tidyTyDecl td)
tidyTyDecl (MatrixD r c td) = liftM3 MatrixD  (tidyLExpr r)
                                              (tidyLExpr c)
                                              (tidyTyDecl td)
tidyTyDecl (TySynD v)       = liftM TySynD    (tidyLVar v)
tidyTyDecl d                = return d

tidyArg :: Arg Var -> TidyM (Arg Var)
tidyArg (Arg v td) 
    = liftM2 Arg (tidyLVar v) (tidyTyDecl td)
tidyArg (ArgConst v td i) 
    = liftM3 ArgConst (tidyLVar v) (tidyTyDecl td) (T.mapM tidyLExpr i)

tidyFld :: (Located Var, TyDecl Var) -> TidyM (Located Var, TyDecl Var)
tidyFld (v, td) = liftM2 (,) (tidyLVar v) (tidyTyDecl td)

tidyLStmt :: LStmt Var -> TidyM (LStmt Var)
tidyLStmt = tidyLoc tidyStmt

tidyStmt :: Stmt Var -> TidyM (Stmt Var)
tidyStmt (VDecl  vd)     
    = liftM  VDecl (tidyVarDecl vd)
tidyStmt (CDecl  cd)     
    = liftM  CDecl (tidyConstDecl cd)
tidyStmt (Assign lvs es) 
    = liftM2 Assign (mapM tidyLVal lvs) (mapM tidyTLExpr es)
tidyStmt (FCallS v es)   
    = liftM2 FCallS (tidyVar v) (mapM tidyTLExpr es)
tidyStmt (Ret es)        
    = liftM  Ret (mapM tidyTLExpr es)
tidyStmt (Ite e ss mss)  
    = liftM3 Ite (tidyTLExpr e)
                 (mapM tidyLStmt ss)
                 (T.mapM (mapM tidyLStmt) mss)
tidyStmt (Seq iter ss)   
    = liftM2 Seq (tidySeqIter iter) (mapM tidyLStmt ss)
tidyStmt (While e ss)    
    = liftM2 While (tidyTLExpr e) (mapM tidyLStmt ss)
tidyStmt (Nop a)         
    = return (Nop a)

tidyLVal :: LVal Var -> TidyM (LVal Var)
tidyLVal (LVVar v)          = liftM  LVVar (tidyLVar v)
tidyLVal (LVStruct lv fi)   = liftM2 LVStruct (tidyLVal lv) (tidyVar fi)
tidyLVal (LVCont ty lv pat) = liftM2 (LVCont ty) (tidyLVal lv) (tidyAPat pat)

tidySeqIter :: SeqIter Var -> TidyM (SeqIter Var)
tidySeqIter (SeqIter v s e mb is)
    = liftM4 (\v' s' e' mb' -> SeqIter v' s' e' mb' is)
        (tidyVar v) (tidyLExpr s) (tidyLExpr e) (T.mapM tidyLExpr mb)

tidyAPat :: APat Var -> TidyM (APat Var)
tidyAPat (VectP rp)   = liftM  VectP (tidyRowAPat rp)
tidyAPat (MatP rp cp) = liftM2 MatP  (tidyRowAPat rp) (tidyRowAPat cp)

tidyRowAPat :: RowAPat Var -> TidyM (RowAPat Var)
tidyRowAPat (CElem e)    = liftM  CElem  (tidyTLExpr e)
tidyRowAPat (CRange i j) = liftM2 CRange (tidyTLExpr i) (tidyTLExpr j)

tidyMod :: Mod Var -> TidyM (Mod Var)
tidyMod (ModNum e)        = liftM ModNum (tidyLExpr e)
tidyMod (ModPol td ind p) = liftM3 ModPol (tidyTyDecl td)
                                          (tidyVar ind)
                                          (tidyPol p)

tidyLit :: Literal Var -> TidyM (Literal Var)
tidyLit (PLit pol) = liftM PLit (tidyPol pol)
tidyLit l          = return l

tidyPol :: Pol Var -> TidyM (Pol Var)
tidyPol (Pol ms) = liftM Pol (mapM tidyMon ms)

tidyMon :: Mon Var -> TidyM (Mon Var)
tidyMon (Mon c b) = liftM2 Mon (tidyMCoef c) (tidyMBase b)

tidyMCoef :: MCoef Var -> TidyM (MCoef Var)
tidyMCoef (CoefP p) = liftM CoefP (tidyPol p)
tidyMCoef c         = return c

tidyMBase :: MBase Var -> TidyM (MBase Var)
tidyMBase (MExpI n e) = liftM (flip MExpI e) (tidyVar n)
tidyMBase b           = return b


tidyLExpr :: LExpr Var -> TidyM (LExpr Var)
tidyLExpr = tidyLoc tidyExpr

tidyTLExpr :: TLExpr Var -> TidyM (TLExpr Var)
tidyTLExpr = tidyLoc (\ (TyE t e) -> liftM (TyE t) (tidyExpr e))

tidyExpr :: Expr Var -> TidyM (Expr Var)
tidyExpr (Var v)             = liftM  Var (tidyVar v)
tidyExpr (Lit lit)           = liftM  Lit (tidyLit lit)
tidyExpr (FunCall v es)      = liftM2 FunCall (tidyLVar v) (mapM tidyTLExpr es)
tidyExpr (StructProj e fi)   = liftM2 StructProj (tidyTLExpr e) (tidyVar fi)
tidyExpr (UnaryOp op e)      = liftM  (UnaryOp op) (tidyTLExpr e)
tidyExpr (BinaryOp op e1 e2) = liftM2 (BinaryOp op) (tidyTLExpr e1)
                                                    (tidyTLExpr e2)
tidyExpr (Access e pat)      = liftM2 Access (tidyTLExpr e) (tidyAPat pat)
tidyExpr (Cast b td e)       = liftM2 (Cast b) (mapM tidyLTyDecl td)
                                                    (tidyTLExpr e)

tidyLVar :: Located Var -> TidyM (Located Var)
tidyLVar = tidyLoc tidyVar

tidyVar :: Var -> TidyM Var
tidyVar v
    | isLocal v = do
        vars <- gets seen
        case Map.lookup v vars of
            Nothing -> do
                v' <- newSymbol v
                t'  <- tidyType $ varType v'
                let  v'' = setType t' v'
                modify (\s -> s { seen = Map.insert v v'' (seen s) })
                return v''
            Just v' -> return v'
    | isGlobalInit v = return v
    | nsTyVar v      = return v
    | isCCast v      = return v
    | isCFunction v  = return v
    | otherwise      = do
        t' <- tidyType $ varType v
        return $ setSymbol (addPrefix $ getSymbol v) (setType t' v)

newSymbol :: Var -> TidyM Var
newSymbol v = do
    sc <- gets symbolCount
    let vs = getSymbol v
    case Map.lookup vs sc of
        Nothing -> do
            modify (\s -> s { symbolCount = Map.insert vs 1 (symbolCount s) })
            return $ setSymbol (addPrefix vs) v
        Just i  -> do
            modify (\s -> s { symbolCount = Map.adjust (+1) vs (symbolCount s) })
            return $ setSymbol (addPrefix $ vs ++ '_' : show i) v

tidyType :: Type Var -> TidyM (Type Var)
tidyType (Bits s sz) = 
    liftM (Bits s) $ tidyIExpr sz
tidyType (Vector n t) = 
    liftM2 Vector (tidyIExpr n) (tidyType t)
tidyType (Matrix n m t) = 
    liftM3 Matrix (tidyIExpr n) (tidyIExpr m) (tidyType t)
tidyType (Mod Nothing Nothing (Pol [Mon (CoefI i) EZero])) = do
    i' <- tidyIExpr i
    return $ Mod Nothing Nothing (Pol [Mon (CoefI i') EZero])
tidyType (Mod (Just im@(Mod Nothing Nothing (Pol [Mon (CoefI _) EZero]))) 
         (Just i) 
         (Pol pol)) = do
    im'  <- tidyType im
    pol' <- mapM aux pol
    return $ Mod (Just im') (Just i) (Pol pol')
    where
    aux (Mon (CoefI co) e) = do
        co' <- tidyIExpr co
        return $ Mon (CoefI co') e
    aux _ = error "<tidyType>: not expected case"
tidyType (TySyn v t) = do
    t' <- tidyType t
    let tct    = TySyn newvar t'
        newvar = setType tct v
    return tct
tidyType (Struct s flds) = do
    fldtys' <- mapM tidyFld' flds
    let tct    = Struct newvar flds'
        newvar = setType tct s
        flds'  = map(\(v, ty) -> (setType (SField newvar ty) v, ty)) fldtys'
    return tct
    where
    tidyFld' (a, sf) = tidyType sf >>= \ sf' -> return (a, sf')

tidyType e = return e

tidyIExpr :: IExpr Var -> TidyM (IExpr Var)
tidyIExpr (IInd v)          = liftM  IInd $ tidyVar v
tidyIExpr (ISum l)          = liftM  ISum $ mapM tidyIExpr l
tidyIExpr (IArith op e1 e2) = liftM2 (IArith op) (tidyIExpr e1) (tidyIExpr e2)
tidyIExpr (ISym e)          = liftM  ISym $ tidyIExpr e
tidyIExpr n@(IInt _)        = return n

{-# INLINE addPrefix #-}
addPrefix :: String -> String
addPrefix = ("c_" ++)

{-# INLINE tidyLoc #-}
tidyLoc :: (a -> TidyM a) -> Located a -> TidyM (Located a)
tidyLoc f (L l a) = liftM (L l) (f a)