ghc/compiler/coreSyn/CoreTidy.lhs

   1 %
   2 % (c) The AQUA Project, Glasgow University, 1996-1998
   3 %
   4
   5 \begin{code}
   6 module CoreTidy (
   7         tidyExpr, tidyVarOcc, tidyRule, tidyRules
   8     ) where
   9
  10 #include "HsVersions.h"
  11
  12 import CoreSyn
  13 import CoreUtils        ( exprArity )
  14 import Unify            ( coreRefineTys )
  15 import DataCon          ( DataCon, isVanillaDataCon )
  16 import Id               ( Id, mkUserLocal, idInfo, setIdInfo, idUnique,
  17                           idType, setIdType )
  18 import IdInfo           ( setArityInfo, vanillaIdInfo,
  19                           newStrictnessInfo, setAllStrictnessInfo,
  20                           newDemandInfo, setNewDemandInfo )
  21 import Type             ( Type, tidyType, tidyTyVarBndr, substTy, mkOpenTvSubst )
  22 import Var              ( Var, TyVar, varName )
  23 import VarEnv
  24 import UniqFM           ( lookupUFM )
  25 import Name             ( Name, getOccName )
  26 import OccName          ( tidyOccName )
  27 import SrcLoc           ( noSrcLoc )
  28 import Maybes           ( orElse )
  29 import Outputable
  30 import Util             ( mapAccumL )
  31 \end{code}
  32
  33
  34 This module contains "tidying" code for *nested* expressions, bindings, rules.
  35 The code for *top-level* bindings is in TidyPgm.
  36
  37 %************************************************************************
  38 %*                                                                      *
  39 \subsection{Tidying expressions, rules}
  40 %*                                                                      *
  41 %************************************************************************
  42
  43 \begin{code}
  44 tidyBind :: TidyEnv
  45          -> CoreBind
  46          ->  (TidyEnv, CoreBind)
  47
  48 tidyBind env (NonRec bndr rhs)
  49   = tidyLetBndr env (bndr,rhs) =: \ (env', bndr') ->
  50     (env', NonRec bndr' (tidyExpr env' rhs))
  51
  52 tidyBind env (Rec prs)
  53   = mapAccumL tidyLetBndr  env prs      =: \ (env', bndrs') ->
  54     map (tidyExpr env') (map snd prs)   =: \ rhss' ->
  55     (env', Rec (zip bndrs' rhss'))
  56
  57
  58 ------------  Expressions  --------------
  59 tidyExpr :: TidyEnv -> CoreExpr -> CoreExpr
  60 tidyExpr env (Var v)    =  Var (tidyVarOcc env v)
  61 tidyExpr env (Type ty)  =  Type (tidyType env ty)
  62 tidyExpr env (Lit lit)  =  Lit lit
  63 tidyExpr env (App f a)  =  App (tidyExpr env f) (tidyExpr env a)
  64 tidyExpr env (Note n e) =  Note (tidyNote env n) (tidyExpr env e)
  65
  66 tidyExpr env (Let b e)
  67   = tidyBind env b      =: \ (env', b') ->
  68     Let b' (tidyExpr env' e)
  69
  70 tidyExpr env (Case e b ty alts)
  71   = tidyBndr env b      =: \ (env', b) ->
  72     Case (tidyExpr env e) b (tidyType env ty)
  73          (map (tidyAlt b env') alts)
  74
  75 tidyExpr env (Lam b e)
  76   = tidyBndr env b      =: \ (env', b) ->
  77     Lam b (tidyExpr env' e)
  78
  79 ------------  Case alternatives  --------------
  80 tidyAlt case_bndr env (DataAlt con, vs, rhs)
  81   | not (isVanillaDataCon con)  -- GADT case
  82   = tidyBndrs env tvs           =: \ (env1, tvs') ->
  83     let
  84         env2 = refineTidyEnv env con tvs' scrut_ty
  85     in
  86     tidyBndrs env2 ids  =: \ (env3, ids') ->
  87     (DataAlt con, tvs' ++ ids', tidyExpr env3 rhs)
  88   where
  89     (tvs, ids) = span isTyVar vs
  90     scrut_ty = idType case_bndr
  91
  92 tidyAlt case_bndr env (con, vs, rhs)
  93   = tidyBndrs env vs    =: \ (env', vs) ->
  94     (con, vs, tidyExpr env' rhs)
  95
  96 refineTidyEnv :: TidyEnv -> DataCon -> [TyVar] -> Type -> TidyEnv
  97 -- Refine the TidyEnv in the light of the type refinement from coreRefineTys
  98 refineTidyEnv tidy_env@(occ_env, var_env)  con tvs scrut_ty
  99   = case coreRefineTys con tvs scrut_ty of
 100         Nothing -> tidy_env
 101         Just (tv_subst, all_bound_here)
 102             | all_bound_here    -- Local type refinement only
 103             -> tidy_env
 104             | otherwise         -- Apply the refining subst to the tidy env
 105                                 -- This ensures that occurences have the most refined type
 106                                 -- And that means that exprType will work right everywhere
 107             -> (occ_env, mapVarEnv (refine subst) var_env)
 108             where
 109               subst = mkOpenTvSubst tv_subst
 110   where
 111     refine subst var | isId var  = setIdType var (substTy subst (idType var))
 112                      | otherwise = var
 113
 114 ------------  Notes  --------------
 115 tidyNote env (Coerce t1 t2)  = Coerce (tidyType env t1) (tidyType env t2)
 116 tidyNote env note            = note
 117
 118 ------------  Rules  --------------
 119 tidyRules :: TidyEnv -> [CoreRule] -> [CoreRule]
 120 tidyRules env [] = []
 121 tidyRules env (rule : rules)
 122   = tidyRule env rule           =: \ rule ->
 123     tidyRules env rules         =: \ rules ->
 124     (rule : rules)
 125
 126 tidyRule :: TidyEnv -> CoreRule -> CoreRule
 127 tidyRule env rule@(BuiltinRule {}) = rule
 128 tidyRule env rule@(Rule { ru_bndrs = bndrs, ru_args = args, ru_rhs = rhs,
 129                           ru_fn = fn, ru_rough = mb_ns })
 130   = tidyBndrs env bndrs         =: \ (env', bndrs) ->
 131     map (tidyExpr env') args    =: \ args ->
 132     rule { ru_bndrs = bndrs, ru_args = args,
 133            ru_rhs   = tidyExpr env' rhs,
 134            ru_fn    = tidyNameOcc env fn,
 135            ru_rough = map (fmap (tidyNameOcc env')) mb_ns }
 136 \end{code}
 137
 138
 139 %************************************************************************
 140 %*                                                                      *
 141 \subsection{Tidying non-top-level binders}
 142 %*                                                                      *
 143 %************************************************************************
 144
 145 \begin{code}
 146 tidyNameOcc :: TidyEnv -> Name -> Name
 147 -- In rules and instances, we have Names, and we must tidy them too
 148 -- Fortunately, we can lookup in the VarEnv with a name
 149 tidyNameOcc (_, var_env) n = case lookupUFM var_env n of
 150                                 Nothing -> n
 151                                 Just v  -> varName v
 152
 153 tidyVarOcc :: TidyEnv -> Var -> Var
 154 tidyVarOcc (_, var_env) v = lookupVarEnv var_env v `orElse` v
 155
 156 -- tidyBndr is used for lambda and case binders
 157 tidyBndr :: TidyEnv -> Var -> (TidyEnv, Var)
 158 tidyBndr env var
 159   | isTyVar var = tidyTyVarBndr env var
 160   | otherwise   = tidyIdBndr env var
 161
 162 tidyBndrs :: TidyEnv -> [Var] -> (TidyEnv, [Var])
 163 tidyBndrs env vars = mapAccumL tidyBndr env vars
 164
 165 tidyLetBndr :: TidyEnv -> (Id, CoreExpr) -> (TidyEnv, Var)
 166 -- Used for local (non-top-level) let(rec)s
 167 tidyLetBndr env (id,rhs)
 168   = ((tidy_env,new_var_env), final_id)
 169   where
 170     ((tidy_env,var_env), new_id) = tidyIdBndr env id
 171
 172         -- We need to keep around any interesting strictness and
 173         -- demand info because later on we may need to use it when
 174         -- converting to A-normal form.
 175         -- eg.
 176         --      f (g x),  where f is strict in its argument, will be converted
 177         --      into  case (g x) of z -> f z  by CorePrep, but only if f still
 178         --      has its strictness info.
 179         --
 180         -- Similarly for the demand info - on a let binder, this tells
 181         -- CorePrep to turn the let into a case.
 182         --
 183         -- Similarly arity info for eta expansion in CorePrep
 184         --
 185     final_id = new_id `setIdInfo` new_info
 186     idinfo   = idInfo id
 187     new_info = vanillaIdInfo
 188                 `setArityInfo`          exprArity rhs
 189                 `setAllStrictnessInfo`  newStrictnessInfo idinfo
 190                 `setNewDemandInfo`      newDemandInfo idinfo
 191
 192     -- Override the env we get back from tidyId with the new IdInfo
 193     -- so it gets propagated to the usage sites.
 194     new_var_env = extendVarEnv var_env id final_id
 195
 196 -- Non-top-level variables
 197 tidyIdBndr :: TidyEnv -> Id -> (TidyEnv, Id)
 198 tidyIdBndr env@(tidy_env, var_env) id
 199   = -- do this pattern match strictly, otherwise we end up holding on to
 200     -- stuff in the OccName.
 201     case tidyOccName tidy_env (getOccName id) of { (tidy_env', occ') ->
 202     let
 203         -- Give the Id a fresh print-name, *and* rename its type
 204         -- The SrcLoc isn't important now,
 205         -- though we could extract it from the Id
 206         --
 207         -- All nested Ids now have the same IdInfo, namely vanillaIdInfo,
 208         -- which should save some space.
 209         -- But note that tidyLetBndr puts some of it back.
 210         ty'               = tidyType env (idType id)
 211         id'               = mkUserLocal occ' (idUnique id) ty' noSrcLoc
 212                                 `setIdInfo` vanillaIdInfo
 213         var_env'          = extendVarEnv var_env id id'
 214     in
 215      ((tidy_env', var_env'), id')
 216    }
 217 \end{code}
 218
 219 \begin{code}
 220 m =: k = m `seq` k m
 221 \end{code}