{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DeriveFunctor #-}

-- | A bit like 'Fence', but not thread safe and optimised for avoiding taking the fence
module General.Wait(
    Wait(Now,Later), runWait, quickly, fromLater,
    firstJustWaitUnordered, firstLeftWaitUnordered
    ) where

import Control.Monad.Extra
import Control.Monad.IO.Class
import Data.IORef.Extra
import Data.List.Extra
import Data.Primitive.Array
import GHC.Exts(RealWorld)
import Control.Monad.Fail
import Prelude


runWait :: Monad m => Wait m a -> m (Wait m a)
runWait :: forall (m :: * -> *) a. Monad m => Wait m a -> m (Wait m a)
runWait (Lift m (Wait m a)
x) = forall (m :: * -> *) a. Monad m => Wait m a -> m (Wait m a)
runWait forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< m (Wait m a)
x
runWait Wait m a
x = forall (f :: * -> *) a. Applicative f => a -> f a
pure Wait m a
x

fromLater :: Monad m => Wait m a -> (a -> m ()) -> m ()
fromLater :: forall (m :: * -> *) a. Monad m => Wait m a -> (a -> m ()) -> m ()
fromLater (Lift m (Wait m a)
x) a -> m ()
f = do Wait m a
x <- m (Wait m a)
x; forall (m :: * -> *) a. Monad m => Wait m a -> (a -> m ()) -> m ()
fromLater Wait m a
x a -> m ()
f
fromLater (Now a
x) a -> m ()
f = a -> m ()
f a
x
fromLater (Later (a -> m ()) -> m ()
x) a -> m ()
f = (a -> m ()) -> m ()
x a -> m ()
f

quickly :: Functor m => m a -> Wait m a
quickly :: forall (m :: * -> *) a. Functor m => m a -> Wait m a
quickly = forall (m :: * -> *) a. m (Wait m a) -> Wait m a
Lift forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall (m :: * -> *) a. a -> Wait m a
Now

data Wait m a = Now a
              | Lift (m (Wait m a))
              | Later ((a -> m ()) -> m ())
                deriving forall a b. a -> Wait m b -> Wait m a
forall a b. (a -> b) -> Wait m a -> Wait m b
forall (m :: * -> *) a b. Functor m => a -> Wait m b -> Wait m a
forall (m :: * -> *) a b.
Functor m =>
(a -> b) -> Wait m a -> Wait m b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
<$ :: forall a b. a -> Wait m b -> Wait m a
$c<$ :: forall (m :: * -> *) a b. Functor m => a -> Wait m b -> Wait m a
fmap :: forall a b. (a -> b) -> Wait m a -> Wait m b
$cfmap :: forall (m :: * -> *) a b.
Functor m =>
(a -> b) -> Wait m a -> Wait m b
Functor

instance (Monad m, Applicative m) => Applicative (Wait m) where
    pure :: forall a. a -> Wait m a
pure = forall (m :: * -> *) a. a -> Wait m a
Now
    Now a -> b
x <*> :: forall a b. Wait m (a -> b) -> Wait m a -> Wait m b
<*> Wait m a
y = a -> b
x forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Wait m a
y
    Lift m (Wait m (a -> b))
x <*> Wait m a
y = forall (m :: * -> *) a. m (Wait m a) -> Wait m a
Lift forall a b. (a -> b) -> a -> b
$ (forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Wait m a
y) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> m (Wait m (a -> b))
x
    Later ((a -> b) -> m ()) -> m ()
x <*> Now a
y = forall (m :: * -> *) a. ((a -> m ()) -> m ()) -> Wait m a
Later forall a b. (a -> b) -> a -> b
$ \b -> m ()
c -> ((a -> b) -> m ()) -> m ()
x forall a b. (a -> b) -> a -> b
$ \a -> b
x -> b -> m ()
c forall a b. (a -> b) -> a -> b
$ a -> b
x a
y
    -- Note: We pull the Lift from the right BEFORE the Later, to enable parallelism
    Later ((a -> b) -> m ()) -> m ()
x <*> Lift m (Wait m a)
y = forall (m :: * -> *) a. m (Wait m a) -> Wait m a
Lift forall a b. (a -> b) -> a -> b
$ do Wait m a
y <- m (Wait m a)
y; forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. ((a -> m ()) -> m ()) -> Wait m a
Later ((a -> b) -> m ()) -> m ()
x forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Wait m a
y
    Later ((a -> b) -> m ()) -> m ()
x <*> Later (a -> m ()) -> m ()
y = forall (m :: * -> *) a. ((a -> m ()) -> m ()) -> Wait m a
Later forall a b. (a -> b) -> a -> b
$ \b -> m ()
c -> ((a -> b) -> m ()) -> m ()
x forall a b. (a -> b) -> a -> b
$ \a -> b
x -> (a -> m ()) -> m ()
y forall a b. (a -> b) -> a -> b
$ \a
y -> b -> m ()
c forall a b. (a -> b) -> a -> b
$ a -> b
x a
y

instance (Monad m, Applicative m) => Monad (Wait m) where
    return :: forall a. a -> Wait m a
return = forall (f :: * -> *) a. Applicative f => a -> f a
pure
    >> :: forall a b. Wait m a -> Wait m b -> Wait m b
(>>) = forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
(*>)
    Now a
x >>= :: forall a b. Wait m a -> (a -> Wait m b) -> Wait m b
>>= a -> Wait m b
f = a -> Wait m b
f a
x
    Lift m (Wait m a)
x >>= a -> Wait m b
f = forall (m :: * -> *) a. m (Wait m a) -> Wait m a
Lift forall a b. (a -> b) -> a -> b
$ do Wait m a
x <- m (Wait m a)
x; forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ Wait m a
x forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= a -> Wait m b
f
    Later (a -> m ()) -> m ()
x >>= a -> Wait m b
f = forall (m :: * -> *) a. ((a -> m ()) -> m ()) -> Wait m a
Later forall a b. (a -> b) -> a -> b
$ \b -> m ()
c -> (a -> m ()) -> m ()
x forall a b. (a -> b) -> a -> b
$ \a
x -> do
        Wait m b
x <- forall (m :: * -> *) a. Monad m => Wait m a -> m (Wait m a)
runWait forall a b. (a -> b) -> a -> b
$ a -> Wait m b
f a
x
        case Wait m b
x of
            Now b
x -> b -> m ()
c b
x
            Wait m b
_ -> forall (m :: * -> *) a. Monad m => Wait m a -> (a -> m ()) -> m ()
fromLater Wait m b
x b -> m ()
c

instance (MonadIO m,  Applicative m) => MonadIO (Wait m) where
    liftIO :: forall a. IO a -> Wait m a
liftIO = forall (m :: * -> *) a. m (Wait m a) -> Wait m a
Lift forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall (m :: * -> *) a. a -> Wait m a
Now

instance MonadFail m => MonadFail (Wait m) where
    fail :: forall a. String -> Wait m a
fail = forall (m :: * -> *) a. m (Wait m a) -> Wait m a
Lift forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) a. MonadFail m => String -> m a
Control.Monad.Fail.fail

firstJustWaitUnordered :: MonadIO m => (a -> Wait m (Maybe b)) -> [a] -> Wait m (Maybe b)
firstJustWaitUnordered :: forall (m :: * -> *) a b.
MonadIO m =>
(a -> Wait m (Maybe b)) -> [a] -> Wait m (Maybe b)
firstJustWaitUnordered a -> Wait m (Maybe b)
f = forall (m :: * -> *) a.
MonadIO m =>
Int
-> [(Maybe a -> m ()) -> m ()]
-> [Wait m (Maybe a)]
-> Wait m (Maybe a)
go Int
0 [] forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map a -> Wait m (Maybe b)
f
    where
        -- keep a list of those things we might visit later, and ask for each we see in turn
        go :: MonadIO m => Int -> [(Maybe a -> m ()) -> m ()] -> [Wait m (Maybe a)] -> Wait m (Maybe a)
        go :: forall (m :: * -> *) a.
MonadIO m =>
Int
-> [(Maybe a -> m ()) -> m ()]
-> [Wait m (Maybe a)]
-> Wait m (Maybe a)
go !Int
nlater [(Maybe a -> m ()) -> m ()]
later (Wait m (Maybe a)
x:[Wait m (Maybe a)]
xs) = case Wait m (Maybe a)
x of
            Now (Just a
a) -> forall (m :: * -> *) a. a -> Wait m a
Now forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a
Just a
a
            Now Maybe a
Nothing -> forall (m :: * -> *) a.
MonadIO m =>
Int
-> [(Maybe a -> m ()) -> m ()]
-> [Wait m (Maybe a)]
-> Wait m (Maybe a)
go Int
nlater [(Maybe a -> m ()) -> m ()]
later [Wait m (Maybe a)]
xs
            Later (Maybe a -> m ()) -> m ()
l -> forall (m :: * -> *) a.
MonadIO m =>
Int
-> [(Maybe a -> m ()) -> m ()]
-> [Wait m (Maybe a)]
-> Wait m (Maybe a)
go (forall a. Enum a => a -> a
succ Int
nlater) ((Maybe a -> m ()) -> m ()
lforall a. a -> [a] -> [a]
:[(Maybe a -> m ()) -> m ()]
later) [Wait m (Maybe a)]
xs
            Lift m (Wait m (Maybe a))
x -> forall (m :: * -> *) a. m (Wait m a) -> Wait m a
Lift forall a b. (a -> b) -> a -> b
$ do
                Wait m (Maybe a)
x <- m (Wait m (Maybe a))
x
                forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a.
MonadIO m =>
Int
-> [(Maybe a -> m ()) -> m ()]
-> [Wait m (Maybe a)]
-> Wait m (Maybe a)
go Int
nlater [(Maybe a -> m ()) -> m ()]
later (Wait m (Maybe a)
xforall a. a -> [a] -> [a]
:[Wait m (Maybe a)]
xs)
        go Int
_ [] [] = forall (m :: * -> *) a. a -> Wait m a
Now forall a. Maybe a
Nothing
        go Int
_ [(Maybe a -> m ()) -> m ()
l] [] = forall (m :: * -> *) a. ((a -> m ()) -> m ()) -> Wait m a
Later (Maybe a -> m ()) -> m ()
l
        go Int
nls [(Maybe a -> m ()) -> m ()]
ls [] = forall (m :: * -> *) a. ((a -> m ()) -> m ()) -> Wait m a
Later forall a b. (a -> b) -> a -> b
$ \Maybe a -> m ()
callback -> do
            IORef Int
ref <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. a -> IO (IORef a)
newIORef Int
nls
            forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [(Maybe a -> m ()) -> m ()]
ls forall a b. (a -> b) -> a -> b
$ \(Maybe a -> m ()) -> m ()
l -> (Maybe a -> m ()) -> m ()
l forall a b. (a -> b) -> a -> b
$ \Maybe a
r -> do
                Int
old <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. IORef a -> IO a
readIORef IORef Int
ref
                forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
old forall a. Ord a => a -> a -> Bool
> Int
0) forall a b. (a -> b) -> a -> b
$ case Maybe a
r of
                    Just a
a -> do
                        forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. IORef a -> a -> IO ()
writeIORef' IORef Int
ref Int
0
                        Maybe a -> m ()
callback forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a
Just a
a
                    Maybe a
Nothing -> do
                        forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. IORef a -> a -> IO ()
writeIORef' IORef Int
ref forall a b. (a -> b) -> a -> b
$ Int
oldforall a. Num a => a -> a -> a
-Int
1
                        forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
old forall a. Eq a => a -> a -> Bool
== Int
1) forall a b. (a -> b) -> a -> b
$ Maybe a -> m ()
callback forall a. Maybe a
Nothing


firstLeftWaitUnordered :: MonadIO m => (a -> Wait m (Either e b)) -> [a] -> Wait m (Either e [b])
firstLeftWaitUnordered :: forall (m :: * -> *) a e b.
MonadIO m =>
(a -> Wait m (Either e b)) -> [a] -> Wait m (Either e [b])
firstLeftWaitUnordered a -> Wait m (Either e b)
f [a]
xs = do
        let n :: Int
n = forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs
        MutableArray RealWorld b
mut <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (MutableArray (PrimState m) a)
newArray Int
n forall a. HasCallStack => a
undefined
        Maybe e
res <- forall (m :: * -> *) b e.
MonadIO m =>
MutableArray RealWorld b
-> [(Int, (Either e b -> m ()) -> m ())]
-> [(Int, Wait m (Either e b))]
-> Wait m (Maybe e)
go MutableArray RealWorld b
mut [] forall a b. (a -> b) -> a -> b
$ forall a b. Enum a => a -> [b] -> [(a, b)]
zipFrom Int
0 forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map a -> Wait m (Either e b)
f [a]
xs
        case Maybe e
res of
            Just e
e -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall a b. a -> Either a b
Left e
e
            Maybe e
Nothing -> forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a b. b -> Either a b
Right forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> Int -> m a
readArray MutableArray RealWorld b
mut) [Int
0..Int
nforall a. Num a => a -> a -> a
-Int
1]
    where
        -- keep a list of those things we might visit later, and ask for each we see in turn
        go :: MonadIO m => MutableArray RealWorld b -> [(Int, (Either e b -> m ()) -> m ())] -> [(Int, Wait m (Either e b))] -> Wait m (Maybe e)
        go :: forall (m :: * -> *) b e.
MonadIO m =>
MutableArray RealWorld b
-> [(Int, (Either e b -> m ()) -> m ())]
-> [(Int, Wait m (Either e b))]
-> Wait m (Maybe e)
go MutableArray RealWorld b
mut [(Int, (Either e b -> m ()) -> m ())]
later ((Int
i,Wait m (Either e b)
x):[(Int, Wait m (Either e b))]
xs) = case Wait m (Either e b)
x of
            Now (Left e
e) -> forall (m :: * -> *) a. a -> Wait m a
Now forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a
Just e
e
            Now (Right b
b) -> do
                forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> Int -> a -> m ()
writeArray MutableArray RealWorld b
mut Int
i b
b
                forall (m :: * -> *) b e.
MonadIO m =>
MutableArray RealWorld b
-> [(Int, (Either e b -> m ()) -> m ())]
-> [(Int, Wait m (Either e b))]
-> Wait m (Maybe e)
go MutableArray RealWorld b
mut [(Int, (Either e b -> m ()) -> m ())]
later [(Int, Wait m (Either e b))]
xs
            Later (Either e b -> m ()) -> m ()
l -> forall (m :: * -> *) b e.
MonadIO m =>
MutableArray RealWorld b
-> [(Int, (Either e b -> m ()) -> m ())]
-> [(Int, Wait m (Either e b))]
-> Wait m (Maybe e)
go MutableArray RealWorld b
mut ((Int
i,(Either e b -> m ()) -> m ()
l)forall a. a -> [a] -> [a]
:[(Int, (Either e b -> m ()) -> m ())]
later) [(Int, Wait m (Either e b))]
xs
            Lift m (Wait m (Either e b))
x -> forall (m :: * -> *) a. m (Wait m a) -> Wait m a
Lift forall a b. (a -> b) -> a -> b
$ do
                Wait m (Either e b)
x <- m (Wait m (Either e b))
x
                forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) b e.
MonadIO m =>
MutableArray RealWorld b
-> [(Int, (Either e b -> m ()) -> m ())]
-> [(Int, Wait m (Either e b))]
-> Wait m (Maybe e)
go MutableArray RealWorld b
mut [(Int, (Either e b -> m ()) -> m ())]
later ((Int
i,Wait m (Either e b)
x)forall a. a -> [a] -> [a]
:[(Int, Wait m (Either e b))]
xs)
        go MutableArray RealWorld b
_ [] [] = forall (m :: * -> *) a. a -> Wait m a
Now forall a. Maybe a
Nothing
        go MutableArray RealWorld b
mut [(Int, (Either e b -> m ()) -> m ())]
ls [] = forall (m :: * -> *) a. ((a -> m ()) -> m ()) -> Wait m a
Later forall a b. (a -> b) -> a -> b
$ \Maybe e -> m ()
callback -> do
            IORef Int
ref <- 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 [(Int, (Either e b -> m ()) -> m ())]
ls
            forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [(Int, (Either e b -> m ()) -> m ())]
ls forall a b. (a -> b) -> a -> b
$ \(Int
i,(Either e b -> m ()) -> m ()
l) -> (Either e b -> m ()) -> m ()
l forall a b. (a -> b) -> a -> b
$ \Either e b
r -> do
                Int
old <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. IORef a -> IO a
readIORef IORef Int
ref
                forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
old forall a. Ord a => a -> a -> Bool
> Int
0) forall a b. (a -> b) -> a -> b
$ case Either e b
r of
                    Left e
a -> do
                        forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. IORef a -> a -> IO ()
writeIORef' IORef Int
ref Int
0
                        Maybe e -> m ()
callback forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a
Just e
a
                    Right b
v -> do
                        forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> Int -> a -> m ()
writeArray MutableArray RealWorld b
mut Int
i b
v
                        forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. IORef a -> a -> IO ()
writeIORef' IORef Int
ref forall a b. (a -> b) -> a -> b
$ Int
oldforall a. Num a => a -> a -> a
-Int
1
                        forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
old forall a. Eq a => a -> a -> Bool
== Int
1) forall a b. (a -> b) -> a -> b
$ Maybe e -> m ()
callback forall a. Maybe a
Nothing