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