<commit_before>module Async (
    Async,
    async,
    withAsync,
    cancel,
    waitCatch,
    waitCatchSTM,
    waitSTM,
    wait,
    waitEither,
    waitBoth,
    concurrently
  ) where

import Control.Concurrent.STM
import Control.Exception
import Control.Concurrent (ThreadId, forkIO)

-- <<Async
data Async a = Async ThreadId (STM (Either SomeException a))
-- >>

forkFinally :: IO a -> (Either SomeException a -> IO ()) -> IO ThreadId
forkFinally action fun =
  mask $ \restore ->
    forkIO (do r <- try (restore action); fun r)

-- <<async
async :: IO a -> IO (Async a)
async action = do
  var <- newEmptyTMVarIO
  t <- forkFinally action (atomically . putTMVar var)
  return (Async t (readTMVar var))
-- >>

-- <<Functor
instance Functor Async where
  fmap f (Async t stm) = Async t stm'
    where stm' = do
            r <- stm
            case r of
              Left e  -> return (Left e)
              Right a -> return (Right (f a))
-- >>

--- <<watchCatch
waitCatch :: Async a -> IO (Either SomeException a)
waitCatch = atomically . waitCatchSTM
-- >>

-- <<waitCatchSTM
waitCatchSTM :: Async a -> STM (Either SomeException a)
waitCatchSTM (Async _ stm) = stm
-- >>

-- <<waitSTM
waitSTM :: Async a -> STM a
waitSTM a = do
  r <- waitCatchSTM a
  case r of
    Left e  -> throwSTM e
    Right a -> return a
-- >>

-- <<wait
wait :: Async a -> IO a
wait = atomically . waitSTM
-- >>

-- <<cancel
cancel :: Async a -> IO ()
cancel (Async t _) = throwTo t ThreadKilled
-- >>

-- <<waitEither
waitEither :: Async a -> Async b -> IO (Either a b)
waitEither a b = atomically $
  fmap Left (waitSTM a)
    `orElse`
  fmap Right (waitSTM b)
-- >>

-- <<waitAny
waitAny :: [Async a] -> IO a
waitAny asyncs =
  atomically $ foldr orElse retry $ map waitSTM asyncs
-- >>

-- <<withAsync
withAsync :: IO a -> (Async a -> IO b) -> IO b
withAsync io operation = bracket (async io) cancel operation
-- >>

-- <<waitBoth
waitBoth :: Async a -> Async b -> IO (a,b)
waitBoth a1 a2 =
  atomically $ do
    r1 <- waitSTM a1 `orElse` (do waitSTM a2; retry) -- <1>
    r2 <- waitSTM a2
    return (r1,r2)
-- >>

-- <<concurrently
concurrently :: IO a -> IO b -> IO (a,b)
concurrently ioa iob =
  withAsync ioa $ \a ->
  withAsync iob $ \b ->
    waitBoth a b
-- >>

-- <<race
race :: IO a -> IO b -> IO (Either a b)
race ioa iob =
  withAsync ioa $ \a ->
  withAsync iob $ \b ->
    waitEither a b
-- >>
<commit_msg>waitAny should be exported from Async.<commit_after>module Async (
    Async,
    async,
    withAsync,
    cancel,
    waitCatch,
    waitCatchSTM,
    waitSTM,
    wait,
    waitEither,
    waitAny,
    waitBoth,
    concurrently
  ) where

import Control.Concurrent.STM
import Control.Exception
import Control.Concurrent (ThreadId, forkIO)

-- <<Async
data Async a = Async ThreadId (STM (Either SomeException a))
-- >>

forkFinally :: IO a -> (Either SomeException a -> IO ()) -> IO ThreadId
forkFinally action fun =
  mask $ \restore ->
    forkIO (do r <- try (restore action); fun r)

-- <<async
async :: IO a -> IO (Async a)
async action = do
  var <- newEmptyTMVarIO
  t <- forkFinally action (atomically . putTMVar var)
  return (Async t (readTMVar var))
-- >>

-- <<Functor
instance Functor Async where
  fmap f (Async t stm) = Async t stm'
    where stm' = do
            r <- stm
            case r of
              Left e  -> return (Left e)
              Right a -> return (Right (f a))
-- >>

--- <<watchCatch
waitCatch :: Async a -> IO (Either SomeException a)
waitCatch = atomically . waitCatchSTM
-- >>

-- <<waitCatchSTM
waitCatchSTM :: Async a -> STM (Either SomeException a)
waitCatchSTM (Async _ stm) = stm
-- >>

-- <<waitSTM
waitSTM :: Async a -> STM a
waitSTM a = do
  r <- waitCatchSTM a
  case r of
    Left e  -> throwSTM e
    Right a -> return a
-- >>

-- <<wait
wait :: Async a -> IO a
wait = atomically . waitSTM
-- >>

-- <<cancel
cancel :: Async a -> IO ()
cancel (Async t _) = throwTo t ThreadKilled
-- >>

-- <<waitEither
waitEither :: Async a -> Async b -> IO (Either a b)
waitEither a b = atomically $
  fmap Left (waitSTM a)
    `orElse`
  fmap Right (waitSTM b)
-- >>

-- <<waitAny
waitAny :: [Async a] -> IO a
waitAny asyncs =
  atomically $ foldr orElse retry $ map waitSTM asyncs
-- >>

-- <<withAsync
withAsync :: IO a -> (Async a -> IO b) -> IO b
withAsync io operation = bracket (async io) cancel operation
-- >>

-- <<waitBoth
waitBoth :: Async a -> Async b -> IO (a,b)
waitBoth a1 a2 =
  atomically $ do
    r1 <- waitSTM a1 `orElse` (do waitSTM a2; retry) -- <1>
    r2 <- waitSTM a2
    return (r1,r2)
-- >>

-- <<concurrently
concurrently :: IO a -> IO b -> IO (a,b)
concurrently ioa iob =
  withAsync ioa $ \a ->
  withAsync iob $ \b ->
    waitBoth a b
-- >>

-- <<race
race :: IO a -> IO b -> IO (Either a b)
race ioa iob =
  withAsync ioa $ \a ->
  withAsync iob $ \b ->
    waitEither a b
-- >>

{'fraction_non_alphanumeric': 0.14144736842105263, 'fraction_numerical': 0.004111842105263158, 'mean_word_length': 2.345942228335626, 'pattern_counts': {'":': 0, '<': 75, '<?xml version=': 0, '>': 121, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1703008', 'n_tokens_mistral': 1751, 'n_tokens_neox': 1677, 'n_words': 606}