From 1fe65f674d36b657fc4225ec3d908f64cbfcfbb9 Mon Sep 17 00:00:00 2001
From: Marcin Szamotulski <profunctor@pm.me>
Date: Sat, 10 Aug 2019 08:22:52 +0200
Subject: [PATCH] Added cancelWith to the MonadAsync type class

---
 io-sim-classes/src/Control/Monad/Class/MonadAsync.hs | 10 ++++++----
 io-sim/src/Control/Monad/IOSim.hs                    |  3 ++-
 2 files changed, 8 insertions(+), 5 deletions(-)

diff --git a/io-sim-classes/src/Control/Monad/Class/MonadAsync.hs b/io-sim-classes/src/Control/Monad/Class/MonadAsync.hs
index f67127552b1..19db40a773c 100644
--- a/io-sim-classes/src/Control/Monad/Class/MonadAsync.hs
+++ b/io-sim-classes/src/Control/Monad/Class/MonadAsync.hs
@@ -17,7 +17,7 @@ import           Control.Concurrent.Async (AsyncCancelled(..))
 
 class MonadSTM m => MonadAsync m where
 
-  {-# MINIMAL async, cancel, waitCatchSTM, pollSTM #-}
+  {-# MINIMAL async, cancel, cancelWith, waitCatchSTM, pollSTM #-}
 
   -- | An asynchronous action
   type Async m :: * -> *
@@ -26,9 +26,10 @@ class MonadSTM m => MonadAsync m where
   withAsync             :: m a -> (Async m a -> m b) -> m b
 
   wait                  :: Async m a -> m a
-  poll                  :: Async m a -> m (Maybe (Either SomeException a)) 
+  poll                  :: Async m a -> m (Maybe (Either SomeException a))
   waitCatch             :: Async m a -> m (Either SomeException a)
   cancel                :: Async m a -> m ()
+  cancelWith            :: Exception e => Async m a -> e -> m ()
   uninterruptibleCancel :: Async m a -> m ()
 
   waitSTM               :: Async m a -> STM m a
@@ -37,7 +38,7 @@ class MonadSTM m => MonadAsync m where
 
   waitAny               :: [Async m a] -> m (Async m a, a)
   waitAnyCatch          :: [Async m a] -> m (Async m a, Either SomeException a)
-  waitAnyCancel         :: [Async m a] -> m (Async m a, a) 
+  waitAnyCancel         :: [Async m a] -> m (Async m a, a)
   waitAnyCatchCancel    :: [Async m a] -> m (Async m a, Either SomeException a)
   waitEither            :: Async m a -> Async m b -> m (Either a b)
 
@@ -70,7 +71,7 @@ class MonadSTM m => MonadAsync m where
   default uninterruptibleCancel
                         :: MonadMask m => Async m a -> m ()
   default waitSTM       :: MonadThrow (STM m) => Async m a -> STM m a
-  default waitAnyCancel         :: MonadThrow m => [Async m a] -> m (Async m a, a) 
+  default waitAnyCancel         :: MonadThrow m => [Async m a] -> m (Async m a, a)
   default waitAnyCatchCancel    :: MonadThrow m => [Async m a]
                                 -> m (Async m a, Either SomeException a)
   default waitEitherCancel      :: MonadThrow m => Async m a -> Async m b
@@ -227,6 +228,7 @@ instance MonadAsync IO where
   poll                  = Async.poll
   waitCatch             = Async.waitCatch
   cancel                = Async.cancel
+  cancelWith            = Async.cancelWith
   uninterruptibleCancel = Async.uninterruptibleCancel
 
   waitSTM               = Async.waitSTM
diff --git a/io-sim/src/Control/Monad/IOSim.hs b/io-sim/src/Control/Monad/IOSim.hs
index c2cf4e3724b..f5284136608 100644
--- a/io-sim/src/Control/Monad/IOSim.hs
+++ b/io-sim/src/Control/Monad/IOSim.hs
@@ -333,6 +333,7 @@ instance MonadAsync (SimM s) where
     return (Async tid var)
 
   cancel a@(Async tid _) = throwTo tid AsyncCancelled <* waitCatch a
+  cancelWith a@(Async tid _) e = throwTo tid e <* waitCatch a
 
   waitCatchSTM (Async _ var) = readTMVar var
   pollSTM      (Async _ var) = tryReadTMVar var
@@ -643,7 +644,7 @@ schedule thread@Thread{
 
     Catch action' handler k -> do
       -- push the failure and success continuations onto the control stack
-      let thread' = thread { threadControl = ThreadControl action' 
+      let thread' = thread { threadControl = ThreadControl action'
                                                (CatchFrame handler k ctl) }
       schedule thread' simstate