Add "afterFirst" run an action after first event

src/Asyncly.hs

@@ -36,6 +36,8 @@ module Asyncly
     , (*>>)
     , thenDiscard
     , (>>*)
+    , afterFirst
+    , (>>|)
 
     , Log
     , Loggable

src/Asyncly/AsyncT.hs

@@ -34,6 +34,8 @@ module Asyncly.AsyncT
     , (*>>)
     , thenDiscard
     , (>>*)
+    , afterFirst
+    , (>>|)
 
     -- internal
     , dbg
@@ -603,31 +605,22 @@ instance (Num a, Monad (AsyncT m)) => Num (AsyncT m a) where
 -- Special compositions
 ------------------------------------------------------------------------------
 
-infixr 1 >>*, *>>
+infixr 1 >>*, *>>, >>|
 
-{-
+-- XXX This can be moved to utility functions as it is purely app level
 -- | Run @b@ once, discarding its result when the first task in task set @a@
 -- has finished. Useful to start a singleton task after the first task has been
 -- setup.
-afterFirst :: MonadIO m => AsyncT m a -> AsyncT m b -> AsyncT m a
-afterFirst ma mb = AsyncT $ do
-  fs  <- getContinuations
-  ref <- liftIO $ newIORef False
-  setContinuation ma (cont ref)  fs
-  r   <- runAsyncT ma
-  restoreStack fs
-  return  r
-  where cont ref x = AsyncT $ do
-          n <- liftIO $ readIORef ref
-          if n == True
-            then return $ Just x
-            else do liftIO $ writeIORef ref True
-                    runAsyncT mb
-                    return $ Just x
+afterFirst :: MonadAsync m => AsyncT m a -> AsyncT m b -> AsyncT m a
+afterFirst ma mb = do
+    ref <- liftIO $ newIORef False
+    x <- ma
+    done <- liftIO $ readIORef ref
+    when (not done) $ (liftIO $ writeIORef ref True) >>* mb
+    return x
 
-(<|) :: MonadIO m => AsyncT m a -> AsyncT m b -> AsyncT m a
-(<|) = afterFirst
-                    -}
+(>>|) :: MonadAsync m => AsyncT m a -> AsyncT m b -> AsyncT m a
+(>>|) = afterFirst
 
 -- | Run 'm a' in "isolation" and discard its result, and then run 'm b' and
 -- return its result.  Isolation means that any alternative actions inside 'm

test/Main.hs

@@ -101,6 +101,16 @@ main = hspec $ do
         )
         `shouldReturn` ([2] :: [Int])
 
+    -- 2 and 3 should be printed once and only once
+    it ">>| works as expected" $
+        ((wait $ ((async (liftIO (putStrLn "0") >> return 0)
+                    <|> (liftIO (putStrLn "1") >> return 1))
+                 >>| (async (liftIO (putStrLn "2") >> return 2)
+                        <|> (liftIO (putStrLn "3") >> return 3))
+                )
+        ) >>= return . sort)
+        `shouldReturn` ([0, 1] :: [Int])
+
 generalExample :: AsyncT IO Int
 generalExample = do
     liftIO $ return ()