cλash/CLasH/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 CLasH.Normalize (getNormalized, normalizeExpr, splitNormalized) where
   8
   9 -- Standard modules
  10 import Debug.Trace
  11 import qualified Maybe
  12 import qualified List
  13 import qualified "transformers" Control.Monad.Trans as Trans
  14 import qualified Control.Monad as Monad
  15 import qualified Control.Monad.Trans.Writer as Writer
  16 import qualified Data.Map as Map
  17 import qualified Data.Monoid as Monoid
  18 import Data.Accessor
  19
  20 -- GHC API
  21 import CoreSyn
  22 import qualified UniqSupply
  23 import qualified CoreUtils
  24 import qualified Type
  25 import qualified TcType
  26 import qualified Name
  27 import qualified Id
  28 import qualified Var
  29 import qualified VarSet
  30 import qualified NameSet
  31 import qualified CoreFVs
  32 import qualified CoreUtils
  33 import qualified MkCore
  34 import qualified HscTypes
  35 import Outputable ( showSDoc, ppr, nest )
  36
  37 -- Local imports
  38 import CLasH.Normalize.NormalizeTypes
  39 import CLasH.Translator.TranslatorTypes
  40 import CLasH.Normalize.NormalizeTools
  41 import CLasH.VHDL.VHDLTypes
  42 import qualified CLasH.Utils as Utils
  43 import CLasH.Utils.Core.CoreTools
  44 import CLasH.Utils.Core.BinderTools
  45 import CLasH.Utils.Pretty
  46
  47 --------------------------------
  48 -- Start of transformations
  49 --------------------------------
  50
  51 --------------------------------
  52 -- η abstraction
  53 --------------------------------
  54 eta, etatop :: Transform
  55 eta expr | is_fun expr && not (is_lam expr) = do
  56   let arg_ty = (fst . Type.splitFunTy . CoreUtils.exprType) expr
  57   id <- Trans.lift $ mkInternalVar "param" arg_ty
  58   change (Lam id (App expr (Var id)))
  59 -- Leave all other expressions unchanged
  60 eta e = return e
  61 etatop = notappargs ("eta", eta)
  62
  63 --------------------------------
  64 -- β-reduction
  65 --------------------------------
  66 beta, betatop :: Transform
  67 -- Substitute arg for x in expr
  68 beta (App (Lam x expr) arg) = change $ substitute [(x, arg)] expr
  69 -- Propagate the application into the let
  70 beta (App (Let binds expr) arg) = change $ Let binds (App expr arg)
  71 -- Propagate the application into each of the alternatives
  72 beta (App (Case scrut b ty alts) arg) = change $ Case scrut b ty' alts'
  73   where
  74     alts' = map (\(con, bndrs, expr) -> (con, bndrs, (App expr arg))) alts
  75     ty' = CoreUtils.applyTypeToArg ty arg
  76 -- Leave all other expressions unchanged
  77 beta expr = return expr
  78 -- Perform this transform everywhere
  79 betatop = everywhere ("beta", beta)
  80
  81 --------------------------------
  82 -- Cast propagation
  83 --------------------------------
  84 -- Try to move casts as much downward as possible.
  85 castprop, castproptop :: Transform
  86 castprop (Cast (Let binds expr) ty) = change $ Let binds (Cast expr ty)
  87 castprop expr@(Cast (Case scrut b _ alts) ty) = change (Case scrut b ty alts')
  88   where
  89     alts' = map (\(con, bndrs, expr) -> (con, bndrs, (Cast expr ty))) alts
  90 -- Leave all other expressions unchanged
  91 castprop expr = return expr
  92 -- Perform this transform everywhere
  93 castproptop = everywhere ("castprop", castprop)
  94
  95 --------------------------------
  96 -- Cast simplification. Mostly useful for state packing and unpacking, but
  97 -- perhaps for others as well.
  98 --------------------------------
  99 castsimpl, castsimpltop :: Transform
 100 castsimpl expr@(Cast val ty) = do
 101   -- Don't extract values that are already simpl
 102   local_var <- Trans.lift $ is_local_var val
 103   -- Don't extract values that are not representable, to prevent loops with
 104   -- inlinenonrep
 105   repr <- isRepr val
 106   if (not local_var) && repr
 107     then do
 108       -- Generate a binder for the expression
 109       id <- Trans.lift $ mkBinderFor val "castval"
 110       -- Extract the expression
 111       change $ Let (NonRec id val) (Cast (Var id) ty)
 112     else
 113       return expr
 114 -- Leave all other expressions unchanged
 115 castsimpl expr = return expr
 116 -- Perform this transform everywhere
 117 castsimpltop = everywhere ("castsimpl", castsimpl)
 118
 119
 120 --------------------------------
 121 -- Lambda simplication
 122 --------------------------------
 123 -- Ensure that a lambda always evaluates to a let expressions or a simple
 124 -- variable reference.
 125 lambdasimpl, lambdasimpltop :: Transform
 126 -- Don't simplify a lambda that evaluates to let, since this is already
 127 -- normal form (and would cause infinite loops).
 128 lambdasimpl expr@(Lam _ (Let _ _)) = return expr
 129 -- Put the of a lambda in its own binding, but not when the expression is
 130 -- already a local variable, or not representable (to prevent loops with
 131 -- inlinenonrep).
 132 lambdasimpl expr@(Lam bndr res) = do
 133   repr <- isRepr res
 134   local_var <- Trans.lift $ is_local_var res
 135   if not local_var && repr
 136     then do
 137       id <- Trans.lift $ mkBinderFor res "res"
 138       change $ Lam bndr (Let (NonRec id res) (Var id))
 139     else
 140       -- If the result is already a local var or not representable, don't
 141       -- extract it.
 142       return expr
 143
 144 -- Leave all other expressions unchanged
 145 lambdasimpl expr = return expr
 146 -- Perform this transform everywhere
 147 lambdasimpltop = everywhere ("lambdasimpl", lambdasimpl)
 148
 149 --------------------------------
 150 -- let derecursification
 151 --------------------------------
 152 letderec, letderectop :: Transform
 153 letderec expr@(Let (Rec binds) res) = case liftable of
 154   -- Nothing is liftable, just return
 155   [] -> return expr
 156   -- Something can be lifted, generate a new let expression
 157   _ -> change $ mkNonRecLets liftable (Let (Rec nonliftable) res)
 158   where
 159     -- Make a list of all the binders bound in this recursive let
 160     bndrs = map fst binds
 161     -- See which bindings are liftable
 162     (liftable, nonliftable) = List.partition canlift binds
 163     -- Any expression that does not use any of the binders in this recursive let
 164     -- can be lifted into a nonrec let. It can't use its own binder either,
 165     -- since that would mean the binding is self-recursive and should be in a
 166     -- single bind recursive let.
 167     canlift (bndr, e) = not $ expr_uses_binders bndrs e
 168 -- Leave all other expressions unchanged
 169 letderec expr = return expr
 170 -- Perform this transform everywhere
 171 letderectop = everywhere ("letderec", letderec)
 172
 173 --------------------------------
 174 -- let simplification
 175 --------------------------------
 176 letsimpl, letsimpltop :: Transform
 177 -- Don't simplify a let that evaluates to another let, since this is already
 178 -- normal form (and would cause infinite loops with letflat below).
 179 letsimpl expr@(Let _ (Let _ _)) = return expr
 180 -- Put the "in ..." value of a let in its own binding, but not when the
 181 -- expression is already a local variable, or not representable (to prevent loops with inlinenonrep).
 182 letsimpl expr@(Let binds res) = do
 183   repr <- isRepr res
 184   local_var <- Trans.lift $ is_local_var res
 185   if not local_var && repr
 186     then do
 187       -- If the result is not a local var already (to prevent loops with
 188       -- ourselves), extract it.
 189       id <- Trans.lift $ mkBinderFor res "foo"
 190       change $ Let binds (Let (NonRec id  res) (Var id))
 191     else
 192       -- If the result is already a local var, don't extract it.
 193       return expr
 194
 195 -- Leave all other expressions unchanged
 196 letsimpl expr = return expr
 197 -- Perform this transform everywhere
 198 letsimpltop = everywhere ("letsimpl", letsimpl)
 199
 200 --------------------------------
 201 -- let flattening
 202 --------------------------------
 203 -- Takes a let that binds another let, and turns that into two nested lets.
 204 -- e.g., from:
 205 -- let b = (let b' = expr' in res') in res
 206 -- to:
 207 -- let b' = expr' in (let b = res' in res)
 208 letflat, letflattop :: Transform
 209 -- Turn a nonrec let that binds a let into two nested lets.
 210 letflat (Let (NonRec b (Let binds  res')) res) =
 211   change $ Let binds (Let (NonRec b res') res)
 212 letflat (Let (Rec binds) expr) = do
 213   -- Flatten each binding.
 214   binds' <- Utils.concatM $ Monad.mapM flatbind binds
 215   -- Return the new let. We don't use change here, since possibly nothing has
 216   -- changed. If anything has changed, flatbind has already flagged that
 217   -- change.
 218   return $ Let (Rec binds') expr
 219   where
 220     -- Turns a binding of a let into a multiple bindings, or any other binding
 221     -- into a list with just that binding
 222     flatbind :: (CoreBndr, CoreExpr) -> TransformMonad [(CoreBndr, CoreExpr)]
 223     flatbind (b, Let (Rec binds) expr) = change ((b, expr):binds)
 224     flatbind (b, Let (NonRec b' expr') expr) = change [(b, expr), (b', expr')]
 225     flatbind (b, expr) = return [(b, expr)]
 226 -- Leave all other expressions unchanged
 227 letflat expr = return expr
 228 -- Perform this transform everywhere
 229 letflattop = everywhere ("letflat", letflat)
 230
 231 --------------------------------
 232 -- empty let removal
 233 --------------------------------
 234 -- Remove empty (recursive) lets
 235 letremove, letremovetop :: Transform
 236 letremove (Let (Rec []) res) = change $ res
 237 -- Leave all other expressions unchanged
 238 letremove expr = return expr
 239 -- Perform this transform everywhere
 240 letremovetop = everywhere ("letremove", letremove)
 241
 242 --------------------------------
 243 -- Simple let binding removal
 244 --------------------------------
 245 -- Remove a = b bindings from let expressions everywhere
 246 letremovesimpletop :: Transform
 247 letremovesimpletop = everywhere ("letremovesimple", inlinebind (\(b, e) -> Trans.lift $ is_local_var e))
 248
 249 --------------------------------
 250 -- Unused let binding removal
 251 --------------------------------
 252 letremoveunused, letremoveunusedtop :: Transform
 253 letremoveunused expr@(Let (NonRec b bound) res) = do
 254   let used = expr_uses_binders [b] res
 255   if used
 256     then return expr
 257     else change res
 258 letremoveunused expr@(Let (Rec binds) res) = do
 259   -- Filter out all unused binds.
 260   let binds' = filter dobind binds
 261   -- Only set the changed flag if binds got removed
 262   changeif (length binds' /= length binds) (Let (Rec binds') res)
 263     where
 264       bound_exprs = map snd binds
 265       -- For each bind check if the bind is used by res or any of the bound
 266       -- expressions
 267       dobind (bndr, _) = any (expr_uses_binders [bndr]) (res:bound_exprs)
 268 -- Leave all other expressions unchanged
 269 letremoveunused expr = return expr
 270 letremoveunusedtop = everywhere ("letremoveunused", letremoveunused)
 271
 272 {-
 273 --------------------------------
 274 -- Identical let binding merging
 275 --------------------------------
 276 -- Merge two bindings in a let if they are identical
 277 -- TODO: We would very much like to use GHC's CSE module for this, but that
 278 -- doesn't track if something changed or not, so we can't use it properly.
 279 letmerge, letmergetop :: Transform
 280 letmerge expr@(Let _ _) = do
 281   let (binds, res) = flattenLets expr
 282   binds' <- domerge binds
 283   return $ mkNonRecLets binds' res
 284   where
 285     domerge :: [(CoreBndr, CoreExpr)] -> TransformMonad [(CoreBndr, CoreExpr)]
 286     domerge [] = return []
 287     domerge (e:es) = do
 288       es' <- mapM (mergebinds e) es
 289       es'' <- domerge es'
 290       return (e:es'')
 291
 292     -- Uses the second bind to simplify the second bind, if applicable.
 293     mergebinds :: (CoreBndr, CoreExpr) -> (CoreBndr, CoreExpr) -> TransformMonad (CoreBndr, CoreExpr)
 294     mergebinds (b1, e1) (b2, e2)
 295       -- Identical expressions? Replace the second binding with a reference to
 296       -- the first binder.
 297       | CoreUtils.cheapEqExpr e1 e2 = change $ (b2, Var b1)
 298       -- Different expressions? Don't change
 299       | otherwise = return (b2, e2)
 300 -- Leave all other expressions unchanged
 301 letmerge expr = return expr
 302 letmergetop = everywhere ("letmerge", letmerge)
 303 -}
 304
 305 --------------------------------
 306 -- Function inlining
 307 --------------------------------
 308 -- Remove a = B bindings, with B :: a -> b, or B :: forall x . T, from let
 309 -- expressions everywhere. This means that any value that still needs to be
 310 -- applied to something else (polymorphic values need to be applied to a
 311 -- Type) will be inlined, and will eventually be applied to all their
 312 -- arguments.
 313 --
 314 -- This is a tricky function, which is prone to create loops in the
 315 -- transformations. To fix this, we make sure that no transformation will
 316 -- create a new let binding with a function type. These other transformations
 317 -- will just not work on those function-typed values at first, but the other
 318 -- transformations (in particular β-reduction) should make sure that the type
 319 -- of those values eventually becomes primitive.
 320 inlinenonreptop :: Transform
 321 inlinenonreptop = everywhere ("inlinenonrep", inlinebind ((Monad.liftM not) . isRepr . snd))
 322
 323 inlinetoplevel, inlinetopleveltop :: Transform
 324 -- Any system name is candidate for inlining. Never inline user-defined
 325 -- functions, to preserver structure.
 326 inlinetoplevel expr@(Var f) | not $ isUserDefined f = do
 327   norm <- isNormalizeable f
 328   -- See if this is a top level binding for which we have a body
 329   body_maybe <- Trans.lift $ getGlobalBind f
 330   if norm && Maybe.isJust body_maybe
 331     then do
 332       -- Get the normalized version
 333       norm <- Trans.lift $ getNormalized f
 334       if needsInline norm
 335         then do
 336           -- Regenerate all uniques in the to-be-inlined expression
 337           norm_uniqued <- Trans.lift $ genUniques norm
 338           change norm_uniqued
 339         else
 340           return expr
 341     else
 342       -- No body or not normalizeable.
 343       return expr
 344 -- Leave all other expressions unchanged
 345 inlinetoplevel expr = return expr
 346 inlinetopleveltop = everywhere ("inlinetoplevel", inlinetoplevel)
 347
 348 needsInline :: CoreExpr -> Bool
 349 needsInline expr = case splitNormalized expr of
 350   -- Inline any function that only has a single definition, it is probably
 351   -- simple enough. This might inline some stuff that it shouldn't though it
 352   -- will never inline user-defined functions (inlinetoplevel only tries
 353   -- system names) and inlining should never break things.
 354   (args, [bind], res) -> True
 355   _ -> False
 356
 357 --------------------------------
 358 -- Scrutinee simplification
 359 --------------------------------
 360 scrutsimpl,scrutsimpltop :: Transform
 361 -- Don't touch scrutinees that are already simple
 362 scrutsimpl expr@(Case (Var _) _ _ _) = return expr
 363 -- Replace all other cases with a let that binds the scrutinee and a new
 364 -- simple scrutinee, but only when the scrutinee is representable (to prevent
 365 -- loops with inlinenonrep, though I don't think a non-representable scrutinee
 366 -- will be supported anyway...)
 367 scrutsimpl expr@(Case scrut b ty alts) = do
 368   repr <- isRepr scrut
 369   if repr
 370     then do
 371       id <- Trans.lift $ mkBinderFor scrut "scrut"
 372       change $ Let (NonRec id scrut) (Case (Var id) b ty alts)
 373     else
 374       return expr
 375 -- Leave all other expressions unchanged
 376 scrutsimpl expr = return expr
 377 -- Perform this transform everywhere
 378 scrutsimpltop = everywhere ("scrutsimpl", scrutsimpl)
 379
 380 --------------------------------
 381 -- Case binder wildening
 382 --------------------------------
 383 casesimpl, casesimpltop :: Transform
 384 -- This is already a selector case (or, if x does not appear in bndrs, a very
 385 -- simple case statement that will be removed by caseremove below). Just leave
 386 -- it be.
 387 casesimpl expr@(Case scrut b ty [(con, bndrs, Var x)]) = return expr
 388 -- Make sure that all case alternatives have only wild binders and simple
 389 -- expressions.
 390 -- This is done by creating a new let binding for each non-wild binder, which
 391 -- is bound to a new simple selector case statement and for each complex
 392 -- expression. We do this only for representable types, to prevent loops with
 393 -- inlinenonrep.
 394 casesimpl expr@(Case scrut b ty alts) = do
 395   (bindingss, alts') <- (Monad.liftM unzip) $ mapM doalt alts
 396   let bindings = concat bindingss
 397   -- Replace the case with a let with bindings and a case
 398   let newlet = mkNonRecLets bindings (Case scrut b ty alts')
 399   -- If there are no non-wild binders, or this case is already a simple
 400   -- selector (i.e., a single alt with exactly one binding), already a simple
 401   -- selector altan no bindings (i.e., no wild binders in the original case),
 402   -- don't change anything, otherwise, replace the case.
 403   if null bindings then return expr else change newlet
 404   where
 405   -- Generate a single wild binder, since they are all the same
 406   wild = MkCore.mkWildBinder
 407   -- Wilden the binders of one alt, producing a list of bindings as a
 408   -- sideeffect.
 409   doalt :: CoreAlt -> TransformMonad ([(CoreBndr, CoreExpr)], CoreAlt)
 410   doalt (con, bndrs, expr) = do
 411     -- Make each binder wild, if possible
 412     bndrs_res <- Monad.zipWithM dobndr bndrs [0..]
 413     let (newbndrs, bindings_maybe) = unzip bndrs_res
 414     -- Extract a complex expression, if possible. For this we check if any of
 415     -- the new list of bndrs are used by expr. We can't use free_vars here,
 416     -- since that looks at the old bndrs.
 417     let uses_bndrs = not $ VarSet.isEmptyVarSet $ CoreFVs.exprSomeFreeVars (`elem` newbndrs) $ expr
 418     (exprbinding_maybe, expr') <- doexpr expr uses_bndrs
 419     -- Create a new alternative
 420     let newalt = (con, newbndrs, expr')
 421     let bindings = Maybe.catMaybes (bindings_maybe ++ [exprbinding_maybe])
 422     return (bindings, newalt)
 423     where
 424       -- Make wild alternatives for each binder
 425       wildbndrs = map (\bndr -> MkCore.mkWildBinder (Id.idType bndr)) bndrs
 426       -- A set of all the binders that are used by the expression
 427       free_vars = CoreFVs.exprSomeFreeVars (`elem` bndrs) expr
 428       -- Look at the ith binder in the case alternative. Return a new binder
 429       -- for it (either the same one, or a wild one) and optionally a let
 430       -- binding containing a case expression.
 431       dobndr :: CoreBndr -> Int -> TransformMonad (CoreBndr, Maybe (CoreBndr, CoreExpr))
 432       dobndr b i = do
 433         repr <- isRepr b
 434         -- Is b wild (e.g., not a free var of expr. Since b is only in scope
 435         -- in expr, this means that b is unused if expr does not use it.)
 436         let wild = not (VarSet.elemVarSet b free_vars)
 437         -- Create a new binding for any representable binder that is not
 438         -- already wild and is representable (to prevent loops with
 439         -- inlinenonrep).
 440         if (not wild) && repr
 441           then do
 442             -- Create on new binder that will actually capture a value in this
 443             -- case statement, and return it.
 444             let bty = (Id.idType b)
 445             id <- Trans.lift $ mkInternalVar "sel" bty
 446             let binders = take i wildbndrs ++ [id] ++ drop (i+1) wildbndrs
 447             let caseexpr = Case scrut b bty [(con, binders, Var id)]
 448             return (wildbndrs!!i, Just (b, caseexpr))
 449           else
 450             -- Just leave the original binder in place, and don't generate an
 451             -- extra selector case.
 452             return (b, Nothing)
 453       -- Process the expression of a case alternative. Accepts an expression
 454       -- and whether this expression uses any of the binders in the
 455       -- alternative. Returns an optional new binding and a new expression.
 456       doexpr :: CoreExpr -> Bool -> TransformMonad (Maybe (CoreBndr, CoreExpr), CoreExpr)
 457       doexpr expr uses_bndrs = do
 458         local_var <- Trans.lift $ is_local_var expr
 459         repr <- isRepr expr
 460         -- Extract any expressions that do not use any binders from this
 461         -- alternative, is not a local var already and is representable (to
 462         -- prevent loops with inlinenonrep).
 463         if (not uses_bndrs) && (not local_var) && repr
 464           then do
 465             id <- Trans.lift $ mkBinderFor expr "caseval"
 466             -- We don't flag a change here, since casevalsimpl will do that above
 467             -- based on Just we return here.
 468             return $ (Just (id, expr), Var id)
 469           else
 470             -- Don't simplify anything else
 471             return (Nothing, expr)
 472 -- Leave all other expressions unchanged
 473 casesimpl expr = return expr
 474 -- Perform this transform everywhere
 475 casesimpltop = everywhere ("casesimpl", casesimpl)
 476
 477 --------------------------------
 478 -- Case removal
 479 --------------------------------
 480 -- Remove case statements that have only a single alternative and only wild
 481 -- binders.
 482 caseremove, caseremovetop :: Transform
 483 -- Replace a useless case by the value of its single alternative
 484 caseremove (Case scrut b ty [(con, bndrs, expr)]) | not usesvars = change expr
 485     -- Find if any of the binders are used by expr
 486     where usesvars = (not . VarSet.isEmptyVarSet . (CoreFVs.exprSomeFreeVars (`elem` bndrs))) expr
 487 -- Leave all other expressions unchanged
 488 caseremove expr = return expr
 489 -- Perform this transform everywhere
 490 caseremovetop = everywhere ("caseremove", caseremove)
 491
 492 --------------------------------
 493 -- Argument extraction
 494 --------------------------------
 495 -- Make sure that all arguments of a representable type are simple variables.
 496 appsimpl, appsimpltop :: Transform
 497 -- Simplify all representable arguments. Do this by introducing a new Let
 498 -- that binds the argument and passing the new binder in the application.
 499 appsimpl expr@(App f arg) = do
 500   -- Check runtime representability
 501   repr <- isRepr arg
 502   local_var <- Trans.lift $ is_local_var arg
 503   if repr && not local_var
 504     then do -- Extract representable arguments
 505       id <- Trans.lift $ mkBinderFor arg "arg"
 506       change $ Let (NonRec id arg) (App f (Var id))
 507     else -- Leave non-representable arguments unchanged
 508       return expr
 509 -- Leave all other expressions unchanged
 510 appsimpl expr = return expr
 511 -- Perform this transform everywhere
 512 appsimpltop = everywhere ("appsimpl", appsimpl)
 513
 514 --------------------------------
 515 -- Function-typed argument propagation
 516 --------------------------------
 517 -- Remove all applications to function-typed arguments, by duplication the
 518 -- function called with the function-typed parameter replaced by the free
 519 -- variables of the argument passed in.
 520 argprop, argproptop :: Transform
 521 -- Transform any application of a named function (i.e., skip applications of
 522 -- lambda's). Also skip applications that have arguments with free type
 523 -- variables, since we can't inline those.
 524 argprop expr@(App _ _) | is_var fexpr = do
 525   -- Find the body of the function called
 526   body_maybe <- Trans.lift $ getGlobalBind f
 527   case body_maybe of
 528     Just body -> do
 529       -- Process each of the arguments in turn
 530       (args', changed) <- Writer.listen $ mapM doarg args
 531       -- See if any of the arguments changed
 532       case Monoid.getAny changed of
 533         True -> do
 534           let (newargs', newparams', oldargs) = unzip3 args'
 535           let newargs = concat newargs'
 536           let newparams = concat newparams'
 537           -- Create a new body that consists of a lambda for all new arguments and
 538           -- the old body applied to some arguments.
 539           let newbody = MkCore.mkCoreLams newparams (MkCore.mkCoreApps body oldargs)
 540           -- Create a new function with the same name but a new body
 541           newf <- Trans.lift $ mkFunction f newbody
 542           -- Replace the original application with one of the new function to the
 543           -- new arguments.
 544           change $ MkCore.mkCoreApps (Var newf) newargs
 545         False ->
 546           -- Don't change the expression if none of the arguments changed
 547           return expr
 548
 549     -- If we don't have a body for the function called, leave it unchanged (it
 550     -- should be a primitive function then).
 551     Nothing -> return expr
 552   where
 553     -- Find the function called and the arguments
 554     (fexpr, args) = collectArgs expr
 555     Var f = fexpr
 556
 557     -- Process a single argument and return (args, bndrs, arg), where args are
 558     -- the arguments to replace the given argument in the original
 559     -- application, bndrs are the binders to include in the top-level lambda
 560     -- in the new function body, and arg is the argument to apply to the old
 561     -- function body.
 562     doarg :: CoreExpr -> TransformMonad ([CoreExpr], [CoreBndr], CoreExpr)
 563     doarg arg = do
 564       repr <- isRepr arg
 565       bndrs <- Trans.lift getGlobalBinders
 566       let interesting var = Var.isLocalVar var && (not $ var `elem` bndrs)
 567       if not repr && not (is_var arg && interesting (exprToVar arg)) && not (has_free_tyvars arg)
 568         then do
 569           -- Propagate all complex arguments that are not representable, but not
 570           -- arguments with free type variables (since those would require types
 571           -- not known yet, which will always be known eventually).
 572           -- Find interesting free variables, each of which should be passed to
 573           -- the new function instead of the original function argument.
 574           --
 575           -- Interesting vars are those that are local, but not available from the
 576           -- top level scope (functions from this module are defined as local, but
 577           -- they're not local to this function, so we can freely move references
 578           -- to them into another function).
 579           let free_vars = VarSet.varSetElems $ CoreFVs.exprSomeFreeVars interesting arg
 580           -- Mark the current expression as changed
 581           setChanged
 582           return (map Var free_vars, free_vars, arg)
 583         else do
 584           -- Representable types will not be propagated, and arguments with free
 585           -- type variables will be propagated later.
 586           -- TODO: preserve original naming?
 587           id <- Trans.lift $ mkBinderFor arg "param"
 588           -- Just pass the original argument to the new function, which binds it
 589           -- to a new id and just pass that new id to the old function body.
 590           return ([arg], [id], mkReferenceTo id)
 591 -- Leave all other expressions unchanged
 592 argprop expr = return expr
 593 -- Perform this transform everywhere
 594 argproptop = everywhere ("argprop", argprop)
 595
 596 --------------------------------
 597 -- Function-typed argument extraction
 598 --------------------------------
 599 -- This transform takes any function-typed argument that cannot be propagated
 600 -- (because the function that is applied to it is a builtin function), and
 601 -- puts it in a brand new top level binder. This allows us to for example
 602 -- apply map to a lambda expression This will not conflict with inlinenonrep,
 603 -- since that only inlines local let bindings, not top level bindings.
 604 funextract, funextracttop :: Transform
 605 funextract expr@(App _ _) | is_var fexpr = do
 606   body_maybe <- Trans.lift $ getGlobalBind f
 607   case body_maybe of
 608     -- We don't have a function body for f, so we can perform this transform.
 609     Nothing -> do
 610       -- Find the new arguments
 611       args' <- mapM doarg args
 612       -- And update the arguments. We use return instead of changed, so the
 613       -- changed flag doesn't get set if none of the args got changed.
 614       return $ MkCore.mkCoreApps fexpr args'
 615     -- We have a function body for f, leave this application to funprop
 616     Just _ -> return expr
 617   where
 618     -- Find the function called and the arguments
 619     (fexpr, args) = collectArgs expr
 620     Var f = fexpr
 621     -- Change any arguments that have a function type, but are not simple yet
 622     -- (ie, a variable or application). This means to create a new function
 623     -- for map (\f -> ...) b, but not for map (foo a) b.
 624     --
 625     -- We could use is_applicable here instead of is_fun, but I think
 626     -- arguments to functions could only have forall typing when existential
 627     -- typing is enabled. Not sure, though.
 628     doarg arg | not (is_simple arg) && is_fun arg = do
 629       -- Create a new top level binding that binds the argument. Its body will
 630       -- be extended with lambda expressions, to take any free variables used
 631       -- by the argument expression.
 632       let free_vars = VarSet.varSetElems $ CoreFVs.exprFreeVars arg
 633       let body = MkCore.mkCoreLams free_vars arg
 634       id <- Trans.lift $ mkBinderFor body "fun"
 635       Trans.lift $ addGlobalBind id body
 636       -- Replace the argument with a reference to the new function, applied to
 637       -- all vars it uses.
 638       change $ MkCore.mkCoreApps (Var id) (map Var free_vars)
 639     -- Leave all other arguments untouched
 640     doarg arg = return arg
 641
 642 -- Leave all other expressions unchanged
 643 funextract expr = return expr
 644 -- Perform this transform everywhere
 645 funextracttop = everywhere ("funextract", funextract)
 646
 647 --------------------------------
 648 -- Ensure that a function that just returns another function (or rather,
 649 -- another top-level binder) is still properly normalized. This is a temporary
 650 -- solution, we should probably integrate this pass with lambdasimpl and
 651 -- letsimpl instead.
 652 --------------------------------
 653 simplrestop expr@(Lam _ _) = return expr
 654 simplrestop expr@(Let _ _) = return expr
 655 simplrestop expr = do
 656   local_var <- Trans.lift $ is_local_var expr
 657   -- Don't extract values that are not representable, to prevent loops with
 658   -- inlinenonrep
 659   repr <- isRepr expr
 660   if local_var || not repr
 661     then
 662       return expr
 663     else do
 664       id <- Trans.lift $ mkBinderFor expr "res"
 665       change $ Let (NonRec id expr) (Var id)
 666 --------------------------------
 667 -- End of transformations
 668 --------------------------------
 669
 670
 671
 672
 673 -- What transforms to run?
 674 transforms = [inlinetopleveltop, argproptop, funextracttop, etatop, betatop, castproptop, letremovesimpletop, letderectop, letremovetop, letsimpltop, letflattop, scrutsimpltop, casesimpltop, caseremovetop, inlinenonreptop, appsimpltop, letremoveunusedtop, castsimpltop, lambdasimpltop, simplrestop]
 675
 676 -- | Returns the normalized version of the given function.
 677 getNormalized ::
 678   CoreBndr -- ^ The function to get
 679   -> TranslatorSession CoreExpr -- The normalized function body
 680
 681 getNormalized bndr = Utils.makeCached bndr tsNormalized $ do
 682   if is_poly (Var bndr)
 683     then
 684       -- This should really only happen at the top level... TODO: Give
 685       -- a different error if this happens down in the recursion.
 686       error $ "\nNormalize.normalizeBind: Function " ++ show bndr ++ " is polymorphic, can't normalize"
 687     else do
 688       expr <- getBinding bndr
 689       normalizeExpr (show bndr) expr
 690
 691 -- | Normalize an expression
 692 normalizeExpr ::
 693   String -- ^ What are we normalizing? For debug output only.
 694   -> CoreSyn.CoreExpr -- ^ The expression to normalize
 695   -> TranslatorSession CoreSyn.CoreExpr -- ^ The normalized expression
 696
 697 normalizeExpr what expr = do
 698       expr_uniqued <- genUniques expr
 699       -- Normalize this expression
 700       trace (what ++ " before normalization:\n\n" ++ showSDoc ( ppr expr_uniqued ) ++ "\n") $ return ()
 701       expr' <- dotransforms transforms expr_uniqued
 702       trace ("\n" ++ what ++ " after normalization:\n\n" ++ showSDoc ( ppr expr')) $ return ()
 703       return expr'
 704
 705 -- | Get the value that is bound to the given binder at top level. Fails when
 706 --   there is no such binding.
 707 getBinding ::
 708   CoreBndr -- ^ The binder to get the expression for
 709   -> TranslatorSession CoreExpr -- ^ The value bound to the binder
 710
 711 getBinding bndr = Utils.makeCached bndr tsBindings $ do
 712   -- If the binding isn't in the "cache" (bindings map), then we can't create
 713   -- it out of thin air, so return an error.
 714   error $ "Normalize.getBinding: Unknown function requested: " ++ show bndr
 715
 716 -- | Split a normalized expression into the argument binders, top level
 717 --   bindings and the result binder.
 718 splitNormalized ::
 719   CoreExpr -- ^ The normalized expression
 720   -> ([CoreBndr], [Binding], CoreBndr)
 721 splitNormalized expr = (args, binds, res)
 722   where
 723     (args, letexpr) = CoreSyn.collectBinders expr
 724     (binds, resexpr) = flattenLets letexpr
 725     res = case resexpr of
 726       (Var x) -> x
 727       _ -> error $ "Normalize.splitNormalized: Not in normal form: " ++ pprString expr ++ "\n"