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
 336           change norm
 337         else
 338           return expr
 339     else
 340       -- No body or not normalizeable.
 341       return expr
 342 -- Leave all other expressions unchanged
 343 inlinetoplevel expr = return expr
 344 inlinetopleveltop = everywhere ("inlinetoplevel", inlinetoplevel)
 345
 346 needsInline :: CoreExpr -> Bool
 347 needsInline expr = case splitNormalized expr of
 348   -- Inline any function that only has a single definition, it is probably
 349   -- simple enough. This might inline some stuff that it shouldn't though it
 350   -- will never inline user-defined functions (inlinetoplevel only tries
 351   -- system names) and inlining should never break things.
 352   (args, [bind], res) -> True
 353   _ -> False
 354
 355 --------------------------------
 356 -- Scrutinee simplification
 357 --------------------------------
 358 scrutsimpl,scrutsimpltop :: Transform
 359 -- Don't touch scrutinees that are already simple
 360 scrutsimpl expr@(Case (Var _) _ _ _) = return expr
 361 -- Replace all other cases with a let that binds the scrutinee and a new
 362 -- simple scrutinee, but only when the scrutinee is representable (to prevent
 363 -- loops with inlinenonrep, though I don't think a non-representable scrutinee
 364 -- will be supported anyway...)
 365 scrutsimpl expr@(Case scrut b ty alts) = do
 366   repr <- isRepr scrut
 367   if repr
 368     then do
 369       id <- Trans.lift $ mkBinderFor scrut "scrut"
 370       change $ Let (NonRec id scrut) (Case (Var id) b ty alts)
 371     else
 372       return expr
 373 -- Leave all other expressions unchanged
 374 scrutsimpl expr = return expr
 375 -- Perform this transform everywhere
 376 scrutsimpltop = everywhere ("scrutsimpl", scrutsimpl)
 377
 378 --------------------------------
 379 -- Case binder wildening
 380 --------------------------------
 381 casesimpl, casesimpltop :: Transform
 382 -- This is already a selector case (or, if x does not appear in bndrs, a very
 383 -- simple case statement that will be removed by caseremove below). Just leave
 384 -- it be.
 385 casesimpl expr@(Case scrut b ty [(con, bndrs, Var x)]) = return expr
 386 -- Make sure that all case alternatives have only wild binders and simple
 387 -- expressions.
 388 -- This is done by creating a new let binding for each non-wild binder, which
 389 -- is bound to a new simple selector case statement and for each complex
 390 -- expression. We do this only for representable types, to prevent loops with
 391 -- inlinenonrep.
 392 casesimpl expr@(Case scrut b ty alts) = do
 393   (bindingss, alts') <- (Monad.liftM unzip) $ mapM doalt alts
 394   let bindings = concat bindingss
 395   -- Replace the case with a let with bindings and a case
 396   let newlet = mkNonRecLets bindings (Case scrut b ty alts')
 397   -- If there are no non-wild binders, or this case is already a simple
 398   -- selector (i.e., a single alt with exactly one binding), already a simple
 399   -- selector altan no bindings (i.e., no wild binders in the original case),
 400   -- don't change anything, otherwise, replace the case.
 401   if null bindings then return expr else change newlet
 402   where
 403   -- Generate a single wild binder, since they are all the same
 404   wild = MkCore.mkWildBinder
 405   -- Wilden the binders of one alt, producing a list of bindings as a
 406   -- sideeffect.
 407   doalt :: CoreAlt -> TransformMonad ([(CoreBndr, CoreExpr)], CoreAlt)
 408   doalt (con, bndrs, expr) = do
 409     -- Make each binder wild, if possible
 410     bndrs_res <- Monad.zipWithM dobndr bndrs [0..]
 411     let (newbndrs, bindings_maybe) = unzip bndrs_res
 412     -- Extract a complex expression, if possible. For this we check if any of
 413     -- the new list of bndrs are used by expr. We can't use free_vars here,
 414     -- since that looks at the old bndrs.
 415     let uses_bndrs = not $ VarSet.isEmptyVarSet $ CoreFVs.exprSomeFreeVars (`elem` newbndrs) $ expr
 416     (exprbinding_maybe, expr') <- doexpr expr uses_bndrs
 417     -- Create a new alternative
 418     let newalt = (con, newbndrs, expr')
 419     let bindings = Maybe.catMaybes (bindings_maybe ++ [exprbinding_maybe])
 420     return (bindings, newalt)
 421     where
 422       -- Make wild alternatives for each binder
 423       wildbndrs = map (\bndr -> MkCore.mkWildBinder (Id.idType bndr)) bndrs
 424       -- A set of all the binders that are used by the expression
 425       free_vars = CoreFVs.exprSomeFreeVars (`elem` bndrs) expr
 426       -- Look at the ith binder in the case alternative. Return a new binder
 427       -- for it (either the same one, or a wild one) and optionally a let
 428       -- binding containing a case expression.
 429       dobndr :: CoreBndr -> Int -> TransformMonad (CoreBndr, Maybe (CoreBndr, CoreExpr))
 430       dobndr b i = do
 431         repr <- isRepr b
 432         -- Is b wild (e.g., not a free var of expr. Since b is only in scope
 433         -- in expr, this means that b is unused if expr does not use it.)
 434         let wild = not (VarSet.elemVarSet b free_vars)
 435         -- Create a new binding for any representable binder that is not
 436         -- already wild and is representable (to prevent loops with
 437         -- inlinenonrep).
 438         if (not wild) && repr
 439           then do
 440             -- Create on new binder that will actually capture a value in this
 441             -- case statement, and return it.
 442             let bty = (Id.idType b)
 443             id <- Trans.lift $ mkInternalVar "sel" bty
 444             let binders = take i wildbndrs ++ [id] ++ drop (i+1) wildbndrs
 445             let caseexpr = Case scrut b bty [(con, binders, Var id)]
 446             return (wildbndrs!!i, Just (b, caseexpr))
 447           else
 448             -- Just leave the original binder in place, and don't generate an
 449             -- extra selector case.
 450             return (b, Nothing)
 451       -- Process the expression of a case alternative. Accepts an expression
 452       -- and whether this expression uses any of the binders in the
 453       -- alternative. Returns an optional new binding and a new expression.
 454       doexpr :: CoreExpr -> Bool -> TransformMonad (Maybe (CoreBndr, CoreExpr), CoreExpr)
 455       doexpr expr uses_bndrs = do
 456         local_var <- Trans.lift $ is_local_var expr
 457         repr <- isRepr expr
 458         -- Extract any expressions that do not use any binders from this
 459         -- alternative, is not a local var already and is representable (to
 460         -- prevent loops with inlinenonrep).
 461         if (not uses_bndrs) && (not local_var) && repr
 462           then do
 463             id <- Trans.lift $ mkBinderFor expr "caseval"
 464             -- We don't flag a change here, since casevalsimpl will do that above
 465             -- based on Just we return here.
 466             return $ (Just (id, expr), Var id)
 467           else
 468             -- Don't simplify anything else
 469             return (Nothing, expr)
 470 -- Leave all other expressions unchanged
 471 casesimpl expr = return expr
 472 -- Perform this transform everywhere
 473 casesimpltop = everywhere ("casesimpl", casesimpl)
 474
 475 --------------------------------
 476 -- Case removal
 477 --------------------------------
 478 -- Remove case statements that have only a single alternative and only wild
 479 -- binders.
 480 caseremove, caseremovetop :: Transform
 481 -- Replace a useless case by the value of its single alternative
 482 caseremove (Case scrut b ty [(con, bndrs, expr)]) | not usesvars = change expr
 483     -- Find if any of the binders are used by expr
 484     where usesvars = (not . VarSet.isEmptyVarSet . (CoreFVs.exprSomeFreeVars (`elem` bndrs))) expr
 485 -- Leave all other expressions unchanged
 486 caseremove expr = return expr
 487 -- Perform this transform everywhere
 488 caseremovetop = everywhere ("caseremove", caseremove)
 489
 490 --------------------------------
 491 -- Argument extraction
 492 --------------------------------
 493 -- Make sure that all arguments of a representable type are simple variables.
 494 appsimpl, appsimpltop :: Transform
 495 -- Simplify all representable arguments. Do this by introducing a new Let
 496 -- that binds the argument and passing the new binder in the application.
 497 appsimpl expr@(App f arg) = do
 498   -- Check runtime representability
 499   repr <- isRepr arg
 500   local_var <- Trans.lift $ is_local_var arg
 501   if repr && not local_var
 502     then do -- Extract representable arguments
 503       id <- Trans.lift $ mkBinderFor arg "arg"
 504       change $ Let (NonRec id arg) (App f (Var id))
 505     else -- Leave non-representable arguments unchanged
 506       return expr
 507 -- Leave all other expressions unchanged
 508 appsimpl expr = return expr
 509 -- Perform this transform everywhere
 510 appsimpltop = everywhere ("appsimpl", appsimpl)
 511
 512 --------------------------------
 513 -- Function-typed argument propagation
 514 --------------------------------
 515 -- Remove all applications to function-typed arguments, by duplication the
 516 -- function called with the function-typed parameter replaced by the free
 517 -- variables of the argument passed in.
 518 argprop, argproptop :: Transform
 519 -- Transform any application of a named function (i.e., skip applications of
 520 -- lambda's). Also skip applications that have arguments with free type
 521 -- variables, since we can't inline those.
 522 argprop expr@(App _ _) | is_var fexpr = do
 523   -- Find the body of the function called
 524   body_maybe <- Trans.lift $ getGlobalBind f
 525   case body_maybe of
 526     Just body -> do
 527       -- Process each of the arguments in turn
 528       (args', changed) <- Writer.listen $ mapM doarg args
 529       -- See if any of the arguments changed
 530       case Monoid.getAny changed of
 531         True -> do
 532           let (newargs', newparams', oldargs) = unzip3 args'
 533           let newargs = concat newargs'
 534           let newparams = concat newparams'
 535           -- Create a new body that consists of a lambda for all new arguments and
 536           -- the old body applied to some arguments.
 537           let newbody = MkCore.mkCoreLams newparams (MkCore.mkCoreApps body oldargs)
 538           -- Create a new function with the same name but a new body
 539           newf <- Trans.lift $ mkFunction f newbody
 540           -- Replace the original application with one of the new function to the
 541           -- new arguments.
 542           change $ MkCore.mkCoreApps (Var newf) newargs
 543         False ->
 544           -- Don't change the expression if none of the arguments changed
 545           return expr
 546
 547     -- If we don't have a body for the function called, leave it unchanged (it
 548     -- should be a primitive function then).
 549     Nothing -> return expr
 550   where
 551     -- Find the function called and the arguments
 552     (fexpr, args) = collectArgs expr
 553     Var f = fexpr
 554
 555     -- Process a single argument and return (args, bndrs, arg), where args are
 556     -- the arguments to replace the given argument in the original
 557     -- application, bndrs are the binders to include in the top-level lambda
 558     -- in the new function body, and arg is the argument to apply to the old
 559     -- function body.
 560     doarg :: CoreExpr -> TransformMonad ([CoreExpr], [CoreBndr], CoreExpr)
 561     doarg arg = do
 562       repr <- isRepr arg
 563       bndrs <- Trans.lift getGlobalBinders
 564       let interesting var = Var.isLocalVar var && (not $ var `elem` bndrs)
 565       if not repr && not (is_var arg && interesting (exprToVar arg)) && not (has_free_tyvars arg)
 566         then do
 567           -- Propagate all complex arguments that are not representable, but not
 568           -- arguments with free type variables (since those would require types
 569           -- not known yet, which will always be known eventually).
 570           -- Find interesting free variables, each of which should be passed to
 571           -- the new function instead of the original function argument.
 572           --
 573           -- Interesting vars are those that are local, but not available from the
 574           -- top level scope (functions from this module are defined as local, but
 575           -- they're not local to this function, so we can freely move references
 576           -- to them into another function).
 577           let free_vars = VarSet.varSetElems $ CoreFVs.exprSomeFreeVars interesting arg
 578           -- Mark the current expression as changed
 579           setChanged
 580           return (map Var free_vars, free_vars, arg)
 581         else do
 582           -- Representable types will not be propagated, and arguments with free
 583           -- type variables will be propagated later.
 584           -- TODO: preserve original naming?
 585           id <- Trans.lift $ mkBinderFor arg "param"
 586           -- Just pass the original argument to the new function, which binds it
 587           -- to a new id and just pass that new id to the old function body.
 588           return ([arg], [id], mkReferenceTo id)
 589 -- Leave all other expressions unchanged
 590 argprop expr = return expr
 591 -- Perform this transform everywhere
 592 argproptop = everywhere ("argprop", argprop)
 593
 594 --------------------------------
 595 -- Function-typed argument extraction
 596 --------------------------------
 597 -- This transform takes any function-typed argument that cannot be propagated
 598 -- (because the function that is applied to it is a builtin function), and
 599 -- puts it in a brand new top level binder. This allows us to for example
 600 -- apply map to a lambda expression This will not conflict with inlinenonrep,
 601 -- since that only inlines local let bindings, not top level bindings.
 602 funextract, funextracttop :: Transform
 603 funextract expr@(App _ _) | is_var fexpr = do
 604   body_maybe <- Trans.lift $ getGlobalBind f
 605   case body_maybe of
 606     -- We don't have a function body for f, so we can perform this transform.
 607     Nothing -> do
 608       -- Find the new arguments
 609       args' <- mapM doarg args
 610       -- And update the arguments. We use return instead of changed, so the
 611       -- changed flag doesn't get set if none of the args got changed.
 612       return $ MkCore.mkCoreApps fexpr args'
 613     -- We have a function body for f, leave this application to funprop
 614     Just _ -> return expr
 615   where
 616     -- Find the function called and the arguments
 617     (fexpr, args) = collectArgs expr
 618     Var f = fexpr
 619     -- Change any arguments that have a function type, but are not simple yet
 620     -- (ie, a variable or application). This means to create a new function
 621     -- for map (\f -> ...) b, but not for map (foo a) b.
 622     --
 623     -- We could use is_applicable here instead of is_fun, but I think
 624     -- arguments to functions could only have forall typing when existential
 625     -- typing is enabled. Not sure, though.
 626     doarg arg | not (is_simple arg) && is_fun arg = do
 627       -- Create a new top level binding that binds the argument. Its body will
 628       -- be extended with lambda expressions, to take any free variables used
 629       -- by the argument expression.
 630       let free_vars = VarSet.varSetElems $ CoreFVs.exprFreeVars arg
 631       let body = MkCore.mkCoreLams free_vars arg
 632       id <- Trans.lift $ mkBinderFor body "fun"
 633       Trans.lift $ addGlobalBind id body
 634       -- Replace the argument with a reference to the new function, applied to
 635       -- all vars it uses.
 636       change $ MkCore.mkCoreApps (Var id) (map Var free_vars)
 637     -- Leave all other arguments untouched
 638     doarg arg = return arg
 639
 640 -- Leave all other expressions unchanged
 641 funextract expr = return expr
 642 -- Perform this transform everywhere
 643 funextracttop = everywhere ("funextract", funextract)
 644
 645 --------------------------------
 646 -- Ensure that a function that just returns another function (or rather,
 647 -- another top-level binder) is still properly normalized. This is a temporary
 648 -- solution, we should probably integrate this pass with lambdasimpl and
 649 -- letsimpl instead.
 650 --------------------------------
 651 simplrestop expr@(Lam _ _) = return expr
 652 simplrestop expr@(Let _ _) = return expr
 653 simplrestop expr = do
 654   local_var <- Trans.lift $ is_local_var expr
 655   -- Don't extract values that are not representable, to prevent loops with
 656   -- inlinenonrep
 657   repr <- isRepr expr
 658   if local_var || not repr
 659     then
 660       return expr
 661     else do
 662       id <- Trans.lift $ mkBinderFor expr "res"
 663       change $ Let (NonRec id expr) (Var id)
 664 --------------------------------
 665 -- End of transformations
 666 --------------------------------
 667
 668
 669
 670
 671 -- What transforms to run?
 672 transforms = [inlinetopleveltop, argproptop, funextracttop, etatop, betatop, castproptop, letremovesimpletop, letderectop, letremovetop, letsimpltop, letflattop, scrutsimpltop, casesimpltop, caseremovetop, inlinenonreptop, appsimpltop, letremoveunusedtop, castsimpltop, lambdasimpltop, simplrestop]
 673
 674 -- | Returns the normalized version of the given function.
 675 getNormalized ::
 676   CoreBndr -- ^ The function to get
 677   -> TranslatorSession CoreExpr -- The normalized function body
 678
 679 getNormalized bndr = Utils.makeCached bndr tsNormalized $ do
 680   if is_poly (Var bndr)
 681     then
 682       -- This should really only happen at the top level... TODO: Give
 683       -- a different error if this happens down in the recursion.
 684       error $ "\nNormalize.normalizeBind: Function " ++ show bndr ++ " is polymorphic, can't normalize"
 685     else do
 686       expr <- getBinding bndr
 687       normalizeExpr (show bndr) expr
 688
 689 -- | Normalize an expression
 690 normalizeExpr ::
 691   String -- ^ What are we normalizing? For debug output only.
 692   -> CoreSyn.CoreExpr -- ^ The expression to normalize
 693   -> TranslatorSession CoreSyn.CoreExpr -- ^ The normalized expression
 694
 695 normalizeExpr what expr = do
 696       expr_uniqued <- genUniques expr
 697       -- Normalize this expression
 698       trace (what ++ " before normalization:\n\n" ++ showSDoc ( ppr expr_uniqued ) ++ "\n") $ return ()
 699       expr' <- dotransforms transforms expr_uniqued
 700       trace ("\n" ++ what ++ " after normalization:\n\n" ++ showSDoc ( ppr expr')) $ return ()
 701       return expr'
 702
 703 -- | Get the value that is bound to the given binder at top level. Fails when
 704 --   there is no such binding.
 705 getBinding ::
 706   CoreBndr -- ^ The binder to get the expression for
 707   -> TranslatorSession CoreExpr -- ^ The value bound to the binder
 708
 709 getBinding bndr = Utils.makeCached bndr tsBindings $ do
 710   -- If the binding isn't in the "cache" (bindings map), then we can't create
 711   -- it out of thin air, so return an error.
 712   error $ "Normalize.getBinding: Unknown function requested: " ++ show bndr
 713
 714 -- | Split a normalized expression into the argument binders, top level
 715 --   bindings and the result binder.
 716 splitNormalized ::
 717   CoreExpr -- ^ The normalized expression
 718   -> ([CoreBndr], [Binding], CoreBndr)
 719 splitNormalized expr = (args, binds, res)
 720   where
 721     (args, letexpr) = CoreSyn.collectBinders expr
 722     (binds, resexpr) = flattenLets letexpr
 723     res = case resexpr of
 724       (Var x) -> x
 725       _ -> error $ "Normalize.splitNormalized: Not in normal form: " ++ pprString expr ++ "\n"