Fix recently broken withLiftMapIO and add regression tests

Broken in e778e23.

effectful-core/src/Effectful/Dispatch/Dynamic.hs

@@ -902,9 +902,10 @@ withLiftMapIO
   -> ((forall a b. (IO a -> IO b) -> Eff localEs a -> Eff localEs b) -> Eff es r)
   -- ^ Continuation with the lifting function in scope.
   -> Eff es r
-withLiftMapIO (LocalEnv les) k = k $ \mapIO m -> unsafeEff $ \es -> do
+withLiftMapIO (LocalEnv les) k = unsafeEff $ \es -> do
   requireMatchingStorages es les
-  seqUnliftIO es $ \unlift -> mapIO $ unlift m
+  (`unEff` es) $ k $ \mapIO m -> unsafeEff $ \localEs -> do
+    seqUnliftIO localEs $ \unlift -> mapIO $ unlift m
 {-# INLINE withLiftMapIO #-}
 
 ----------------------------------------

effectful/tests/UnliftTests.hs

@@ -5,6 +5,9 @@ import Test.Tasty.HUnit
 import UnliftIO.Async qualified as A
 
 import Effectful
+import Effectful.Concurrent.Async qualified as E
+import Effectful.Dispatch.Dynamic
+import Effectful.State.Static.Local
 import Utils qualified as U
 
 unliftTests :: TestTree
@@ -21,6 +24,7 @@ unliftTests = testGroup "Unlift"
     , testCase "Uses in same thread" test_persistentSameThread
     , testCase "Uses in multiple threads" test_persistentMultipleThreads
     ]
+  , testCase "Unlifting functions work correctly" test_unliftingFunctions
   ]
 
 test_threadStrategy :: Assertion
@@ -74,6 +78,78 @@ test_persistentMultipleThreads = runEff $ do
       inThread $ runInIO $ return ()
       inThread $ runInIO $ return ()
 
+test_unliftingFunctions :: Assertion
+test_unliftingFunctions = runEff . E.runConcurrent $ do
+  testFork "runFork1" runFork1
+  testFork "runFork2" runFork2
+  testFork "runFork3" runFork3
+  testFork "runFork4" runFork4
+  testFork "runFork5" runFork5
+  where
+    testFork description runFork = do
+      a <- runFork . send $ ForkWithUnmask $ \unmask -> do
+        evalState @Int 0 $ raiseWith SeqUnlift $ \unlift -> do
+          unlift $ modify @Int (+1)
+          unmask . unlift $ modify @Int (+2)
+          unlift $ modify @Int (+4)
+          unmask . unlift $ modify @Int (+8)
+          unlift $ U.assertEqual (description ++ ": correct state") 15 =<< get @Int
+      E.waitCatch a >>= \case
+        Right () -> pure ()
+        Left err -> U.assertFailure $ description ++ ": " ++ show err
+
+data Fork :: Effect where
+  ForkWithUnmask :: ((forall a. m a -> m a) -> m r) -> Fork m (A.Async r)
+type instance DispatchOf Fork = Dynamic
+
+-- | Uses 'localUnliftIO' and 'withLiftMapIO'.
+runFork1 :: IOE :> es => Eff (Fork : es) a -> Eff es a
+runFork1 = interpret $ \env -> \case
+  ForkWithUnmask m -> do
+    withLiftMapIO env $ \liftMap -> do
+      localUnliftIO env strategy $ \unlift -> do
+        A.asyncWithUnmask $ \unmask -> unlift $ m $ liftMap unmask
+  where
+    strategy = ConcUnlift Ephemeral $ Limited 1
+
+-- | Uses 'localUnlift' and 'withLiftMap'.
+runFork2 :: (IOE :> es, E.Concurrent :> es) => Eff (Fork : es) a -> Eff es a
+runFork2 = interpret $ \env -> \case
+  ForkWithUnmask m -> do
+    withLiftMap env $ \liftMap -> do
+      localUnlift env strategy $ \unlift -> do
+        E.asyncWithUnmask $ \unmask -> unlift $ m $ liftMap unmask
+  where
+    strategy = ConcUnlift Ephemeral $ Limited 1
+
+-- | Uses 'localLiftUnliftIO'.
+runFork3 :: IOE :> es => Eff (Fork : es) a -> Eff es a
+runFork3 = interpret $ \env -> \case
+  ForkWithUnmask m -> do
+    localLiftUnliftIO env strategy $ \lift unlift -> do
+      A.asyncWithUnmask $ \unmask -> unlift $ m $ lift . unmask . unlift
+  where
+    strategy = ConcUnlift Persistent $ Limited 1
+
+-- | Uses 'localLiftUnlift'.
+runFork4 :: (IOE :> es, E.Concurrent :> es) => Eff (Fork : es) a -> Eff es a
+runFork4 = interpret $ \env -> \case
+  ForkWithUnmask m -> do
+    localLiftUnlift env strategy $ \lift unlift -> do
+      E.asyncWithUnmask $ \unmask -> unlift $ m $ lift . unmask . unlift
+  where
+    strategy = ConcUnlift Persistent $ Limited 1
+
+-- | Uses 'localLift' and 'localUnlift'.
+runFork5 :: (IOE :> es, E.Concurrent :> es) => Eff (Fork : es) a -> Eff es a
+runFork5 = interpret $ \env -> \case
+  ForkWithUnmask m -> do
+    localLift env strategy $ \lift -> do
+      localUnlift env strategy $ \unlift -> do
+        E.asyncWithUnmask $ \unmask -> unlift $ m $ lift . unmask . unlift
+  where
+    strategy = ConcUnlift Persistent $ Limited 1
+
 ----------------------------------------
 -- Helpers