{-# LANGUAGE LambdaCase #-}

module General.Fence(
    Fence, newFence, signalFence, waitFence, testFence,
    exceptFence
    ) where

import Control.Monad
import Control.Monad.IO.Class
import Control.Exception.Extra
import Development.Shake.Internal.Errors
import Data.Maybe
import Data.Either.Extra
import Data.IORef


---------------------------------------------------------------------
-- FENCE

-- | Like a barrier, but based on callbacks
newtype Fence m a = Fence (IORef (Either (a -> m ()) a))
instance Show (Fence m a) where show :: Fence m a -> String
show Fence m a
_ = String
"Fence"

newFence :: MonadIO m => IO (Fence m a)
newFence :: IO (Fence m a)
newFence = IORef (Either (a -> m ()) a) -> Fence m a
forall (m :: * -> *) a. IORef (Either (a -> m ()) a) -> Fence m a
Fence (IORef (Either (a -> m ()) a) -> Fence m a)
-> IO (IORef (Either (a -> m ()) a)) -> IO (Fence m a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Either (a -> m ()) a -> IO (IORef (Either (a -> m ()) a))
forall a. a -> IO (IORef a)
newIORef ((a -> m ()) -> Either (a -> m ()) a
forall a b. a -> Either a b
Left ((a -> m ()) -> Either (a -> m ()) a)
-> (a -> m ()) -> Either (a -> m ()) a
forall a b. (a -> b) -> a -> b
$ m () -> a -> m ()
forall a b. a -> b -> a
const (m () -> a -> m ()) -> m () -> a -> m ()
forall a b. (a -> b) -> a -> b
$ () -> m ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ())

signalFence :: (Partial, MonadIO m) => Fence m a -> a -> m ()
signalFence :: Fence m a -> a -> m ()
signalFence (Fence IORef (Either (a -> m ()) a)
ref) a
v = m (m ()) -> m ()
forall (m :: * -> *) a. Monad m => m (m a) -> m a
join (m (m ()) -> m ()) -> m (m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ IO (m ()) -> m (m ())
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (m ()) -> m (m ())) -> IO (m ()) -> m (m ())
forall a b. (a -> b) -> a -> b
$ IORef (Either (a -> m ()) a)
-> (Either (a -> m ()) a -> (Either (a -> m ()) a, m ()))
-> IO (m ())
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' IORef (Either (a -> m ()) a)
ref ((Either (a -> m ()) a -> (Either (a -> m ()) a, m ()))
 -> IO (m ()))
-> (Either (a -> m ()) a -> (Either (a -> m ()) a, m ()))
-> IO (m ())
forall a b. (a -> b) -> a -> b
$ \case
    Left a -> m ()
queue -> (a -> Either (a -> m ()) a
forall a b. b -> Either a b
Right a
v, a -> m ()
queue a
v)
    Right a
_ -> SomeException -> (Either (a -> m ()) a, m ())
forall a. SomeException -> a
throwImpure (SomeException -> (Either (a -> m ()) a, m ()))
-> SomeException -> (Either (a -> m ()) a, m ())
forall a b. (a -> b) -> a -> b
$ Partial => String -> SomeException
String -> SomeException
errorInternal String
"signalFence called twice on one Fence"

waitFence :: MonadIO m => Fence m a -> (a -> m ()) -> m ()
waitFence :: Fence m a -> (a -> m ()) -> m ()
waitFence (Fence IORef (Either (a -> m ()) a)
ref) a -> m ()
call = m (m ()) -> m ()
forall (m :: * -> *) a. Monad m => m (m a) -> m a
join (m (m ()) -> m ()) -> m (m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ IO (m ()) -> m (m ())
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (m ()) -> m (m ())) -> IO (m ()) -> m (m ())
forall a b. (a -> b) -> a -> b
$ IORef (Either (a -> m ()) a)
-> (Either (a -> m ()) a -> (Either (a -> m ()) a, m ()))
-> IO (m ())
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' IORef (Either (a -> m ()) a)
ref ((Either (a -> m ()) a -> (Either (a -> m ()) a, m ()))
 -> IO (m ()))
-> (Either (a -> m ()) a -> (Either (a -> m ()) a, m ()))
-> IO (m ())
forall a b. (a -> b) -> a -> b
$ \case
    Left a -> m ()
queue -> ((a -> m ()) -> Either (a -> m ()) a
forall a b. a -> Either a b
Left (\a
a -> a -> m ()
queue a
a m () -> m () -> m ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> a -> m ()
call a
a), () -> m ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ())
    Right a
v -> (a -> Either (a -> m ()) a
forall a b. b -> Either a b
Right a
v, a -> m ()
call a
v)

testFence :: Fence m a -> IO (Maybe a)
testFence :: Fence m a -> IO (Maybe a)
testFence (Fence IORef (Either (a -> m ()) a)
x) = Either (a -> m ()) a -> Maybe a
forall a b. Either a b -> Maybe b
eitherToMaybe (Either (a -> m ()) a -> Maybe a)
-> IO (Either (a -> m ()) a) -> IO (Maybe a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IORef (Either (a -> m ()) a) -> IO (Either (a -> m ()) a)
forall a. IORef a -> IO a
readIORef IORef (Either (a -> m ()) a)
x


---------------------------------------------------------------------
-- FENCE COMPOSITES

exceptFence :: MonadIO m => [Fence m (Either e r)] -> m (Fence m (Either e [r]))
exceptFence :: [Fence m (Either e r)] -> m (Fence m (Either e [r]))
exceptFence [Fence m (Either e r)]
xs = do
    -- number of items still to complete, becomes negative after it has triggered
    IORef Int
todo <- IO (IORef Int) -> m (IORef Int)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (IORef Int) -> m (IORef Int))
-> IO (IORef Int) -> m (IORef Int)
forall a b. (a -> b) -> a -> b
$ Int -> IO (IORef Int)
forall a. a -> IO (IORef a)
newIORef (Int -> IO (IORef Int)) -> Int -> IO (IORef Int)
forall a b. (a -> b) -> a -> b
$ [Fence m (Either e r)] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Fence m (Either e r)]
xs
    Fence m (Either e [r])
fence <- IO (Fence m (Either e [r])) -> m (Fence m (Either e [r]))
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO IO (Fence m (Either e [r]))
forall (m :: * -> *) a. MonadIO m => IO (Fence m a)
newFence

    [Fence m (Either e r)] -> (Fence m (Either e r) -> m ()) -> m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Fence m (Either e r)]
xs ((Fence m (Either e r) -> m ()) -> m ())
-> (Fence m (Either e r) -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \Fence m (Either e r)
x -> Fence m (Either e r) -> (Either e r -> m ()) -> m ()
forall (m :: * -> *) a.
MonadIO m =>
Fence m a -> (a -> m ()) -> m ()
waitFence Fence m (Either e r)
x ((Either e r -> m ()) -> m ()) -> (Either e r -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \Either e r
res ->
        m (m ()) -> m ()
forall (m :: * -> *) a. Monad m => m (m a) -> m a
join (m (m ()) -> m ()) -> m (m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ IO (m ()) -> m (m ())
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (m ()) -> m (m ())) -> IO (m ()) -> m (m ())
forall a b. (a -> b) -> a -> b
$ IORef Int -> (Int -> (Int, m ())) -> IO (m ())
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' IORef Int
todo ((Int -> (Int, m ())) -> IO (m ()))
-> (Int -> (Int, m ())) -> IO (m ())
forall a b. (a -> b) -> a -> b
$ \Int
i -> case Either e r
res of
            Left e
e | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 -> (-Int
1, Fence m (Either e [r]) -> Either e [r] -> m ()
forall (m :: * -> *) a.
(Partial, MonadIO m) =>
Fence m a -> a -> m ()
signalFence Fence m (Either e [r])
fence (Either e [r] -> m ()) -> Either e [r] -> m ()
forall a b. (a -> b) -> a -> b
$ e -> Either e [r]
forall a b. a -> Either a b
Left e
e)
            Either e r
_ | Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1 -> (-Int
1, Fence m (Either e [r]) -> Either e [r] -> m ()
forall (m :: * -> *) a.
(Partial, MonadIO m) =>
Fence m a -> a -> m ()
signalFence Fence m (Either e [r])
fence (Either e [r] -> m ()) -> ([r] -> Either e [r]) -> [r] -> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [r] -> Either e [r]
forall a b. b -> Either a b
Right ([r] -> m ()) -> m [r] -> m ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< IO [r] -> m [r]
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO ((Fence m (Either e r) -> IO r) -> [Fence m (Either e r)] -> IO [r]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ((Maybe (Either e r) -> r) -> IO (Maybe (Either e r)) -> IO r
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Either e r -> r
forall l r. Partial => Either l r -> r
fromRight' (Either e r -> r)
-> (Maybe (Either e r) -> Either e r) -> Maybe (Either e r) -> r
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Maybe (Either e r) -> Either e r
forall a. Partial => Maybe a -> a
fromJust) (IO (Maybe (Either e r)) -> IO r)
-> (Fence m (Either e r) -> IO (Maybe (Either e r)))
-> Fence m (Either e r)
-> IO r
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Fence m (Either e r) -> IO (Maybe (Either e r))
forall (m :: * -> *) a. Fence m a -> IO (Maybe a)
testFence) [Fence m (Either e r)]
xs))
              | Bool
otherwise -> (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1, () -> m ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ())
    Fence m (Either e [r]) -> m (Fence m (Either e [r]))
forall (f :: * -> *) a. Applicative f => a -> f a
pure Fence m (Either e [r])
fence