@@ -104,6 +104,7 @@ import Control.Monad.Exception
104104import Control.Monad.Primitive
105105import Control.Monad.Reader
106106import Control.Monad.Ref
107+import Control.Monad.State
107108#ifndef USE_TEMPLATE_HASKELL
108109import Data.Functor.Contravariant (phantom)
109110#endif
@@ -217,12 +218,11 @@ data HydrationDomBuilderEnv t m
217218 , _hydrationDomBuilderEnv_commitAction :: !(JSM ()) -- Action to take when all children are ready --TODO: we should probably get rid of this once we invoke it
218219 , _hydrationDomBuilderEnv_hydrationMode :: {-# UNPACK #-} !(IORef HydrationMode) -- In hydration mode?
219220 , _hydrationDomBuilderEnv_hydrationNode :: {-# UNPACK #-} !(IORef (Maybe Node)) -- Hydration node index
220− , _hydrationDomBuilderEnv_hydrationResult :: {-# UNPACK #-} !(IORef [Behavior t (m ())])
221221 , _hydrationDomBuilderEnv_prerenderDepth :: {-# UNPACK #-} !(IORef Word)
222222 , _hydrationDomBuilderEnv_eventChan :: {-# UNPACK #-} !(Chan [DSum (EventTriggerRef t) TriggerInvocation])
223223 }
224224
225−newtype HydrationDomBuilderT t m a = HydrationDomBuilderT { unHydrationDomBuilderT :: ReaderT (HydrationDomBuilderEnv t m) (RequesterT t JSM Identity (TriggerEventT t m)) a }
225+newtype HydrationDomBuilderT t m a = HydrationDomBuilderT { unHydrationDomBuilderT :: ReaderT (HydrationDomBuilderEnv t m) (StateT [Behavior t (m ())] (RequesterT t JSM Identity (TriggerEventT t m))) a }
226226 deriving (Functor, Applicative, Monad, MonadFix, MonadIO, MonadException
227227#if MIN_VERSION_base(4,9,1)
228228 , MonadAsyncException
@@ -234,14 +234,17 @@ instance PrimMonad m => PrimMonad (HydrationDomBuilderT x m) where
234234 primitive = lift . primitive
235235
236236instance MonadTrans (HydrationDomBuilderT t) where
237− lift = HydrationDomBuilderT . lift . lift . lift
237+ lift = HydrationDomBuilderT . lift . lift . lift . lift
238238
239239instance (Reflex t, MonadFix m) => DomRenderHook t (HydrationDomBuilderT t m) where
240240 withRenderHook hook (HydrationDomBuilderT a) = do
241241 e <- HydrationDomBuilderT ask
242− HydrationDomBuilderT $ lift $ withRequesting $ \rsp -> do
243− (x, req) <- lift $ runRequesterT (runReaderT a e) $ runIdentity <$> rsp
242+ s <- HydrationDomBuilderT get
243+ (a, s') <- HydrationDomBuilderT $ lift $ lift $ withRequesting $ \rsp -> do
244+ (x, req) <- lift $ runRequesterT (runStateT (runReaderT a e) s) $ runIdentity <$> rsp
244245 return (ffor req $ \rm -> hook $ traverseRequesterData (\r -> Identity <$> r) rm, x)
246+ HydrationDomBuilderT $ put s'
247+ pure a
245248 requestDomAction = HydrationDomBuilderT . lift . requestingIdentity
246249 requestDomAction_ = HydrationDomBuilderT . lift . requesting_
247250
@@ -258,10 +261,10 @@ runHydrationDomBuilderT
258261 => HydrationDomBuilderT t m a
259262 -> HydrationDomBuilderEnv t m
260263-- -> Chan [DSum (EventTriggerRef t) TriggerInvocation]
261− -> m a
264+ -> m (a, [Behavior t (m ())])
262265runHydrationDomBuilderT (HydrationDomBuilderT a) env =
263266 flip runTriggerEventT (_hydrationDomBuilderEnv_eventChan env) $ do
264− rec (x, req) <- runRequesterT (runReaderT a env) rsp
267+ rec (x, req) <- runRequesterT (runStateT (runReaderT a env) []) rsp
265268 rsp <- performEventAsync $ ffor req $ \rm f -> liftJSM $ runInAnimationFrame f $
266269 traverseRequesterData (\r -> Identity <$> r) rm
267270 return x
@@ -280,7 +283,7 @@ askParent = HydrationDomBuilderT $ asks _hydrationDomBuilderEnv_parent
280283
281284{-# INLINABLE askEvents #-}
282285askEvents :: Monad m => HydrationDomBuilderT t m (Chan [DSum (EventTriggerRef t) TriggerInvocation])
283−askEvents = HydrationDomBuilderT . lift . lift $ TriggerEventT.askEvents
286+askEvents = HydrationDomBuilderT . lift . lift . lift $ TriggerEventT.askEvents
284287
285288localEnv :: Monad m => (HydrationDomBuilderEnv t m -> HydrationDomBuilderEnv t m) -> HydrationDomBuilderT t m a -> HydrationDomBuilderT t m a
286289localEnv f = HydrationDomBuilderT . local f . unHydrationDomBuilderT
@@ -473,8 +476,7 @@ makeElement elementTag cfg child = do
473476 let beh = pure $ void . flip runHydrationDomBuilderT env $ do
474477 es <- wrap events e $ extractRawElementConfig cfg
475478 liftIO $ fire es
476− hydrationResultRef <- HydrationDomBuilderT $ asks _hydrationDomBuilderEnv_hydrationResult
477− liftIO $ modifyIORef' hydrationResultRef (beh :)
479+ HydrationDomBuilderT $ modify (beh :)
478480 collapsed <- switchHold never $ ffor ese $ \es -> DMap.singleton (WrapArg Click) . Identity <$> (select es (WrapArg Click)) -- TODO! this can't stay here, obviously.
479481 let evts = fan collapsed
480482 return ((Element evts (), result), e)
@@ -729,22 +731,19 @@ instance (Reflex t, Adjustable t m, MonadJSM m, MonadHold t m, MonadFix m, PrimM
729731 HydrationMode_Immediate -> toNode <$> createDocumentFragment doc
730732 unreadyChildren <- liftIO $ newIORef 0
731733 hydrationNode <- liftIO $ newIORef Nothing
732− r <- liftIO $ newIORef []
733− result <- runReaderT (unHydrationDomBuilderT a0) initialEnv
734+ (result, h') <- flip runStateT [] $ runReaderT (unHydrationDomBuilderT a0) initialEnv
734735 { _hydrationDomBuilderEnv_unreadyChildren = unreadyChildren
735736 , _hydrationDomBuilderEnv_commitAction = myCommitAction
736737 , _hydrationDomBuilderEnv_hydrationNode = hydrationNode
737738 , _hydrationDomBuilderEnv_parent = p
738− , _hydrationDomBuilderEnv_hydrationResult = r
739739 }
740− h' <- liftIO $ readIORef r
741740 liftIO $ readIORef unreadyChildren >>= \case
742741 0 -> writeIORef haveEverBeenReady True
743742 _ -> modifyIORef' parentUnreadyChildren succ
744743 liftIO $ putStrLn "drawInitialChild: end"
745744 return (h', result)
746745 a'' <- numberOccurrences a'
747− (result0, child') <- HydrationDomBuilderT $ lift $ runWithReplace drawInitialChild $ ffor a'' $ \(cohortId, child) -> do
746+ (result0, child') <- HydrationDomBuilderT $ lift $ lift $ runWithReplace drawInitialChild $ ffor a'' $ \(cohortId, child) -> do
748747 p <- liftIO (readIORef hydrating) >>= \case
749748 HydrationMode_Hydrating -> pure parent
750749 HydrationMode_Immediate -> toNode <$> createDocumentFragment doc
@@ -758,17 +757,14 @@ instance (Reflex t, Adjustable t m, MonadJSM m, MonadHold t m, MonadFix m, PrimM
758757 liftIO $ writeIORef currentCohort cohortId
759758 myCommitAction
760759 hydrationNode <- liftIO $ newIORef (Just $ toNode before)
761− r <- liftIO $ newIORef []
762− result <- runReaderT (unHydrationDomBuilderT child) $ initialEnv
760+ (result, h') <- flip runStateT [] $ runReaderT (unHydrationDomBuilderT child) $ initialEnv
763761 { _hydrationDomBuilderEnv_unreadyChildren = unreadyChildren
764762 , _hydrationDomBuilderEnv_commitAction = commitAction
765763 , _hydrationDomBuilderEnv_hydrationNode = hydrationNode
766764 , _hydrationDomBuilderEnv_parent = p
767− , _hydrationDomBuilderEnv_hydrationResult = r
768765 }
769766 uc <- liftIO $ readIORef unreadyChildren
770767 h <- liftIO $ readIORef hydrating
771− h' <- liftIO $ readIORef r
772768 let commitActionToRunNow = if uc == 0
773769 then Just $ commitAction
774770 else Nothing -- A child will run it when unreadyChildren is decremented to 0
@@ -776,13 +772,10 @@ instance (Reflex t, Adjustable t m, MonadJSM m, MonadHold t m, MonadFix m, PrimM
776772 hydrationAction = if h == HydrationMode_Hydrating then Just h' else Nothing
777773 return (hydrationAction, commitActionToRunNow', result)
778774 env <- HydrationDomBuilderT ask
779− let f :: RequesterT t JSM Identity (TriggerEventT t m) x -> m ()
780− f = void . flip runHydrationDomBuilderT env . HydrationDomBuilderT . lift
781− g :: [Behavior t (m ())] -> Behavior t (m ())
782− g = fmap (void . sequence) . sequence
783− beh <- hold (g $ fst result0) (fmapMaybe (\(h,_,_) -> g <$> h) child')
784− hydrationResultRef <- HydrationDomBuilderT $ asks _hydrationDomBuilderEnv_hydrationResult
785− liftIO $ modifyIORef' hydrationResultRef (join beh :)
775+ let f :: [Behavior t (m ())] -> Behavior t (m ())
776+ f = fmap (void . sequence) . sequence
777+ beh <- hold (f $ fst result0) (fmapMaybe (\(h,_,_) -> f <$> h) child')
778+ HydrationDomBuilderT $ modify (join beh :)
786779 requestDomAction_ $ fmapMaybe (\(_,c,_) -> c) child'
787780 return (snd result0, fmap (\(_,_,r) -> r) child')
788781
@@ -953,7 +946,7 @@ hoistTraverseIntMapWithKeyWithAdjust base updateChildUnreadiness applyDomUpdate_
953946 liftIO $ writeIORef pendingChange (newUnready, p)
954947 when (IntMap.null newUnready) $ do
955948 applyDomUpdate p
956− (children0, children') <- HydrationDomBuilderT $ lift $ base (\k v -> drawChildUpdateInt initialEnv markChildReady $ f k v) dm0 dm'
949+ (children0, children') <- HydrationDomBuilderT $ lift $ lift $ base (\k v -> drawChildUpdateInt initialEnv markChildReady $ f k v) dm0 dm'
957950 let processChild k (_, _, sRef, _) = do
958951 readIORef sRef >>= \case
959952 ChildReadyStateInt_Ready -> return Nothing
@@ -1043,7 +1036,7 @@ hoistTraverseWithKeyWithAdjust base mapPatch updateChildUnreadiness applyDomUpda
10431036 liftIO $ writeIORef pendingChange (newUnready, p)
10441037 when (DMap.null newUnready) $ do
10451038 applyDomUpdate p
1046− (children0, children') <- HydrationDomBuilderT $ lift $ base (\k v -> drawChildUpdate initialEnv markChildReady $ f k v) dm0 dm'
1039+ (children0, children') <- HydrationDomBuilderT $ lift $ lift $ base (\k v -> drawChildUpdate initialEnv markChildReady $ f k v) dm0 dm'
10471040 let processChild k (Compose (_, _, sRef, _)) = ComposeMaybe <$> do
10481041 readIORef sRef >>= \case
10491042 ChildReadyState_Ready -> return Nothing
@@ -1081,7 +1074,7 @@ drawChildUpdate initialEnv markReady child = do
10811074 childReadyState <- liftIO $ newIORef $ ChildReadyState_Unready Nothing
10821075 unreadyChildren <- liftIO $ newIORef 0
10831076 df <- createDocumentFragment $ _hydrationDomBuilderEnv_document initialEnv
1084− (placeholder, result) <- runReaderT (unHydrationDomBuilderT $ (,) <$> textNodeInternal ("" :: Text) <*> child) $ initialEnv
1077+ (placeholder, result) <- flip evalStateT [] $ runReaderT (unHydrationDomBuilderT $ (,) <$> textNodeInternal ("" :: Text) <*> child) $ initialEnv
10851078 { _hydrationDomBuilderEnv_parent = toNode df
10861079 , _hydrationDomBuilderEnv_unreadyChildren = unreadyChildren
10871080 , _hydrationDomBuilderEnv_commitAction = markReady childReadyState
@@ -1100,7 +1093,7 @@ drawChildUpdateInt initialEnv markReady child = do
11001093 childReadyState <- liftIO $ newIORef $ ChildReadyStateInt_Unready Nothing
11011094 unreadyChildren <- liftIO $ newIORef 0
11021095 df <- createDocumentFragment $ _hydrationDomBuilderEnv_document initialEnv
1103− (placeholder, result) <- runReaderT (unHydrationDomBuilderT $ (,) <$> textNodeInternal ("" :: Text) <*> child) $ initialEnv
1096+ (placeholder, result) <- flip evalStateT [] $ runReaderT (unHydrationDomBuilderT $ (,) <$> textNodeInternal ("" :: Text) <*> child) $ initialEnv
11041097 { _hydrationDomBuilderEnv_parent = toNode df
11051098 , _hydrationDomBuilderEnv_unreadyChildren = unreadyChildren
11061099 , _hydrationDomBuilderEnv_commitAction = markReady childReadyState