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 env (App fun arg)       = App (libCase env fun) (libCase env arg)
 203 libCase env (Note note body)    = Note note (libCase env body)
 204 libCase env (Cast e co)         = Cast (libCase env e) co
 205
 206 libCase env (Lam binder body)
 207   = Lam binder (libCase (addBinders env [binder]) body)
 208
 209 libCase env (Let bind body)
 210   = Let bind' (libCase env_body body)
 211   where
 212     (env_body, bind') = libCaseBind env bind
 213
 214 libCase env (Case scrut bndr ty alts)
 215   = Case (libCase env scrut) bndr ty (map (libCaseAlt env_alts) alts)
 216   where
 217     env_alts = addBinders (mk_alt_env scrut) [bndr]
 218     mk_alt_env (Var scrut_var) = addScrutedVar env scrut_var
 219     mk_alt_env (Cast scrut _)  = mk_alt_env scrut       -- Note [Scrutinee with cast]
 220     mk_alt_env _               = env
 221
 222 libCaseAlt :: LibCaseEnv -> (AltCon, [CoreBndr], CoreExpr)
 223                          -> (AltCon, [CoreBndr], CoreExpr)
 224 libCaseAlt env (con,args,rhs) = (con, args, libCase (addBinders env args) rhs)
 225 \end{code}
 226
 227
 228 Ids
 229 ~~~
 230 \begin{code}
 231 libCaseId :: LibCaseEnv -> Id -> CoreExpr
 232 libCaseId env v
 233   | Just the_bind <- lookupRecId env v  -- It's a use of a recursive thing
 234   , notNull free_scruts                 -- with free vars scrutinised in RHS
 235   = Let the_bind (Var v)
 236
 237   | otherwise
 238   = Var v
 239
 240   where
 241     rec_id_level = lookupLevel env v
 242     free_scruts  = freeScruts env rec_id_level
 243
 244 freeScruts :: LibCaseEnv
 245            -> LibCaseLevel      -- Level of the recursive Id
 246            -> [Id]              -- Ids that are scrutinised between the binding
 247                                 -- of the recursive Id and here
 248 freeScruts env rec_bind_lvl
 249   = [v | (v, scrut_bind_lvl, scrut_at_lvl) <- lc_scruts env
 250        , scrut_bind_lvl <= rec_bind_lvl
 251        , scrut_at_lvl > rec_bind_lvl]
 252         -- Note [When to specialise]
 253         -- Note [Avoiding fruitless liberate-case]
 254 \end{code}
 255
 256 Note [When to specialise]
 257 ~~~~~~~~~~~~~~~~~~~~~~~~~
 258 Consider
 259   f = \x. letrec g = \y. case x of
 260                            True  -> ... (f a) ...
 261                            False -> ... (g b) ...
 262
 263 We get the following levels
 264           f  0
 265           x  1
 266           g  1
 267           y  2
 268
 269 Then 'x' is being scrutinised at a deeper level than its binding, so
 270 it's added to lc_sruts:  [(x,1)]
 271
 272 We do *not* want to specialise the call to 'f', becuase 'x' is not free
 273 in 'f'.  So here the bind-level of 'x' (=1) is not <= the bind-level of 'f' (=0).
 274
 275 We *do* want to specialise the call to 'g', because 'x' is free in g.
 276 Here the bind-level of 'x' (=1) is <= the bind-level of 'g' (=1).
 277
 278 Note [Avoiding fruitless liberate-case]
 279 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 280 Consider also:
 281   f = \x. case top_lvl_thing of
 282                 I# _ -> let g = \y. ... g ...
 283                         in ...
 284
 285 Here, top_lvl_thing is scrutinised at a level (1) deeper than its
 286 binding site (0).  Nevertheless, we do NOT want to specialise the call
 287 to 'g' because all the structure in its free variables is already
 288 visible at the definition site for g.  Hence, when considering specialising
 289 an occurrence of 'g', we want to check that there's a scruted-var v st
 290
 291    a) v's binding site is *outside* g
 292    b) v's scrutinisation site is *inside* g
 293
 294
 295 %************************************************************************
 296 %*                                                                      *
 297         Utility functions
 298 %*                                                                      *
 299 %************************************************************************
 300
 301 \begin{code}
 302 addBinders :: LibCaseEnv -> [CoreBndr] -> LibCaseEnv
 303 addBinders env@(LibCaseEnv { lc_lvl = lvl, lc_lvl_env = lvl_env }) binders
 304   = env { lc_lvl_env = lvl_env' }
 305   where
 306     lvl_env' = extendVarEnvList lvl_env (binders `zip` repeat lvl)
 307
 308 addRecBinds :: LibCaseEnv -> [(Id,CoreExpr)] -> LibCaseEnv
 309 addRecBinds env@(LibCaseEnv {lc_lvl = lvl, lc_lvl_env = lvl_env,
 310                              lc_rec_env = rec_env}) pairs
 311   = env { lc_lvl = lvl', lc_lvl_env = lvl_env', lc_rec_env = rec_env' }
 312   where
 313     lvl'     = lvl + 1
 314     lvl_env' = extendVarEnvList lvl_env [(binder,lvl) | (binder,_) <- pairs]
 315     rec_env' = extendVarEnvList rec_env [(binder, Rec pairs) | (binder,_) <- pairs]
 316
 317 addScrutedVar :: LibCaseEnv
 318               -> Id             -- This Id is being scrutinised by a case expression
 319               -> LibCaseEnv
 320
 321 addScrutedVar env@(LibCaseEnv { lc_lvl = lvl, lc_lvl_env = lvl_env,
 322                                 lc_scruts = scruts }) scrut_var
 323   | bind_lvl < lvl
 324   = env { lc_scruts = scruts' }
 325         -- Add to scruts iff the scrut_var is being scrutinised at
 326         -- a deeper level than its defn
 327
 328   | otherwise = env
 329   where
 330     scruts'  = (scrut_var, bind_lvl, lvl) : scruts
 331     bind_lvl = case lookupVarEnv lvl_env scrut_var of
 332                  Just lvl -> lvl
 333                  Nothing  -> topLevel
 334
 335 lookupRecId :: LibCaseEnv -> Id -> Maybe CoreBind
 336 lookupRecId env id = lookupVarEnv (lc_rec_env env) id
 337
 338 lookupLevel :: LibCaseEnv -> Id -> LibCaseLevel
 339 lookupLevel env id
 340   = case lookupVarEnv (lc_lvl_env env) id of
 341       Just lvl -> lvl
 342       Nothing  -> topLevel
 343 \end{code}
 344
 345 %************************************************************************
 346 %*                                                                      *
 347          The environment
 348 %*                                                                      *
 349 %************************************************************************
 350
 351 \begin{code}
 352 type LibCaseLevel = Int
 353
 354 topLevel :: LibCaseLevel
 355 topLevel = 0
 356 \end{code}
 357
 358 \begin{code}
 359 data LibCaseEnv
 360   = LibCaseEnv {
 361         lc_size :: Maybe Int,   -- Bomb-out size for deciding if
 362                                 -- potential liberatees are too big.
 363                                 -- (passed in from cmd-line args)
 364
 365         lc_lvl :: LibCaseLevel, -- Current level
 366                 -- The level is incremented when (and only when) going
 367                 -- inside the RHS of a (sufficiently small) recursive
 368                 -- function.
 369
 370         lc_lvl_env :: IdEnv LibCaseLevel,
 371                 -- Binds all non-top-level in-scope Ids (top-level and
 372                 -- imported things have a level of zero)
 373
 374         lc_rec_env :: IdEnv CoreBind,
 375                 -- Binds *only* recursively defined ids, to their own
 376                 -- binding group, and *only* in their own RHSs
 377
 378         lc_scruts :: [(Id, LibCaseLevel, LibCaseLevel)]
 379                 -- Each of these Ids was scrutinised by an enclosing
 380                 -- case expression, at a level deeper than its binding
 381                 -- level.
 382                 --
 383                 -- The first LibCaseLevel is the *binding level* of
 384                 --   the scrutinised Id,
 385                 -- The second is the level *at which it was scrutinised*.
 386                 --   (see Note [Avoiding fruitless liberate-case])
 387                 -- The former is a bit redundant, since you could always
 388                 -- look it up in lc_lvl_env, but it's just cached here
 389                 --
 390                 -- The order is insignificant; it's a bag really
 391                 --
 392                 -- There's one element per scrutinisation;
 393                 --    in principle the same Id may appear multiple times,
 394                 --    although that'd be unusual:
 395                 --       case x of { (a,b) -> ....(case x of ...) .. }
 396         }
 397
 398 initEnv :: DynFlags -> LibCaseEnv
 399 initEnv dflags
 400   = LibCaseEnv { lc_size = liberateCaseThreshold dflags,
 401                  lc_lvl = 0,
 402                  lc_lvl_env = emptyVarEnv,
 403                  lc_rec_env = emptyVarEnv,
 404                  lc_scruts = [] }
 405
 406 bombOutSize :: LibCaseEnv -> Maybe Int
 407 bombOutSize = lc_size
 408 \end{code}
 409