ghc/compiler/simplCore/LiberateCase.lhs

   1 %
   2 % (c) The AQUA Project, Glasgow University, 1994
   3 %
   4 \section[LiberateCase]{Unroll recursion to allow evals to be lifted from a loop}
   5
   6
   7 \begin{code}
   8 #include "HsVersions.h"
   9
  10 module LiberateCase ( liberateCase ) where
  11
  12 IMPORT_Trace
  13
  14 import Id               ( localiseId, toplevelishId{-debugging-} )
  15 import IdEnv
  16 import Maybes
  17 import Outputable
  18 import PlainCore
  19 import Pretty
  20 import SimplEnv         ( UnfoldingGuidance(..) )
  21 import Util
  22 \end{code}
  23
  24 This module walks over @Core@, and looks for @case@ on free variables.
  25 The criterion is:
  26         if there is case on a free on the route to the recursive call,
  27         then the recursive call is replaced with an unfolding.
  28
  29 Example
  30
  31 \begin{verbatim}
  32 f = \ t -> case v of
  33                V a b -> a : f t
  34 \end{verbatim}
  35
  36 => the inner f is replaced.
  37
  38 \begin{verbatim}
  39 f = \ t -> case v of
  40                V a b -> a : (letrec
  41                                 f =  \ t -> case v of
  42                                                V a b -> a : f t
  43                              in f) t
  44 \end{verbatim}
  45 (note the NEED for shadowing)
  46
  47 => Run Andr\'e's wonder pass ...
  48 \begin{verbatim}
  49 f = \ t -> case v of
  50                V a b -> a : (letrec
  51                                 f = \ t -> a : f t
  52                              in f t)
  53 \begin{verbatim}
  54 Better code, because 'a' is  free inside the inner letrec, rather
  55 than needing projection from v.
  56
  57
  58 To think about (Apr 94)
  59 ~~~~~~~~~~~~~~
  60
  61 Main worry: duplicating code excessively.  At the moment we duplicate
  62 the entire binding group once at each recursive call.  But there may
  63 be a group of recursive calls which share a common set of evaluated
  64 free variables, in which case the duplication is a plain waste.
  65
  66 Another thing we could consider adding is some unfold-threshold thing,
  67 so that we'll only duplicate if the size of the group rhss isn't too
  68 big.
  69
  70 Data types
  71 ~~~~~~~~~~
  72
  73 The ``level'' of a binder tells how many
  74 recursive defns lexically enclose the binding
  75 A recursive defn "encloses" its RHS, not its
  76 scope.  For example:
  77 \begin{verbatim}
  78         letrec f = let g = ... in ...
  79         in
  80         let h = ...
  81         in ...
  82 \end{verbatim}
  83 Here, the level of @f@ is zero, the level of @g@ is one,
  84 and the level of @h@ is zero (NB not one).
  85
  86 \begin{code}
  87 type LibCaseLevel = Int
  88
  89 topLevel :: LibCaseLevel
  90 topLevel = 0
  91 \end{code}
  92
  93 \begin{code}
  94 data LibCaseEnv
  95   = LibCaseEnv
  96         Int                     -- Bomb-out size for deciding if
  97                                 -- potential liberatees are too big.
  98                                 -- (passed in from cmd-line args)
  99
 100         LibCaseLevel            -- Current level
 101
 102         (IdEnv LibCaseLevel)    -- Binds all non-top-level in-scope Ids
 103                                 -- (top-level and imported things have
 104                                 -- a level of zero)
 105
 106         (IdEnv PlainCoreBinding)-- Binds *only* recursively defined
 107                                 -- Ids, to their own binding group,
 108                                 -- and *only* in their own RHSs
 109
 110         [(Id,LibCaseLevel)]     -- Each of these Ids was scrutinised by an
 111                                 -- enclosing case expression, with the
 112                                 -- specified number of enclosing
 113                                 -- recursive bindings; furthermore,
 114                                 -- the Id is bound at a lower level
 115                                 -- than the case expression.  The
 116                                 -- order is insignificant; it's a bag
 117                                 -- really
 118
 119 initEnv :: Int -> LibCaseEnv
 120 initEnv bomb_size = LibCaseEnv bomb_size 0 nullIdEnv nullIdEnv []
 121
 122 bombOutSize (LibCaseEnv bomb_size _ _ _ _) = bomb_size
 123 \end{code}
 124
 125
 126 Programs
 127 ~~~~~~~~
 128 \begin{code}
 129 liberateCase :: Int -> [PlainCoreBinding] -> [PlainCoreBinding]
 130 liberateCase bomb_size prog
 131   = do_prog (initEnv bomb_size) prog
 132   where
 133     do_prog env [] = []
 134     do_prog env (bind:binds) = bind' : do_prog env' binds
 135                              where
 136                                (env', bind') = libCaseBind env bind
 137 \end{code}
 138
 139 Bindings
 140 ~~~~~~~~
 141
 142 \begin{code}
 143 libCaseBind :: LibCaseEnv -> PlainCoreBinding -> (LibCaseEnv, PlainCoreBinding)
 144
 145 libCaseBind env (CoNonRec binder rhs)
 146   = (addBinders env [binder], CoNonRec binder (libCase env rhs))
 147
 148 libCaseBind env (CoRec pairs)
 149   = (env_body, CoRec pairs')
 150   where
 151     (binders, rhss) = unzip pairs
 152
 153     env_body = addBinders env binders
 154
 155     pairs' = [(binder, libCase env_rhs rhs) | (binder,rhs) <- pairs]
 156
 157     env_rhs = if all rhs_small_enough rhss then extended_env else env
 158
 159         -- We extend the rec-env by binding each Id to its rhs, first
 160         -- processing the rhs with an *un-extended* environment, so
 161         -- that the same process doesn't occur for ever!
 162
 163     extended_env
 164       = addRecBinds env [ (localiseId binder, libCase env_body rhs)
 165                         | (binder, rhs) <- pairs ]
 166
 167         -- Why "localiseId" above?  Because we're creating a new local
 168         -- copy of the original binding.  In particular, the original
 169         -- binding might have been for a TopLevId, and this copy clearly
 170         -- will not be top-level!
 171
 172         -- It is enough to change just the binder, because subsequent
 173         -- simplification will propagate the right info from the binder.
 174
 175         -- Why does it matter?  Because the codeGen keeps a separate
 176         -- environment for top-level Ids, and it is disastrous for it
 177         -- to think that something is top-level when it isn't.
 178
 179     rhs_small_enough rhs
 180       = case (calcUnfoldingGuidance True{-sccs OK-} lIBERATE_BOMB_SIZE rhs) of
 181           UnfoldNever -> False
 182           _           -> True   -- we didn't BOMB, so it must be OK
 183
 184     lIBERATE_BOMB_SIZE = bombOutSize env
 185 \end{code}
 186
 187
 188 Expressions
 189 ~~~~~~~~~~~
 190
 191 \begin{code}
 192 libCase :: LibCaseEnv
 193         -> PlainCoreExpr
 194         -> PlainCoreExpr
 195
 196 libCase env (CoLit lit)          = CoLit lit
 197 libCase env (CoVar v)            = mkCoLetsNoUnboxed (libCaseId env v) (CoVar v)
 198 libCase env (CoApp fun arg)      = mkCoLetsNoUnboxed (libCaseAtom env arg) (CoApp (libCase env fun) arg)
 199 libCase env (CoTyApp fun ty)     = CoTyApp (libCase env fun) ty
 200 libCase env (CoCon con tys args) = mkCoLetsNoUnboxed (libCaseAtoms env args) (CoCon con tys args)
 201 libCase env (CoPrim op tys args) = mkCoLetsNoUnboxed (libCaseAtoms env args) (CoPrim op tys args)
 202 libCase env (CoTyLam tyvar body) = CoTyLam tyvar (libCase env body)
 203 libCase env (CoSCC cc body)      = CoSCC cc (libCase env body)
 204
 205 libCase env (CoLam binders body)
 206   = CoLam binders (libCase env' body)
 207   where
 208     env' = addBinders env binders
 209
 210 libCase env (CoLet bind body)
 211   = CoLet bind' (libCase env_body body)
 212   where
 213     (env_body, bind') = libCaseBind env bind
 214
 215 libCase env (CoCase scrut alts)
 216   = CoCase (libCase env scrut) (libCaseAlts env_alts alts)
 217   where
 218     env_alts = case scrut of
 219                   CoVar scrut_var -> addScrutedVar env scrut_var
 220                   other           -> env
 221 \end{code}
 222
 223
 224 Case alternatives
 225 ~~~~~~~~~~~~~~~~~
 226
 227 \begin{code}
 228 libCaseAlts env (CoAlgAlts alts deflt)
 229   = CoAlgAlts (map do_alt alts) (libCaseDeflt env deflt)
 230   where
 231     do_alt (con,args,rhs) = (con, args, libCase (addBinders env args) rhs)
 232
 233 libCaseAlts env (CoPrimAlts alts deflt)
 234   = CoPrimAlts (map do_alt alts) (libCaseDeflt env deflt)
 235   where
 236     do_alt (lit,rhs) = (lit, libCase env rhs)
 237
 238 libCaseDeflt env CoNoDefault
 239    = CoNoDefault
 240 libCaseDeflt env (CoBindDefault binder rhs)
 241    = CoBindDefault binder (libCase (addBinders env [binder]) rhs)
 242 \end{code}
 243
 244 Atoms and Ids
 245 ~~~~~~~~~~~~~
 246 \begin{code}
 247 libCaseAtoms :: LibCaseEnv -> [PlainCoreAtom] -> [PlainCoreBinding]
 248 libCaseAtoms env atoms = concat [libCaseAtom env atom | atom <- atoms]
 249
 250 libCaseAtom :: LibCaseEnv -> PlainCoreAtom -> [PlainCoreBinding]
 251 libCaseAtom env (CoVarAtom arg_id) = libCaseId env arg_id
 252 libCaseAtom env (CoLitAtom lit)    = []
 253
 254 libCaseId :: LibCaseEnv -> Id -> [PlainCoreBinding]
 255 libCaseId env v
 256   | maybeToBool maybe_rec_bind &&       -- It's a use of a recursive thing
 257     there_are_free_scruts               -- with free vars scrutinised in RHS
 258   = [the_bind]
 259
 260   | otherwise
 261   = []
 262
 263   where
 264     maybe_rec_bind :: Maybe PlainCoreBinding    -- The binding of the recursive thingy
 265     maybe_rec_bind = lookupRecId env v
 266     Just the_bind = maybe_rec_bind
 267
 268     rec_id_level = lookupLevel env v
 269
 270     there_are_free_scruts = freeScruts env rec_id_level
 271 \end{code}
 272
 273
 274
 275 Utility functions
 276 ~~~~~~~~~~~~~~~~~
 277 \begin{code}
 278 addBinders :: LibCaseEnv -> [Id] -> LibCaseEnv
 279 addBinders (LibCaseEnv bomb lvl lvl_env rec_env scruts) binders
 280   = LibCaseEnv bomb lvl lvl_env' rec_env scruts
 281   where
 282     lvl_env' = growIdEnvList lvl_env (binders `zip` repeat lvl)
 283
 284 addRecBinds :: LibCaseEnv -> [(Id,PlainCoreExpr)] -> LibCaseEnv
 285 addRecBinds (LibCaseEnv bomb lvl lvl_env rec_env scruts) pairs
 286   = LibCaseEnv bomb lvl' lvl_env' rec_env' scruts
 287   where
 288     lvl'     = lvl + 1
 289     lvl_env' = growIdEnvList lvl_env [(binder,lvl) | (binder,_) <- pairs]
 290     rec_env' = growIdEnvList rec_env [(binder, CoRec pairs) | (binder,_) <- pairs]
 291
 292 addScrutedVar :: LibCaseEnv
 293               -> Id             -- This Id is being scrutinised by a case expression
 294               -> LibCaseEnv
 295
 296 addScrutedVar env@(LibCaseEnv bomb lvl lvl_env rec_env scruts) scrut_var
 297   | bind_lvl < lvl
 298   = LibCaseEnv bomb lvl lvl_env rec_env scruts'
 299         -- Add to scruts iff the scrut_var is being scrutinised at
 300         -- a deeper level than its defn
 301
 302   | otherwise = env
 303   where
 304     scruts'  = (scrut_var, lvl) : scruts
 305     bind_lvl = case lookupIdEnv lvl_env scrut_var of
 306                  Just lvl -> lvl
 307                  Nothing  -> --false: ASSERT(toplevelishId scrut_var)
 308                              topLevel
 309
 310 lookupRecId :: LibCaseEnv -> Id -> Maybe PlainCoreBinding
 311 lookupRecId (LibCaseEnv bomb lvl lvl_env rec_env scruts) id
 312 #ifndef DEBUG
 313   = lookupIdEnv rec_env id
 314 #else
 315   = case (lookupIdEnv rec_env id) of
 316       xxx@(Just _) -> xxx
 317       xxx          -> --false: ASSERT(toplevelishId id)
 318                       xxx
 319 #endif
 320
 321 lookupLevel :: LibCaseEnv -> Id -> LibCaseLevel
 322 lookupLevel (LibCaseEnv bomb lvl lvl_env rec_env scruts) id
 323   = case lookupIdEnv lvl_env id of
 324       Just lvl -> lvl
 325       Nothing  -> ASSERT(toplevelishId id)
 326                   topLevel
 327
 328 freeScruts :: LibCaseEnv
 329            -> LibCaseLevel      -- Level of the recursive Id
 330            -> Bool              -- True <=> there is an enclosing case of a variable
 331                                 -- bound outside (ie level <=) the recursive Id.
 332 freeScruts (LibCaseEnv bomb lvl lvl_env rec_env scruts) rec_bind_lvl
 333   = not (null free_scruts)
 334   where
 335     free_scruts = [v | (v,lvl) <- scruts, lvl > rec_bind_lvl]
 336 \end{code}