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