Fix polysemy-research#268

src/Polysemy/Async.hs

@@ -67,7 +67,7 @@ sequenceConcurrently t = traverse async t >>= traverse await
 --
 -- @since 1.0.0.0
 asyncToIO
-    :: Member (Embed IO) r
+    :: (Member (Embed IO) r, Member (Final IO) r)
     => Sem (Async ': r) a
     -> Sem r a
 asyncToIO m = withLowerToIO $ \lower _ -> lower $

src/Polysemy/IO.hs

@@ -56,6 +56,7 @@ embedToMonadIO = runEmbedded $ liftIO @m
 lowerEmbedded
     :: ( MonadIO m
        , Member (Embed IO) r
+       , Member (Final IO) r
        )
     => (forall x. m x -> IO x)  -- ^ The means of running this monad.
     -> Sem (Embed m ': r) a

src/Polysemy/Internal/Forklift.hs

@@ -11,6 +11,7 @@ import           Control.Concurrent.Chan.Unagi
 import           Control.Concurrent.MVar
 import           Control.Exception
 import           Polysemy.Internal
+import           Polysemy.Final
 import           Polysemy.Internal.Union
 
 
@@ -19,7 +20,7 @@ import           Polysemy.Internal.Union
 --
 -- @since 0.5.0.0
 data Forklift r = forall a. Forklift
-  { responseMVar :: MVar a
+  { responseMVar :: MVar (Either SomeException a)
   , request      :: Union r (Sem r) a
   }
 
@@ -30,18 +31,24 @@ data Forklift r = forall a. Forklift
 --
 -- @since 0.5.0.0
 runViaForklift
-    :: Member (Embed IO) r
+    :: (Member (Embed IO) r, Member (Final IO) r)
     => InChan (Forklift r)
     -> Sem r a
     -> IO a
 runViaForklift chan = usingSem $ \u -> do
   case prj u of
     Just (Weaving (Embed m) s _ ex _) ->
       ex . (<$ s) <$> m
-    _ -> do
-      mvar <- newEmptyMVar
-      writeChan chan $ Forklift mvar u
-      takeMVar mvar
+    _ -> case prj u of
+      Just (Weaving (WithWeavingToFinal wav) s wv ex ins) ->
+        ex <$> wav s (runViaForklift chan . wv) ins
+      _ -> do
+        mvar <- newEmptyMVar
+        writeChan chan $ Forklift mvar u
+        res <- takeMVar mvar
+        case res of
+          Right a -> pure a
+          Left  e -> throwIO e
 {-# INLINE runViaForklift #-}
 
 
@@ -54,7 +61,7 @@ runViaForklift chan = usingSem $ \u -> do
 --
 -- @since 0.5.0.0
 withLowerToIO
-    :: Member (Embed IO) r
+    :: (Member (Embed IO) r, Member (Final IO) r)
     => ((forall x. Sem r x -> IO x) -> IO () -> IO a)
        -- ^ A lambda that takes the lowering function, and a finalizing 'IO'
        -- action to mark a the forked thread as being complete. The finalizing
@@ -75,7 +82,9 @@ withLowerToIO action = do
         case raced of
           Left () -> embed $ A.wait res
           Right (Forklift mvar req) -> do
-            resp <- liftSem req
+            resp <- withWeavingToFinal $ \s wv _ ->
+                     wv (liftSem (fmap Right req) <$ s)
+              `catch` \e -> pure (Left e <$ s)
             embed $ putMVar mvar $ resp
             me_b
       {-# INLINE me #-}

test/AsyncSpec.hs

@@ -14,7 +14,8 @@ import Test.Hspec
 spec :: Spec
 spec = describe "async" $ do
   it "should thread state and not lock" $ do
-    (ts, (s, r)) <- runM
+    (ts, (s, r)) <- runFinal
+                  . embedToFinal @IO
                   . runTraceList
                   . runState "hello"
                   . asyncToIO $ do

test/BracketSpec.hs

@@ -240,9 +240,10 @@ runTest = pure
         . runError @()
 
 runTest2
-  :: Sem '[Error (), Resource, State [Char], Trace, Output String, Embed IO] a
+  :: Sem '[Error (), Resource, State [Char], Trace, Output String, Embed IO, Final IO] a
   -> IO ([String], ([Char], Either () a))
-runTest2 = runM
+runTest2 = runFinal
+         . embedToFinal @IO
          . ignoreOutput
          . runTraceList
          . runState ""
@@ -284,17 +285,15 @@ testAllThree name k m = do
 testTheIOTwo
     :: String
     -> (([String], ([Char], Either () a)) -> Expectation)
-    -> (Sem '[Error (), Resource, State [Char], Trace, Output String, Embed IO] a)
+    -> (Sem '[Error (), Resource, State [Char], Trace, Output String, Embed IO, Final IO] a)
     -> Spec
 testTheIOTwo name k m = do
   describe name $ do
     it "via resourceToIO" $ do
       z <- runTest2 m
       k z
-    -- NOTE(sandy): This unsafeCoerces are safe, because we're just weakening
-    -- the end of the union
     it "via resourceToIOFinal" $ do
-      z <- runTest3 $ unsafeCoerce m
+      z <- runTest3 $ m
       k z