ghc/compiler/coreSyn/CoreUtils.lhs

   1 %
   2 % (c) The GRASP/AQUA Project, Glasgow University, 1992-1996
   3 %
   4 \section[CoreUtils]{Utility functions on @Core@ syntax}
   5
   6 \begin{code}
   7 module CoreUtils (
   8         coreExprType, coreAltsType, coreExprCc,
   9
  10         mkCoreIfThenElse,
  11         argToExpr,
  12         unTagBinders, unTagBindersAlts,
  13
  14         maybeErrorApp,
  15         nonErrorRHSs,
  16         squashableDictishCcExpr,
  17         idSpecVars
  18     ) where
  19
  20 #include "HsVersions.h"
  21
  22 import CoreSyn
  23
  24 import CostCentre       ( isDictCC, CostCentre, noCostCentre )
  25 import MkId             ( mkSysLocal )
  26 import Id               ( idType, isBottomingId, getIdSpecialisation,
  27                           mkIdWithNewUniq,
  28                           dataConRepType,
  29                           addOneToIdEnv, growIdEnvList, lookupIdEnv,
  30                           isNullIdEnv, IdEnv, Id
  31                         )
  32 import Literal          ( literalType, Literal(..) )
  33 import Maybes           ( catMaybes, maybeToBool )
  34 import PprCore
  35 import PrimOp           ( primOpType, PrimOp(..) )
  36 import SpecEnv          ( specEnvValues )
  37 import SrcLoc           ( noSrcLoc )
  38 import TyVar            ( cloneTyVar,
  39                           isEmptyTyVarEnv, addToTyVarEnv, TyVarEnv,
  40                           TyVar, GenTyVar
  41                         )
  42 import Type             ( mkFunTy, mkForAllTy, mkTyVarTy,
  43                           splitFunTy_maybe, applyTys, isUnpointedType,
  44                           splitSigmaTy, splitFunTys, instantiateTy,
  45                           Type
  46                         )
  47 import TysWiredIn       ( trueDataCon, falseDataCon )
  48 import Unique           ( Unique )
  49 import BasicTypes       ( Unused )
  50 import UniqSupply       ( returnUs, thenUs,
  51                           mapUs, mapAndUnzipUs, getUnique,
  52                           UniqSM, UniqSupply
  53                         )
  54 import Util             ( zipEqual )
  55 import Outputable
  56
  57 type TypeEnv = TyVarEnv Type
  58 \end{code}
  59
  60 %************************************************************************
  61 %*                                                                      *
  62 \subsection{Find the type of a Core atom/expression}
  63 %*                                                                      *
  64 %************************************************************************
  65
  66 \begin{code}
  67 coreExprType :: CoreExpr -> Type
  68
  69 coreExprType (Var var) = idType   var
  70 coreExprType (Lit lit) = literalType lit
  71
  72 coreExprType (Let _ body)       = coreExprType body
  73 coreExprType (Case _ alts)      = coreAltsType alts
  74
  75 coreExprType (Note (Coerce ty _) e) = ty
  76 coreExprType (Note other_note e)    = coreExprType e
  77
  78 -- a Con is a fully-saturated application of a data constructor
  79 -- a Prim is <ditto> of a PrimOp
  80
  81 coreExprType (Con con args) =
  82 --                            pprTrace "appTyArgs" (hsep [ppr con, semi,
  83 --                                                         ppr con_ty, semi,
  84 --                                                         ppr args]) $
  85                               applyTypeToArgs con_ty args
  86                             where
  87                                 con_ty = dataConRepType con
  88
  89 coreExprType (Prim op args) = applyTypeToArgs (primOpType op) args
  90
  91 coreExprType (Lam (ValBinder binder) expr)
  92   = idType binder `mkFunTy` coreExprType expr
  93
  94 coreExprType (Lam (TyBinder tyvar) expr)
  95   = mkForAllTy tyvar (coreExprType expr)
  96
  97 coreExprType (App expr (TyArg ty))
  98   =     -- Gather type args; more efficient to instantiate the type all at once
  99     go expr [ty]
 100   where
 101     go (App expr (TyArg ty)) tys = go expr (ty:tys)
 102     go expr                  tys = applyTys (coreExprType expr) tys
 103
 104 coreExprType (App expr val_arg)
 105   = ASSERT(isValArg val_arg)
 106     let
 107         fun_ty = coreExprType expr
 108     in
 109     case (splitFunTy_maybe fun_ty) of
 110           Just (_, result_ty) -> result_ty
 111 #ifdef DEBUG
 112           Nothing -> pprPanic "coreExprType:\n"
 113                         (vcat [ppr fun_ty,  ppr (App expr val_arg)])
 114 #endif
 115 \end{code}
 116
 117 \begin{code}
 118 coreAltsType :: CoreCaseAlts -> Type
 119
 120 coreAltsType (AlgAlts [] deflt)         = default_ty deflt
 121 coreAltsType (AlgAlts ((_,_,rhs1):_) _) = coreExprType rhs1
 122
 123 coreAltsType (PrimAlts [] deflt)       = default_ty deflt
 124 coreAltsType (PrimAlts ((_,rhs1):_) _) = coreExprType rhs1
 125
 126 default_ty NoDefault           = panic "coreExprType:Case:default_ty"
 127 default_ty (BindDefault _ rhs) = coreExprType rhs
 128 \end{code}
 129
 130 \begin{code}
 131 applyTypeToArgs op_ty (TyArg ty : args)
 132   =     -- Accumulate type arguments so we can instantiate all at once
 133     applyTypeToArgs (applyTys op_ty tys) rest_args
 134   where
 135     (tys, rest_args)         = go [ty] args
 136     go tys (TyArg ty : args) = go (ty:tys) args
 137     go tys rest_args         = (reverse tys, rest_args)
 138
 139 applyTypeToArgs op_ty (val_or_lit_arg:args)
 140   = case (splitFunTy_maybe op_ty) of
 141         Just (_, res_ty) -> applyTypeToArgs res_ty args
 142
 143 applyTypeToArgs op_ty [] = op_ty
 144 \end{code}
 145
 146 coreExprCc gets the cost centre enclosing an expression, if any.
 147 It looks inside lambdas because (scc "foo" \x.e) = \x.scc "foo" e
 148
 149 \begin{code}
 150 coreExprCc :: GenCoreExpr val_bdr val_occ flexi -> CostCentre
 151 coreExprCc (Note (SCC cc) e)   = cc
 152 coreExprCc (Note other_note e) = coreExprCc e
 153 coreExprCc (Lam _ e)           = coreExprCc e
 154 coreExprCc other               = noCostCentre
 155 \end{code}
 156
 157 %************************************************************************
 158 %*                                                                      *
 159 \subsection{Routines to manufacture bits of @CoreExpr@}
 160 %*                                                                      *
 161 %************************************************************************
 162
 163 \begin{code}
 164 mkCoreIfThenElse (Var bool) then_expr else_expr
 165     | bool == trueDataCon   = then_expr
 166     | bool == falseDataCon  = else_expr
 167
 168 mkCoreIfThenElse guard then_expr else_expr
 169   = Case guard
 170       (AlgAlts [ (trueDataCon,  [], then_expr),
 171                  (falseDataCon, [], else_expr) ]
 172        NoDefault )
 173 \end{code}
 174
 175 For making @Apps@ and @Lets@, we must take appropriate evasive
 176 action if the thing being bound has unboxed type.  @mkCoApp@ requires
 177 a name supply to do its work.
 178
 179 @mkCoApps@, @mkCoCon@ and @mkCoPrim@ also handle the
 180 arguments-must-be-atoms constraint.
 181
 182 \begin{code}
 183 data CoreArgOrExpr
 184   = AnArg   CoreArg
 185   | AnExpr  CoreExpr
 186
 187 mkCoApps :: CoreExpr -> [CoreArgOrExpr] -> UniqSM CoreExpr
 188 mkCoCon  :: Id       -> [CoreArgOrExpr] -> UniqSM CoreExpr
 189 mkCoPrim :: PrimOp   -> [CoreArgOrExpr] -> UniqSM CoreExpr
 190
 191 mkCoApps fun args = co_thing (mkGenApp fun) args
 192 mkCoCon  con args = co_thing (Con  con)     args
 193 mkCoPrim  op args = co_thing (Prim op)      args
 194
 195 co_thing :: ([CoreArg] -> CoreExpr)
 196          -> [CoreArgOrExpr]
 197          -> UniqSM CoreExpr
 198
 199 co_thing thing arg_exprs
 200   = mapAndUnzipUs expr_to_arg arg_exprs `thenUs` \ (args, maybe_binds) ->
 201     returnUs (mkCoLetsUnboxedToCase (catMaybes maybe_binds) (thing args))
 202   where
 203     expr_to_arg :: CoreArgOrExpr
 204                 -> UniqSM (CoreArg, Maybe CoreBinding)
 205
 206     expr_to_arg (AnArg  arg)     = returnUs (arg,      Nothing)
 207     expr_to_arg (AnExpr (Var v)) = returnUs (VarArg v, Nothing)
 208     expr_to_arg (AnExpr (Lit l)) = returnUs (LitArg l, Nothing)
 209     expr_to_arg (AnExpr other_expr)
 210       = let
 211             e_ty = coreExprType other_expr
 212         in
 213         getUnique `thenUs` \ uniq ->
 214         let
 215             new_var  = mkSysLocal SLIT("a") uniq e_ty noSrcLoc
 216         in
 217         returnUs (VarArg new_var, Just (NonRec new_var other_expr))
 218 \end{code}
 219
 220 \begin{code}
 221 argToExpr ::
 222   GenCoreArg val_occ flexi -> GenCoreExpr val_bdr val_occ flexi
 223
 224 argToExpr (VarArg v)   = Var v
 225 argToExpr (LitArg lit) = Lit lit
 226 \end{code}
 227
 228 All the following functions operate on binders, perform a uniform
 229 transformation on them; ie. the function @(\ x -> (x,False))@
 230 annotates all binders with False.
 231
 232 \begin{code}
 233 unTagBinders :: GenCoreExpr (Id,tag) bdee flexi -> GenCoreExpr Id bdee flexi
 234 unTagBinders expr = bop_expr fst expr
 235
 236 unTagBindersAlts :: GenCoreCaseAlts (Id,tag) bdee flexi -> GenCoreCaseAlts Id bdee flexi
 237 unTagBindersAlts alts = bop_alts fst alts
 238 \end{code}
 239
 240 \begin{code}
 241 bop_expr  :: (a -> b) -> GenCoreExpr a bdee flexi -> GenCoreExpr b bdee flexi
 242
 243 bop_expr f (Var b)           = Var b
 244 bop_expr f (Lit lit)         = Lit lit
 245 bop_expr f (Con con args)    = Con con args
 246 bop_expr f (Prim op args)    = Prim op args
 247 bop_expr f (Lam binder expr) = Lam  (bop_binder f binder) (bop_expr f expr)
 248 bop_expr f (App expr arg)    = App  (bop_expr f expr) arg
 249 bop_expr f (Note note expr)  = Note note (bop_expr f expr)
 250 bop_expr f (Let bind expr)   = Let  (bop_bind f bind) (bop_expr f expr)
 251 bop_expr f (Case expr alts)  = Case (bop_expr f expr) (bop_alts f alts)
 252
 253 bop_binder f (ValBinder   v) = ValBinder (f v)
 254 bop_binder f (TyBinder    t) = TyBinder    t
 255
 256 bop_bind f (NonRec b e) = NonRec (f b) (bop_expr f e)
 257 bop_bind f (Rec pairs)  = Rec [(f b, bop_expr f e) | (b, e) <- pairs]
 258
 259 bop_alts f (AlgAlts alts deflt)
 260   = AlgAlts  [ (con, [f b | b <- binders], bop_expr f e)
 261              | (con, binders, e) <- alts ]
 262              (bop_deflt f deflt)
 263
 264 bop_alts f (PrimAlts alts deflt)
 265   = PrimAlts [ (lit, bop_expr f e) | (lit, e) <- alts ]
 266              (bop_deflt f deflt)
 267
 268 bop_deflt f (NoDefault)          = NoDefault
 269 bop_deflt f (BindDefault b expr) = BindDefault (f b) (bop_expr f expr)
 270 \end{code}
 271
 272 OLD (but left here because of the nice example): @singleAlt@ checks
 273 whether a bunch of case alternatives is actually just one alternative.
 274 It specifically {\em ignores} alternatives which consist of just a
 275 call to @error@, because they won't result in any code duplication.
 276
 277 Example:
 278 \begin{verbatim}
 279         case (case <something> of
 280                 True  -> <rhs>
 281                 False -> error "Foo") of
 282         <alts>
 283
 284 ===>
 285
 286         case <something> of
 287            True ->  case <rhs> of
 288                     <alts>
 289            False -> case error "Foo" of
 290                     <alts>
 291
 292 ===>
 293
 294         case <something> of
 295            True ->  case <rhs> of
 296                     <alts>
 297            False -> error "Foo"
 298 \end{verbatim}
 299 Notice that the \tr{<alts>} don't get duplicated.
 300
 301 \begin{code}
 302 nonErrorRHSs :: GenCoreCaseAlts a Id Unused -> [GenCoreExpr a Id Unused]
 303
 304 nonErrorRHSs alts
 305   = filter not_error_app (find_rhss alts)
 306   where
 307     find_rhss (AlgAlts  as deflt) = [rhs | (_,_,rhs) <- as] ++ deflt_rhs deflt
 308     find_rhss (PrimAlts as deflt) = [rhs | (_,rhs)   <- as] ++ deflt_rhs deflt
 309
 310     deflt_rhs NoDefault           = []
 311     deflt_rhs (BindDefault _ rhs) = [rhs]
 312
 313     not_error_app rhs
 314       = case (maybeErrorApp rhs Nothing) of
 315           Just _  -> False
 316           Nothing -> True
 317 \end{code}
 318
 319 maybeErrorApp checks whether an expression is of the form
 320
 321         error ty args
 322
 323 If so, it returns
 324
 325         Just (error ty' args)
 326
 327 where ty' is supplied as an argument to maybeErrorApp.
 328
 329 Here's where it is useful:
 330
 331                 case (error ty "Foo" e1 e2) of <alts>
 332  ===>
 333                 error ty' "Foo"
 334
 335 where ty' is the type of any of the alternatives.  You might think
 336 this never occurs, but see the comments on the definition of
 337 @singleAlt@.
 338
 339 Note: we *avoid* the case where ty' might end up as a primitive type:
 340 this is very uncool (totally wrong).
 341
 342 NOTICE: in the example above we threw away e1 and e2, but not the
 343 string "Foo".  How did we know to do that?
 344
 345 Answer: for now anyway, we only handle the case of a function whose
 346 type is of form
 347
 348         bottomingFn :: forall a. t1 -> ... -> tn -> a
 349                               ^---------------------^ NB!
 350
 351 Furthermore, we only count a bottomingApp if the function is applied
 352 to more than n args.  If so, we transform:
 353
 354         bottomingFn ty e1 ... en en+1 ... em
 355 to
 356         bottomingFn ty' e1 ... en
 357
 358 That is, we discard en+1 .. em
 359
 360 \begin{code}
 361 maybeErrorApp
 362         :: GenCoreExpr a Id Unused      -- Expr to look at
 363         -> Maybe Type                   -- Just ty => a result type *already cloned*;
 364                                         -- Nothing => don't know result ty; we
 365                                         -- *pretend* that the result ty won't be
 366                                         -- primitive -- somebody later must
 367                                         -- ensure this.
 368         -> Maybe (GenCoreExpr b Id Unused)
 369
 370 maybeErrorApp expr result_ty_maybe
 371   = case (collectArgs expr) of
 372       (Var fun, [ty], other_args)
 373         | isBottomingId fun
 374         && maybeToBool result_ty_maybe -- we *know* the result type
 375                                        -- (otherwise: live a fairy-tale existence...)
 376         && not (isUnpointedType result_ty) ->
 377
 378         case (splitSigmaTy (idType fun)) of
 379           ([tyvar], [], tau_ty) ->
 380               case (splitFunTys tau_ty) of { (arg_tys, res_ty) ->
 381               let
 382                   n_args_to_keep = length arg_tys
 383                   args_to_keep   = take n_args_to_keep other_args
 384               in
 385               if  (res_ty == mkTyVarTy tyvar)
 386                && n_args_to_keep <= length other_args
 387               then
 388                     -- Phew!  We're in business
 389                   Just (mkGenApp (Var fun) (TyArg result_ty : args_to_keep))
 390               else
 391                   Nothing
 392               }
 393
 394           other -> Nothing  -- Function type wrong shape
 395       other -> Nothing
 396   where
 397     Just result_ty = result_ty_maybe
 398 \end{code}
 399
 400 \begin{code}
 401 squashableDictishCcExpr :: CostCentre -> GenCoreExpr a b c -> Bool
 402
 403 squashableDictishCcExpr cc expr
 404   = if not (isDictCC cc) then
 405         False -- that was easy...
 406     else
 407         squashable expr -- note: quite like the "atomic_rhs" stuff in simplifier
 408   where
 409     squashable (Var _)      = True
 410     squashable (Con  _ _)   = True -- I think so... WDP 94/09
 411     squashable (Prim _ _)   = True -- ditto
 412     squashable (App f a)
 413       | notValArg a         = squashable f
 414     squashable other        = False
 415 \end{code}
 416
 417
 418 Given an Id, idSpecVars returns all its specialisations.
 419 We extract these from its SpecEnv.
 420 This is used by the occurrence analyser and free-var finder;
 421 we regard an Id's specialisations as free in the Id's definition.
 422
 423 \begin{code}
 424 idSpecVars :: Id -> [Id]
 425 idSpecVars id
 426   = map get_spec (specEnvValues (getIdSpecialisation id))
 427   where
 428     -- get_spec is another cheapo function like dictRhsFVs
 429     -- It knows what these specialisation temlates look like,
 430     -- and just goes for the jugular
 431     get_spec (App f _) = get_spec f
 432     get_spec (Lam _ b) = get_spec b
 433     get_spec (Var v)   = v
 434 \end{code}