Normalize.hs

   1 {-# LANGUAGE PackageImports #-}
   2 --
   3 -- Functions to bring a Core expression in normal form. This module provides a
   4 -- top level function "normalize", and defines the actual transformation passes that
   5 -- are performed.
   6 --
   7 module Normalize (normalizeModule) where
   8
   9 -- Standard modules
  10 import Debug.Trace
  11 import qualified Maybe
  12 import qualified "transformers" Control.Monad.Trans as Trans
  13 import qualified Control.Monad as Monad
  14 import qualified Control.Monad.Trans.Writer as Writer
  15 import qualified Data.Map as Map
  16 import qualified Data.Monoid as Monoid
  17 import Data.Accessor
  18
  19 -- GHC API
  20 import CoreSyn
  21 import qualified UniqSupply
  22 import qualified CoreUtils
  23 import qualified Type
  24 import qualified Id
  25 import qualified Var
  26 import qualified VarSet
  27 import qualified CoreFVs
  28 import qualified CoreUtils
  29 import qualified MkCore
  30 import Outputable ( showSDoc, ppr, nest )
  31
  32 -- Local imports
  33 import NormalizeTypes
  34 import NormalizeTools
  35 import CoreTools
  36 import Pretty
  37
  38 --------------------------------
  39 -- Start of transformations
  40 --------------------------------
  41
  42 --------------------------------
  43 -- η abstraction
  44 --------------------------------
  45 eta, etatop :: Transform
  46 eta expr | is_fun expr && not (is_lam expr) = do
  47 eta expr | is_fun expr && not (is_lam expr) = do
  48   let arg_ty = (fst . Type.splitFunTy . CoreUtils.exprType) expr
  49   id <- mkInternalVar "param" arg_ty
  50   change (Lam id (App expr (Var id)))
  51 -- Leave all other expressions unchanged
  52 eta e = return e
  53 etatop = notappargs ("eta", eta)
  54
  55 --------------------------------
  56 -- β-reduction
  57 --------------------------------
  58 beta, betatop :: Transform
  59 -- Substitute arg for x in expr
  60 beta (App (Lam x expr) arg) = change $ substitute [(x, arg)] expr
  61 -- Propagate the application into the let
  62 beta (App (Let binds expr) arg) = change $ Let binds (App expr arg)
  63 -- Propagate the application into each of the alternatives
  64 beta (App (Case scrut b ty alts) arg) = change $ Case scrut b ty' alts'
  65   where
  66     alts' = map (\(con, bndrs, expr) -> (con, bndrs, (App expr arg))) alts
  67     ty' = CoreUtils.applyTypeToArg ty arg
  68 -- Leave all other expressions unchanged
  69 beta expr = return expr
  70 -- Perform this transform everywhere
  71 betatop = everywhere ("beta", beta)
  72
  73 --------------------------------
  74 -- let recursification
  75 --------------------------------
  76 letrec, letrectop :: Transform
  77 letrec (Let (NonRec b expr) res) = change $ Let (Rec [(b, expr)]) res
  78 -- Leave all other expressions unchanged
  79 letrec expr = return expr
  80 -- Perform this transform everywhere
  81 letrectop = everywhere ("letrec", letrec)
  82
  83 --------------------------------
  84 -- let simplification
  85 --------------------------------
  86 letsimpl, letsimpltop :: Transform
  87 -- Don't simplifiy lets that are already simple
  88 letsimpl expr@(Let _ (Var _)) = return expr
  89 -- Put the "in ..." value of a let in its own binding, but not when the
  90 -- expression is applicable (to prevent loops with inlinefun).
  91 letsimpl (Let (Rec binds) expr) | not $ is_applicable expr = do
  92   id <- mkInternalVar "foo" (CoreUtils.exprType expr)
  93   let bind = (id, expr)
  94   change $ Let (Rec (bind:binds)) (Var id)
  95 -- Leave all other expressions unchanged
  96 letsimpl expr = return expr
  97 -- Perform this transform everywhere
  98 letsimpltop = everywhere ("letsimpl", letsimpl)
  99
 100 --------------------------------
 101 -- let flattening
 102 --------------------------------
 103 letflat, letflattop :: Transform
 104 letflat (Let (Rec binds) expr) = do
 105   -- Turn each binding into a list of bindings (possibly containing just one
 106   -- element, of course)
 107   bindss <- Monad.mapM flatbind binds
 108   -- Concat all the bindings
 109   let binds' = concat bindss
 110   -- Return the new let. We don't use change here, since possibly nothing has
 111   -- changed. If anything has changed, flatbind has already flagged that
 112   -- change.
 113   return $ Let (Rec binds') expr
 114   where
 115     -- Turns a binding of a let into a multiple bindings, or any other binding
 116     -- into a list with just that binding
 117     flatbind :: (CoreBndr, CoreExpr) -> TransformMonad [(CoreBndr, CoreExpr)]
 118     flatbind (b, Let (Rec binds) expr) = change ((b, expr):binds)
 119     flatbind (b, expr) = return [(b, expr)]
 120 -- Leave all other expressions unchanged
 121 letflat expr = return expr
 122 -- Perform this transform everywhere
 123 letflattop = everywhere ("letflat", letflat)
 124
 125 --------------------------------
 126 -- Simple let binding removal
 127 --------------------------------
 128 -- Remove a = b bindings from let expressions everywhere
 129 letremovetop :: Transform
 130 letremovetop = everywhere ("letremove", inlinebind (\(b, e) -> case e of (Var v) -> True; otherwise -> False))
 131
 132 --------------------------------
 133 -- Function inlining
 134 --------------------------------
 135 -- Remove a = B bindings, with B :: a -> b, or B :: forall x . T, from let
 136 -- expressions everywhere. This means that any value that still needs to be
 137 -- applied to something else (polymorphic values need to be applied to a
 138 -- Type) will be inlined, and will eventually be applied to all their
 139 -- arguments.
 140 --
 141 -- This is a tricky function, which is prone to create loops in the
 142 -- transformations. To fix this, we make sure that no transformation will
 143 -- create a new let binding with a function type. These other transformations
 144 -- will just not work on those function-typed values at first, but the other
 145 -- transformations (in particular β-reduction) should make sure that the type
 146 -- of those values eventually becomes primitive.
 147 inlinefuntop :: Transform
 148 inlinefuntop = everywhere ("inlinefun", inlinebind (is_applicable . snd))
 149
 150 --------------------------------
 151 -- Scrutinee simplification
 152 --------------------------------
 153 scrutsimpl,scrutsimpltop :: Transform
 154 -- Don't touch scrutinees that are already simple
 155 scrutsimpl expr@(Case (Var _) _ _ _) = return expr
 156 -- Replace all other cases with a let that binds the scrutinee and a new
 157 -- simple scrutinee, but not when the scrutinee is applicable (to prevent
 158 -- loops with inlinefun, though I don't think a scrutinee can be
 159 -- applicable...)
 160 scrutsimpl (Case scrut b ty alts) | not $ is_applicable scrut = do
 161   id <- mkInternalVar "scrut" (CoreUtils.exprType scrut)
 162   change $ Let (Rec [(id, scrut)]) (Case (Var id) b ty alts)
 163 -- Leave all other expressions unchanged
 164 scrutsimpl expr = return expr
 165 -- Perform this transform everywhere
 166 scrutsimpltop = everywhere ("scrutsimpl", scrutsimpl)
 167
 168 --------------------------------
 169 -- Case binder wildening
 170 --------------------------------
 171 casewild, casewildtop :: Transform
 172 casewild expr@(Case scrut b ty alts) = do
 173   (bindingss, alts') <- (Monad.liftM unzip) $ mapM doalt alts
 174   let bindings = concat bindingss
 175   -- Replace the case with a let with bindings and a case
 176   let newlet = (Let (Rec bindings) (Case scrut b ty alts'))
 177   -- If there are no non-wild binders, or this case is already a simple
 178   -- selector (i.e., a single alt with exactly one binding), already a simple
 179   -- selector altan no bindings (i.e., no wild binders in the original case),
 180   -- don't change anything, otherwise, replace the case.
 181   if null bindings || length alts == 1 && length bindings == 1 then return expr else change newlet
 182   where
 183   -- Generate a single wild binder, since they are all the same
 184   wild = Id.mkWildId
 185   -- Wilden the binders of one alt, producing a list of bindings as a
 186   -- sideeffect.
 187   doalt :: CoreAlt -> TransformMonad ([(CoreBndr, CoreExpr)], CoreAlt)
 188   doalt (con, bndrs, expr) = do
 189     bindings_maybe <- Monad.zipWithM mkextracts bndrs [0..]
 190     let bindings = Maybe.catMaybes bindings_maybe
 191     -- We replace the binders with wild binders only. We can leave expr
 192     -- unchanged, since the new bindings bind the same vars as the original
 193     -- did.
 194     let newalt = (con, wildbndrs, expr)
 195     return (bindings, newalt)
 196     where
 197       -- Make all binders wild
 198       wildbndrs = map (\bndr -> Id.mkWildId (Id.idType bndr)) bndrs
 199       -- Creates a case statement to retrieve the ith element from the scrutinee
 200       -- and binds that to b.
 201       mkextracts :: CoreBndr -> Int -> TransformMonad (Maybe (CoreBndr, CoreExpr))
 202       mkextracts b i =
 203         if is_wild b || Type.isFunTy (Id.idType b)
 204           -- Don't create extra bindings for binders that are already wild, or
 205           -- for binders that bind function types (to prevent loops with
 206           -- inlinefun).
 207           then return Nothing
 208           else do
 209             -- Create on new binder that will actually capture a value in this
 210             -- case statement, and return it
 211             let bty = (Id.idType b)
 212             id <- mkInternalVar "sel" bty
 213             let binders = take i wildbndrs ++ [id] ++ drop (i+1) wildbndrs
 214             return $ Just (b, Case scrut b bty [(con, binders, Var id)])
 215 -- Leave all other expressions unchanged
 216 casewild expr = return expr
 217 -- Perform this transform everywhere
 218 casewildtop = everywhere ("casewild", casewild)
 219
 220 --------------------------------
 221 -- Case value simplification
 222 --------------------------------
 223 casevalsimpl, casevalsimpltop :: Transform
 224 casevalsimpl expr@(Case scrut b ty alts) = do
 225   -- Try to simplify each alternative, resulting in an optional binding and a
 226   -- new alternative.
 227   (bindings_maybe, alts') <- (Monad.liftM unzip) $ mapM doalt alts
 228   let bindings = Maybe.catMaybes bindings_maybe
 229   -- Create a new let around the case, that binds of the cases values.
 230   let newlet = Let (Rec bindings) (Case scrut b ty alts')
 231   -- If there were no values that needed and allowed simplification, don't
 232   -- change the case.
 233   if null bindings then return expr else change newlet
 234   where
 235     doalt :: CoreAlt -> TransformMonad (Maybe (CoreBndr, CoreExpr), CoreAlt)
 236     -- Don't simplify values that are already simple
 237     doalt alt@(con, bndrs, Var _) = return (Nothing, alt)
 238     -- Simplify each alt by creating a new id, binding the case value to it and
 239     -- replacing the case value with that id. Only do this when the case value
 240     -- does not use any of the binders bound by this alternative, for that would
 241     -- cause those binders to become unbound when moving the value outside of
 242     -- the case statement. Also, don't create a binding for applicable
 243     -- expressions, to prevent loops with inlinefun.
 244     doalt (con, bndrs, expr) | (not usesvars) && (not $ is_applicable expr) = do
 245       id <- mkInternalVar "caseval" (CoreUtils.exprType expr)
 246       -- We don't flag a change here, since casevalsimpl will do that above
 247       -- based on Just we return here.
 248       return $ (Just (id, expr), (con, bndrs, Var id))
 249       -- Find if any of the binders are used by expr
 250       where usesvars = (not . VarSet.isEmptyVarSet . (CoreFVs.exprSomeFreeVars (`elem` bndrs))) expr
 251     -- Don't simplify anything else
 252     doalt alt = return (Nothing, alt)
 253 -- Leave all other expressions unchanged
 254 casevalsimpl expr = return expr
 255 -- Perform this transform everywhere
 256 casevalsimpltop = everywhere ("casevalsimpl", casevalsimpl)
 257
 258 --------------------------------
 259 -- Case removal
 260 --------------------------------
 261 -- Remove case statements that have only a single alternative and only wild
 262 -- binders.
 263 caseremove, caseremovetop :: Transform
 264 -- Replace a useless case by the value of its single alternative
 265 caseremove (Case scrut b ty [(con, bndrs, expr)]) | not usesvars = change expr
 266     -- Find if any of the binders are used by expr
 267     where usesvars = (not . VarSet.isEmptyVarSet . (CoreFVs.exprSomeFreeVars (`elem` bndrs))) expr
 268 -- Leave all other expressions unchanged
 269 caseremove expr = return expr
 270 -- Perform this transform everywhere
 271 caseremovetop = everywhere ("caseremove", caseremove)
 272
 273 --------------------------------
 274 -- Application simplification
 275 --------------------------------
 276 -- Make sure that all arguments in an application are simple variables.
 277 appsimpl, appsimpltop :: Transform
 278 -- Don't simplify arguments that are already simple
 279 appsimpl expr@(App f (Var _)) = return expr
 280 -- Simplify all non-applicable (to prevent loops with inlinefun) arguments,
 281 -- except for type arguments (since a let can't bind type vars, only a lambda
 282 -- can). Do this by introducing a new Let that binds the argument and passing
 283 -- the new binder in the application.
 284 appsimpl (App f expr) | (not $ is_applicable expr) && (not $ CoreSyn.isTypeArg expr) = do
 285   id <- mkInternalVar "arg" (CoreUtils.exprType expr)
 286   change $ Let (Rec [(id, expr)]) (App f (Var id))
 287 -- Leave all other expressions unchanged
 288 appsimpl expr = return expr
 289 -- Perform this transform everywhere
 290 appsimpltop = everywhere ("appsimpl", appsimpl)
 291
 292
 293 --------------------------------
 294 -- Type argument propagation
 295 --------------------------------
 296 -- Remove all applications to type arguments, by duplicating the function
 297 -- called with the type application in its new definition. We leave
 298 -- dictionaries that might be associated with the type untouched, the funprop
 299 -- transform should propagate these later on.
 300 typeprop, typeproptop :: Transform
 301 -- Transform any function that is applied to a type argument. Since type
 302 -- arguments are always the first ones to apply and we'll remove all type
 303 -- arguments, we can simply do them one by one. We only propagate type
 304 -- arguments without any free tyvars, since tyvars those wouldn't be in scope
 305 -- in the new function.
 306 typeprop expr@(App (Var f) arg@(Type ty)) | not $ has_free_tyvars arg = do
 307   body_maybe <- Trans.lift $ getGlobalBind f
 308   case body_maybe of
 309     Just body -> do
 310       let newbody = App body (Type ty)
 311       -- Create a new function with the same name but a new body
 312       newf <- mkFunction f newbody
 313       -- Replace the application with this new function
 314       change (Var newf)
 315     -- If we don't have a body for the function called, leave it unchanged (it
 316     -- should be a primitive function then).
 317     Nothing -> return expr
 318 -- Leave all other expressions unchanged
 319 typeprop expr = return expr
 320 -- Perform this transform everywhere
 321 typeproptop = everywhere ("typeprop", typeprop)
 322
 323
 324 --------------------------------
 325 -- Function-typed argument propagation
 326 --------------------------------
 327 -- Remove all applications to function-typed arguments, by duplication the
 328 -- function called with the function-typed parameter replaced by the free
 329 -- variables of the argument passed in.
 330 funprop, funproptop :: Transform
 331 -- Transform any application of a named function (i.e., skip applications of
 332 -- lambda's). Also skip applications that have arguments with free type
 333 -- variables, since we can't inline those.
 334 funprop expr@(App _ _) | is_var fexpr && not (any has_free_tyvars args) = do
 335   -- Find the body of the function called
 336   body_maybe <- Trans.lift $ getGlobalBind f
 337   case body_maybe of
 338     Just body -> do
 339       -- Process each of the arguments in turn
 340       (args', changed) <- Writer.listen $ mapM doarg args
 341       -- See if any of the arguments changed
 342       case Monoid.getAny changed of
 343         True -> do
 344           let (newargs', newparams', oldargs) = unzip3 args'
 345           let newargs = concat newargs'
 346           let newparams = concat newparams'
 347           -- Create a new body that consists of a lambda for all new arguments and
 348           -- the old body applied to some arguments.
 349           let newbody = MkCore.mkCoreLams newparams (MkCore.mkCoreApps body oldargs)
 350           -- Create a new function with the same name but a new body
 351           newf <- mkFunction f newbody
 352           -- Replace the original application with one of the new function to the
 353           -- new arguments.
 354           change $ MkCore.mkCoreApps (Var newf) newargs
 355         False ->
 356           -- Don't change the expression if none of the arguments changed
 357           return expr
 358
 359     -- If we don't have a body for the function called, leave it unchanged (it
 360     -- should be a primitive function then).
 361     Nothing -> return expr
 362   where
 363     -- Find the function called and the arguments
 364     (fexpr, args) = collectArgs expr
 365     Var f = fexpr
 366
 367     -- Process a single argument and return (args, bndrs, arg), where args are
 368     -- the arguments to replace the given argument in the original
 369     -- application, bndrs are the binders to include in the top-level lambda
 370     -- in the new function body, and arg is the argument to apply to the old
 371     -- function body.
 372     doarg :: CoreExpr -> TransformMonad ([CoreExpr], [CoreBndr], CoreExpr)
 373     doarg arg | is_fun arg = do
 374       bndrs <- Trans.lift getGlobalBinders
 375       -- Find interesting free variables, each of which should be passed to
 376       -- the new function instead of the original function argument.
 377       --
 378       -- Interesting vars are those that are local, but not available from the
 379       -- top level scope (functions from this module are defined as local, but
 380       -- they're not local to this function, so we can freely move references
 381       -- to them into another function).
 382       let interesting var = Var.isLocalVar var && (not $ var `elem` bndrs)
 383       let free_vars = VarSet.varSetElems $ CoreFVs.exprSomeFreeVars interesting arg
 384       -- Mark the current expression as changed
 385       setChanged
 386       return (map Var free_vars, free_vars, arg)
 387     -- Non-functiontyped arguments can be unchanged. Note that this handles
 388     -- both values and types.
 389     doarg arg = do
 390       -- TODO: preserve original naming?
 391       id <- mkBinderFor arg "param"
 392       -- Just pass the original argument to the new function, which binds it
 393       -- to a new id and just pass that new id to the old function body.
 394       return ([arg], [id], mkReferenceTo id)
 395 -- Leave all other expressions unchanged
 396 funprop expr = return expr
 397 -- Perform this transform everywhere
 398 funproptop = everywhere ("funprop", funprop)
 399
 400
 401 -- TODO: introduce top level let if needed?
 402
 403 --------------------------------
 404 -- End of transformations
 405 --------------------------------
 406
 407
 408
 409
 410 -- What transforms to run?
 411 transforms = [typeproptop, funproptop, etatop, betatop, letremovetop, letrectop, letsimpltop, letflattop, casewildtop, scrutsimpltop, casevalsimpltop, caseremovetop, inlinefuntop, appsimpltop]
 412
 413 -- Turns the given bind into VHDL
 414 normalizeModule ::
 415   UniqSupply.UniqSupply -- ^ A UniqSupply we can use
 416   -> [(CoreBndr, CoreExpr)]  -- ^ All bindings we know (i.e., in the current module)
 417   -> [CoreBndr]  -- ^ The bindings to generate VHDL for (i.e., the top level bindings)
 418   -> [Bool] -- ^ For each of the bindings to generate VHDL for, if it is stateful
 419   -> [(CoreBndr, CoreExpr)] -- ^ The resulting VHDL
 420
 421 normalizeModule uniqsupply bindings generate_for statefuls = runTransformSession uniqsupply $ do
 422   -- Put all the bindings in this module in the tsBindings map
 423   putA tsBindings (Map.fromList bindings)
 424   -- (Recursively) normalize each of the requested bindings
 425   mapM normalizeBind generate_for
 426   -- Get all initial bindings and the ones we produced
 427   bindings_map <- getA tsBindings
 428   let bindings = Map.assocs bindings_map
 429   normalized_bindings <- getA tsNormalized
 430   -- But return only the normalized bindings
 431   return $ filter ((flip VarSet.elemVarSet normalized_bindings) . fst) bindings
 432
 433 normalizeBind :: CoreBndr -> TransformSession ()
 434 normalizeBind bndr =
 435   -- Don't normalize global variables, these should be either builtin
 436   -- functions or data constructors.
 437   Monad.when (Var.isLocalIdVar bndr) $ do
 438     -- Skip binders that have a polymorphic type, since it's impossible to
 439     -- create polymorphic hardware.
 440     if is_poly (Var bndr)
 441       then
 442         -- This should really only happen at the top level... TODO: Give
 443         -- a different error if this happens down in the recursion.
 444         error $ "Function " ++ show bndr ++ " is polymorphic, can't normalize"
 445       else do
 446         normalized_funcs <- getA tsNormalized
 447         -- See if this function was normalized already
 448         if VarSet.elemVarSet bndr normalized_funcs
 449           then
 450             -- Yup, don't do it again
 451             return ()
 452           else do
 453             -- Nope, note that it has been and do it.
 454             modA tsNormalized (flip VarSet.extendVarSet bndr)
 455             expr_maybe <- getGlobalBind bndr
 456             case expr_maybe of
 457               Just expr -> do
 458                 -- Introduce an empty Let at the top level, so there will always be
 459                 -- a let in the expression (none of the transformations will remove
 460                 -- the last let).
 461                 let expr' = Let (Rec []) expr
 462                 -- Normalize this expression
 463                 trace ("Transforming " ++ (show bndr) ++ "\nBefore:\n\n" ++ showSDoc ( ppr expr' ) ++ "\n") $ return ()
 464                 expr' <- dotransforms transforms expr'
 465                 trace ("\nAfter:\n\n" ++ showSDoc ( ppr expr')) $ return ()
 466                 -- And store the normalized version in the session
 467                 modA tsBindings (Map.insert bndr expr')
 468                 -- Find all vars used with a function type. All of these should be global
 469                 -- binders (i.e., functions used), since any local binders with a function
 470                 -- type should have been inlined already.
 471                 let used_funcs_set = CoreFVs.exprSomeFreeVars (\v -> (Type.isFunTy . snd . Type.splitForAllTys . Id.idType) v) expr'
 472                 let used_funcs = VarSet.varSetElems used_funcs_set
 473                 -- Process each of the used functions recursively
 474                 mapM normalizeBind used_funcs
 475                 return ()
 476               -- We don't have a value for this binder. This really shouldn't
 477               -- happen for local id's...
 478               Nothing -> error $ "No value found for binder " ++ pprString bndr ++ "? This should not happen!"