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