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