{-
  Copyright (c) Meta Platforms, Inc. and affiliates.
  All rights reserved.

  This source code is licensed under the BSD-style license found in the
  LICENSE file in the root directory of this source tree.
-}

{-# OPTIONS_GHC -fno-warn-missing-signatures #-}
{-# LANGUAGE CPP #-}
module Util.Concurrent
  ( concurrently3
  , concurrently4
  , concurrently5
  , concurrently6
  , concurrently7
  -- * ThreadLock
  , ThreadLock
  , newThreadLock
  , withThreadLock
  -- * Caching
  , cacheSuccess
  ) where

import Control.Concurrent
import Control.Concurrent.Async
import Control.Monad (void)
import Control.Monad.IO.Class

#if __GLASGOW_HASKELL__ >= 902
import Control.Monad.Catch
#else
import Exception (ExceptionMonad(..))
#endif

c :: IO a -> IO b -> IO (a, b)
c = concurrently

concurrently3 :: IO t1 -> IO t2 -> IO t3 -> IO (t1, t2, t3)
concurrently3 a1 a2 a3 = flatten3 <$> a1 `c` a2 `c` a3

concurrently4 :: IO t1 -> IO t2 -> IO t3 -> IO t4 -> IO (t1, t2, t3, t4)
concurrently4 a1 a2 a3 a4 = flatten4 <$> a1 `c` a2 `c` a3 `c` a4

concurrently5 :: IO t1 -> IO t2 -> IO t3 -> IO t4 -> IO t5
              -> IO (t1, t2, t3, t4, t5)
concurrently5 a1 a2 a3 a4 a5 = flatten5 <$> a1 `c` a2 `c` a3 `c` a4 `c` a5

concurrently6 :: IO t1 -> IO t2 -> IO t3 -> IO t4 -> IO t5 -> IO t6
              -> IO (t1, t2, t3, t4, t5, t6)
concurrently6 a1 a2 a3 a4 a5 a6 =
  flatten6 <$> a1 `c` a2 `c` a3 `c` a4 `c` a5 `c` a6

concurrently7 :: IO t1 -> IO t2 -> IO t3 -> IO t4 -> IO t5 -> IO t6 -> IO t7
              -> IO (t1, t2, t3, t4, t5, t6, t7)
concurrently7 a1 a2 a3 a4 a5 a6 a7 =
  flatten7 <$> a1 `c` a2 `c` a3 `c` a4 `c` a5 `c` a6 `c` a7

flatten3 ((a1, a2), a3) = (a1, a2, a3)
flatten4 (((a1, a2), a3), a4) = (a1, a2, a3, a4)
flatten5 ((((a1, a2), a3), a4), a5) = (a1, a2, a3, a4, a5)
flatten6 (((((a1, a2), a3), a4), a5), a6) = (a1, a2, a3, a4, a5, a6)
flatten7 ((((((a1, a2), a3), a4), a5), a6), a7) = (a1, a2, a3, a4, a5, a6, a7)

newtype ThreadLock = ThreadLock (IO (IO ()))

newThreadLock :: IO ThreadLock
newThreadLock = do
  mvar <- newEmptyMVar
  return $ ThreadLock $ do
    ours <- myThreadId
    theirs <- tryReadMVar mvar
    if theirs == Just ours
    then return $ return () -- we are holding the lock, so this is noop
    else do
      putMVar mvar ours
      return $ void $ takeMVar mvar

-- | withThreadLock guarantees that only one thread holds the lock at the same
-- time. It can be nested so it is safe to use it in a recursive function.
#if __GLASGOW_HASKELL__ >= 902
withThreadLock :: (MonadIO m, MonadMask m) => ThreadLock -> m a -> m a
withThreadLock (ThreadLock takeIO) action = do
  fmap fst $ generalBracket (liftIO takeIO) (\releaseIO _ -> liftIO releaseIO) $ const action
#else
withThreadLock :: ExceptionMonad m => ThreadLock -> m a -> m a
withThreadLock (ThreadLock takeIO) action = do
  gbracket (liftIO takeIO) (\releaseIO -> liftIO releaseIO) $ const action
#endif

-- | Build an IO action that will cache its result the first time it
-- runs successfullly to completion. If it fails with an exception,
-- the exception is re-raised, but the IO action will be attempted
-- again the next time it is invoked.
--
-- Like `Control.Concurrent.Extra.once`, but without caching exceptions.
--
cacheSuccess :: IO a -> IO (IO a)
cacheSuccess act = do
  m <- newMVar Nothing
  return $ modifyMVar m $ \x -> case x of
    Nothing -> do a <- act; return (Just a, a)
    Just a -> return (Just a, a)