module Mischief.ECS.World.Par
  ( par,
    par',
    parIterList,
    parFor,
    parFor_,
    parFor',
    parFor'_,
  )
where

import Control.Concurrent.Async (async, wait)
import Control.Monad
import Control.Monad.IO.Class
import Control.Monad.Reader
import Data.Foldable
import Data.IORef (modifyIORef', newIORef, readIORef)
import Data.Traversable
import GHC.Conc (numCapabilities)
import Mischief.ECS.Hidden
import Mischief.ECS.World
import Mischief.ECS.World.Defer

-- | Receives a list of @ParSystems@ and runs them in parallel. If ran in 'System', flushes the deferred queue when done!
par :: (MonadSystem w m) => [ParSystem ()] -> m ()
par :: forall w (m :: * -> *). MonadSystem w m => [ParSystem ()] -> m ()
par ![ParSystem ()]
x =
  System () -> ParSystem () -> m ()
forall a. System a -> ParSystem a -> m a
forall (m :: * -> *) a.
SystemAny m =>
System a -> ParSystem a -> m a
systemAny
    ( do
        [ParSystem ()] -> System ()
forall w (m :: * -> *). MonadSystem w m => [ParSystem ()] -> m ()
par' [ParSystem ()]
x
        System ()
flush
    )
    ([ParSystem ()] -> ParSystem ()
forall w (m :: * -> *). MonadSystem w m => [ParSystem ()] -> m ()
par' [ParSystem ()]
x)

-- | Receives a list of @ParSystems@ and runs them in parallel. Doesn't flush the deferred queue when done!
par' :: (MonadSystem w m) => [ParSystem ()] -> m ()
par' :: forall w (m :: * -> *). MonadSystem w m => [ParSystem ()] -> m ()
par' ![ParSystem ()]
parSystems = do
  !world <- m World
forall w (m :: * -> *). MonadSystem w m => m World
unsafeGetWorld

  !x <- forM parSystems $ \(ParSystem !ReaderT ParWorld IO ()
p) -> do
    !systems <- IO (IORef [System ()]) -> m (IORef [System ()])
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (IORef [System ()]) -> m (IORef [System ()]))
-> IO (IORef [System ()]) -> m (IORef [System ()])
forall a b. (a -> b) -> a -> b
$ [System ()] -> IO (IORef [System ()])
forall a. a -> IO (IORef a)
newIORef []
    !id <- liftIO $ async $ runReaderT p ParWorld {world = Hidden world, parDeferred = systems}
    pure (id, systems)

  for_ x $ \(!Async ()
id, !IORef [System ()]
systems) -> do
    IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ Async () -> IO ()
forall a. Async a -> IO a
wait Async ()
id
    !systems <- IO [System ()] -> m [System ()]
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO [System ()] -> m [System ()])
-> IO [System ()] -> m [System ()]
forall a b. (a -> b) -> a -> b
$ IORef [System ()] -> IO [System ()]
forall a. IORef a -> IO a
readIORef IORef [System ()]
systems
    liftIO $ modifyIORef' world.deferred (++ systems)

parIterList :: (MonadSystem w m, Foldable t) => t a -> ([a] -> ParSystem b) -> m [b]
parIterList :: forall w (m :: * -> *) (t :: * -> *) a b.
(MonadSystem w m, Foldable t) =>
t a -> ([a] -> ParSystem b) -> m [b]
parIterList !t a
list ![a] -> ParSystem b
s = do
  !world <- m World
forall w (m :: * -> *). MonadSystem w m => m World
unsafeGetWorld

  let !n = Int
numCapabilities
  let !len = t a -> Int
forall a. t a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length t a
list

  let !chunks = Int -> [a] -> [[a]]
forall a. Int -> [a] -> [[a]]
group (Int
len Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
n) (t a -> [a]
forall a. t a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList t a
list)

  !x <- forM chunks $ \(![a]
chunk) -> do
    !systems <- IO (IORef [System ()]) -> m (IORef [System ()])
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (IORef [System ()]) -> m (IORef [System ()]))
-> IO (IORef [System ()]) -> m (IORef [System ()])
forall a b. (a -> b) -> a -> b
$ [System ()] -> IO (IORef [System ()])
forall a. a -> IO (IORef a)
newIORef []
    let ParSystem !p = s chunk
    !id <- liftIO $ async $ runReaderT p ParWorld {world = Hidden world, parDeferred = systems}
    pure (id, systems)

  for x $ \(!Async b
id, !IORef [System ()]
systems) -> do
    !a <- IO b -> m b
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO b -> m b) -> IO b -> m b
forall a b. (a -> b) -> a -> b
$ Async b -> IO b
forall a. Async a -> IO a
wait Async b
id
    !systems <- liftIO $ readIORef systems
    liftIO $ modifyIORef' world.deferred (++ systems)
    pure a

group :: Int -> [a] -> [[a]]
group :: forall a. Int -> [a] -> [[a]]
group Int
_ [] = []
group Int
0 ![a]
l = [[a]
l]
group !Int
n ![a]
l = Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
take Int
n [a]
l [a] -> [[a]] -> [[a]]
forall a. a -> [a] -> [a]
: Int -> [a] -> [[a]]
forall a. Int -> [a] -> [[a]]
group Int
n (Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
drop Int
n [a]
l)

-- | Iterates over a list in parallel. If ran from 'System', flushes the deferred queue when done.
parFor :: (Foldable t, MonadSystem w m) => t a -> (a -> ParSystem b) -> m [b]
parFor :: forall (t :: * -> *) w (m :: * -> *) a b.
(Foldable t, MonadSystem w m) =>
t a -> (a -> ParSystem b) -> m [b]
parFor !t a
x !a -> ParSystem b
s =
  System [b] -> ParSystem [b] -> m [b]
forall a. System a -> ParSystem a -> m a
forall (m :: * -> *) a.
SystemAny m =>
System a -> ParSystem a -> m a
systemAny
    ( do
        !a <- t a -> (a -> ParSystem b) -> System [b]
forall w (m :: * -> *) (t :: * -> *) a b.
(MonadSystem w m, Foldable t) =>
t a -> (a -> ParSystem b) -> m [b]
parFor' t a
x a -> ParSystem b
s
        flush
        pure a
    )
    (t a -> (a -> ParSystem b) -> ParSystem [b]
forall w (m :: * -> *) (t :: * -> *) a b.
(MonadSystem w m, Foldable t) =>
t a -> (a -> ParSystem b) -> m [b]
parFor' t a
x a -> ParSystem b
s)

-- | Iterates over a list in parallel and discards the results. If ran from 'System', flushes the deffered queue when done.
parFor_ :: (Foldable t) => t a -> (a -> ParSystem b) -> System ()
parFor_ :: forall (t :: * -> *) a b.
Foldable t =>
t a -> (a -> ParSystem b) -> System ()
parFor_ !t a
x !a -> ParSystem b
s = do
  System [b] -> System ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (System [b] -> System ()) -> System [b] -> System ()
forall a b. (a -> b) -> a -> b
$ t a -> (a -> ParSystem b) -> System [b]
forall w (m :: * -> *) (t :: * -> *) a b.
(MonadSystem w m, Foldable t) =>
t a -> (a -> ParSystem b) -> m [b]
parFor' t a
x a -> ParSystem b
s
  System ()
flush

-- | Iterates over a list in parallel. Doesn't flush the deferred queue when done.
parFor' :: (MonadSystem w m, Foldable t) => t a -> (a -> ParSystem b) -> m [b]
parFor' :: forall w (m :: * -> *) (t :: * -> *) a b.
(MonadSystem w m, Foldable t) =>
t a -> (a -> ParSystem b) -> m [b]
parFor' !t a
x !a -> ParSystem b
s = [[b]] -> [b]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[b]] -> [b]) -> m [[b]] -> m [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> t a -> ([a] -> ParSystem [b]) -> m [[b]]
forall w (m :: * -> *) (t :: * -> *) a b.
(MonadSystem w m, Foldable t) =>
t a -> ([a] -> ParSystem b) -> m [b]
parIterList t a
x ([a] -> (a -> ParSystem b) -> ParSystem [b]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
t a -> (a -> f b) -> f (t b)
`for` a -> ParSystem b
s)

-- | Iterates over a list in parallel and discards the results. Doesn't flush the deffered queue when done.
parFor'_ :: (MonadSystem w m, Foldable t) => t a -> (a -> ParSystem b) -> m ()
parFor'_ :: forall w (m :: * -> *) (t :: * -> *) a b.
(MonadSystem w m, Foldable t) =>
t a -> (a -> ParSystem b) -> m ()
parFor'_ !t a
x !a -> ParSystem b
s = m [()] -> m ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (t a -> ([a] -> ParSystem ()) -> m [()]
forall w (m :: * -> *) (t :: * -> *) a b.
(MonadSystem w m, Foldable t) =>
t a -> ([a] -> ParSystem b) -> m [b]
parIterList t a
x ([a] -> (a -> ParSystem b) -> ParSystem ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
`for_` a -> ParSystem b
s))