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.Accessor.Monad.Trans.State as MonadState
  17 import qualified Data.Monoid as Monoid
  18 import qualified Data.Map as Map
  19
  20 -- GHC API
  21 import CoreSyn
  22 import qualified CoreUtils
  23 import qualified Type
  24 import qualified Id
  25 import qualified Var
  26 import qualified Name
  27 import qualified VarSet
  28 import qualified CoreFVs
  29 import qualified Class
  30 import qualified MkCore
  31 import Outputable ( showSDoc, ppr, nest )
  32
  33 -- Local imports
  34 import CLasH.Normalize.NormalizeTypes
  35 import CLasH.Translator.TranslatorTypes
  36 import CLasH.Normalize.NormalizeTools
  37 import CLasH.VHDL.Constants (builtinIds)
  38 import qualified CLasH.Utils as Utils
  39 import CLasH.Utils.Core.CoreTools
  40 import CLasH.Utils.Core.BinderTools
  41 import CLasH.Utils.Pretty
  42
  43 --------------------------------
  44 -- Start of transformations
  45 --------------------------------
  46
  47 --------------------------------
  48 -- η expansion
  49 --------------------------------
  50 -- Make sure all parameters to the normalized functions are named by top
  51 -- level lambda expressions. For this we apply η expansion to the
  52 -- function body (possibly enclosed in some lambda abstractions) while
  53 -- it has a function type. Eventually this will result in a function
  54 -- body consisting of a bunch of nested lambdas containing a
  55 -- non-function value (e.g., a complete application).
  56 eta, etatop :: Transform
  57 eta c expr | is_fun expr && not (is_lam expr) && all (== LambdaBody) c = do
  58   let arg_ty = (fst . Type.splitFunTy . CoreUtils.exprType) expr
  59   id <- Trans.lift $ mkInternalVar "param" arg_ty
  60   change (Lam id (App expr (Var id)))
  61 -- Leave all other expressions unchanged
  62 eta c e = return e
  63 etatop = everywhere ("eta", eta)
  64
  65 --------------------------------
  66 -- β-reduction
  67 --------------------------------
  68 beta, betatop :: Transform
  69 -- Substitute arg for x in expr. For value lambda's, also clone before
  70 -- substitution.
  71 beta c (App (Lam x expr) arg) | CoreSyn.isTyVar x = setChanged >> substitute x arg c expr
  72                               | otherwise         = setChanged >> substitute_clone x arg c expr
  73 -- Propagate the application into the let
  74 beta c (App (Let binds expr) arg) = change $ Let binds (App expr arg)
  75 -- Propagate the application into each of the alternatives
  76 beta c (App (Case scrut b ty alts) arg) = change $ Case scrut b ty' alts'
  77   where
  78     alts' = map (\(con, bndrs, expr) -> (con, bndrs, (App expr arg))) alts
  79     ty' = CoreUtils.applyTypeToArg ty arg
  80 -- Leave all other expressions unchanged
  81 beta c expr = return expr
  82 -- Perform this transform everywhere
  83 betatop = everywhere ("beta", beta)
  84
  85 --------------------------------
  86 -- Cast propagation
  87 --------------------------------
  88 -- Try to move casts as much downward as possible.
  89 castprop, castproptop :: Transform
  90 castprop c (Cast (Let binds expr) ty) = change $ Let binds (Cast expr ty)
  91 castprop c expr@(Cast (Case scrut b _ alts) ty) = change (Case scrut b ty alts')
  92   where
  93     alts' = map (\(con, bndrs, expr) -> (con, bndrs, (Cast expr ty))) alts
  94 -- Leave all other expressions unchanged
  95 castprop c expr = return expr
  96 -- Perform this transform everywhere
  97 castproptop = everywhere ("castprop", castprop)
  98
  99 --------------------------------
 100 -- Cast simplification. Mostly useful for state packing and unpacking, but
 101 -- perhaps for others as well.
 102 --------------------------------
 103 castsimpl, castsimpltop :: Transform
 104 castsimpl c expr@(Cast val ty) = do
 105   -- Don't extract values that are already simpl
 106   local_var <- Trans.lift $ is_local_var val
 107   -- Don't extract values that are not representable, to prevent loops with
 108   -- inlinenonrep
 109   repr <- isRepr val
 110   if (not local_var) && repr
 111     then do
 112       -- Generate a binder for the expression
 113       id <- Trans.lift $ mkBinderFor val "castval"
 114       -- Extract the expression
 115       change $ Let (NonRec id val) (Cast (Var id) ty)
 116     else
 117       return expr
 118 -- Leave all other expressions unchanged
 119 castsimpl c expr = return expr
 120 -- Perform this transform everywhere
 121 castsimpltop = everywhere ("castsimpl", castsimpl)
 122
 123 --------------------------------
 124 -- Return value simplification
 125 --------------------------------
 126 -- Ensure the return value of a function follows proper normal form. eta
 127 -- expansion ensures the body starts with lambda abstractions, this
 128 -- transformation ensures that the lambda abstractions always contain a
 129 -- recursive let and that, when the return value is representable, the
 130 -- let contains a local variable reference in its body.
 131 retvalsimpl c expr | all (== LambdaBody) c && not (is_lam expr) && not (is_let expr) = do
 132   local_var <- Trans.lift $ is_local_var expr
 133   repr <- isRepr expr
 134   if not local_var && repr
 135     then do
 136       id <- Trans.lift $ mkBinderFor expr "res"
 137       change $ Let (Rec [(id, expr)]) (Var id)
 138     else
 139       return expr
 140
 141 retvalsimpl c expr@(Let (Rec binds) body) | all (== LambdaBody) c = do
 142   -- Don't extract values that are already a local variable, to prevent
 143   -- loops with ourselves.
 144   local_var <- Trans.lift $ is_local_var body
 145   -- Don't extract values that are not representable, to prevent loops with
 146   -- inlinenonrep
 147   repr <- isRepr body
 148   if not local_var && repr
 149     then do
 150       id <- Trans.lift $ mkBinderFor body "res"
 151       change $ Let (Rec ((id, body):binds)) (Var id)
 152     else
 153       return expr
 154
 155
 156 -- Leave all other expressions unchanged
 157 retvalsimpl c expr = return expr
 158 -- Perform this transform everywhere
 159 retvalsimpltop = everywhere ("retvalsimpl", retvalsimpl)
 160
 161 --------------------------------
 162 -- let derecursification
 163 --------------------------------
 164 letrec, letrectop :: Transform
 165 letrec c expr@(Let (NonRec bndr val) res) =
 166   change $ Let (Rec [(bndr, val)]) res
 167
 168 -- Leave all other expressions unchanged
 169 letrec c expr = return expr
 170 -- Perform this transform everywhere
 171 letrectop = everywhere ("letrec", letrec)
 172
 173 --------------------------------
 174 -- let flattening
 175 --------------------------------
 176 -- Takes a let that binds another let, and turns that into two nested lets.
 177 -- e.g., from:
 178 -- let b = (let b' = expr' in res') in res
 179 -- to:
 180 -- let b' = expr' in (let b = res' in res)
 181 letflat, letflattop :: Transform
 182 -- Turn a nonrec let that binds a let into two nested lets.
 183 letflat c (Let (NonRec b (Let binds  res')) res) =
 184   change $ Let binds (Let (NonRec b res') res)
 185 letflat c (Let (Rec binds) expr) = do
 186   -- Flatten each binding.
 187   binds' <- Utils.concatM $ Monad.mapM flatbind binds
 188   -- Return the new let. We don't use change here, since possibly nothing has
 189   -- changed. If anything has changed, flatbind has already flagged that
 190   -- change.
 191   return $ Let (Rec binds') expr
 192   where
 193     -- Turns a binding of a let into a multiple bindings, or any other binding
 194     -- into a list with just that binding
 195     flatbind :: (CoreBndr, CoreExpr) -> TransformMonad [(CoreBndr, CoreExpr)]
 196     flatbind (b, Let (Rec binds) expr) = change ((b, expr):binds)
 197     flatbind (b, Let (NonRec b' expr') expr) = change [(b, expr), (b', expr')]
 198     flatbind (b, expr) = return [(b, expr)]
 199 -- Leave all other expressions unchanged
 200 letflat c expr = return expr
 201 -- Perform this transform everywhere
 202 letflattop = everywhere ("letflat", letflat)
 203
 204 --------------------------------
 205 -- empty let removal
 206 --------------------------------
 207 -- Remove empty (recursive) lets
 208 letremove, letremovetop :: Transform
 209 letremove c (Let (Rec []) res) = change res
 210 -- Leave all other expressions unchanged
 211 letremove c expr = return expr
 212 -- Perform this transform everywhere
 213 letremovetop = everywhere ("letremove", letremove)
 214
 215 --------------------------------
 216 -- Simple let binding removal
 217 --------------------------------
 218 -- Remove a = b bindings from let expressions everywhere
 219 letremovesimpletop :: Transform
 220 letremovesimpletop = everywhere ("letremovesimple", inlinebind (\(b, e) -> Trans.lift $ is_local_var e))
 221
 222 --------------------------------
 223 -- Unused let binding removal
 224 --------------------------------
 225 letremoveunused, letremoveunusedtop :: Transform
 226 letremoveunused c expr@(Let (NonRec b bound) res) = do
 227   let used = expr_uses_binders [b] res
 228   if used
 229     then return expr
 230     else change res
 231 letremoveunused c expr@(Let (Rec binds) res) = do
 232   -- Filter out all unused binds.
 233   let binds' = filter dobind binds
 234   -- Only set the changed flag if binds got removed
 235   changeif (length binds' /= length binds) (Let (Rec binds') res)
 236     where
 237       bound_exprs = map snd binds
 238       -- For each bind check if the bind is used by res or any of the bound
 239       -- expressions
 240       dobind (bndr, _) = any (expr_uses_binders [bndr]) (res:bound_exprs)
 241 -- Leave all other expressions unchanged
 242 letremoveunused c expr = return expr
 243 letremoveunusedtop = everywhere ("letremoveunused", letremoveunused)
 244
 245 {-
 246 --------------------------------
 247 -- Identical let binding merging
 248 --------------------------------
 249 -- Merge two bindings in a let if they are identical
 250 -- TODO: We would very much like to use GHC's CSE module for this, but that
 251 -- doesn't track if something changed or not, so we can't use it properly.
 252 letmerge, letmergetop :: Transform
 253 letmerge c expr@(Let _ _) = do
 254   let (binds, res) = flattenLets expr
 255   binds' <- domerge binds
 256   return $ mkNonRecLets binds' res
 257   where
 258     domerge :: [(CoreBndr, CoreExpr)] -> TransformMonad [(CoreBndr, CoreExpr)]
 259     domerge [] = return []
 260     domerge (e:es) = do
 261       es' <- mapM (mergebinds e) es
 262       es'' <- domerge es'
 263       return (e:es'')
 264
 265     -- Uses the second bind to simplify the second bind, if applicable.
 266     mergebinds :: (CoreBndr, CoreExpr) -> (CoreBndr, CoreExpr) -> TransformMonad (CoreBndr, CoreExpr)
 267     mergebinds (b1, e1) (b2, e2)
 268       -- Identical expressions? Replace the second binding with a reference to
 269       -- the first binder.
 270       | CoreUtils.cheapEqExpr e1 e2 = change $ (b2, Var b1)
 271       -- Different expressions? Don't change
 272       | otherwise = return (b2, e2)
 273 -- Leave all other expressions unchanged
 274 letmerge c expr = return expr
 275 letmergetop = everywhere ("letmerge", letmerge)
 276 -}
 277
 278 --------------------------------
 279 -- Non-representable binding inlining
 280 --------------------------------
 281 -- Remove a = B bindings, with B of a non-representable type, from let
 282 -- expressions everywhere. This means that any value that we can't generate a
 283 -- signal for, will be inlined and hopefully turned into something we can
 284 -- represent.
 285 --
 286 -- This is a tricky function, which is prone to create loops in the
 287 -- transformations. To fix this, we make sure that no transformation will
 288 -- create a new let binding with a non-representable type. These other
 289 -- transformations will just not work on those function-typed values at first,
 290 -- but the other transformations (in particular β-reduction) should make sure
 291 -- that the type of those values eventually becomes representable.
 292 inlinenonreptop :: Transform
 293 inlinenonreptop = everywhere ("inlinenonrep", inlinebind ((Monad.liftM not) . isRepr . snd))
 294
 295 --------------------------------
 296 -- Top level function inlining
 297 --------------------------------
 298 -- This transformation inlines simple top level bindings. Simple
 299 -- currently means that the body is only a single application (though
 300 -- the complexity of the arguments is not currently checked) or that the
 301 -- normalized form only contains a single binding. This should catch most of the
 302 -- cases where a top level function is created that simply calls a type class
 303 -- method with a type and dictionary argument, e.g.
 304 --   fromInteger = GHC.Num.fromInteger (SizedWord D8) $dNum
 305 -- which is later called using simply
 306 --   fromInteger (smallInteger 10)
 307 --
 308 -- These useless wrappers are created by GHC automatically. If we don't
 309 -- inline them, we get loads of useless components cluttering the
 310 -- generated VHDL.
 311 --
 312 -- Note that the inlining could also inline simple functions defined by
 313 -- the user, not just GHC generated functions. It turns out to be near
 314 -- impossible to reliably determine what functions are generated and
 315 -- what functions are user-defined. Instead of guessing (which will
 316 -- inline less than we want) we will just inline all simple functions.
 317 --
 318 -- Only functions that are actually completely applied and bound by a
 319 -- variable in a let expression are inlined. These are the expressions
 320 -- that will eventually generate instantiations of trivial components.
 321 -- By not inlining any other reference, we also prevent looping problems
 322 -- with funextract and inlinedict.
 323 inlinetoplevel, inlinetopleveltop :: Transform
 324 inlinetoplevel (LetBinding:_) expr | not (is_fun expr) =
 325   case collectArgs expr of
 326         (Var f, args) -> do
 327           body_maybe <- needsInline f
 328           case body_maybe of
 329                 Just body -> do
 330                         -- Regenerate all uniques in the to-be-inlined expression
 331                         body_uniqued <- Trans.lift $ genUniques body
 332                         -- And replace the variable reference with the unique'd body.
 333                         change (mkApps body_uniqued args)
 334                         -- No need to inline
 335                 Nothing -> return expr
 336         -- This is not an application of a binder, leave it unchanged.
 337         _ -> return expr
 338
 339 -- Leave all other expressions unchanged
 340 inlinetoplevel c expr = return expr
 341 inlinetopleveltop = everywhere ("inlinetoplevel", inlinetoplevel)
 342
 343 -- | Does the given binder need to be inlined? If so, return the body to
 344 -- be used for inlining.
 345 needsInline :: CoreBndr -> TransformMonad (Maybe CoreExpr)
 346 needsInline f = do
 347   body_maybe <- Trans.lift $ getGlobalBind f
 348   case body_maybe of
 349     -- No body available?
 350     Nothing -> return Nothing
 351     Just body -> case CoreSyn.collectArgs body of
 352       -- The body is some (top level) binder applied to 0 or more
 353       -- arguments. That should be simple enough to inline.
 354       (Var f, args) -> return $ Just body
 355       -- Body is more complicated, try normalizing it
 356       _ -> do
 357         norm_maybe <- Trans.lift $ getNormalized_maybe f
 358         case norm_maybe of
 359           -- Noth normalizeable
 360           Nothing -> return Nothing
 361           Just norm -> case splitNormalizedNonRep norm of
 362             -- The function has just a single binding, so that's simple
 363             -- enough to inline.
 364             (args, [bind], Var res) -> return $ Just norm
 365             -- More complicated function, don't inline
 366             _ -> return Nothing
 367
 368 --------------------------------
 369 -- Dictionary inlining
 370 --------------------------------
 371 -- Inline all top level dictionaries, that are in a position where
 372 -- classopresolution can actually resolve them. This makes this
 373 -- transformation look similar to classoperesolution below, but we'll
 374 -- keep them separated for clarity. By not inlining other dictionaries,
 375 -- we prevent expression sizes exploding when huge type level integer
 376 -- dictionaries are inlined which can never be expanded (in casts, for
 377 -- example).
 378 inlinedict c expr@(App (App (Var sel) ty) (Var dict)) | not is_builtin && is_classop = do
 379   body_maybe <- Trans.lift $ getGlobalBind dict
 380   case body_maybe of
 381     -- No body available (no source available, or a local variable /
 382     -- argument)
 383     Nothing -> return expr
 384     Just body -> change (App (App (Var sel) ty) body)
 385   where
 386     -- Is this a builtin function / method?
 387     is_builtin = elem (Name.getOccString sel) builtinIds
 388     -- Are we dealing with a class operation selector?
 389     is_classop = Maybe.isJust (Id.isClassOpId_maybe sel)
 390
 391 -- Leave all other expressions unchanged
 392 inlinedict c expr = return expr
 393 inlinedicttop = everywhere ("inlinedict", inlinedict)
 394
 395 --------------------------------
 396 -- ClassOp resolution
 397 --------------------------------
 398 -- Resolves any class operation to the actual operation whenever
 399 -- possible. Class methods (as well as parent dictionary selectors) are
 400 -- special "functions" that take a type and a dictionary and evaluate to
 401 -- the corresponding method. A dictionary is nothing more than a
 402 -- special dataconstructor applied to the type the dictionary is for,
 403 -- each of the superclasses and all of the class method definitions for
 404 -- that particular type. Since dictionaries all always inlined (top
 405 -- levels dictionaries are inlined by inlinedict, local dictionaries are
 406 -- inlined by inlinenonrep), we will eventually have something like:
 407 --
 408 --   baz
 409 --     @ CLasH.HardwareTypes.Bit
 410 --     (D:Baz @ CLasH.HardwareTypes.Bit bitbaz)
 411 --
 412 -- Here, baz is the method selector for the baz method, while
 413 -- D:Baz is the dictionary constructor for the Baz and bitbaz is the baz
 414 -- method defined in the Baz Bit instance declaration.
 415 --
 416 -- To resolve this, we can look at the ClassOp IdInfo from the baz Id,
 417 -- which contains the Class it is defined for. From the Class, we can
 418 -- get a list of all selectors (both parent class selectors as well as
 419 -- method selectors). Since the arguments to D:Baz (after the type
 420 -- argument) correspond exactly to this list, we then look up baz in
 421 -- that list and replace the entire expression by the corresponding
 422 -- argument to D:Baz.
 423 --
 424 -- We don't resolve methods that have a builtin translation (such as
 425 -- ==), since the actual implementation is not always (easily)
 426 -- translateable. For example, when deriving ==, GHC generates code
 427 -- using $con2tag functions to translate a datacon to an int and compare
 428 -- that with GHC.Prim.==# . Better to avoid that for now.
 429 classopresolution, classopresolutiontop :: Transform
 430 classopresolution c expr@(App (App (Var sel) ty) dict) | not is_builtin =
 431   case Id.isClassOpId_maybe sel of
 432     -- Not a class op selector
 433     Nothing -> return expr
 434     Just cls -> case collectArgs dict of
 435       (_, []) -> return expr -- Dict is not an application (e.g., not inlined yet)
 436       (Var dictdc, (ty':selectors)) | not (Maybe.isJust (Id.isDataConId_maybe dictdc)) -> return expr -- Dictionary is not a datacon yet (but e.g., a top level binder)
 437                                 | tyargs_neq ty ty' -> error $ "Normalize.classopresolution: Applying class selector to dictionary without matching type?\n" ++ pprString expr
 438                                 | otherwise ->
 439         let selector_ids = Class.classSelIds cls in
 440         -- Find the selector used in the class' list of selectors
 441         case List.elemIndex sel selector_ids of
 442           Nothing -> error $ "Normalize.classopresolution: Selector not found in class' selector list? This should not happen!\nExpression: " ++ pprString expr ++ "\nClass: " ++ show cls ++ "\nSelectors: " ++ show selector_ids
 443           -- Get the corresponding argument from the dictionary
 444           Just n -> change (selectors!!n)
 445       (_, _) -> return expr -- Not applying a variable? Don't touch
 446   where
 447     -- Compare two type arguments, returning True if they are _not_
 448     -- equal
 449     tyargs_neq (Type ty1) (Type ty2) = not $ Type.coreEqType ty1 ty2
 450     tyargs_neq _ _ = True
 451     -- Is this a builtin function / method?
 452     is_builtin = elem (Name.getOccString sel) builtinIds
 453
 454 -- Leave all other expressions unchanged
 455 classopresolution c expr = return expr
 456 -- Perform this transform everywhere
 457 classopresolutiontop = everywhere ("classopresolution", classopresolution)
 458
 459 --------------------------------
 460 -- Scrutinee simplification
 461 --------------------------------
 462 scrutsimpl,scrutsimpltop :: Transform
 463 -- Don't touch scrutinees that are already simple
 464 scrutsimpl c expr@(Case (Var _) _ _ _) = return expr
 465 -- Replace all other cases with a let that binds the scrutinee and a new
 466 -- simple scrutinee, but only when the scrutinee is representable (to prevent
 467 -- loops with inlinenonrep, though I don't think a non-representable scrutinee
 468 -- will be supported anyway...)
 469 scrutsimpl c expr@(Case scrut b ty alts) = do
 470   repr <- isRepr scrut
 471   if repr
 472     then do
 473       id <- Trans.lift $ mkBinderFor scrut "scrut"
 474       change $ Let (NonRec id scrut) (Case (Var id) b ty alts)
 475     else
 476       return expr
 477 -- Leave all other expressions unchanged
 478 scrutsimpl c expr = return expr
 479 -- Perform this transform everywhere
 480 scrutsimpltop = everywhere ("scrutsimpl", scrutsimpl)
 481
 482 --------------------------------
 483 -- Scrutinee binder removal
 484 --------------------------------
 485 -- A case expression can have an extra binder, to which the scrutinee is bound
 486 -- after bringing it to WHNF. This is used for forcing evaluation of strict
 487 -- arguments. Since strictness does not matter for us (rather, everything is
 488 -- sort of strict), this binder is ignored when generating VHDL, and must thus
 489 -- be wild in the normal form.
 490 scrutbndrremove, scrutbndrremovetop :: Transform
 491 -- If the scrutinee is already simple, and the bndr is not wild yet, replace
 492 -- all occurences of the binder with the scrutinee variable.
 493 scrutbndrremove c (Case (Var scrut) bndr ty alts) | bndr_used = do
 494     alts' <- mapM subs_bndr alts
 495     change $ Case (Var scrut) wild ty alts'
 496   where
 497     is_used (_, _, expr) = expr_uses_binders [bndr] expr
 498     bndr_used = or $ map is_used alts
 499     subs_bndr (con, bndrs, expr) = do
 500       expr' <- substitute bndr (Var scrut) c expr
 501       return (con, bndrs, expr')
 502     wild = MkCore.mkWildBinder (Id.idType bndr)
 503 -- Leave all other expressions unchanged
 504 scrutbndrremove c expr = return expr
 505 scrutbndrremovetop = everywhere ("scrutbndrremove", scrutbndrremove)
 506
 507 --------------------------------
 508 -- Case binder wildening
 509 --------------------------------
 510 casesimpl, casesimpltop :: Transform
 511 -- This is already a selector case (or, if x does not appear in bndrs, a very
 512 -- simple case statement that will be removed by caseremove below). Just leave
 513 -- it be.
 514 casesimpl c expr@(Case scrut b ty [(con, bndrs, Var x)]) = return expr
 515 -- Make sure that all case alternatives have only wild binders and simple
 516 -- expressions.
 517 -- This is done by creating a new let binding for each non-wild binder, which
 518 -- is bound to a new simple selector case statement and for each complex
 519 -- expression. We do this only for representable types, to prevent loops with
 520 -- inlinenonrep.
 521 casesimpl c expr@(Case scrut bndr ty alts) | not bndr_used = do
 522   (bindingss, alts') <- (Monad.liftM unzip) $ mapM doalt alts
 523   let bindings = concat bindingss
 524   -- Replace the case with a let with bindings and a case
 525   let newlet = mkNonRecLets bindings (Case scrut bndr ty alts')
 526   -- If there are no non-wild binders, or this case is already a simple
 527   -- selector (i.e., a single alt with exactly one binding), already a simple
 528   -- selector altan no bindings (i.e., no wild binders in the original case),
 529   -- don't change anything, otherwise, replace the case.
 530   if null bindings then return expr else change newlet
 531   where
 532   -- Check if the scrutinee binder is used
 533   is_used (_, _, expr) = expr_uses_binders [bndr] expr
 534   bndr_used = or $ map is_used alts
 535   -- Generate a single wild binder, since they are all the same
 536   wild = MkCore.mkWildBinder
 537   -- Wilden the binders of one alt, producing a list of bindings as a
 538   -- sideeffect.
 539   doalt :: CoreAlt -> TransformMonad ([(CoreBndr, CoreExpr)], CoreAlt)
 540   doalt (con, bndrs, expr) = do
 541     -- Make each binder wild, if possible
 542     bndrs_res <- Monad.zipWithM dobndr bndrs [0..]
 543     let (newbndrs, bindings_maybe) = unzip bndrs_res
 544     -- Extract a complex expression, if possible. For this we check if any of
 545     -- the new list of bndrs are used by expr. We can't use free_vars here,
 546     -- since that looks at the old bndrs.
 547     let uses_bndrs = not $ VarSet.isEmptyVarSet $ CoreFVs.exprSomeFreeVars (`elem` newbndrs) expr
 548     (exprbinding_maybe, expr') <- doexpr expr uses_bndrs
 549     -- Create a new alternative
 550     let newalt = (con, newbndrs, expr')
 551     let bindings = Maybe.catMaybes (bindings_maybe ++ [exprbinding_maybe])
 552     return (bindings, newalt)
 553     where
 554       -- Make wild alternatives for each binder
 555       wildbndrs = map (\bndr -> MkCore.mkWildBinder (Id.idType bndr)) bndrs
 556       -- A set of all the binders that are used by the expression
 557       free_vars = CoreFVs.exprSomeFreeVars (`elem` bndrs) expr
 558       -- Look at the ith binder in the case alternative. Return a new binder
 559       -- for it (either the same one, or a wild one) and optionally a let
 560       -- binding containing a case expression.
 561       dobndr :: CoreBndr -> Int -> TransformMonad (CoreBndr, Maybe (CoreBndr, CoreExpr))
 562       dobndr b i = do
 563         repr <- isRepr b
 564         -- Is b wild (e.g., not a free var of expr. Since b is only in scope
 565         -- in expr, this means that b is unused if expr does not use it.)
 566         let wild = not (VarSet.elemVarSet b free_vars)
 567         -- Create a new binding for any representable binder that is not
 568         -- already wild and is representable (to prevent loops with
 569         -- inlinenonrep).
 570         if (not wild) && repr
 571           then do
 572             -- Create on new binder that will actually capture a value in this
 573             -- case statement, and return it.
 574             let bty = (Id.idType b)
 575             id <- Trans.lift $ mkInternalVar "sel" bty
 576             let binders = take i wildbndrs ++ [id] ++ drop (i+1) wildbndrs
 577             let caseexpr = Case scrut b bty [(con, binders, Var id)]
 578             return (wildbndrs!!i, Just (b, caseexpr))
 579           else
 580             -- Just leave the original binder in place, and don't generate an
 581             -- extra selector case.
 582             return (b, Nothing)
 583       -- Process the expression of a case alternative. Accepts an expression
 584       -- and whether this expression uses any of the binders in the
 585       -- alternative. Returns an optional new binding and a new expression.
 586       doexpr :: CoreExpr -> Bool -> TransformMonad (Maybe (CoreBndr, CoreExpr), CoreExpr)
 587       doexpr expr uses_bndrs = do
 588         local_var <- Trans.lift $ is_local_var expr
 589         repr <- isRepr expr
 590         -- Extract any expressions that do not use any binders from this
 591         -- alternative, is not a local var already and is representable (to
 592         -- prevent loops with inlinenonrep).
 593         if (not uses_bndrs) && (not local_var) && repr
 594           then do
 595             id <- Trans.lift $ mkBinderFor expr "caseval"
 596             -- We don't flag a change here, since casevalsimpl will do that above
 597             -- based on Just we return here.
 598             return (Just (id, expr), Var id)
 599           else
 600             -- Don't simplify anything else
 601             return (Nothing, expr)
 602 -- Leave all other expressions unchanged
 603 casesimpl c expr = return expr
 604 -- Perform this transform everywhere
 605 casesimpltop = everywhere ("casesimpl", casesimpl)
 606
 607 --------------------------------
 608 -- Case removal
 609 --------------------------------
 610 -- Remove case statements that have only a single alternative and only wild
 611 -- binders.
 612 caseremove, caseremovetop :: Transform
 613 -- Replace a useless case by the value of its single alternative
 614 caseremove c (Case scrut b ty [(con, bndrs, expr)]) | not usesvars = change expr
 615     -- Find if any of the binders are used by expr
 616     where usesvars = (not . VarSet.isEmptyVarSet . (CoreFVs.exprSomeFreeVars (`elem` b:bndrs))) expr
 617 -- Leave all other expressions unchanged
 618 caseremove c expr = return expr
 619 -- Perform this transform everywhere
 620 caseremovetop = everywhere ("caseremove", caseremove)
 621
 622 --------------------------------
 623 -- Argument extraction
 624 --------------------------------
 625 -- Make sure that all arguments of a representable type are simple variables.
 626 appsimpl, appsimpltop :: Transform
 627 -- Simplify all representable arguments. Do this by introducing a new Let
 628 -- that binds the argument and passing the new binder in the application.
 629 appsimpl c expr@(App f arg) = do
 630   -- Check runtime representability
 631   repr <- isRepr arg
 632   local_var <- Trans.lift $ is_local_var arg
 633   if repr && not local_var
 634     then do -- Extract representable arguments
 635       id <- Trans.lift $ mkBinderFor arg "arg"
 636       change $ Let (NonRec id arg) (App f (Var id))
 637     else -- Leave non-representable arguments unchanged
 638       return expr
 639 -- Leave all other expressions unchanged
 640 appsimpl c expr = return expr
 641 -- Perform this transform everywhere
 642 appsimpltop = everywhere ("appsimpl", appsimpl)
 643
 644 --------------------------------
 645 -- Function-typed argument propagation
 646 --------------------------------
 647 -- Remove all applications to function-typed arguments, by duplication the
 648 -- function called with the function-typed parameter replaced by the free
 649 -- variables of the argument passed in.
 650 argprop, argproptop :: Transform
 651 -- Transform any application of a named function (i.e., skip applications of
 652 -- lambda's). Also skip applications that have arguments with free type
 653 -- variables, since we can't inline those.
 654 argprop c expr@(App _ _) | is_var fexpr = do
 655   -- Find the body of the function called
 656   body_maybe <- Trans.lift $ getGlobalBind f
 657   case body_maybe of
 658     Just body -> do
 659       -- Process each of the arguments in turn
 660       (args', changed) <- Writer.listen $ mapM doarg args
 661       -- See if any of the arguments changed
 662       case Monoid.getAny changed of
 663         True -> do
 664           let (newargs', newparams', oldargs) = unzip3 args'
 665           let newargs = concat newargs'
 666           let newparams = concat newparams'
 667           -- Create a new body that consists of a lambda for all new arguments and
 668           -- the old body applied to some arguments.
 669           let newbody = MkCore.mkCoreLams newparams (MkCore.mkCoreApps body oldargs)
 670           -- Create a new function with the same name but a new body
 671           newf <- Trans.lift $ mkFunction f newbody
 672
 673           Trans.lift $ MonadState.modify tsInitStates (\ismap ->
 674             let init_state_maybe = Map.lookup f ismap in
 675             case init_state_maybe of
 676               Nothing -> ismap
 677               Just init_state -> Map.insert newf init_state ismap)
 678           -- Replace the original application with one of the new function to the
 679           -- new arguments.
 680           change $ MkCore.mkCoreApps (Var newf) newargs
 681         False ->
 682           -- Don't change the expression if none of the arguments changed
 683           return expr
 684
 685     -- If we don't have a body for the function called, leave it unchanged (it
 686     -- should be a primitive function then).
 687     Nothing -> return expr
 688   where
 689     -- Find the function called and the arguments
 690     (fexpr, args) = collectArgs expr
 691     Var f = fexpr
 692
 693     -- Process a single argument and return (args, bndrs, arg), where args are
 694     -- the arguments to replace the given argument in the original
 695     -- application, bndrs are the binders to include in the top-level lambda
 696     -- in the new function body, and arg is the argument to apply to the old
 697     -- function body.
 698     doarg :: CoreExpr -> TransformMonad ([CoreExpr], [CoreBndr], CoreExpr)
 699     doarg arg = do
 700       repr <- isRepr arg
 701       bndrs <- Trans.lift getGlobalBinders
 702       let interesting var = Var.isLocalVar var && (var `notElem` bndrs)
 703       if not repr && not (is_var arg && interesting (exprToVar arg)) && not (has_free_tyvars arg)
 704         then do
 705           -- Propagate all complex arguments that are not representable, but not
 706           -- arguments with free type variables (since those would require types
 707           -- not known yet, which will always be known eventually).
 708           -- Find interesting free variables, each of which should be passed to
 709           -- the new function instead of the original function argument.
 710           --
 711           -- Interesting vars are those that are local, but not available from the
 712           -- top level scope (functions from this module are defined as local, but
 713           -- they're not local to this function, so we can freely move references
 714           -- to them into another function).
 715           let free_vars = VarSet.varSetElems $ CoreFVs.exprSomeFreeVars interesting arg
 716           -- Mark the current expression as changed
 717           setChanged
 718           -- TODO: Clone the free_vars (and update references in arg), since
 719           -- this might cause conflicts if two arguments that are propagated
 720           -- share a free variable. Also, we are now introducing new variables
 721           -- into a function that are not fresh, which violates the binder
 722           -- uniqueness invariant.
 723           return (map Var free_vars, free_vars, arg)
 724         else do
 725           -- Representable types will not be propagated, and arguments with free
 726           -- type variables will be propagated later.
 727           -- Note that we implicitly remove any type variables in the type of
 728           -- the original argument by using the type of the actual argument
 729           -- for the new formal parameter.
 730           -- TODO: preserve original naming?
 731           id <- Trans.lift $ mkBinderFor arg "param"
 732           -- Just pass the original argument to the new function, which binds it
 733           -- to a new id and just pass that new id to the old function body.
 734           return ([arg], [id], mkReferenceTo id)
 735 -- Leave all other expressions unchanged
 736 argprop c expr = return expr
 737 -- Perform this transform everywhere
 738 argproptop = everywhere ("argprop", argprop)
 739
 740 --------------------------------
 741 -- Function-typed argument extraction
 742 --------------------------------
 743 -- This transform takes any function-typed argument that cannot be propagated
 744 -- (because the function that is applied to it is a builtin function), and
 745 -- puts it in a brand new top level binder. This allows us to for example
 746 -- apply map to a lambda expression This will not conflict with inlinenonrep,
 747 -- since that only inlines local let bindings, not top level bindings.
 748 funextract, funextracttop :: Transform
 749 funextract c expr@(App _ _) | is_var fexpr = do
 750   body_maybe <- Trans.lift $ getGlobalBind f
 751   case body_maybe of
 752     -- We don't have a function body for f, so we can perform this transform.
 753     Nothing -> do
 754       -- Find the new arguments
 755       args' <- mapM doarg args
 756       -- And update the arguments. We use return instead of changed, so the
 757       -- changed flag doesn't get set if none of the args got changed.
 758       return $ MkCore.mkCoreApps fexpr args'
 759     -- We have a function body for f, leave this application to funprop
 760     Just _ -> return expr
 761   where
 762     -- Find the function called and the arguments
 763     (fexpr, args) = collectArgs expr
 764     Var f = fexpr
 765     -- Change any arguments that have a function type, but are not simple yet
 766     -- (ie, a variable or application). This means to create a new function
 767     -- for map (\f -> ...) b, but not for map (foo a) b.
 768     --
 769     -- We could use is_applicable here instead of is_fun, but I think
 770     -- arguments to functions could only have forall typing when existential
 771     -- typing is enabled. Not sure, though.
 772     doarg arg | not (is_simple arg) && is_fun arg = do
 773       -- Create a new top level binding that binds the argument. Its body will
 774       -- be extended with lambda expressions, to take any free variables used
 775       -- by the argument expression.
 776       let free_vars = VarSet.varSetElems $ CoreFVs.exprFreeVars arg
 777       let body = MkCore.mkCoreLams free_vars arg
 778       id <- Trans.lift $ mkBinderFor body "fun"
 779       Trans.lift $ addGlobalBind id body
 780       -- Replace the argument with a reference to the new function, applied to
 781       -- all vars it uses.
 782       change $ MkCore.mkCoreApps (Var id) (map Var free_vars)
 783     -- Leave all other arguments untouched
 784     doarg arg = return arg
 785
 786 -- Leave all other expressions unchanged
 787 funextract c expr = return expr
 788 -- Perform this transform everywhere
 789 funextracttop = everywhere ("funextract", funextract)
 790
 791 --------------------------------
 792 -- End of transformations
 793 --------------------------------
 794
 795
 796
 797
 798 -- What transforms to run?
 799 transforms = [inlinedicttop, inlinetopleveltop, classopresolutiontop, argproptop, funextracttop, etatop, betatop, castproptop, letremovesimpletop, letrectop, letremovetop, retvalsimpltop, letflattop, scrutsimpltop, scrutbndrremovetop, casesimpltop, caseremovetop, inlinenonreptop, appsimpltop, letremoveunusedtop, castsimpltop]
 800
 801 -- | Returns the normalized version of the given function, or an error
 802 -- if it is not a known global binder.
 803 getNormalized ::
 804   CoreBndr -- ^ The function to get
 805   -> TranslatorSession CoreExpr -- The normalized function body
 806 getNormalized bndr = do
 807   norm <- getNormalized_maybe bndr
 808   return $ Maybe.fromMaybe
 809     (error $ "Normalize.getNormalized: Unknown or non-representable function requested: " ++ show bndr)
 810     norm
 811
 812 -- | Returns the normalized version of the given function, or Nothing
 813 -- when the binder is not a known global binder or is not normalizeable.
 814 getNormalized_maybe ::
 815   CoreBndr -- ^ The function to get
 816   -> TranslatorSession (Maybe CoreExpr) -- The normalized function body
 817
 818 getNormalized_maybe bndr = do
 819     expr_maybe <- getGlobalBind bndr
 820     normalizeable <- isNormalizeable' bndr
 821     if not normalizeable || Maybe.isNothing expr_maybe
 822       then
 823         -- Binder not normalizeable or not found
 824         return Nothing
 825       else if is_poly (Var bndr)
 826         then
 827           -- This should really only happen at the top level... TODO: Give
 828           -- a different error if this happens down in the recursion.
 829           error $ "\nNormalize.normalizeBind: Function " ++ show bndr ++ " is polymorphic, can't normalize"
 830         else do
 831           -- Binder found and is monomorphic. Normalize the expression
 832           -- and cache the result.
 833           normalized <- Utils.makeCached bndr tsNormalized $
 834             normalizeExpr (show bndr) (Maybe.fromJust expr_maybe)
 835           return (Just normalized)
 836
 837 -- | Normalize an expression
 838 normalizeExpr ::
 839   String -- ^ What are we normalizing? For debug output only.
 840   -> CoreSyn.CoreExpr -- ^ The expression to normalize
 841   -> TranslatorSession CoreSyn.CoreExpr -- ^ The normalized expression
 842
 843 normalizeExpr what expr = do
 844       startcount <- MonadState.get tsTransformCounter
 845       expr_uniqued <- genUniques expr
 846       -- Normalize this expression
 847       trace (what ++ " before normalization:\n\n" ++ showSDoc ( ppr expr_uniqued ) ++ "\n") $ return ()
 848       expr' <- dotransforms transforms expr_uniqued
 849       endcount <- MonadState.get tsTransformCounter
 850       trace ("\n" ++ what ++ " after normalization:\n\n" ++ showSDoc ( ppr expr')
 851              ++ "\nNeeded " ++ show (endcount - startcount) ++ " transformations to normalize " ++ what) $
 852        return expr'
 853
 854 -- | Split a normalized expression into the argument binders, top level
 855 --   bindings and the result binder. This function returns an error if
 856 --   the type of the expression is not representable.
 857 splitNormalized ::
 858   CoreExpr -- ^ The normalized expression
 859   -> ([CoreBndr], [Binding], CoreBndr)
 860 splitNormalized expr =
 861   case splitNormalizedNonRep expr of
 862     (args, binds, Var res) -> (args, binds, res)
 863     _ -> error $ "Normalize.splitNormalized: Not in normal form: " ++ pprString expr ++ "\n"
 864
 865 -- Split a normalized expression, whose type can be unrepresentable.
 866 splitNormalizedNonRep::
 867   CoreExpr -- ^ The normalized expression
 868   -> ([CoreBndr], [Binding], CoreExpr)
 869 splitNormalizedNonRep expr = (args, binds, resexpr)
 870   where
 871     (args, letexpr) = CoreSyn.collectBinders expr
 872     (binds, resexpr) = flattenLets letexpr