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