    {-# LANGUAGE ExistentialQuantification #-}

    module Development.Shake.Internal.Core.Rendezvous(
        Waiting, newWaiting, afterWaiting,
        Answer(..), Compute(..),
        rendezvous
        ) where

    import Control.Monad
    import Data.IORef.Extra
    import Data.Primitive.Array
    import Development.Shake.Internal.Errors

    -- | Given a sequence of 'Answer' values the sequence stops
    --   when there is a single 'Abort' or all values end up as 'Continue'.
    data Answer a c
        = Abort a
        | Continue c

    -- | A compuation that either has a result available immediate,
    --   or has a result that can be collected later.
    data Compute a
        = Now a
        | Later (Waiting a)

    partitionAnswer :: [Answer a c] -> ([a], [c])
    partitionAnswer = foldr f ([],[])
        where f (Abort    a) ~(as,cs) = (a:as,cs)
              f (Continue c) ~(as,cs) = (as,c:cs)

    partitionCompute :: [Compute a] -> ([a], [Waiting a])
    partitionCompute = foldr f ([],[])
        where f (Now   x) ~(xs,ws) = (x:xs,ws)
              f (Later w) ~(xs,ws) = (xs,w:ws)

    -- | A type representing someone waiting for a result.
    data Waiting a = forall b . Waiting (b -> a) (IORef (b -> IO ()))
        -- Contains a functor value to apply, along with somewhere to register callbacks

    instance Functor Waiting where
        fmap f (Waiting op ref) = Waiting (f . op) ref

    instance Show (Waiting a) where
        show _ = "Waiting"

    newWaiting :: IO (Waiting a, a -> IO ())
    newWaiting = do
        ref <- newIORef $ \_ -> return ()
        let run x = ($ x) =<< readIORef ref
        return (Waiting id ref, run)

    afterWaiting :: Waiting a -> (a -> IO ()) -> IO ()
    afterWaiting (Waiting op ref) act = modifyIORef' ref (\a s -> a s >> act (op s))

    rendezvous :: [Compute (Answer a c)] -> IO (Compute (Either a [c]))
    rendezvous xs = do
        let (now, later) = partitionCompute xs
        let (abort, continue) = partitionAnswer now
        if not $ null abort then
            return $ Now $ Left $ head abort
         else if null later then
            return $ Now $ Right continue
         else do
            (waiting, run) <- newWaiting
            let n = length xs
            result <- newArray n $ errorInternal "rendezvous"
            todo <- newIORef $ length later
            forM_ (zip [0..] xs) $ \(i,x) -> case x of
                Now (Continue c) -> writeArray result i c
                Later w -> afterWaiting w $ \v -> do
                    t <- readIORef todo
                    case v of
                        _ | t == 0 -> return () -- must have already aborted
                        Abort a -> do
                            writeIORef todo 0
                            run $ Left a
                        Continue c -> do
                            writeArray result i c
                            writeIORef' todo $ t-1
                            when (t == 1) $ do
                                rs <- unsafeFreezeArray result
                                run $ Right $ map (indexArray rs) [0..n-1]
            return $ Later waiting