Add TODO regarding removal of is_wild.
[matthijs/master-project/cλash.git] / Normalize.hs
index 653cd68cfc1f599fd99cf7e69086154399ff9c19..0004ee3a9c01a90060173d6978abc5f7891873b6 100644 (file)
@@ -1,3 +1,4 @@
+{-# LANGUAGE PackageImports #-}
 --
 -- Functions to bring a Core expression in normal form. This module provides a
 -- top level function "normalize", and defines the actual transformation passes that
@@ -8,8 +9,11 @@ module Normalize (normalizeModule) where
 -- Standard modules
 import Debug.Trace
 import qualified Maybe
+import qualified "transformers" Control.Monad.Trans as Trans
 import qualified Control.Monad as Monad
+import qualified Control.Monad.Trans.Writer as Writer
 import qualified Data.Map as Map
+import qualified Data.Monoid as Monoid
 import Data.Accessor
 
 -- GHC API
@@ -18,14 +22,18 @@ import qualified UniqSupply
 import qualified CoreUtils
 import qualified Type
 import qualified Id
+import qualified Var
 import qualified VarSet
 import qualified CoreFVs
+import qualified CoreUtils
+import qualified MkCore
 import Outputable ( showSDoc, ppr, nest )
 
 -- Local imports
 import NormalizeTypes
 import NormalizeTools
 import CoreTools
+import Pretty
 
 --------------------------------
 -- Start of transformations
@@ -41,7 +49,7 @@ eta expr | is_fun expr && not (is_lam expr) = do
   change (Lam id (App expr (Var id)))
 -- Leave all other expressions unchanged
 eta e = return e
-etatop = notapplied ("eta", eta)
+etatop = notappargs ("eta", eta)
 
 --------------------------------
 -- β-reduction
@@ -55,12 +63,26 @@ beta (App (Let binds expr) arg) = change $ Let binds (App expr arg)
 beta (App (Case scrut b ty alts) arg) = change $ Case scrut b ty' alts'
   where 
     alts' = map (\(con, bndrs, expr) -> (con, bndrs, (App expr arg))) alts
-    (_, ty') = Type.splitFunTy ty
+    ty' = CoreUtils.applyTypeToArg ty arg
 -- Leave all other expressions unchanged
 beta expr = return expr
 -- Perform this transform everywhere
 betatop = everywhere ("beta", beta)
 
+--------------------------------
+-- Cast propagation
+--------------------------------
+-- Try to move casts as much downward as possible.
+castprop, castproptop :: Transform
+castprop (Cast (Let binds expr) ty) = change $ Let binds (Cast expr ty)
+castprop expr@(Cast (Case scrut b _ alts) ty) = change (Case scrut b ty alts')
+  where
+    alts' = map (\(con, bndrs, expr) -> (con, bndrs, (Cast expr ty))) alts
+-- Leave all other expressions unchanged
+castprop expr = return expr
+-- Perform this transform everywhere
+castproptop = everywhere ("castprop", castprop)
+
 --------------------------------
 -- let recursification
 --------------------------------
@@ -78,8 +100,8 @@ letsimpl, letsimpltop :: Transform
 -- Don't simplifiy lets that are already simple
 letsimpl expr@(Let _ (Var _)) = return expr
 -- Put the "in ..." value of a let in its own binding, but not when the
--- expression has a function type (to prevent loops with inlinefun).
-letsimpl (Let (Rec binds) expr) | not $ is_fun expr = do
+-- expression is applicable (to prevent loops with inlinefun).
+letsimpl (Let (Rec binds) expr) | not $ is_applicable expr = do
   id <- mkInternalVar "foo" (CoreUtils.exprType expr)
   let bind = (id, expr)
   change $ Let (Rec (bind:binds)) (Var id)
@@ -118,12 +140,17 @@ letflattop = everywhere ("letflat", letflat)
 --------------------------------
 -- Remove a = b bindings from let expressions everywhere
 letremovetop :: Transform
-letremovetop = everywhere ("letremove", inlinebind (\(b, e) -> case e of (Var v) -> True; otherwise -> False))
+letremovetop = everywhere ("letremove", inlinebind (\(b, e) -> case e of (Var v) | not $ Id.isDataConWorkId v -> True; otherwise -> False))
 
 --------------------------------
 -- Function inlining
 --------------------------------
--- Remove a = B bindings, with B :: a -> b, from let expressions everywhere.
+-- Remove a = B bindings, with B :: a -> b, or B :: forall x . T, from let
+-- expressions everywhere. This means that any value that still needs to be
+-- applied to something else (polymorphic values need to be applied to a
+-- Type) will be inlined, and will eventually be applied to all their
+-- arguments.
+--
 -- This is a tricky function, which is prone to create loops in the
 -- transformations. To fix this, we make sure that no transformation will
 -- create a new let binding with a function type. These other transformations
@@ -131,7 +158,7 @@ letremovetop = everywhere ("letremove", inlinebind (\(b, e) -> case e of (Var v)
 -- transformations (in particular β-reduction) should make sure that the type
 -- of those values eventually becomes primitive.
 inlinefuntop :: Transform
-inlinefuntop = everywhere ("inlinefun", inlinebind (Type.isFunTy . CoreUtils.exprType . snd))
+inlinefuntop = everywhere ("inlinefun", inlinebind (is_applicable . snd))
 
 --------------------------------
 -- Scrutinee simplification
@@ -140,10 +167,10 @@ scrutsimpl,scrutsimpltop :: Transform
 -- Don't touch scrutinees that are already simple
 scrutsimpl expr@(Case (Var _) _ _ _) = return expr
 -- Replace all other cases with a let that binds the scrutinee and a new
--- simple scrutinee, but not when the scrutinee is a function type (to prevent
--- loops with inlinefun, though I don't think a scrutinee can have a function
--- type...)
-scrutsimpl (Case scrut b ty alts) | not $ is_fun scrut = do
+-- simple scrutinee, but not when the scrutinee is applicable (to prevent
+-- loops with inlinefun, though I don't think a scrutinee can be
+-- applicable...)
+scrutsimpl (Case scrut b ty alts) | not $ is_applicable scrut = do
   id <- mkInternalVar "scrut" (CoreUtils.exprType scrut)
   change $ Let (Rec [(id, scrut)]) (Case (Var id) b ty alts)
 -- Leave all other expressions unchanged
@@ -186,6 +213,7 @@ casewild expr@(Case scrut b ty alts) = do
       -- and binds that to b.
       mkextracts :: CoreBndr -> Int -> TransformMonad (Maybe (CoreBndr, CoreExpr))
       mkextracts b i =
+        -- TODO: Use free variables instead of is_wild. is_wild is a hack.
         if is_wild b || Type.isFunTy (Id.idType b) 
           -- Don't create extra bindings for binders that are already wild, or
           -- for binders that bind function types (to prevent loops with
@@ -225,9 +253,9 @@ casevalsimpl expr@(Case scrut b ty alts) = do
     -- replacing the case value with that id. Only do this when the case value
     -- does not use any of the binders bound by this alternative, for that would
     -- cause those binders to become unbound when moving the value outside of
-    -- the case statement. Also, don't create a binding for function-typed
+    -- the case statement. Also, don't create a binding for applicable
     -- expressions, to prevent loops with inlinefun.
-    doalt (con, bndrs, expr) | (not usesvars) && (not $ is_fun expr) = do
+    doalt (con, bndrs, expr) | (not usesvars) && (not $ is_applicable expr) = do
       id <- mkInternalVar "caseval" (CoreUtils.exprType expr)
       -- We don't flag a change here, since casevalsimpl will do that above
       -- based on Just we return here.
@@ -261,12 +289,14 @@ caseremovetop = everywhere ("caseremove", caseremove)
 --------------------------------
 -- Make sure that all arguments in an application are simple variables.
 appsimpl, appsimpltop :: Transform
--- Don't simplify arguments that are already simple
-appsimpl expr@(App f (Var _)) = return expr
--- Simplify all arguments that do not have a function type (to prevent loops
--- with inlinefun) and is not a type argument. Do this by introducing a new
--- Let that binds the argument and passing the new binder in the application.
-appsimpl (App f expr) | (not $ is_fun expr) && (not $ CoreSyn.isTypeArg expr) = do
+-- Don't simplify arguments that are already simple. Do simplify datacons,
+-- however, since we can't portmap literals.
+appsimpl expr@(App f (Var v)) | not $ Id.isDataConWorkId v = return expr
+-- Simplify all non-applicable (to prevent loops with inlinefun) arguments,
+-- except for type arguments (since a let can't bind type vars, only a lambda
+-- can). Do this by introducing a new Let that binds the argument and passing
+-- the new binder in the application.
+appsimpl (App f expr) | (not $ is_applicable expr) && (not $ CoreSyn.isTypeArg expr) = do
   id <- mkInternalVar "arg" (CoreUtils.exprType expr)
   change $ Let (Rec [(id, expr)]) (App f (Var id))
 -- Leave all other expressions unchanged
@@ -274,7 +304,164 @@ appsimpl expr = return expr
 -- Perform this transform everywhere
 appsimpltop = everywhere ("appsimpl", appsimpl)
 
--- TODO: introduce top level let if needed?
+
+--------------------------------
+-- Type argument propagation
+--------------------------------
+-- Remove all applications to type arguments, by duplicating the function
+-- called with the type application in its new definition. We leave
+-- dictionaries that might be associated with the type untouched, the funprop
+-- transform should propagate these later on.
+typeprop, typeproptop :: Transform
+-- Transform any function that is applied to a type argument. Since type
+-- arguments are always the first ones to apply and we'll remove all type
+-- arguments, we can simply do them one by one. We only propagate type
+-- arguments without any free tyvars, since tyvars those wouldn't be in scope
+-- in the new function.
+typeprop expr@(App (Var f) arg@(Type ty)) | not $ has_free_tyvars arg = do
+  body_maybe <- Trans.lift $ getGlobalBind f
+  case body_maybe of
+    Just body -> do
+      let newbody = App body (Type ty)
+      -- Create a new function with the same name but a new body
+      newf <- mkFunction f newbody
+      -- Replace the application with this new function
+      change (Var newf)
+    -- If we don't have a body for the function called, leave it unchanged (it
+    -- should be a primitive function then).
+    Nothing -> return expr
+-- Leave all other expressions unchanged
+typeprop expr = return expr
+-- Perform this transform everywhere
+typeproptop = everywhere ("typeprop", typeprop)
+
+
+--------------------------------
+-- Function-typed argument propagation
+--------------------------------
+-- Remove all applications to function-typed arguments, by duplication the
+-- function called with the function-typed parameter replaced by the free
+-- variables of the argument passed in.
+funprop, funproptop :: Transform
+-- Transform any application of a named function (i.e., skip applications of
+-- lambda's). Also skip applications that have arguments with free type
+-- variables, since we can't inline those.
+funprop expr@(App _ _) | is_var fexpr && not (any has_free_tyvars args) = do
+  -- Find the body of the function called
+  body_maybe <- Trans.lift $ getGlobalBind f
+  case body_maybe of
+    Just body -> do
+      -- Process each of the arguments in turn
+      (args', changed) <- Writer.listen $ mapM doarg args
+      -- See if any of the arguments changed
+      case Monoid.getAny changed of
+        True -> do
+          let (newargs', newparams', oldargs) = unzip3 args'
+          let newargs = concat newargs'
+          let newparams = concat newparams'
+          -- Create a new body that consists of a lambda for all new arguments and
+          -- the old body applied to some arguments.
+          let newbody = MkCore.mkCoreLams newparams (MkCore.mkCoreApps body oldargs)
+          -- Create a new function with the same name but a new body
+          newf <- mkFunction f newbody
+          -- Replace the original application with one of the new function to the
+          -- new arguments.
+          change $ MkCore.mkCoreApps (Var newf) newargs
+        False ->
+          -- Don't change the expression if none of the arguments changed
+          return expr
+      
+    -- If we don't have a body for the function called, leave it unchanged (it
+    -- should be a primitive function then).
+    Nothing -> return expr
+  where
+    -- Find the function called and the arguments
+    (fexpr, args) = collectArgs expr
+    Var f = fexpr
+
+    -- Process a single argument and return (args, bndrs, arg), where args are
+    -- the arguments to replace the given argument in the original
+    -- application, bndrs are the binders to include in the top-level lambda
+    -- in the new function body, and arg is the argument to apply to the old
+    -- function body.
+    doarg :: CoreExpr -> TransformMonad ([CoreExpr], [CoreBndr], CoreExpr)
+    doarg arg | is_fun arg = do
+      bndrs <- Trans.lift getGlobalBinders
+      -- Find interesting free variables, each of which should be passed to
+      -- the new function instead of the original function argument.
+      -- 
+      -- Interesting vars are those that are local, but not available from the
+      -- top level scope (functions from this module are defined as local, but
+      -- they're not local to this function, so we can freely move references
+      -- to them into another function).
+      let interesting var = Var.isLocalVar var && (not $ var `elem` bndrs)
+      let free_vars = VarSet.varSetElems $ CoreFVs.exprSomeFreeVars interesting arg
+      -- Mark the current expression as changed
+      setChanged
+      return (map Var free_vars, free_vars, arg)
+    -- Non-functiontyped arguments can be unchanged. Note that this handles
+    -- both values and types.
+    doarg arg = do
+      -- TODO: preserve original naming?
+      id <- mkBinderFor arg "param"
+      -- Just pass the original argument to the new function, which binds it
+      -- to a new id and just pass that new id to the old function body.
+      return ([arg], [id], mkReferenceTo id) 
+-- Leave all other expressions unchanged
+funprop expr = return expr
+-- Perform this transform everywhere
+funproptop = everywhere ("funprop", funprop)
+
+--------------------------------
+-- Function-typed argument extraction
+--------------------------------
+-- This transform takes any function-typed argument that cannot be propagated
+-- (because the function that is applied to it is a builtin function), and
+-- puts it in a brand new top level binder. This allows us to for example
+-- apply map to a lambda expression This will not conflict with inlinefun,
+-- since that only inlines local let bindings, not top level bindings.
+funextract, funextracttop :: Transform
+funextract expr@(App _ _) | is_var fexpr = do
+  body_maybe <- Trans.lift $ getGlobalBind f
+  case body_maybe of
+    -- We don't have a function body for f, so we can perform this transform.
+    Nothing -> do
+      -- Find the new arguments
+      args' <- mapM doarg args
+      -- And update the arguments. We use return instead of changed, so the
+      -- changed flag doesn't get set if none of the args got changed.
+      return $ MkCore.mkCoreApps fexpr args'
+    -- We have a function body for f, leave this application to funprop
+    Just _ -> return expr
+  where
+    -- Find the function called and the arguments
+    (fexpr, args) = collectArgs expr
+    Var f = fexpr
+    -- Change any arguments that have a function type, but are not simple yet
+    -- (ie, a variable or application). This means to create a new function
+    -- for map (\f -> ...) b, but not for map (foo a) b.
+    --
+    -- We could use is_applicable here instead of is_fun, but I think
+    -- arguments to functions could only have forall typing when existential
+    -- typing is enabled. Not sure, though.
+    doarg arg | not (is_simple arg) && is_fun arg = do
+      -- Create a new top level binding that binds the argument. Its body will
+      -- be extended with lambda expressions, to take any free variables used
+      -- by the argument expression.
+      let free_vars = VarSet.varSetElems $ CoreFVs.exprFreeVars arg
+      let body = MkCore.mkCoreLams free_vars arg
+      id <- mkBinderFor body "fun"
+      Trans.lift $ addGlobalBind id body
+      -- Replace the argument with a reference to the new function, applied to
+      -- all vars it uses.
+      change $ MkCore.mkCoreApps (Var id) (map Var free_vars)
+    -- Leave all other arguments untouched
+    doarg arg = return arg
+
+-- Leave all other expressions unchanged
+funextract expr = return expr
+-- Perform this transform everywhere
+funextracttop = everywhere ("funextract", funextract)
 
 --------------------------------
 -- End of transformations
@@ -284,7 +471,7 @@ appsimpltop = everywhere ("appsimpl", appsimpl)
 
 
 -- What transforms to run?
-transforms = [etatop, betatop, letremovetop, letrectop, letsimpltop, letflattop, casewildtop, scrutsimpltop, casevalsimpltop, caseremovetop, inlinefuntop, appsimpltop]
+transforms = [typeproptop, funproptop, funextracttop, etatop, betatop, castproptop, letremovetop, letrectop, letsimpltop, letflattop, casewildtop, scrutsimpltop, casevalsimpltop, caseremovetop, inlinefuntop, appsimpltop]
 
 -- Turns the given bind into VHDL
 normalizeModule :: 
@@ -307,33 +494,48 @@ normalizeModule uniqsupply bindings generate_for statefuls = runTransformSession
   return $ filter ((flip VarSet.elemVarSet normalized_bindings) . fst) bindings
 
 normalizeBind :: CoreBndr -> TransformSession ()
-normalizeBind bndr = do
-  normalized_funcs <- getA tsNormalized
-  -- See if this function was normalized already
-  if VarSet.elemVarSet bndr normalized_funcs
-    then
-      -- Yup, don't do it again
-      return ()
-    else do
-      -- Nope, note that it has been and do it.
-      modA tsNormalized (flip VarSet.extendVarSet bndr)
-      expr_maybe <- getGlobalBind bndr
-      case expr_maybe of 
-        Just expr -> do
-          -- Normalize this expression
-          expr' <- dotransforms transforms expr
-          let expr'' = trace ("Before:\n\n" ++ showSDoc ( ppr expr ) ++ "\n\nAfter:\n\n" ++ showSDoc ( ppr expr')) expr'
-          -- And store the normalized version in the session
-          modA tsBindings (Map.insert bndr expr'')
-          -- Find all vars used with a function type. All of these should be global
-          -- binders (i.e., functions used), since any local binders with a function
-          -- type should have been inlined already.
-          let used_funcs_set = CoreFVs.exprSomeFreeVars (\v -> trace (showSDoc $ ppr $ Id.idType v) ((Type.isFunTy . snd . Type.splitForAllTys . Id.idType)v)) expr''
-          let used_funcs = VarSet.varSetElems used_funcs_set
-          -- Process each of the used functions recursively
-          mapM normalizeBind (trace (show used_funcs) used_funcs)
-          return ()
-        -- We don't have a value for this binder, let's assume this is a builtin
-        -- function. This might need some extra checking and a nice error
-        -- message).
-        Nothing -> return ()
+normalizeBind bndr =
+  -- Don't normalize global variables, these should be either builtin
+  -- functions or data constructors.
+  Monad.when (Var.isLocalIdVar bndr) $ do
+    -- Skip binders that have a polymorphic type, since it's impossible to
+    -- create polymorphic hardware.
+    if is_poly (Var bndr)
+      then
+        -- This should really only happen at the top level... TODO: Give
+        -- a different error if this happens down in the recursion.
+        error $ "\nNormalize.normalizeBind: Function " ++ show bndr ++ " is polymorphic, can't normalize"
+      else do
+        normalized_funcs <- getA tsNormalized
+        -- See if this function was normalized already
+        if VarSet.elemVarSet bndr normalized_funcs
+          then
+            -- Yup, don't do it again
+            return ()
+          else do
+            -- Nope, note that it has been and do it.
+            modA tsNormalized (flip VarSet.extendVarSet bndr)
+            expr_maybe <- getGlobalBind bndr
+            case expr_maybe of 
+              Just expr -> do
+                -- Introduce an empty Let at the top level, so there will always be
+                -- a let in the expression (none of the transformations will remove
+                -- the last let).
+                let expr' = Let (Rec []) expr
+                -- Normalize this expression
+                trace ("Transforming " ++ (show bndr) ++ "\nBefore:\n\n" ++ showSDoc ( ppr expr' ) ++ "\n") $ return ()
+                expr' <- dotransforms transforms expr'
+                trace ("\nAfter:\n\n" ++ showSDoc ( ppr expr')) $ return ()
+                -- And store the normalized version in the session
+                modA tsBindings (Map.insert bndr expr')
+                -- Find all vars used with a function type. All of these should be global
+                -- binders (i.e., functions used), since any local binders with a function
+                -- type should have been inlined already.
+                let used_funcs_set = CoreFVs.exprSomeFreeVars (\v -> (Type.isFunTy . snd . Type.splitForAllTys . Id.idType) v) expr'
+                let used_funcs = VarSet.varSetElems used_funcs_set
+                -- Process each of the used functions recursively
+                mapM normalizeBind used_funcs
+                return ()
+              -- We don't have a value for this binder. This really shouldn't
+              -- happen for local id's...
+              Nothing -> error $ "\nNormalize.normalizeBind: No value found for binder " ++ pprString bndr ++ "? This should not happen!"