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
-- 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
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