ghc/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 CmdLineOpts      ( DynFlags, DynFlag(..), opt_LiberateCaseThreshold )
  12 import CoreLint         ( showPass, endPass )
  13 import CoreSyn
  14 import CoreUnfold       ( couldBeSmallEnoughToInline )
  15 import Var              ( Id )
  16 import VarEnv
  17 import UniqFM           ( ufmToList )
  18 import Outputable
  19 \end{code}
  20
  21 This module walks over @Core@, and looks for @case@ on free variables.
  22 The criterion is:
  23         if there is case on a free on the route to the recursive call,
  24         then the recursive call is replaced with an unfolding.
  25
  26 Example
  27
  28 \begin{verbatim}
  29 f = \ t -> case v of
  30                V a b -> a : f t
  31 \end{verbatim}
  32
  33 => the inner f is replaced.
  34
  35 \begin{verbatim}
  36 f = \ t -> case v of
  37                V a b -> a : (letrec
  38                                 f =  \ t -> case v of
  39                                                V a b -> a : f t
  40                              in f) t
  41 \end{verbatim}
  42 (note the NEED for shadowing)
  43
  44 => Simplify
  45
  46 \begin{verbatim}
  47 f = \ t -> case v of
  48                V a b -> a : (letrec
  49                                 f = \ t -> a : f t
  50                              in f t)
  51 \begin{verbatim}
  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
  74 Similarly drop:
  75
  76         drop n [] = []
  77         drop 0 xs = xs
  78         drop n (x:xs) = drop (n-1) xs
  79
  80 Would like to pass n along unboxed.
  81
  82
  83 To think about (Apr 94)
  84 ~~~~~~~~~~~~~~
  85
  86 Main worry: duplicating code excessively.  At the moment we duplicate
  87 the entire binding group once at each recursive call.  But there may
  88 be a group of recursive calls which share a common set of evaluated
  89 free variables, in which case the duplication is a plain waste.
  90
  91 Another thing we could consider adding is some unfold-threshold thing,
  92 so that we'll only duplicate if the size of the group rhss isn't too
  93 big.
  94
  95 Data types
  96 ~~~~~~~~~~
  97
  98 The ``level'' of a binder tells how many
  99 recursive defns lexically enclose the binding
 100 A recursive defn "encloses" its RHS, not its
 101 scope.  For example:
 102 \begin{verbatim}
 103         letrec f = let g = ... in ...
 104         in
 105         let h = ...
 106         in ...
 107 \end{verbatim}
 108 Here, the level of @f@ is zero, the level of @g@ is one,
 109 and the level of @h@ is zero (NB not one).
 110
 111 \begin{code}
 112 type LibCaseLevel = Int
 113
 114 topLevel :: LibCaseLevel
 115 topLevel = 0
 116 \end{code}
 117
 118 \begin{code}
 119 data LibCaseEnv
 120   = LibCaseEnv
 121         Int                     -- Bomb-out size for deciding if
 122                                 -- potential liberatees are too big.
 123                                 -- (passed in from cmd-line args)
 124
 125         LibCaseLevel            -- Current level
 126
 127         (IdEnv LibCaseLevel)    -- Binds all non-top-level in-scope Ids
 128                                 -- (top-level and imported things have
 129                                 -- a level of zero)
 130
 131         (IdEnv CoreBind)        -- Binds *only* recursively defined
 132                                 -- Ids, to their own binding group,
 133                                 -- and *only* in their own RHSs
 134
 135         [(Id,LibCaseLevel)]     -- Each of these Ids was scrutinised by an
 136                                 -- enclosing case expression, with the
 137                                 -- specified number of enclosing
 138                                 -- recursive bindings; furthermore,
 139                                 -- the Id is bound at a lower level
 140                                 -- than the case expression.  The
 141                                 -- order is insignificant; it's a bag
 142                                 -- really
 143
 144 initEnv :: Int -> LibCaseEnv
 145 initEnv bomb_size = LibCaseEnv bomb_size 0 emptyVarEnv emptyVarEnv []
 146
 147 pprEnv :: LibCaseEnv -> SDoc
 148 pprEnv (LibCaseEnv _ lvl lvl_env _ scruts)
 149   = vcat [text "LibCaseEnv" <+> int lvl,
 150           fsep (map ppr (ufmToList lvl_env)),
 151           fsep (map ppr scruts)]
 152
 153 bombOutSize (LibCaseEnv bomb_size _ _ _ _) = bomb_size
 154 \end{code}
 155
 156
 157 Programs
 158 ~~~~~~~~
 159 \begin{code}
 160 liberateCase :: DynFlags -> [CoreBind] -> IO [CoreBind]
 161 liberateCase dflags binds
 162   = do {
 163         showPass dflags "Liberate case" ;
 164         let { binds' = do_prog (initEnv opt_LiberateCaseThreshold) binds } ;
 165         endPass dflags "Liberate case" Opt_D_verbose_core2core binds'
 166                                 {- no specific flag for dumping -}
 167     }
 168   where
 169     do_prog env [] = []
 170     do_prog env (bind:binds) = bind' : do_prog env' binds
 171                              where
 172                                (env', bind') = libCaseBind env bind
 173 \end{code}
 174
 175 Bindings
 176 ~~~~~~~~
 177
 178 \begin{code}
 179 libCaseBind :: LibCaseEnv -> CoreBind -> (LibCaseEnv, CoreBind)
 180
 181 libCaseBind env (NonRec binder rhs)
 182   = (addBinders env [binder], NonRec binder (libCase env rhs))
 183
 184 libCaseBind env (Rec pairs)
 185   = (env_body, Rec pairs')
 186   where
 187     (binders, rhss) = unzip pairs
 188
 189     env_body = addBinders env binders
 190
 191     pairs' = [(binder, libCase env_rhs rhs) | (binder,rhs) <- pairs]
 192
 193     env_rhs = if all rhs_small_enough rhss then extended_env else env
 194
 195         -- We extend the rec-env by binding each Id to its rhs, first
 196         -- processing the rhs with an *un-extended* environment, so
 197         -- that the same process doesn't occur for ever!
 198
 199     extended_env = addRecBinds env [ (binder, libCase env_body rhs)
 200                                    | (binder, rhs) <- pairs ]
 201
 202     rhs_small_enough rhs = couldBeSmallEnoughToInline lIBERATE_BOMB_SIZE rhs
 203     lIBERATE_BOMB_SIZE   = bombOutSize env
 204 \end{code}
 205
 206
 207 Expressions
 208 ~~~~~~~~~~~
 209
 210 \begin{code}
 211 libCase :: LibCaseEnv
 212         -> CoreExpr
 213         -> CoreExpr
 214
 215 libCase env (Var v)             = libCaseId env v
 216 libCase env (Lit lit)           = Lit lit
 217 libCase env (Type ty)           = Type ty
 218 libCase env (App fun arg)       = App (libCase env fun) (libCase env arg)
 219 libCase env (Note note body)    = Note note (libCase env body)
 220
 221 libCase env (Lam binder body)
 222   = Lam binder (libCase (addBinders env [binder]) body)
 223
 224 libCase env (Let bind body)
 225   = Let bind' (libCase env_body body)
 226   where
 227     (env_body, bind') = libCaseBind env bind
 228
 229 libCase env (Case scrut bndr alts)
 230   = Case (libCase env scrut) bndr (map (libCaseAlt env_alts) alts)
 231   where
 232     env_alts = addBinders env_with_scrut [bndr]
 233     env_with_scrut = case scrut of
 234                         Var scrut_var -> addScrutedVar env scrut_var
 235                         other         -> env
 236
 237 libCaseAlt env (con,args,rhs) = (con, args, libCase (addBinders env args) rhs)
 238 \end{code}
 239
 240 Ids
 241 ~~~
 242 \begin{code}
 243 libCaseId :: LibCaseEnv -> Id -> CoreExpr
 244 libCaseId env v
 245   | Just the_bind <- lookupRecId env v  -- It's a use of a recursive thing
 246   , not (null free_scruts)              -- with free vars scrutinised in RHS
 247   = Let the_bind (Var v)
 248
 249   | otherwise
 250   = Var v
 251
 252   where
 253     rec_id_level = lookupLevel env v
 254     free_scruts  = freeScruts env rec_id_level
 255 \end{code}
 256
 257
 258
 259 Utility functions
 260 ~~~~~~~~~~~~~~~~~
 261 \begin{code}
 262 addBinders :: LibCaseEnv -> [CoreBndr] -> LibCaseEnv
 263 addBinders (LibCaseEnv bomb lvl lvl_env rec_env scruts) binders
 264   = LibCaseEnv bomb lvl lvl_env' rec_env scruts
 265   where
 266     lvl_env' = extendVarEnvList lvl_env (binders `zip` repeat lvl)
 267
 268 addRecBinds :: LibCaseEnv -> [(Id,CoreExpr)] -> LibCaseEnv
 269 addRecBinds (LibCaseEnv bomb lvl lvl_env rec_env scruts) pairs
 270   = LibCaseEnv bomb lvl' lvl_env' rec_env' scruts
 271   where
 272     lvl'     = lvl + 1
 273     lvl_env' = extendVarEnvList lvl_env [(binder,lvl) | (binder,_) <- pairs]
 274     rec_env' = extendVarEnvList rec_env [(binder, Rec pairs) | (binder,_) <- pairs]
 275
 276 addScrutedVar :: LibCaseEnv
 277               -> Id             -- This Id is being scrutinised by a case expression
 278               -> LibCaseEnv
 279
 280 addScrutedVar env@(LibCaseEnv bomb lvl lvl_env rec_env scruts) scrut_var
 281   | bind_lvl < lvl
 282   = LibCaseEnv bomb lvl lvl_env rec_env scruts'
 283         -- Add to scruts iff the scrut_var is being scrutinised at
 284         -- a deeper level than its defn
 285
 286   | otherwise = env
 287   where
 288     scruts'  = (scrut_var, lvl) : scruts
 289     bind_lvl = case lookupVarEnv lvl_env scrut_var of
 290                  Just lvl -> lvl
 291                  Nothing  -> topLevel
 292
 293 lookupRecId :: LibCaseEnv -> Id -> Maybe CoreBind
 294 lookupRecId (LibCaseEnv bomb lvl lvl_env rec_env scruts) id
 295   = lookupVarEnv rec_env id
 296
 297 lookupLevel :: LibCaseEnv -> Id -> LibCaseLevel
 298 lookupLevel (LibCaseEnv bomb lvl lvl_env rec_env scruts) id
 299   = case lookupVarEnv lvl_env id of
 300       Just lvl -> lvl
 301       Nothing  -> topLevel
 302
 303 freeScruts :: LibCaseEnv
 304            -> LibCaseLevel      -- Level of the recursive Id
 305            -> [Id]              -- Ids that are scrutinised between the binding
 306                                 -- of the recursive Id and here
 307 freeScruts (LibCaseEnv bomb lvl lvl_env rec_env scruts) rec_bind_lvl
 308   = [v | (v,scrut_lvl) <- scruts, scrut_lvl > rec_bind_lvl]
 309 \end{code}