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