cao-0.1.1: src/Language/CAO/Syntax/Tidy.hs
{- 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)