@@ -6,7 +6,8 @@ import Control.Monad (void)
66import Control.Monad.Catch (throwM, MonadThrow)
77import Control.Monad.Error.Class (throwError, MonadError)
88import Control.Monad.IO.Class (MonadIO(liftIO))
9−import Control.Monad.State (MonadState, get, modify, runStateT, StateT)
9+import Control.Monad.State (MonadState, get, modify, runStateT
10+ , StateT(..), evalStateT)
1011import Control.Monad.Trans.Except (runExceptT, ExceptT)
1112import Data.Foldable (forM_, fold)
1213import Data.Function ((&))
@@ -436,6 +437,21 @@ parallelTests = H.parallel $ do
436437
437438 it "foldrM is lazy enough" checkFoldrLaziness
438439
440+ ---------------------------------------------------------------------------
441+ -- Monadic state transfer in concurrent tasks
442+ ---------------------------------------------------------------------------
443+
444+ -- XXX Can we write better test cases to hit every case?
445+
446+ it "async: state is saved and used if the work is partially enqueued"
447+ (checkMonadicStateTransfer async)
448+
449+ it "wAsync: state is saved and used if the work is partially enqueued"
450+ (checkMonadicStateTransfer wAsync)
451+
452+ it "ahead: state is saved and used if the work is partially enqueued"
453+ (checkMonadicStateTransfer ahead)
454+
439455 ---------------------------------------------------------------------------
440456 -- Monadic state snapshot in concurrent tasks
441457 ---------------------------------------------------------------------------
@@ -561,6 +577,27 @@ takeCombined n t = do
561577 S.take n (constr ([] :: [Int]) <> constr ([] :: [Int]))
562578 r `shouldBe` []
563579
580+checkMonadicStateTransfer
581+ :: (IsStream t1, IsStream t2)
582+ => ( t1 (StateT Int IO) ()
583+ -> t2 (StateT Int IO) ()
584+ -> SerialT (StateT Int IO) a3 )
585+ -> IO ()
586+checkMonadicStateTransfer op = evalStateT str (0 :: Int)
587+ where
588+ str =
589+ S.drain $
590+ maxBuffer 1 $
591+ (serially $ S.mapM snapshoti $ S.fromList [1..10]) `op`
592+ (serially $ S.mapM snapshoti $ S.fromList [1..10])
593+ snapshoti y = do
594+ modify (+ 1)
595+ x <- get
596+ lift1 $ x `shouldBe` y
597+ lift1 m = StateT $ \s -> do
598+ a <- m
599+ return (a, s)
600+
564601checkFoldrLaziness :: IO ()
565602checkFoldrLaziness = do
566603 S.foldrM (\x xs -> if odd x then return True else xs)