ghc/compiler/coreSyn/CoreTidy.lhs

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