ghc/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 CmdLineOpts      ( DynFlags, DynFlag(..), dopt, opt_LiberateCaseThreshold )
  12 import CoreLint         ( showPass, endPass )
  13 import CoreSyn
  14 import CoreUnfold       ( couldBeSmallEnoughToInline )
  15 import Var              ( Id )
  16 import VarEnv
  17 import Maybes
  18 \end{code}
  19
  20 This module walks over @Core@, and looks for @case@ on free variables.
  21 The criterion is:
  22         if there is case on a free on the route to the recursive call,
  23         then the recursive call is replaced with an unfolding.
  24
  25 Example
  26
  27 \begin{verbatim}
  28 f = \ t -> case v of
  29                V a b -> a : f t
  30 \end{verbatim}
  31
  32 => the inner f is replaced.
  33
  34 \begin{verbatim}
  35 f = \ t -> case v of
  36                V a b -> a : (letrec
  37                                 f =  \ t -> case v of
  38                                                V a b -> a : f t
  39                              in f) t
  40 \end{verbatim}
  41 (note the NEED for shadowing)
  42
  43 => Run Andr\'e's wonder pass ...
  44 \begin{verbatim}
  45 f = \ t -> case v of
  46                V a b -> a : (letrec
  47                                 f = \ t -> a : f t
  48                              in f t)
  49 \begin{verbatim}
  50 Better code, because 'a' is  free inside the inner letrec, rather
  51 than needing projection from v.
  52
  53 Other examples we'd like to catch with this kind of transformation
  54
  55         last []     = error
  56         last (x:[]) = x
  57         last (x:xs) = last xs
  58
  59 We'd like to avoid the redundant pattern match, transforming to
  60
  61         last [] = error
  62         last (x:[]) = x
  63         last (x:(y:ys)) = last' y ys
  64                 where
  65                   last' y []     = y
  66                   last' _ (y:ys) = last' y ys
  67
  68         (is this necessarily an improvement)
  69
  70
  71 Similarly drop:
  72
  73         drop n [] = []
  74         drop 0 xs = xs
  75         drop n (x:xs) = drop (n-1) xs
  76
  77 Would like to pass n along unboxed.
  78
  79
  80 To think about (Apr 94)
  81 ~~~~~~~~~~~~~~
  82
  83 Main worry: duplicating code excessively.  At the moment we duplicate
  84 the entire binding group once at each recursive call.  But there may
  85 be a group of recursive calls which share a common set of evaluated
  86 free variables, in which case the duplication is a plain waste.
  87
  88 Another thing we could consider adding is some unfold-threshold thing,
  89 so that we'll only duplicate if the size of the group rhss isn't too
  90 big.
  91
  92 Data types
  93 ~~~~~~~~~~
  94
  95 The ``level'' of a binder tells how many
  96 recursive defns lexically enclose the binding
  97 A recursive defn "encloses" its RHS, not its
  98 scope.  For example:
  99 \begin{verbatim}
 100         letrec f = let g = ... in ...
 101         in
 102         let h = ...
 103         in ...
 104 \end{verbatim}
 105 Here, the level of @f@ is zero, the level of @g@ is one,
 106 and the level of @h@ is zero (NB not one).
 107
 108 \begin{code}
 109 type LibCaseLevel = Int
 110
 111 topLevel :: LibCaseLevel
 112 topLevel = 0
 113 \end{code}
 114
 115 \begin{code}
 116 data LibCaseEnv
 117   = LibCaseEnv
 118         Int                     -- Bomb-out size for deciding if
 119                                 -- potential liberatees are too big.
 120                                 -- (passed in from cmd-line args)
 121
 122         LibCaseLevel            -- Current level
 123
 124         (IdEnv LibCaseLevel)    -- Binds all non-top-level in-scope Ids
 125                                 -- (top-level and imported things have
 126                                 -- a level of zero)
 127
 128         (IdEnv CoreBind)-- Binds *only* recursively defined
 129                                 -- Ids, to their own binding group,
 130                                 -- and *only* in their own RHSs
 131
 132         [(Id,LibCaseLevel)]     -- Each of these Ids was scrutinised by an
 133                                 -- enclosing case expression, with the
 134                                 -- specified number of enclosing
 135                                 -- recursive bindings; furthermore,
 136                                 -- the Id is bound at a lower level
 137                                 -- than the case expression.  The
 138                                 -- order is insignificant; it's a bag
 139                                 -- really
 140
 141 initEnv :: Int -> LibCaseEnv
 142 initEnv bomb_size = LibCaseEnv bomb_size 0 emptyVarEnv emptyVarEnv []
 143
 144 bombOutSize (LibCaseEnv bomb_size _ _ _ _) = bomb_size
 145 \end{code}
 146
 147
 148 Programs
 149 ~~~~~~~~
 150 \begin{code}
 151 liberateCase :: DynFlags -> [CoreBind] -> IO [CoreBind]
 152 liberateCase dflags binds
 153   = do {
 154         showPass dflags "Liberate case" ;
 155         let { binds' = do_prog (initEnv opt_LiberateCaseThreshold) binds } ;
 156         endPass dflags "Liberate case"
 157                 (dopt Opt_D_verbose_core2core dflags)
 158                                 {- no specific flag for dumping -}
 159                 binds'
 160     }
 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 Bindings
 169 ~~~~~~~~
 170
 171 \begin{code}
 172 libCaseBind :: LibCaseEnv -> CoreBind -> (LibCaseEnv, CoreBind)
 173
 174 libCaseBind env (NonRec binder rhs)
 175   = (addBinders env [binder], NonRec binder (libCase env rhs))
 176
 177 libCaseBind env (Rec pairs)
 178   = (env_body, Rec pairs')
 179   where
 180     (binders, rhss) = unzip pairs
 181
 182     env_body = addBinders env binders
 183
 184     pairs' = [(binder, libCase env_rhs rhs) | (binder,rhs) <- pairs]
 185
 186     env_rhs = if all rhs_small_enough rhss then extended_env else env
 187
 188         -- We extend the rec-env by binding each Id to its rhs, first
 189         -- processing the rhs with an *un-extended* environment, so
 190         -- that the same process doesn't occur for ever!
 191
 192     extended_env
 193       = addRecBinds env [ (binder, libCase env_body rhs)
 194                         | (binder, rhs) <- pairs ]
 195
 196         -- Why "localiseId" above?  Because we're creating a new local
 197         -- copy of the original binding.  In particular, the original
 198         -- binding might have been for a top-level, and this copy clearly
 199         -- will not be top-level!
 200
 201         -- It is enough to change just the binder, because subsequent
 202         -- simplification will propagate the right info from the binder.
 203
 204         -- Why does it matter?  Because the codeGen keeps a separate
 205         -- environment for top-level Ids, and it is disastrous for it
 206         -- to think that something is top-level when it isn't.
 207         --
 208         -- [May 98: all this is now handled by SimplCore.tidyCore]
 209
 210     rhs_small_enough rhs = couldBeSmallEnoughToInline lIBERATE_BOMB_SIZE rhs
 211
 212     lIBERATE_BOMB_SIZE = bombOutSize env
 213 \end{code}
 214
 215
 216 Expressions
 217 ~~~~~~~~~~~
 218
 219 \begin{code}
 220 libCase :: LibCaseEnv
 221         -> CoreExpr
 222         -> CoreExpr
 223
 224 libCase env (Var v)             = libCaseId env v
 225 libCase env (Lit lit)           = Lit lit
 226 libCase env (Type ty)           = Type ty
 227 libCase env (App fun arg)       = App (libCase env fun) (libCase env arg)
 228 libCase env (Note note body)    = Note note (libCase env body)
 229
 230 libCase env (Lam binder body)
 231   = Lam binder (libCase (addBinders env [binder]) body)
 232
 233 libCase env (Let bind body)
 234   = Let bind' (libCase env_body body)
 235   where
 236     (env_body, bind') = libCaseBind env bind
 237
 238 libCase env (Case scrut bndr alts)
 239   = Case (libCase env scrut) bndr (map (libCaseAlt env_alts) alts)
 240   where
 241     env_alts = addBinders env_with_scrut [bndr]
 242     env_with_scrut = case scrut of
 243                         Var scrut_var -> addScrutedVar env scrut_var
 244                         other         -> env
 245
 246 libCaseAlt env (con,args,rhs) = (con, args, libCase (addBinders env args) rhs)
 247 \end{code}
 248
 249 Ids
 250 ~~~
 251 \begin{code}
 252 libCaseId :: LibCaseEnv -> Id -> CoreExpr
 253 libCaseId env v
 254   | maybeToBool maybe_rec_bind &&       -- It's a use of a recursive thing
 255     there_are_free_scruts               -- with free vars scrutinised in RHS
 256   = Let the_bind (Var v)
 257
 258   | otherwise
 259   = Var v
 260
 261   where
 262     maybe_rec_bind :: Maybe CoreBind    -- The binding of the recursive thingy
 263     maybe_rec_bind = lookupRecId env v
 264     Just the_bind  = maybe_rec_bind
 265
 266     rec_id_level = lookupLevel env v
 267
 268     there_are_free_scruts = freeScruts env rec_id_level
 269 \end{code}
 270
 271
 272
 273 Utility functions
 274 ~~~~~~~~~~~~~~~~~
 275 \begin{code}
 276 addBinders :: LibCaseEnv -> [CoreBndr] -> LibCaseEnv
 277 addBinders (LibCaseEnv bomb lvl lvl_env rec_env scruts) binders
 278   = LibCaseEnv bomb lvl lvl_env' rec_env scruts
 279   where
 280     lvl_env' = extendVarEnvList lvl_env (binders `zip` repeat lvl)
 281
 282 addRecBinds :: LibCaseEnv -> [(Id,CoreExpr)] -> LibCaseEnv
 283 addRecBinds (LibCaseEnv bomb lvl lvl_env rec_env scruts) pairs
 284   = LibCaseEnv bomb lvl' lvl_env' rec_env' scruts
 285   where
 286     lvl'     = lvl + 1
 287     lvl_env' = extendVarEnvList lvl_env [(binder,lvl) | (binder,_) <- pairs]
 288     rec_env' = extendVarEnvList rec_env [(binder, Rec pairs) | (binder,_) <- pairs]
 289
 290 addScrutedVar :: LibCaseEnv
 291               -> Id             -- This Id is being scrutinised by a case expression
 292               -> LibCaseEnv
 293
 294 addScrutedVar env@(LibCaseEnv bomb lvl lvl_env rec_env scruts) scrut_var
 295   | bind_lvl < lvl
 296   = LibCaseEnv bomb lvl lvl_env rec_env scruts'
 297         -- Add to scruts iff the scrut_var is being scrutinised at
 298         -- a deeper level than its defn
 299
 300   | otherwise = env
 301   where
 302     scruts'  = (scrut_var, lvl) : scruts
 303     bind_lvl = case lookupVarEnv lvl_env scrut_var of
 304                  Just lvl -> lvl
 305                  Nothing  -> topLevel
 306
 307 lookupRecId :: LibCaseEnv -> Id -> Maybe CoreBind
 308 lookupRecId (LibCaseEnv bomb lvl lvl_env rec_env scruts) id
 309 #ifndef DEBUG
 310   = lookupVarEnv rec_env id
 311 #else
 312   = case (lookupVarEnv rec_env id) of
 313       xxx@(Just _) -> xxx
 314       xxx          -> xxx
 315 #endif
 316
 317 lookupLevel :: LibCaseEnv -> Id -> LibCaseLevel
 318 lookupLevel (LibCaseEnv bomb lvl lvl_env rec_env scruts) id
 319   = case lookupVarEnv lvl_env id of
 320       Just lvl -> lvl
 321       Nothing  -> topLevel
 322
 323 freeScruts :: LibCaseEnv
 324            -> LibCaseLevel      -- Level of the recursive Id
 325            -> Bool              -- True <=> there is an enclosing case of a variable
 326                                 -- bound outside (ie level <=) the recursive Id.
 327 freeScruts (LibCaseEnv bomb lvl lvl_env rec_env scruts) rec_bind_lvl
 328   = not (null free_scruts)
 329   where
 330     free_scruts = [v | (v,lvl) <- scruts, lvl > rec_bind_lvl]
 331 \end{code}