Add getNormalized_maybe that does not fail if a binder was not found.

[matthijs/master-project/cλash.git] / cλash / CLasH / Normalize.hs

diff --git a/cλash/CLasH/Normalize.hs b/cλash/CLasH/Normalize.hs
index 4743824bf2ba85313b3bb4bf1241f00e0f47e2d2..fa14ef65ef7dec1b818e73f86c3da4fdb5842028 100644 (file)
--- a/cλash/CLasH/Normalize.hs
+++ b/cλash/CLasH/Normalize.hs
@@ -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 "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.Monoid as Monoid

+import qualified Data.Map as Map

 
 -- GHC API
 import CoreSyn
 
 -- GHC API
 import CoreSyn


@@ -402,7 +404,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
 -- 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
   where
     is_used (_, _, expr) = expr_uses_binders [bndr] expr
     bndr_used = or $ map is_used alts


@@ -576,6 +578,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
           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
           -- Replace the original application with one of the new function to the
           -- new arguments.
           change $ MkCore.mkCoreApps (Var newf) newargs


@@ -616,6 +624,11 @@ argprop expr@(App _ _) | is_var fexpr = do
           let free_vars = VarSet.varSetElems $ CoreFVs.exprSomeFreeVars interesting arg
           -- Mark the current expression as changed
           setChanged
           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
           return (map Var free_vars, free_vars, arg)
         else do
           -- Representable types will not be propagated, and arguments with free


@@ -713,20 +726,40 @@ 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]
 
 -- 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]
 

--- | 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 ::
   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 function requested: " ++ show bndr)
+    norm
+
+-- | Returns the normalized version of the given function, or Nothing
+-- when the binder is not a known global binder.
+getNormalized_maybe ::
+  CoreBndr -- ^ The function to get
+  -> TranslatorSession (Maybe CoreExpr) -- The normalized function body
+
+getNormalized_maybe bndr = do
+    expr_maybe <- getGlobalBind bndr
+    if Maybe.isNothing expr_maybe
+      then
+        -- Binder 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 ::
 
 -- | Normalize an expression
 normalizeExpr ::


@@ -742,17 +775,6 @@ normalizeExpr what expr = do
       trace ("\n" ++ what ++ " after normalization:\n\n" ++ showSDoc ( ppr expr')) $ return ()
       return expr'
 
       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 ::
 -- | Split a normalized expression into the argument binders, top level
 --   bindings and the result binder.
 splitNormalized ::