From 85015a4a08f358907be6d56320640672514a38d9 Mon Sep 17 00:00:00 2001
From: Andrzej Rybczak <andrzej@rybczak.net>
Date: Thu, 15 Jul 2021 17:00:37 +0200
Subject: [PATCH] Adjustments to Error effect

Prelude for runError2, runError3 etc.
---
 src/Effectful/Async.hs |  2 ++
 src/Effectful/Error.hs | 66 +++++++++++++++++++++++++++---------------
 2 files changed, 45 insertions(+), 23 deletions(-)

diff --git a/src/Effectful/Async.hs b/src/Effectful/Async.hs
index 20d25f5..1a8fda4 100644
--- a/src/Effectful/Async.hs
+++ b/src/Effectful/Async.hs
@@ -88,6 +88,8 @@ import Effectful.Internal.Monad
 --
 -- /Note:/ thread local state changes in 'Eff' operations run asynchronously
 -- will not affect the parent thread.
+--
+-- /TODO:/ write about 'AsyncE' not respecting scoped operations.
 data AsyncE :: Effect where
   AsyncE :: AsyncE m r
 
diff --git a/src/Effectful/Error.hs b/src/Effectful/Error.hs
index ab5c148..0c98e9d 100644
--- a/src/Effectful/Error.hs
+++ b/src/Effectful/Error.hs
@@ -43,42 +43,44 @@ import Effectful.Internal.Env
 import Effectful.Internal.Monad
 
 data Error e :: Effect where
-  Error :: Unique -> Error e m r
+  Error :: ErrorId -> Error e m r
 
--- | TODO: write about possibility of an error escaping the scope of 'runError'
--- when misused with 'AsyncE'.
 runError
   :: forall e es a. Typeable e
   => Eff (Error e : es) a
   -> Eff es (Either (CallStack, e) a)
 runError m = unsafeEff $ \es0 -> mask $ \release -> do
-  -- A unique tag is picked so that different runError handlers for the same
-  -- type don't catch each other's exceptions.
-  tag <- newUnique
+  eid <- newErrorId
   size0 <- sizeEnv es0
-  es <- unsafeConsEnv (IdE (Error @e tag)) noRelinker es0
-  r <- tryErrorIO tag (release $ unEff m es) `onException` unsafeTailEnv size0 es
+  es <- unsafeConsEnv (IdE (Error @e eid)) noRelinker es0
+  r <- tryErrorIO release eid es `onException` unsafeTailEnv size0 es
   _ <- unsafeTailEnv size0 es
   pure r
+  where
+    tryErrorIO release eid es = try (release $ unEff m es) >>= \case
+      Right a -> pure $ Right a
+      Left ex -> tryHandler ex eid (\cs e -> Left (cs, e))
+               $ throwIO ex
 
 throwError
   :: forall e es a. (HasCallStack, Typeable e, Error e :> es)
   => e
   -> Eff es a
-throwError e = readerEffectM @(Error e) $ \(IdE (Error tag)) -> unsafeEff_ $ do
-  throwIO $ ErrorEx tag callStack e
+throwError e = unsafeEff $ \es -> do
+  IdE (Error eid) <- getEnv @(Error e) es
+  throwIO $ ErrorEx eid callStack e
 
 catchError
   :: forall e es a. (Typeable e, Error e :> es)
   => Eff es a
   -> (CallStack -> e -> Eff es a)
   -> Eff es a
-catchError m handler = do
-  readerEffectM @(Error e) $ \(IdE (Error tag)) -> unsafeEff $ \es -> do
-    size <- sizeEnv es
-    catchErrorIO tag (unEff m es) $ \cs e -> do
-      checkSizeEnv size es
-      unEff (handler cs e) es
+catchError m handler = unsafeEff $ \es -> do
+  IdE (Error eid) <- getEnv @(Error e) es
+  size <- sizeEnv es
+  catchErrorIO eid (unEff m es) $ \cs e -> do
+    checkSizeEnv size es
+    unEff (handler cs e) es
 
 tryError
   :: forall e es a. (Typeable e, Error e :> es)
@@ -89,7 +91,28 @@ tryError m = (Right <$> m) `catchError` \es e -> pure $ Left (es, e)
 ----------------------------------------
 -- Helpers
 
-data ErrorEx e = ErrorEx Unique CallStack e
+newtype ErrorId = ErrorId Unique
+  deriving Eq
+
+-- | A unique is picked so that distinct 'Error' handlers for the same type
+-- don't catch each other's exceptions.
+newErrorId :: IO ErrorId
+newErrorId = ErrorId <$> newUnique
+
+tryHandler
+  :: Typeable e
+  => SomeException
+  -> ErrorId
+  -> (CallStack -> e -> r)
+  -> IO r
+  -> IO r
+tryHandler ex eid0 handler next = case fromException ex of
+  Just (ErrorEx eid cs e)
+    | eid0 == eid -> pure $ handler cs e
+    | otherwise   -> next
+  Nothing -> next
+
+data ErrorEx e = ErrorEx ErrorId CallStack e
 instance Typeable e => Show (ErrorEx e) where
   showsPrec p (ErrorEx _ cs e)
     = ("Effectful.Error.ErrorEx (" ++)
@@ -98,12 +121,9 @@ instance Typeable e => Show (ErrorEx e) where
     . showsPrec p cs
 instance Typeable e => Exception (ErrorEx e)
 
-catchErrorIO :: Typeable e => Unique -> IO a -> (CallStack -> e -> IO a) -> IO a
-catchErrorIO tag m handler = do
+catchErrorIO :: Typeable e => ErrorId -> IO a -> (CallStack -> e -> IO a) -> IO a
+catchErrorIO eid m handler = do
   m `catch` \err@(ErrorEx etag e cs) -> do
-    if tag == etag
+    if eid == etag
       then handler e cs
       else throwIO err
-
-tryErrorIO :: Typeable e => Unique -> IO a -> IO (Either (CallStack, e) a)
-tryErrorIO tag m = catchErrorIO tag (Right <$> m) $ \es e -> pure $ Left (es, e)