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