@@ -217,12 +217,19 @@ data HydrationDomBuilderEnv t m
217217 , _hydrationDomBuilderEnv_unreadyChildren :: {-# UNPACK #-} !(IORef Word) -- Number of children who still aren't fully rendered
218218 , _hydrationDomBuilderEnv_commitAction :: !(JSM ()) -- Action to take when all children are ready --TODO: we should probably get rid of this once we invoke it
219219 , _hydrationDomBuilderEnv_hydrationMode :: {-# UNPACK #-} !(IORef HydrationMode) -- In hydration mode?
220− , _hydrationDomBuilderEnv_hydrationNode :: {-# UNPACK #-} !(IORef (Maybe Node)) -- Hydration node index
221220 , _hydrationDomBuilderEnv_prerenderDepth :: {-# UNPACK #-} !(IORef Word)
222221 , _hydrationDomBuilderEnv_eventChan :: {-# UNPACK #-} !(Chan [DSum (EventTriggerRef t) TriggerInvocation])
223222 }
224223
225−newtype HydrationDomBuilderT t m a = HydrationDomBuilderT { unHydrationDomBuilderT :: ReaderT (HydrationDomBuilderEnv t m) (StateT [Behavior t (m ())] (RequesterT t JSM Identity (TriggerEventT t m))) a }
224+data HydrationState t m = HydrationState
225+ { _hydrationState_actions :: [Behavior t (m ())]
226+ , _hydrationState_previousNode :: {-# UNPACK #-} !(Maybe Node)
227+ }
228+
229+addAction :: Monad m => Behavior t (m ()) -> HydrationDomBuilderT t m ()
230+addAction a = HydrationDomBuilderT $ modify $ \(HydrationState as n) -> HydrationState (a : as) n
231+
232+newtype HydrationDomBuilderT t m a = HydrationDomBuilderT { unHydrationDomBuilderT :: ReaderT (HydrationDomBuilderEnv t m) (StateT (HydrationState t m) (RequesterT t JSM Identity (TriggerEventT t m))) a }
226233 deriving (Functor, Applicative, Monad, MonadFix, MonadIO, MonadException
227234#if MIN_VERSION_base(4,9,1)
228235 , MonadAsyncException
@@ -260,14 +267,15 @@ runHydrationDomBuilderT
260267 )
261268 => HydrationDomBuilderT t m a
262269 -> HydrationDomBuilderEnv t m
263−-- -> Chan [DSum (EventTriggerRef t) TriggerInvocation]
270+ -> Maybe Node
264271 -> m (a, [Behavior t (m ())])
265−runHydrationDomBuilderT (HydrationDomBuilderT a) env =
266− flip runTriggerEventT (_hydrationDomBuilderEnv_eventChan env) $ do
267− rec (x, req) <- runRequesterT (runStateT (runReaderT a env) []) rsp
272+runHydrationDomBuilderT (HydrationDomBuilderT a) env prev = do
273+ (a, hs) <- flip runTriggerEventT (_hydrationDomBuilderEnv_eventChan env) $ do
274+ rec (result, req) <- runRequesterT (runStateT (runReaderT a env) (HydrationState [] prev)) rsp
268275 rsp <- performEventAsync $ ffor req $ \rm f -> liftJSM $ runInAnimationFrame f $
269276 traverseRequesterData (\r -> Identity <$> r) rm
270− return x
277+ return result
278+ return (a, _hydrationState_actions hs)
271279 where
272280 runInAnimationFrame f x = void . DOM.inAnimationFrame' $ \_ -> do
273281 v <- synchronously x
@@ -285,8 +293,19 @@ askParent = HydrationDomBuilderT $ asks _hydrationDomBuilderEnv_parent
285293askEvents :: Monad m => HydrationDomBuilderT t m (Chan [DSum (EventTriggerRef t) TriggerInvocation])
286294askEvents = HydrationDomBuilderT . lift . lift . lift $ TriggerEventT.askEvents
287295
288−localEnv :: Monad m => (HydrationDomBuilderEnv t m -> HydrationDomBuilderEnv t m) -> HydrationDomBuilderT t m a -> HydrationDomBuilderT t m a
289−localEnv f = HydrationDomBuilderT . local f . unHydrationDomBuilderT
296+localEnv :: Monad m => (HydrationState t m -> HydrationState t m) -> (HydrationDomBuilderEnv t m -> HydrationDomBuilderEnv t m) -> HydrationDomBuilderT t m a -> HydrationDomBuilderT t m a
297+localEnv f g (HydrationDomBuilderT m) = HydrationDomBuilderT $ do
298+ env <- ask
299+ hs <- get
300+ lift $ lift $ evalStateT (runReaderT m $ g env) (f hs)
301+
302+{-# INLINABLE getPreviousNode #-}
303+getPreviousNode :: Monad m => HydrationDomBuilderT t m (Maybe DOM.Node)
304+getPreviousNode = HydrationDomBuilderT $ gets _hydrationState_previousNode
305+
306+{-# INLINABLE setPreviousNode #-}
307+setPreviousNode :: Monad m => Maybe DOM.Node -> HydrationDomBuilderT t m ()
308+setPreviousNode n = HydrationDomBuilderT $ modify $ \hs -> hs { _hydrationState_previousNode = n }
290309
291310{-# INLINABLE append #-}
292311append :: MonadJSM m => DOM.Node -> HydrationDomBuilderT t m ()
@@ -318,15 +337,15 @@ makeNodeInternal :: (MonadJSM m, IsNode node, Typeable node) => (node -> Hydrati
318337makeNodeInternal check mkNode constructor = do
319338 doc <- askDocument
320339 parent <- askParent
340+ liftIO $ putStr $ show (typeRep constructor) <> ": "
321341 getHydrationMode >>= \case
322342 HydrationMode_Immediate -> do
323− liftIO $ putStrLn $ "Not in hydration mode, creating node: " <> show (typeRep constructor)
343+ liftIO $ putStrLn $ "Not in hydration mode, creating node"
324344 n <- mkNode
325345 append $ toNode n
326346 return n
327347 HydrationMode_Hydrating -> do
328− lastHydrationNodeRef <- HydrationDomBuilderT $ asks _hydrationDomBuilderEnv_hydrationNode
329− lastHydrationNode <- liftIO $ readIORef lastHydrationNodeRef
348+ lastHydrationNode <- getPreviousNode
330349 let go mLastNode = do
331350 node <- maybe (Node.getFirstChildUnchecked parent) Node.getNextSiblingUnchecked mLastNode
332351 DOM.castTo constructor node >>= \case
@@ -341,7 +360,7 @@ makeNodeInternal check mkNode constructor = do
341360 liftIO $ putStrLn $ "Node failed check, skipping"
342361 go (Just node)
343362 n <- go lastHydrationNode
344− liftIO $ modifyIORef' lastHydrationNodeRef (\_ -> Just $ toNode n)
363+ setPreviousNode $ Just $ toNode n
345364 return n
346365
347366-- skipToCommentId "prerender-start" (liftJSM $ createComment doc "rerender-start") $ \new -> \case
@@ -468,15 +487,11 @@ makeElement elementTag cfg child = do
468487 -- TODO: prerender needs completely replacing
469488 e <- makeNodeInternal hasSSRAttribute buildElement DOM.Element
470489 -- Run the child builder with updated parent and previous sibling references
471− childNode <- liftIO $ newIORef Nothing
472− result <- localEnv (\env -> env { _hydrationDomBuilderEnv_parent = toNode e, _hydrationDomBuilderEnv_hydrationNode = childNode }) child
490+ result <- localEnv (\hs -> hs { _hydrationState_previousNode = Nothing }) (\env -> env { _hydrationDomBuilderEnv_parent = toNode e }) child
473491 events <- askEvents
474492 env <- HydrationDomBuilderT ask
475493 (ese, fire) <- newTriggerEvent
476− let beh = pure $ void . flip runHydrationDomBuilderT env $ do
477− es <- wrap events e $ extractRawElementConfig cfg
478− liftIO $ fire es
479− HydrationDomBuilderT $ modify (beh :)
494+ addAction $ pure $ void $ runHydrationDomBuilderT (liftIO . fire <=< wrap events e $ extractRawElementConfig cfg) env Nothing
480495 collapsed <- switchHold never $ ffor ese $ \es -> DMap.singleton (WrapArg Click) . Identity <$> (select es (WrapArg Click)) -- TODO! this can't stay here, obviously.
481496 let evts = fan collapsed
482497 return ((Element evts (), result), e)
@@ -640,9 +655,7 @@ instance SupportsHydrationDomBuilder t m => MountableDomBuilder t (HydrationDomB
640655 type DomFragment (HydrationDomBuilderT t m) = HydrationDomFragment
641656 buildDomFragment w = do
642657 df <- createDocumentFragment =<< askDocument
643− result <- flip localEnv w $ \env -> env
644− { _hydrationDomBuilderEnv_parent = toNode df
645− }
658+ result <- localEnv id (\env -> env { _hydrationDomBuilderEnv_parent = toNode df }) w
646659 state <- liftIO $ newIORef FragmentState_Unmounted
647660 return (HydrationDomFragment df state, result)
648661 mountDomFragment fragment setFragment = do
@@ -663,8 +676,7 @@ skipToCommentId :: MonadJSM m => Text -> HydrationDomBuilderT t m (DOM.Node, Tex
663676skipToCommentId prefix = do
664677 doc <- askDocument
665678 parent <- askParent
666− lastHydrationNodeRef <- HydrationDomBuilderT $ asks _hydrationDomBuilderEnv_hydrationNode
667− lastHydrationNode <- liftIO $ readIORef lastHydrationNodeRef
679+ lastHydrationNode <- getPreviousNode
668680 let go mLastNode = do
669681 node <- maybe (Node.getFirstChildUnchecked parent) Node.getNextSiblingUnchecked mLastNode
670682 DOM.castTo DOM.Comment node >>= \case
@@ -681,7 +693,7 @@ skipToCommentId prefix = do
681693 liftIO $ putStrLn commentText
682694 go (Just node)
683695 (a, key) <- go lastHydrationNode
684− liftIO $ modifyIORef' lastHydrationNodeRef (\_ -> Just $ toNode a)
696+ setPreviousNode $ Just $ toNode a
685697 pure (toNode a, key)
686698
687699skipToReplaceStart :: MonadJSM m => HydrationDomBuilderT t m (DOM.Node, Text)
@@ -730,11 +742,9 @@ instance (Reflex t, Adjustable t m, MonadJSM m, MonadHold t m, MonadFix m, PrimM
730742 HydrationMode_Hydrating -> pure parent
731743 HydrationMode_Immediate -> toNode <$> createDocumentFragment doc
732744 unreadyChildren <- liftIO $ newIORef 0
733− hydrationNode <- liftIO $ newIORef Nothing
734− (result, h') <- flip runStateT [] $ runReaderT (unHydrationDomBuilderT a0) initialEnv
745+ (result, HydrationState h' _) <- flip runStateT (HydrationState [] Nothing) $ runReaderT (unHydrationDomBuilderT a0) initialEnv
735746 { _hydrationDomBuilderEnv_unreadyChildren = unreadyChildren
736747 , _hydrationDomBuilderEnv_commitAction = myCommitAction
737− , _hydrationDomBuilderEnv_hydrationNode = hydrationNode
738748 , _hydrationDomBuilderEnv_parent = p
739749 }
740750 liftIO $ readIORef unreadyChildren >>= \case
@@ -743,7 +753,7 @@ instance (Reflex t, Adjustable t m, MonadJSM m, MonadHold t m, MonadFix m, PrimM
743753 liftIO $ putStrLn "drawInitialChild: end"
744754 return (h', result)
745755 a'' <- numberOccurrences a'
746− (result0, child') <- HydrationDomBuilderT $ lift $ lift $ runWithReplace drawInitialChild $ ffor a'' $ \(cohortId, child) -> do
756+ (result0, evt) <- HydrationDomBuilderT $ lift $ lift $ runWithReplace drawInitialChild $ ffor a'' $ \(cohortId, child) -> do
747757 p <- liftIO (readIORef hydrating) >>= \case
748758 HydrationMode_Hydrating -> pure parent
749759 HydrationMode_Immediate -> toNode <$> createDocumentFragment doc
@@ -756,11 +766,9 @@ instance (Reflex t, Adjustable t m, MonadJSM m, MonadHold t m, MonadFix m, PrimM
756766 insertBefore p after
757767 liftIO $ writeIORef currentCohort cohortId
758768 myCommitAction
759− hydrationNode <- liftIO $ newIORef (Just $ toNode before)
760− (result, h') <- flip runStateT [] $ runReaderT (unHydrationDomBuilderT child) $ initialEnv
769+ (result, HydrationState h' _) <- flip runStateT (HydrationState [] $ Just $ toNode before) $ runReaderT (unHydrationDomBuilderT child) $ initialEnv
761770 { _hydrationDomBuilderEnv_unreadyChildren = unreadyChildren
762771 , _hydrationDomBuilderEnv_commitAction = commitAction
763− , _hydrationDomBuilderEnv_hydrationNode = hydrationNode
764772 , _hydrationDomBuilderEnv_parent = p
765773 }
766774 uc <- liftIO $ readIORef unreadyChildren
@@ -768,16 +776,16 @@ instance (Reflex t, Adjustable t m, MonadJSM m, MonadHold t m, MonadFix m, PrimM
768776 let commitActionToRunNow = if uc == 0
769777 then Just $ commitAction
770778 else Nothing -- A child will run it when unreadyChildren is decremented to 0
771− commitActionToRunNow' = if h == HydrationMode_Hydrating then Nothing else commitActionToRunNow
772− hydrationAction = if h == HydrationMode_Hydrating then Just h' else Nothing
773− return (hydrationAction, commitActionToRunNow', result)
774− env <- HydrationDomBuilderT ask
775− let f :: [Behavior t (m ())] -> Behavior t (m ())
779+ actions = case h of
780+ HydrationMode_Hydrating -> Left h'
781+ HydrationMode_Immediate -> Right commitActionToRunNow
782+ return (actions, result)
783+ let (hydrationAction, commitAction) = fanEither $ fmap fst evt
784+ f :: [Behavior t (m ())] -> Behavior t (m ())
776785 f = fmap (void . sequence) . sequence
777− beh <- hold (f $ fst result0) (fmapMaybe (\(h,_,_) -> f <$> h) child')
778− HydrationDomBuilderT $ modify (join beh :)
779− requestDomAction_ $ fmapMaybe (\(_,c,_) -> c) child'
780− return (snd result0, fmap (\(_,_,r) -> r) child')
786+ addAction . join =<< hold (f $ fst result0) (fmap f hydrationAction)
787+ requestDomAction_ $ fmapMaybe id commitAction
788+ return (snd result0, fmap snd evt)
781789
782790 {-# INLINABLE traverseIntMapWithKeyWithAdjust #-}
783791 traverseIntMapWithKeyWithAdjust = traverseIntMapWithKeyWithAdjust'
@@ -1074,7 +1082,7 @@ drawChildUpdate initialEnv markReady child = do
10741082 childReadyState <- liftIO $ newIORef $ ChildReadyState_Unready Nothing
10751083 unreadyChildren <- liftIO $ newIORef 0
10761084 df <- createDocumentFragment $ _hydrationDomBuilderEnv_document initialEnv
1077− (placeholder, result) <- flip evalStateT [] $ runReaderT (unHydrationDomBuilderT $ (,) <$> textNodeInternal ("" :: Text) <*> child) $ initialEnv
1085+ (placeholder, result) <- flip evalStateT (HydrationState [] Nothing) $ runReaderT (unHydrationDomBuilderT $ (,) <$> textNodeInternal ("" :: Text) <*> child) $ initialEnv
10781086 { _hydrationDomBuilderEnv_parent = toNode df
10791087 , _hydrationDomBuilderEnv_unreadyChildren = unreadyChildren
10801088 , _hydrationDomBuilderEnv_commitAction = markReady childReadyState
@@ -1093,7 +1101,7 @@ drawChildUpdateInt initialEnv markReady child = do
10931101 childReadyState <- liftIO $ newIORef $ ChildReadyStateInt_Unready Nothing
10941102 unreadyChildren <- liftIO $ newIORef 0
10951103 df <- createDocumentFragment $ _hydrationDomBuilderEnv_document initialEnv
1096− (placeholder, result) <- flip evalStateT [] $ runReaderT (unHydrationDomBuilderT $ (,) <$> textNodeInternal ("" :: Text) <*> child) $ initialEnv
1104+ (placeholder, result) <- flip evalStateT (HydrationState [] Nothing) $ runReaderT (unHydrationDomBuilderT $ (,) <$> textNodeInternal ("" :: Text) <*> child) $ initialEnv
10971105 { _hydrationDomBuilderEnv_parent = toNode df
10981106 , _hydrationDomBuilderEnv_unreadyChildren = unreadyChildren
10991107 , _hydrationDomBuilderEnv_commitAction = markReady childReadyState