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      ( opt_D_verbose_core2core, opt_LiberateCaseThreshold )
  12 import CoreLint         ( beginPass, 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 :: [CoreBind] -> IO [CoreBind]
 152 liberateCase binds
 153   = do {
 154         beginPass "Liberate case" ;
 155         let { binds' = do_prog (initEnv opt_LiberateCaseThreshold) binds } ;
 156         endPass "Liberate case"
 157                 opt_D_verbose_core2core         {- no specific flag for dumping -}
 158                 binds'
 159     }
 160   where
 161     do_prog env [] = []
 162     do_prog env (bind:binds) = bind' : do_prog env' binds
 163                              where
 164                                (env', bind') = libCaseBind env bind
 165 \end{code}
 166
 167 Bindings
 168 ~~~~~~~~
 169
 170 \begin{code}
 171 libCaseBind :: LibCaseEnv -> CoreBind -> (LibCaseEnv, CoreBind)
 172
 173 libCaseBind env (NonRec binder rhs)
 174   = (addBinders env [binder], NonRec binder (libCase env rhs))
 175
 176 libCaseBind env (Rec pairs)
 177   = (env_body, Rec pairs')
 178   where
 179     (binders, rhss) = unzip pairs
 180
 181     env_body = addBinders env binders
 182
 183     pairs' = [(binder, libCase env_rhs rhs) | (binder,rhs) <- pairs]
 184
 185     env_rhs = if all rhs_small_enough rhss then extended_env else env
 186
 187         -- We extend the rec-env by binding each Id to its rhs, first
 188         -- processing the rhs with an *un-extended* environment, so
 189         -- that the same process doesn't occur for ever!
 190
 191     extended_env
 192       = addRecBinds env [ (binder, libCase env_body rhs)
 193                         | (binder, rhs) <- pairs ]
 194
 195         -- Why "localiseId" above?  Because we're creating a new local
 196         -- copy of the original binding.  In particular, the original
 197         -- binding might have been for a top-level, and this copy clearly
 198         -- will not be top-level!
 199
 200         -- It is enough to change just the binder, because subsequent
 201         -- simplification will propagate the right info from the binder.
 202
 203         -- Why does it matter?  Because the codeGen keeps a separate
 204         -- environment for top-level Ids, and it is disastrous for it
 205         -- to think that something is top-level when it isn't.
 206         --
 207         -- [May 98: all this is now handled by SimplCore.tidyCore]
 208
 209     rhs_small_enough rhs = couldBeSmallEnoughToInline lIBERATE_BOMB_SIZE rhs
 210
 211     lIBERATE_BOMB_SIZE = bombOutSize env
 212 \end{code}
 213
 214
 215 Expressions
 216 ~~~~~~~~~~~
 217
 218 \begin{code}
 219 libCase :: LibCaseEnv
 220         -> CoreExpr
 221         -> CoreExpr
 222
 223 libCase env (Var v)             = libCaseId env v
 224 libCase env (Lit lit)           = Lit lit
 225 libCase env (Type ty)           = Type ty
 226 libCase env (App fun arg)       = App (libCase env fun) (libCase env arg)
 227 libCase env (Note note body)    = Note note (libCase env body)
 228
 229 libCase env (Lam binder body)
 230   = Lam binder (libCase (addBinders env [binder]) body)
 231
 232 libCase env (Let bind body)
 233   = Let bind' (libCase env_body body)
 234   where
 235     (env_body, bind') = libCaseBind env bind
 236
 237 libCase env (Case scrut bndr alts)
 238   = Case (libCase env scrut) bndr (map (libCaseAlt env_alts) alts)
 239   where
 240     env_alts = addBinders env [bndr]
 241     env_with_scrut = case scrut of
 242                         Var scrut_var -> addScrutedVar env scrut_var
 243                         other             -> env
 244
 245 libCaseAlt env (con,args,rhs) = (con, args, libCase (addBinders env args) rhs)
 246 \end{code}
 247
 248 Ids
 249 ~~~
 250 \begin{code}
 251 libCaseId :: LibCaseEnv -> Id -> CoreExpr
 252 libCaseId env v
 253   | maybeToBool maybe_rec_bind &&       -- It's a use of a recursive thing
 254     there_are_free_scruts               -- with free vars scrutinised in RHS
 255   = Let the_bind (Var v)
 256
 257   | otherwise
 258   = Var v
 259
 260   where
 261     maybe_rec_bind :: Maybe CoreBind    -- The binding of the recursive thingy
 262     maybe_rec_bind = lookupRecId env v
 263     Just the_bind  = maybe_rec_bind
 264
 265     rec_id_level = lookupLevel env v
 266
 267     there_are_free_scruts = freeScruts env rec_id_level
 268 \end{code}
 269
 270
 271
 272 Utility functions
 273 ~~~~~~~~~~~~~~~~~
 274 \begin{code}
 275 addBinders :: LibCaseEnv -> [CoreBndr] -> LibCaseEnv
 276 addBinders (LibCaseEnv bomb lvl lvl_env rec_env scruts) binders
 277   = LibCaseEnv bomb lvl lvl_env' rec_env scruts
 278   where
 279     lvl_env' = extendVarEnvList lvl_env (binders `zip` repeat lvl)
 280
 281 addRecBinds :: LibCaseEnv -> [(Id,CoreExpr)] -> LibCaseEnv
 282 addRecBinds (LibCaseEnv bomb lvl lvl_env rec_env scruts) pairs
 283   = LibCaseEnv bomb lvl' lvl_env' rec_env' scruts
 284   where
 285     lvl'     = lvl + 1
 286     lvl_env' = extendVarEnvList lvl_env [(binder,lvl) | (binder,_) <- pairs]
 287     rec_env' = extendVarEnvList rec_env [(binder, Rec pairs) | (binder,_) <- pairs]
 288
 289 addScrutedVar :: LibCaseEnv
 290               -> Id             -- This Id is being scrutinised by a case expression
 291               -> LibCaseEnv
 292
 293 addScrutedVar env@(LibCaseEnv bomb lvl lvl_env rec_env scruts) scrut_var
 294   | bind_lvl < lvl
 295   = LibCaseEnv bomb lvl lvl_env rec_env scruts'
 296         -- Add to scruts iff the scrut_var is being scrutinised at
 297         -- a deeper level than its defn
 298
 299   | otherwise = env
 300   where
 301     scruts'  = (scrut_var, lvl) : scruts
 302     bind_lvl = case lookupVarEnv lvl_env scrut_var of
 303                  Just lvl -> lvl
 304                  Nothing  -> topLevel
 305
 306 lookupRecId :: LibCaseEnv -> Id -> Maybe CoreBind
 307 lookupRecId (LibCaseEnv bomb lvl lvl_env rec_env scruts) id
 308 #ifndef DEBUG
 309   = lookupVarEnv rec_env id
 310 #else
 311   = case (lookupVarEnv rec_env id) of
 312       xxx@(Just _) -> xxx
 313       xxx          -> xxx
 314 #endif
 315
 316 lookupLevel :: LibCaseEnv -> Id -> LibCaseLevel
 317 lookupLevel (LibCaseEnv bomb lvl lvl_env rec_env scruts) id
 318   = case lookupVarEnv lvl_env id of
 319       Just lvl -> lvl
 320       Nothing  -> topLevel
 321
 322 freeScruts :: LibCaseEnv
 323            -> LibCaseLevel      -- Level of the recursive Id
 324            -> Bool              -- True <=> there is an enclosing case of a variable
 325                                 -- bound outside (ie level <=) the recursive Id.
 326 freeScruts (LibCaseEnv bomb lvl lvl_env rec_env scruts) rec_bind_lvl
 327   = not (null free_scruts)
 328   where
 329     free_scruts = [v | (v,lvl) <- scruts, lvl > rec_bind_lvl]
 330 \end{code}