compiler/simplCore/LiberateCase.lhs

   1 %
   2 % (c) The AQUA Project, Glasgow University, 1994-1998
   3 %
   4 \section[LiberateCase]{Unroll recursion to allow evals to be lifted from a loop}
   5
   6 \begin{code}
   7 module LiberateCase ( liberateCase ) where
   8
   9 #include "HsVersions.h"
  10
  11 import DynFlags
  12 import CoreSyn
  13 import CoreUnfold       ( couldBeSmallEnoughToInline )
  14 import Id
  15 import VarEnv
  16 import Util             ( notNull )
  17 \end{code}
  18
  19 The liberate-case transformation
  20 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
  21 This module walks over @Core@, and looks for @case@ on free variables.
  22 The criterion is:
  23         if there is case on a free on the route to the recursive call,
  24         then the recursive call is replaced with an unfolding.
  25
  26 Example
  27
  28    f = \ t -> case v of
  29                  V a b -> a : f t
  30
  31 => the inner f is replaced.
  32
  33    f = \ t -> case v of
  34                  V a b -> a : (letrec
  35                                 f =  \ t -> case v of
  36                                                V a b -> a : f t
  37                                in f) t
  38 (note the NEED for shadowing)
  39
  40 => Simplify
  41
  42   f = \ t -> case v of
  43                  V a b -> a : (letrec
  44                                 f = \ t -> a : f t
  45                                in f t)
  46
  47 Better code, because 'a' is  free inside the inner letrec, rather
  48 than needing projection from v.
  49
  50 Note that this deals with *free variables*.  SpecConstr deals with
  51 *arguments* that are of known form.  E.g.
  52
  53         last []     = error
  54         last (x:[]) = x
  55         last (x:xs) = last xs
  56
  57
  58 Note [Scrutinee with cast]
  59 ~~~~~~~~~~~~~~~~~~~~~~~~~~
  60 Consider this:
  61     f = \ t -> case (v `cast` co) of
  62                  V a b -> a : f t
  63
  64 Exactly the same optimisation (unrolling one call to f) will work here,
  65 despite the cast.  See mk_alt_env in the Case branch of libCase.
  66
  67
  68 Note [Only functions!]
  69 ~~~~~~~~~~~~~~~~~~~~~~
  70 Consider the following code
  71
  72        f = g (case v of V a b -> a : t f)
  73
  74 where g is expensive. If we aren't careful, liberate case will turn this into
  75
  76        f = g (case v of
  77                V a b -> a : t (letrec f = g (case v of V a b -> a : f t)
  78                                 in f)
  79              )
  80
  81 Yikes! We evaluate g twice. This leads to a O(2^n) explosion
  82 if g calls back to the same code recursively.
  83
  84 Solution: make sure that we only do the liberate-case thing on *functions*
  85
  86 To think about (Apr 94)
  87 ~~~~~~~~~~~~~~
  88 Main worry: duplicating code excessively.  At the moment we duplicate
  89 the entire binding group once at each recursive call.  But there may
  90 be a group of recursive calls which share a common set of evaluated
  91 free variables, in which case the duplication is a plain waste.
  92
  93 Another thing we could consider adding is some unfold-threshold thing,
  94 so that we'll only duplicate if the size of the group rhss isn't too
  95 big.
  96
  97 Data types
  98 ~~~~~~~~~~
  99 The ``level'' of a binder tells how many
 100 recursive defns lexically enclose the binding
 101 A recursive defn "encloses" its RHS, not its
 102 scope.  For example:
 103 \begin{verbatim}
 104         letrec f = let g = ... in ...
 105         in
 106         let h = ...
 107         in ...
 108 \end{verbatim}
 109 Here, the level of @f@ is zero, the level of @g@ is one,
 110 and the level of @h@ is zero (NB not one).
 111
 112
 113 %************************************************************************
 114 %*                                                                      *
 115          Top-level code
 116 %*                                                                      *
 117 %************************************************************************
 118
 119 \begin{code}
 120 liberateCase :: DynFlags -> [CoreBind] -> [CoreBind]
 121 liberateCase dflags binds = do_prog (initEnv dflags) binds
 122   where
 123     do_prog _   [] = []
 124     do_prog env (bind:binds) = bind' : do_prog env' binds
 125                              where
 126                                (env', bind') = libCaseBind env bind
 127 \end{code}
 128
 129
 130 %************************************************************************
 131 %*                                                                      *
 132          Main payload
 133 %*                                                                      *
 134 %************************************************************************
 135
 136 Bindings
 137 ~~~~~~~~
 138 \begin{code}
 139 libCaseBind :: LibCaseEnv -> CoreBind -> (LibCaseEnv, CoreBind)
 140
 141 libCaseBind env (NonRec binder rhs)
 142   = (addBinders env [binder], NonRec binder (libCase env rhs))
 143
 144 libCaseBind env (Rec pairs)
 145   = (env_body, Rec pairs')
 146   where
 147     binders = map fst pairs
 148
 149     env_body = addBinders env binders
 150
 151     pairs' = [(binder, libCase env_rhs rhs) | (binder,rhs) <- pairs]
 152
 153         -- We extend the rec-env by binding each Id to its rhs, first
 154         -- processing the rhs with an *un-extended* environment, so
 155         -- that the same process doesn't occur for ever!
 156     env_rhs = addRecBinds env [ (localiseId binder, libCase env_body rhs)
 157                               | (binder, rhs) <- pairs
 158                               , rhs_small_enough binder rhs ]
 159         -- localiseID : see Note [Need to localiseId in libCaseBind]
 160
 161
 162     rhs_small_enough id rhs     -- Note [Small enough]
 163         =  idArity id > 0       -- Note [Only functions!]
 164         && maybe True (\size -> couldBeSmallEnoughToInline size rhs)
 165                       (bombOutSize env)
 166 \end{code}
 167
 168 Note [Need to localiseId in libCaseBind]
 169 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 170 The call to localiseId is needed for two subtle reasons
 171 (a)  Reset the export flags on the binders so
 172         that we don't get name clashes on exported things if the
 173         local binding floats out to top level.  This is most unlikely
 174         to happen, since the whole point concerns free variables.
 175         But resetting the export flag is right regardless.
 176
 177 (b)  Make the name an Internal one.  External Names should never be
 178         nested; if it were floated to the top level, we'd get a name
 179         clash at code generation time.
 180
 181 Note [Small enough]
 182 ~~~~~~~~~~~~~~~~~~~
 183 Consider
 184   \fv. letrec
 185          f = \x. BIG...(case fv of { (a,b) -> ...g.. })...
 186          g = \y. SMALL...f...
 187 Then we *can* do liberate-case on g (small RHS) but not for f (too big).
 188 But we can choose on a item-by-item basis, and that's what the
 189 rhs_small_enough call in the comprehension for env_rhs does.
 190
 191 Expressions
 192 ~~~~~~~~~~~
 193
 194 \begin{code}
 195 libCase :: LibCaseEnv
 196         -> CoreExpr
 197         -> CoreExpr
 198
 199 libCase env (Var v)             = libCaseId env v
 200 libCase _   (Lit lit)           = Lit lit
 201 libCase _   (Type ty)           = Type ty
 202 libCase _   (Coercion co)       = Coercion co
 203 libCase env (App fun arg)       = App (libCase env fun) (libCase env arg)
 204 libCase env (Note note body)    = Note note (libCase env body)
 205 libCase env (Cast e co)         = Cast (libCase env e) co
 206
 207 libCase env (Lam binder body)
 208   = Lam binder (libCase (addBinders env [binder]) body)
 209
 210 libCase env (Let bind body)
 211   = Let bind' (libCase env_body body)
 212   where
 213     (env_body, bind') = libCaseBind env bind
 214
 215 libCase env (Case scrut bndr ty alts)
 216   = Case (libCase env scrut) bndr ty (map (libCaseAlt env_alts) alts)
 217   where
 218     env_alts = addBinders (mk_alt_env scrut) [bndr]
 219     mk_alt_env (Var scrut_var) = addScrutedVar env scrut_var
 220     mk_alt_env (Cast scrut _)  = mk_alt_env scrut       -- Note [Scrutinee with cast]
 221     mk_alt_env _               = env
 222
 223 libCaseAlt :: LibCaseEnv -> (AltCon, [CoreBndr], CoreExpr)
 224                          -> (AltCon, [CoreBndr], CoreExpr)
 225 libCaseAlt env (con,args,rhs) = (con, args, libCase (addBinders env args) rhs)
 226 \end{code}
 227
 228
 229 Ids
 230 ~~~
 231 \begin{code}
 232 libCaseId :: LibCaseEnv -> Id -> CoreExpr
 233 libCaseId env v
 234   | Just the_bind <- lookupRecId env v  -- It's a use of a recursive thing
 235   , notNull free_scruts                 -- with free vars scrutinised in RHS
 236   = Let the_bind (Var v)
 237
 238   | otherwise
 239   = Var v
 240
 241   where
 242     rec_id_level = lookupLevel env v
 243     free_scruts  = freeScruts env rec_id_level
 244
 245 freeScruts :: LibCaseEnv
 246            -> LibCaseLevel      -- Level of the recursive Id
 247            -> [Id]              -- Ids that are scrutinised between the binding
 248                                 -- of the recursive Id and here
 249 freeScruts env rec_bind_lvl
 250   = [v | (v, scrut_bind_lvl, scrut_at_lvl) <- lc_scruts env
 251        , scrut_bind_lvl <= rec_bind_lvl
 252        , scrut_at_lvl > rec_bind_lvl]
 253         -- Note [When to specialise]
 254         -- Note [Avoiding fruitless liberate-case]
 255 \end{code}
 256
 257 Note [When to specialise]
 258 ~~~~~~~~~~~~~~~~~~~~~~~~~
 259 Consider
 260   f = \x. letrec g = \y. case x of
 261                            True  -> ... (f a) ...
 262                            False -> ... (g b) ...
 263
 264 We get the following levels
 265           f  0
 266           x  1
 267           g  1
 268           y  2
 269
 270 Then 'x' is being scrutinised at a deeper level than its binding, so
 271 it's added to lc_sruts:  [(x,1)]
 272
 273 We do *not* want to specialise the call to 'f', becuase 'x' is not free
 274 in 'f'.  So here the bind-level of 'x' (=1) is not <= the bind-level of 'f' (=0).
 275
 276 We *do* want to specialise the call to 'g', because 'x' is free in g.
 277 Here the bind-level of 'x' (=1) is <= the bind-level of 'g' (=1).
 278
 279 Note [Avoiding fruitless liberate-case]
 280 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 281 Consider also:
 282   f = \x. case top_lvl_thing of
 283                 I# _ -> let g = \y. ... g ...
 284                         in ...
 285
 286 Here, top_lvl_thing is scrutinised at a level (1) deeper than its
 287 binding site (0).  Nevertheless, we do NOT want to specialise the call
 288 to 'g' because all the structure in its free variables is already
 289 visible at the definition site for g.  Hence, when considering specialising
 290 an occurrence of 'g', we want to check that there's a scruted-var v st
 291
 292    a) v's binding site is *outside* g
 293    b) v's scrutinisation site is *inside* g
 294
 295
 296 %************************************************************************
 297 %*                                                                      *
 298         Utility functions
 299 %*                                                                      *
 300 %************************************************************************
 301
 302 \begin{code}
 303 addBinders :: LibCaseEnv -> [CoreBndr] -> LibCaseEnv
 304 addBinders env@(LibCaseEnv { lc_lvl = lvl, lc_lvl_env = lvl_env }) binders
 305   = env { lc_lvl_env = lvl_env' }
 306   where
 307     lvl_env' = extendVarEnvList lvl_env (binders `zip` repeat lvl)
 308
 309 addRecBinds :: LibCaseEnv -> [(Id,CoreExpr)] -> LibCaseEnv
 310 addRecBinds env@(LibCaseEnv {lc_lvl = lvl, lc_lvl_env = lvl_env,
 311                              lc_rec_env = rec_env}) pairs
 312   = env { lc_lvl = lvl', lc_lvl_env = lvl_env', lc_rec_env = rec_env' }
 313   where
 314     lvl'     = lvl + 1
 315     lvl_env' = extendVarEnvList lvl_env [(binder,lvl) | (binder,_) <- pairs]
 316     rec_env' = extendVarEnvList rec_env [(binder, Rec pairs) | (binder,_) <- pairs]
 317
 318 addScrutedVar :: LibCaseEnv
 319               -> Id             -- This Id is being scrutinised by a case expression
 320               -> LibCaseEnv
 321
 322 addScrutedVar env@(LibCaseEnv { lc_lvl = lvl, lc_lvl_env = lvl_env,
 323                                 lc_scruts = scruts }) scrut_var
 324   | bind_lvl < lvl
 325   = env { lc_scruts = scruts' }
 326         -- Add to scruts iff the scrut_var is being scrutinised at
 327         -- a deeper level than its defn
 328
 329   | otherwise = env
 330   where
 331     scruts'  = (scrut_var, bind_lvl, lvl) : scruts
 332     bind_lvl = case lookupVarEnv lvl_env scrut_var of
 333                  Just lvl -> lvl
 334                  Nothing  -> topLevel
 335
 336 lookupRecId :: LibCaseEnv -> Id -> Maybe CoreBind
 337 lookupRecId env id = lookupVarEnv (lc_rec_env env) id
 338
 339 lookupLevel :: LibCaseEnv -> Id -> LibCaseLevel
 340 lookupLevel env id
 341   = case lookupVarEnv (lc_lvl_env env) id of
 342       Just lvl -> lvl
 343       Nothing  -> topLevel
 344 \end{code}
 345
 346 %************************************************************************
 347 %*                                                                      *
 348          The environment
 349 %*                                                                      *
 350 %************************************************************************
 351
 352 \begin{code}
 353 type LibCaseLevel = Int
 354
 355 topLevel :: LibCaseLevel
 356 topLevel = 0
 357 \end{code}
 358
 359 \begin{code}
 360 data LibCaseEnv
 361   = LibCaseEnv {
 362         lc_size :: Maybe Int,   -- Bomb-out size for deciding if
 363                                 -- potential liberatees are too big.
 364                                 -- (passed in from cmd-line args)
 365
 366         lc_lvl :: LibCaseLevel, -- Current level
 367                 -- The level is incremented when (and only when) going
 368                 -- inside the RHS of a (sufficiently small) recursive
 369                 -- function.
 370
 371         lc_lvl_env :: IdEnv LibCaseLevel,
 372                 -- Binds all non-top-level in-scope Ids (top-level and
 373                 -- imported things have a level of zero)
 374
 375         lc_rec_env :: IdEnv CoreBind,
 376                 -- Binds *only* recursively defined ids, to their own
 377                 -- binding group, and *only* in their own RHSs
 378
 379         lc_scruts :: [(Id, LibCaseLevel, LibCaseLevel)]
 380                 -- Each of these Ids was scrutinised by an enclosing
 381                 -- case expression, at a level deeper than its binding
 382                 -- level.
 383                 --
 384                 -- The first LibCaseLevel is the *binding level* of
 385                 --   the scrutinised Id,
 386                 -- The second is the level *at which it was scrutinised*.
 387                 --   (see Note [Avoiding fruitless liberate-case])
 388                 -- The former is a bit redundant, since you could always
 389                 -- look it up in lc_lvl_env, but it's just cached here
 390                 --
 391                 -- The order is insignificant; it's a bag really
 392                 --
 393                 -- There's one element per scrutinisation;
 394                 --    in principle the same Id may appear multiple times,
 395                 --    although that'd be unusual:
 396                 --       case x of { (a,b) -> ....(case x of ...) .. }
 397         }
 398
 399 initEnv :: DynFlags -> LibCaseEnv
 400 initEnv dflags
 401   = LibCaseEnv { lc_size = liberateCaseThreshold dflags,
 402                  lc_lvl = 0,
 403                  lc_lvl_env = emptyVarEnv,
 404                  lc_rec_env = emptyVarEnv,
 405                  lc_scruts = [] }
 406
 407 bombOutSize :: LibCaseEnv -> Maybe Int
 408 bombOutSize = lc_size
 409 \end{code}
 410