compiler/coreSyn/CoreTidy.lhs

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