{-# 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 :: forall (m :: * -> *) a. MonadIO m => IO (Fence m a)
newFence = forall (m :: * -> *) a. IORef (Either (a -> m ()) a) -> Fence m a
Fence forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. a -> IO (IORef a)
newIORef (forall a b. a -> Either a b
Left forall a b. (a -> b) -> a -> b
$ forall a b. a -> b -> a
const forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a. Applicative f => a -> f a
pure ())

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

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

testFence :: Fence m a -> IO (Maybe a)
testFence :: forall (m :: * -> *) a. Fence m a -> IO (Maybe a)
testFence (Fence IORef (Either (a -> m ()) a)
x) = forall a b. Either a b -> Maybe b
eitherToMaybe forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> 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 :: forall (m :: * -> *) e r.
MonadIO m =>
[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 <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. a -> IO (IORef a)
newIORef forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) a. Foldable t => t a -> Int
length [Fence m (Either e r)]
xs
    Fence m (Either e [r])
fence <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall (m :: * -> *) a. MonadIO m => IO (Fence m a)
newFence

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