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