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
  85 To think about (Apr 94)
  86 ~~~~~~~~~~~~~~
  87
  88 Main worry: duplicating code excessively.  At the moment we duplicate
  89 the entire binding group once at each recursive call.  But there may
  90 be a group of recursive calls which share a common set of evaluated
  91 free variables, in which case the duplication is a plain waste.
  92
  93 Another thing we could consider adding is some unfold-threshold thing,
  94 so that we'll only duplicate if the size of the group rhss isn't too
  95 big.
  96
  97 Data types
  98 ~~~~~~~~~~
  99
 100 The ``level'' of a binder tells how many
 101 recursive defns lexically enclose the binding
 102 A recursive defn "encloses" its RHS, not its
 103 scope.  For example:
 104 \begin{verbatim}
 105         letrec f = let g = ... in ...
 106         in
 107         let h = ...
 108         in ...
 109 \end{verbatim}
 110 Here, the level of @f@ is zero, the level of @g@ is one,
 111 and the level of @h@ is zero (NB not one).
 112
 113 \begin{code}
 114 type LibCaseLevel = Int
 115
 116 topLevel :: LibCaseLevel
 117 topLevel = 0
 118 \end{code}
 119
 120 \begin{code}
 121 data LibCaseEnv
 122   = LibCaseEnv
 123         Int                     -- Bomb-out size for deciding if
 124                                 -- potential liberatees are too big.
 125                                 -- (passed in from cmd-line args)
 126
 127         LibCaseLevel            -- Current level
 128
 129         (IdEnv LibCaseLevel)    -- Binds all non-top-level in-scope Ids
 130                                 -- (top-level and imported things have
 131                                 -- a level of zero)
 132
 133         (IdEnv CoreBind)        -- Binds *only* recursively defined
 134                                 -- Ids, to their own binding group,
 135                                 -- and *only* in their own RHSs
 136
 137         [(Id,LibCaseLevel)]     -- Each of these Ids was scrutinised by an
 138                                 -- enclosing case expression, with the
 139                                 -- specified number of enclosing
 140                                 -- recursive bindings; furthermore,
 141                                 -- the Id is bound at a lower level
 142                                 -- than the case expression.  The
 143                                 -- order is insignificant; it's a bag
 144                                 -- really
 145
 146 initEnv :: Int -> LibCaseEnv
 147 initEnv bomb_size = LibCaseEnv bomb_size 0 emptyVarEnv emptyVarEnv []
 148
 149 bombOutSize (LibCaseEnv bomb_size _ _ _ _) = bomb_size
 150 \end{code}
 151
 152
 153 Programs
 154 ~~~~~~~~
 155 \begin{code}
 156 liberateCase :: DynFlags -> [CoreBind] -> IO [CoreBind]
 157 liberateCase dflags binds
 158   = do {
 159         showPass dflags "Liberate case" ;
 160         let { binds' = do_prog (initEnv opt_LiberateCaseThreshold) binds } ;
 161         endPass dflags "Liberate case" Opt_D_verbose_core2core binds'
 162                                 {- no specific flag for dumping -}
 163     }
 164   where
 165     do_prog env [] = []
 166     do_prog env (bind:binds) = bind' : do_prog env' binds
 167                              where
 168                                (env', bind') = libCaseBind env bind
 169 \end{code}
 170
 171 Bindings
 172 ~~~~~~~~
 173
 174 \begin{code}
 175 libCaseBind :: LibCaseEnv -> CoreBind -> (LibCaseEnv, CoreBind)
 176
 177 libCaseBind env (NonRec binder rhs)
 178   = (addBinders env [binder], NonRec binder (libCase env rhs))
 179
 180 libCaseBind env (Rec pairs)
 181   = (env_body, Rec pairs')
 182   where
 183     (binders, rhss) = unzip pairs
 184
 185     env_body = addBinders env binders
 186
 187     pairs' = [(binder, libCase env_rhs rhs) | (binder,rhs) <- pairs]
 188
 189     env_rhs = if all rhs_small_enough rhss then extended_env else env
 190
 191         -- We extend the rec-env by binding each Id to its rhs, first
 192         -- processing the rhs with an *un-extended* environment, so
 193         -- that the same process doesn't occur for ever!
 194         --
 195     extended_env = addRecBinds env [ (adjust binder, libCase env_body rhs)
 196                                    | (binder, rhs) <- pairs ]
 197
 198         -- Two subtle things:
 199         -- (a)  Reset the export flags on the binders so
 200         --      that we don't get name clashes on exported things if the
 201         --      local binding floats out to top level.  This is most unlikely
 202         --      to happen, since the whole point concerns free variables.
 203         --      But resetting the export flag is right regardless.
 204         --
 205         -- (b)  Make the name an Internal one.  External Names should never be
 206         --      nested; if it were floated to the top level, we'd get a name
 207         --      clash at code generation time.
 208     adjust bndr = setIdNotExported (setIdName bndr (localiseName (idName bndr)))
 209
 210     rhs_small_enough rhs = couldBeSmallEnoughToInline lIBERATE_BOMB_SIZE rhs
 211     lIBERATE_BOMB_SIZE   = bombOutSize env
 212 \end{code}
 213
 214
 215 Expressions
 216 ~~~~~~~~~~~
 217
 218 \begin{code}
 219 libCase :: LibCaseEnv
 220         -> CoreExpr
 221         -> CoreExpr
 222
 223 libCase env (Var v)             = libCaseId env v
 224 libCase env (Lit lit)           = Lit lit
 225 libCase env (Type ty)           = Type ty
 226 libCase env (App fun arg)       = App (libCase env fun) (libCase env arg)
 227 libCase env (Note note body)    = Note note (libCase env body)
 228 libCase env (Cast e co)         = Cast (libCase env e) co
 229
 230 libCase env (Lam binder body)
 231   = Lam binder (libCase (addBinders env [binder]) body)
 232
 233 libCase env (Let bind body)
 234   = Let bind' (libCase env_body body)
 235   where
 236     (env_body, bind') = libCaseBind env bind
 237
 238 libCase env (Case scrut bndr ty alts)
 239   = Case (libCase env scrut) bndr ty (map (libCaseAlt env_alts) alts)
 240   where
 241     env_alts = addBinders env_with_scrut [bndr]
 242     env_with_scrut = case scrut of
 243                         Var scrut_var -> addScrutedVar env scrut_var
 244                         other         -> env
 245
 246 libCaseAlt env (con,args,rhs) = (con, args, libCase (addBinders env args) rhs)
 247 \end{code}
 248
 249 Ids
 250 ~~~
 251 \begin{code}
 252 libCaseId :: LibCaseEnv -> Id -> CoreExpr
 253 libCaseId env v
 254   | Just the_bind <- lookupRecId env v  -- It's a use of a recursive thing
 255   , notNull free_scruts                 -- with free vars scrutinised in RHS
 256   = Let the_bind (Var v)
 257
 258   | otherwise
 259   = Var v
 260
 261   where
 262     rec_id_level = lookupLevel env v
 263     free_scruts  = freeScruts env rec_id_level
 264 \end{code}
 265
 266
 267
 268 Utility functions
 269 ~~~~~~~~~~~~~~~~~
 270 \begin{code}
 271 addBinders :: LibCaseEnv -> [CoreBndr] -> LibCaseEnv
 272 addBinders (LibCaseEnv bomb lvl lvl_env rec_env scruts) binders
 273   = LibCaseEnv bomb lvl lvl_env' rec_env scruts
 274   where
 275     lvl_env' = extendVarEnvList lvl_env (binders `zip` repeat lvl)
 276
 277 addRecBinds :: LibCaseEnv -> [(Id,CoreExpr)] -> LibCaseEnv
 278 addRecBinds (LibCaseEnv bomb lvl lvl_env rec_env scruts) pairs
 279   = LibCaseEnv bomb lvl' lvl_env' rec_env' scruts
 280   where
 281     lvl'     = lvl + 1
 282     lvl_env' = extendVarEnvList lvl_env [(binder,lvl) | (binder,_) <- pairs]
 283     rec_env' = extendVarEnvList rec_env [(binder, Rec pairs) | (binder,_) <- pairs]
 284
 285 addScrutedVar :: LibCaseEnv
 286               -> Id             -- This Id is being scrutinised by a case expression
 287               -> LibCaseEnv
 288
 289 addScrutedVar env@(LibCaseEnv bomb lvl lvl_env rec_env scruts) scrut_var
 290   | bind_lvl < lvl
 291   = LibCaseEnv bomb lvl lvl_env rec_env scruts'
 292         -- Add to scruts iff the scrut_var is being scrutinised at
 293         -- a deeper level than its defn
 294
 295   | otherwise = env
 296   where
 297     scruts'  = (scrut_var, lvl) : scruts
 298     bind_lvl = case lookupVarEnv lvl_env scrut_var of
 299                  Just lvl -> lvl
 300                  Nothing  -> topLevel
 301
 302 lookupRecId :: LibCaseEnv -> Id -> Maybe CoreBind
 303 lookupRecId (LibCaseEnv bomb lvl lvl_env rec_env scruts) id
 304   = lookupVarEnv rec_env id
 305
 306 lookupLevel :: LibCaseEnv -> Id -> LibCaseLevel
 307 lookupLevel (LibCaseEnv bomb lvl lvl_env rec_env scruts) id
 308   = case lookupVarEnv lvl_env id of
 309       Just lvl -> lvl
 310       Nothing  -> topLevel
 311
 312 freeScruts :: LibCaseEnv
 313            -> LibCaseLevel      -- Level of the recursive Id
 314            -> [Id]              -- Ids that are scrutinised between the binding
 315                                 -- of the recursive Id and here
 316 freeScruts (LibCaseEnv bomb lvl lvl_env rec_env scruts) rec_bind_lvl
 317   = [v | (v,scrut_lvl) <- scruts, scrut_lvl > rec_bind_lvl]
 318 \end{code}