Copy initial state when specializing functions.

[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..538b60d85c9cd53fc2dd67606c8b4aafc61572e5 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 Data.Accessor.Monad.Trans.State as MonadState
 import qualified Data.Monoid as Monoid
+import qualified Data.Map as Map
 
 -- 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
-    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 +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
+
+          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,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
+          -- 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