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