Normalize.hs

   1 --
   2 -- Functions to bring a Core expression in normal form. This module provides a
   3 -- top level function "normalize", and defines the actual transformation passes that
   4 -- are performed.
   5 --
   6 module Normalize (normalize) where
   7
   8 -- Standard modules
   9 import Debug.Trace
  10 import qualified Maybe
  11 import qualified Control.Monad as Monad
  12
  13 -- GHC API
  14 import CoreSyn
  15 import qualified UniqSupply
  16 import qualified CoreUtils
  17 import qualified Type
  18 import qualified Id
  19 import qualified UniqSet
  20 import qualified CoreFVs
  21 import Outputable ( showSDoc, ppr, nest )
  22
  23 -- Local imports
  24 import NormalizeTypes
  25 import NormalizeTools
  26 import CoreTools
  27
  28 --------------------------------
  29 -- Start of transformations
  30 --------------------------------
  31
  32 --------------------------------
  33 -- η abstraction
  34 --------------------------------
  35 eta, etatop :: Transform
  36 eta expr | is_fun expr && not (is_lam expr) = do
  37   let arg_ty = (fst . Type.splitFunTy . CoreUtils.exprType) expr
  38   id <- mkInternalVar "param" arg_ty
  39   change (Lam id (App expr (Var id)))
  40 -- Leave all other expressions unchanged
  41 eta e = return e
  42 etatop = notapplied ("eta", eta)
  43
  44 --------------------------------
  45 -- β-reduction
  46 --------------------------------
  47 beta, betatop :: Transform
  48 -- Substitute arg for x in expr
  49 beta (App (Lam x expr) arg) = change $ substitute [(x, arg)] expr
  50 -- Propagate the application into the let
  51 beta (App (Let binds expr) arg) = change $ Let binds (App expr arg)
  52 -- Propagate the application into each of the alternatives
  53 beta (App (Case scrut b ty alts) arg) = change $ Case scrut b ty' alts'
  54   where
  55     alts' = map (\(con, bndrs, expr) -> (con, bndrs, (App expr arg))) alts
  56     (_, ty') = Type.splitFunTy ty
  57 -- Leave all other expressions unchanged
  58 beta expr = return expr
  59 -- Perform this transform everywhere
  60 betatop = everywhere ("beta", beta)
  61
  62 --------------------------------
  63 -- let recursification
  64 --------------------------------
  65 letrec, letrectop :: Transform
  66 letrec (Let (NonRec b expr) res) = change $ Let (Rec [(b, expr)]) res
  67 -- Leave all other expressions unchanged
  68 letrec expr = return expr
  69 -- Perform this transform everywhere
  70 letrectop = everywhere ("letrec", letrec)
  71
  72 --------------------------------
  73 -- let simplification
  74 --------------------------------
  75 letsimpl, letsimpltop :: Transform
  76 -- Don't simplifiy lets that are already simple
  77 letsimpl expr@(Let _ (Var _)) = return expr
  78 -- Put the "in ..." value of a let in its own binding, but not when the
  79 -- expression has a function type (to prevent loops with inlinefun).
  80 letsimpl (Let (Rec binds) expr) | not $ is_fun expr = do
  81   id <- mkInternalVar "foo" (CoreUtils.exprType expr)
  82   let bind = (id, expr)
  83   change $ Let (Rec (bind:binds)) (Var id)
  84 -- Leave all other expressions unchanged
  85 letsimpl expr = return expr
  86 -- Perform this transform everywhere
  87 letsimpltop = everywhere ("letsimpl", letsimpl)
  88
  89 --------------------------------
  90 -- let flattening
  91 --------------------------------
  92 letflat, letflattop :: Transform
  93 letflat (Let (Rec binds) expr) = do
  94   -- Turn each binding into a list of bindings (possibly containing just one
  95   -- element, of course)
  96   bindss <- Monad.mapM flatbind binds
  97   -- Concat all the bindings
  98   let binds' = concat bindss
  99   -- Return the new let. We don't use change here, since possibly nothing has
 100   -- changed. If anything has changed, flatbind has already flagged that
 101   -- change.
 102   return $ Let (Rec binds') expr
 103   where
 104     -- Turns a binding of a let into a multiple bindings, or any other binding
 105     -- into a list with just that binding
 106     flatbind :: (CoreBndr, CoreExpr) -> TransformMonad [(CoreBndr, CoreExpr)]
 107     flatbind (b, Let (Rec binds) expr) = change ((b, expr):binds)
 108     flatbind (b, expr) = return [(b, expr)]
 109 -- Leave all other expressions unchanged
 110 letflat expr = return expr
 111 -- Perform this transform everywhere
 112 letflattop = everywhere ("letflat", letflat)
 113
 114 --------------------------------
 115 -- Simple let binding removal
 116 --------------------------------
 117 -- Remove a = b bindings from let expressions everywhere
 118 letremovetop :: Transform
 119 letremovetop = everywhere ("letremove", inlinebind (\(b, e) -> case e of (Var v) -> True; otherwise -> False))
 120
 121 --------------------------------
 122 -- Function inlining
 123 --------------------------------
 124 -- Remove a = B bindings, with B :: a -> b, from let expressions everywhere.
 125 -- This is a tricky function, which is prone to create loops in the
 126 -- transformations. To fix this, we make sure that no transformation will
 127 -- create a new let binding with a function type. These other transformations
 128 -- will just not work on those function-typed values at first, but the other
 129 -- transformations (in particular β-reduction) should make sure that the type
 130 -- of those values eventually becomes primitive.
 131 inlinefuntop :: Transform
 132 inlinefuntop = everywhere ("inlinefun", inlinebind (Type.isFunTy . CoreUtils.exprType . snd))
 133
 134 --------------------------------
 135 -- Scrutinee simplification
 136 --------------------------------
 137 scrutsimpl,scrutsimpltop :: Transform
 138 -- Don't touch scrutinees that are already simple
 139 scrutsimpl expr@(Case (Var _) _ _ _) = return expr
 140 -- Replace all other cases with a let that binds the scrutinee and a new
 141 -- simple scrutinee, but not when the scrutinee is a function type (to prevent
 142 -- loops with inlinefun, though I don't think a scrutinee can have a function
 143 -- type...)
 144 scrutsimpl (Case scrut b ty alts) | not $ is_fun scrut = do
 145   id <- mkInternalVar "scrut" (CoreUtils.exprType scrut)
 146   change $ Let (Rec [(id, scrut)]) (Case (Var id) b ty alts)
 147 -- Leave all other expressions unchanged
 148 scrutsimpl expr = return expr
 149 -- Perform this transform everywhere
 150 scrutsimpltop = everywhere ("scrutsimpl", scrutsimpl)
 151
 152 --------------------------------
 153 -- Case binder wildening
 154 --------------------------------
 155 casewild, casewildtop :: Transform
 156 casewild expr@(Case scrut b ty alts) = do
 157   (bindingss, alts') <- (Monad.liftM unzip) $ mapM doalt alts
 158   let bindings = concat bindingss
 159   -- Replace the case with a let with bindings and a case
 160   let newlet = (Let (Rec bindings) (Case scrut b ty alts'))
 161   -- If there are no non-wild binders, or this case is already a simple
 162   -- selector (i.e., a single alt with exactly one binding), already a simple
 163   -- selector altan no bindings (i.e., no wild binders in the original case),
 164   -- don't change anything, otherwise, replace the case.
 165   if null bindings || length alts == 1 && length bindings == 1 then return expr else change newlet
 166   where
 167   -- Generate a single wild binder, since they are all the same
 168   wild = Id.mkWildId
 169   -- Wilden the binders of one alt, producing a list of bindings as a
 170   -- sideeffect.
 171   doalt :: CoreAlt -> TransformMonad ([(CoreBndr, CoreExpr)], CoreAlt)
 172   doalt (con, bndrs, expr) = do
 173     bindings_maybe <- Monad.zipWithM mkextracts bndrs [0..]
 174     let bindings = Maybe.catMaybes bindings_maybe
 175     -- We replace the binders with wild binders only. We can leave expr
 176     -- unchanged, since the new bindings bind the same vars as the original
 177     -- did.
 178     let newalt = (con, wildbndrs, expr)
 179     return (bindings, newalt)
 180     where
 181       -- Make all binders wild
 182       wildbndrs = map (\bndr -> Id.mkWildId (Id.idType bndr)) bndrs
 183       -- Creates a case statement to retrieve the ith element from the scrutinee
 184       -- and binds that to b.
 185       mkextracts :: CoreBndr -> Int -> TransformMonad (Maybe (CoreBndr, CoreExpr))
 186       mkextracts b i =
 187         if is_wild b || Type.isFunTy (Id.idType b)
 188           -- Don't create extra bindings for binders that are already wild, or
 189           -- for binders that bind function types (to prevent loops with
 190           -- inlinefun).
 191           then return Nothing
 192           else do
 193             -- Create on new binder that will actually capture a value in this
 194             -- case statement, and return it
 195             let bty = (Id.idType b)
 196             id <- mkInternalVar "sel" bty
 197             let binders = take i wildbndrs ++ [id] ++ drop (i+1) wildbndrs
 198             return $ Just (b, Case scrut b bty [(con, binders, Var id)])
 199 -- Leave all other expressions unchanged
 200 casewild expr = return expr
 201 -- Perform this transform everywhere
 202 casewildtop = everywhere ("casewild", casewild)
 203
 204 --------------------------------
 205 -- Case value simplification
 206 --------------------------------
 207 casevalsimpl, casevalsimpltop :: Transform
 208 casevalsimpl expr@(Case scrut b ty alts) = do
 209   -- Try to simplify each alternative, resulting in an optional binding and a
 210   -- new alternative.
 211   (bindings_maybe, alts') <- (Monad.liftM unzip) $ mapM doalt alts
 212   let bindings = Maybe.catMaybes bindings_maybe
 213   -- Create a new let around the case, that binds of the cases values.
 214   let newlet = Let (Rec bindings) (Case scrut b ty alts')
 215   -- If there were no values that needed and allowed simplification, don't
 216   -- change the case.
 217   if null bindings then return expr else change newlet
 218   where
 219     doalt :: CoreAlt -> TransformMonad (Maybe (CoreBndr, CoreExpr), CoreAlt)
 220     -- Don't simplify values that are already simple
 221     doalt alt@(con, bndrs, Var _) = return (Nothing, alt)
 222     -- Simplify each alt by creating a new id, binding the case value to it and
 223     -- replacing the case value with that id. Only do this when the case value
 224     -- does not use any of the binders bound by this alternative, for that would
 225     -- cause those binders to become unbound when moving the value outside of
 226     -- the case statement. Also, don't create a binding for function-typed
 227     -- expressions, to prevent loops with inlinefun.
 228     doalt (con, bndrs, expr) | (not usesvars) && (not $ is_fun expr) = do
 229       id <- mkInternalVar "caseval" (CoreUtils.exprType expr)
 230       -- We don't flag a change here, since casevalsimpl will do that above
 231       -- based on Just we return here.
 232       return $ (Just (id, expr), (con, bndrs, Var id))
 233       -- Find if any of the binders are used by expr
 234       where usesvars = (not . UniqSet.isEmptyUniqSet . (CoreFVs.exprSomeFreeVars (`elem` bndrs))) expr
 235     -- Don't simplify anything else
 236     doalt alt = return (Nothing, alt)
 237 -- Leave all other expressions unchanged
 238 casevalsimpl expr = return expr
 239 -- Perform this transform everywhere
 240 casevalsimpltop = everywhere ("casevalsimpl", casevalsimpl)
 241
 242 --------------------------------
 243 -- Case removal
 244 --------------------------------
 245 -- Remove case statements that have only a single alternative and only wild
 246 -- binders.
 247 caseremove, caseremovetop :: Transform
 248 -- Replace a useless case by the value of its single alternative
 249 caseremove (Case scrut b ty [(con, bndrs, expr)]) | not usesvars = change expr
 250     -- Find if any of the binders are used by expr
 251     where usesvars = (not . UniqSet.isEmptyUniqSet . (CoreFVs.exprSomeFreeVars (`elem` bndrs))) expr
 252 -- Leave all other expressions unchanged
 253 caseremove expr = return expr
 254 -- Perform this transform everywhere
 255 caseremovetop = everywhere ("caseremove", caseremove)
 256
 257 --------------------------------
 258 -- Application simplification
 259 --------------------------------
 260 -- Make sure that all arguments in an application are simple variables.
 261 appsimpl, appsimpltop :: Transform
 262 -- Don't simplify arguments that are already simple
 263 appsimpl expr@(App f (Var _)) = return expr
 264 -- Simplify all arguments that do not have a function type (to prevent loops
 265 -- with inlinefun) and is not a type argument. Do this by introducing a new
 266 -- Let that binds the argument and passing the new binder in the application.
 267 appsimpl (App f expr) | (not $ is_fun expr) && (not $ CoreSyn.isTypeArg expr) = do
 268   id <- mkInternalVar "arg" (CoreUtils.exprType expr)
 269   change $ Let (Rec [(id, expr)]) (App f (Var id))
 270 -- Leave all other expressions unchanged
 271 appsimpl expr = return expr
 272 -- Perform this transform everywhere
 273 appsimpltop = everywhere ("appsimpl", appsimpl)
 274
 275 -- TODO: introduce top level let if needed?
 276
 277 --------------------------------
 278 -- End of transformations
 279 --------------------------------
 280
 281
 282
 283
 284 -- What transforms to run?
 285 transforms = [etatop, betatop, letremovetop, letrectop, letsimpltop, letflattop, casewildtop, scrutsimpltop, casevalsimpltop, caseremovetop, inlinefuntop, appsimpltop]
 286
 287 -- Normalize a core expression by running transforms until none applies
 288 -- anymore. Uses a UniqSupply to generate new identifiers.
 289 normalize :: UniqSupply.UniqSupply -> CoreExpr -> CoreExpr
 290 normalize = dotransforms transforms
 291