Add dictionary inlining transformation.
[matthijs/master-project/cλash.git] / cλash / CLasH / Normalize.hs
index ea1051694f52c405e3fd1d965b578a082389ec27..4c04e8557553d4044944749b0b49a85c699e779a 100644 (file)
@@ -13,7 +13,9 @@ import qualified List
 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.Accessor.Monad.Trans.State as MonadState
 import qualified Data.Monoid as Monoid
+import qualified Data.Map as Map
 
 -- GHC API
 import CoreSyn
@@ -336,23 +338,19 @@ inlinetoplevel, inlinetopleveltop :: Transform
 -- Any system name is candidate for inlining. Never inline user-defined
 -- functions, to preserve structure.
 inlinetoplevel expr@(Var f) | not $ isUserDefined f = do
-  norm <- isNormalizeable f
-  -- See if this is a top level binding for which we have a body
-  body_maybe <- Trans.lift $ getGlobalBind f
-  if norm && Maybe.isJust body_maybe
-    then do
-      -- Get the normalized version
-      norm <- Trans.lift $ getNormalized f
-      if needsInline norm 
-        then do
-          -- Regenerate all uniques in the to-be-inlined expression
-          norm_uniqued <- Trans.lift $ genUniques norm
-          change norm_uniqued
-        else
-          return expr
-    else
+  norm_maybe <- Trans.lift $ getNormalized_maybe f
+  case norm_maybe of
       -- No body or not normalizeable.
-      return expr
+    Nothing -> return expr
+    Just norm -> if needsInline norm then do
+        -- Regenerate all uniques in the to-be-inlined expression
+        norm_uniqued <- Trans.lift $ genUniques norm
+        -- And replace the variable reference with the unique'd body.
+        change norm_uniqued
+      else
+        -- No need to inline
+        return expr
+
 -- Leave all other expressions unchanged
 inlinetoplevel expr = return expr
 inlinetopleveltop = everywhere ("inlinetoplevel", inlinetoplevel)
@@ -366,6 +364,22 @@ needsInline expr = case splitNormalized expr of
   (args, [bind], res) -> True
   _ -> False
 
+
+--------------------------------
+-- Dictionary inlining
+--------------------------------
+-- Inline all top level dictionaries, so we can use them to resolve
+-- class methods based on the dictionary passed. 
+inlinedict expr@(Var f) | Id.isDictId f = do
+  body_maybe <- Trans.lift $ getGlobalBind f
+  case body_maybe of
+    Nothing -> return expr
+    Just body -> change body
+
+-- Leave all other expressions unchanged
+inlinedict expr = return expr
+inlinedicttop = everywhere ("inlinedict", inlinedict)
+
 --------------------------------
 -- Scrutinee simplification
 --------------------------------
@@ -402,7 +416,7 @@ scrutbndrremove, scrutbndrremovetop :: Transform
 -- all occurences of the binder with the scrutinee variable.
 scrutbndrremove (Case (Var scrut) bndr ty alts) | bndr_used = do
     alts' <- mapM subs_bndr alts
-    return $ Case (Var scrut) wild ty alts'
+    change $ Case (Var scrut) wild ty alts'
   where
     is_used (_, _, expr) = expr_uses_binders [bndr] expr
     bndr_used = or $ map is_used alts
@@ -576,6 +590,12 @@ argprop expr@(App _ _) | is_var fexpr = do
           let newbody = MkCore.mkCoreLams newparams (MkCore.mkCoreApps body oldargs)
           -- Create a new function with the same name but a new body
           newf <- Trans.lift $ mkFunction f newbody
+
+          Trans.lift $ MonadState.modify tsInitStates (\ismap ->
+            let init_state_maybe = Map.lookup f ismap in
+            case init_state_maybe of
+              Nothing -> ismap
+              Just init_state -> Map.insert newf init_state ismap)
           -- Replace the original application with one of the new function to the
           -- new arguments.
           change $ MkCore.mkCoreApps (Var newf) newargs
@@ -616,10 +636,18 @@ argprop expr@(App _ _) | is_var fexpr = do
           let free_vars = VarSet.varSetElems $ CoreFVs.exprSomeFreeVars interesting arg
           -- Mark the current expression as changed
           setChanged
+          -- TODO: Clone the free_vars (and update references in arg), since
+          -- this might cause conflicts if two arguments that are propagated
+          -- share a free variable. Also, we are now introducing new variables
+          -- into a function that are not fresh, which violates the binder
+          -- uniqueness invariant.
           return (map Var free_vars, free_vars, arg)
         else do
           -- Representable types will not be propagated, and arguments with free
           -- type variables will be propagated later.
+          -- Note that we implicitly remove any type variables in the type of
+          -- the original argument by using the type of the actual argument
+          -- for the new formal parameter.
           -- TODO: preserve original naming?
           id <- Trans.lift $ mkBinderFor arg "param"
           -- Just pass the original argument to the new function, which binds it
@@ -708,22 +736,43 @@ simplrestop expr = do
 
 
 -- What transforms to run?
-transforms = [inlinetopleveltop, argproptop, funextracttop, etatop, betatop, castproptop, letremovesimpletop, letderectop, letremovetop, letsimpltop, letflattop, scrutsimpltop, scrutbndrremovetop, casesimpltop, caseremovetop, inlinenonreptop, appsimpltop, letremoveunusedtop, castsimpltop, lambdasimpltop, simplrestop]
+transforms = [inlinedicttop, inlinetopleveltop, argproptop, funextracttop, etatop, betatop, castproptop, letremovesimpletop, letderectop, letremovetop, letsimpltop, letflattop, scrutsimpltop, scrutbndrremovetop, casesimpltop, caseremovetop, inlinenonreptop, appsimpltop, letremoveunusedtop, castsimpltop, lambdasimpltop, simplrestop]
 
--- | Returns the normalized version of the given function.
+-- | Returns the normalized version of the given function, or an error
+-- if it is not a known global binder.
 getNormalized ::
   CoreBndr -- ^ The function to get
   -> TranslatorSession CoreExpr -- The normalized function body
-
-getNormalized bndr = Utils.makeCached bndr tsNormalized $
-  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
-      expr <- getBinding bndr
-      normalizeExpr (show bndr) expr
+getNormalized bndr = do
+  norm <- getNormalized_maybe bndr
+  return $ Maybe.fromMaybe
+    (error $ "Normalize.getNormalized: Unknown or non-representable function requested: " ++ show bndr)
+    norm
+
+-- | Returns the normalized version of the given function, or Nothing
+-- when the binder is not a known global binder or is not normalizeable.
+getNormalized_maybe ::
+  CoreBndr -- ^ The function to get
+  -> TranslatorSession (Maybe CoreExpr) -- The normalized function body
+
+getNormalized_maybe bndr = do
+    expr_maybe <- getGlobalBind bndr
+    normalizeable <- isNormalizeable' bndr
+    if not normalizeable || Maybe.isNothing expr_maybe
+      then
+        -- Binder not normalizeable or not found
+        return Nothing
+      else 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
+          -- Binder found and is monomorphic. Normalize the expression
+          -- and cache the result.
+          normalized <- Utils.makeCached bndr tsNormalized $ 
+            normalizeExpr (show bndr) (Maybe.fromJust expr_maybe)
+          return (Just normalized)
 
 -- | Normalize an expression
 normalizeExpr ::
@@ -739,17 +788,6 @@ normalizeExpr what expr = do
       trace ("\n" ++ what ++ " after normalization:\n\n" ++ showSDoc ( ppr expr')) $ return ()
       return expr'
 
--- | Get the value that is bound to the given binder at top level. Fails when
---   there is no such binding.
-getBinding ::
-  CoreBndr -- ^ The binder to get the expression for
-  -> TranslatorSession CoreExpr -- ^ The value bound to the binder
-
-getBinding bndr = Utils.makeCached bndr tsBindings $
-  -- If the binding isn't in the "cache" (bindings map), then we can't create
-  -- it out of thin air, so return an error.
-  error $ "Normalize.getBinding: Unknown function requested: " ++ show bndr
-
 -- | Split a normalized expression into the argument binders, top level
 --   bindings and the result binder.
 splitNormalized ::