{-# LANGUAGE DefaultSignatures #-}

module Multitasking.Core
  ( -- ** Start tasks
    multitask,
    Coordinator,
    start,
    Task (..),
    awaitTask,
    awaitAll,

    -- ** Control thread options
    startWith,
    Ki.ThreadOptions (..),
    Ki.defaultThreadOptions,
    Ki.ThreadAffinity (..),
    Ki.ByteCount,
    Ki.kilobytes,
    Ki.megabytes,
  )
where

import Control.Monad.IO.Class
import Ki qualified
import Multitasking.AsyncOperations
import Multitasking.MonadSTM

-- | Coordinator corresponds to the current multitasking scope.
newtype Coordinator = Coordinator Ki.Scope

-- | Opens up a multitasking scope. No threads launched with the provided 'Coordinator' outlive this scope.
-- Before 'multitask' ends, it will __cancel all threads__. Use 'awaitAll' if you want to wait beforehand.
-- Additionally, exceptions between parent and children are propagated per default,
-- completely shutting down all processes when an exception happens anywhere.
multitask :: (MonadIO m) => (Coordinator -> IO a) -> m a
multitask :: forall (m :: * -> *) a. MonadIO m => (Coordinator -> IO a) -> m a
multitask Coordinator -> IO a
f = IO a -> m a
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO a -> m a) -> IO a -> m a
forall a b. (a -> b) -> a -> b
$ (Scope -> IO a) -> IO a
forall a. (Scope -> IO a) -> IO a
Ki.scoped ((Scope -> IO a) -> IO a) -> (Scope -> IO a) -> IO a
forall a b. (a -> b) -> a -> b
$ Coordinator -> IO a
f (Coordinator -> IO a) -> (Scope -> Coordinator) -> Scope -> IO a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Scope -> Coordinator
Coordinator

-- | A 'Task' is a computation on another thread. Use `await` to wait for the task.
newtype Task a = Task (Ki.Thread a)

-- | 'start' is the main way to spin up new tasks.
-- It will execute the given action in another thread and returns a 'Task'.
start :: (MonadIO m) => Coordinator -> IO a -> m (Task a)
start :: forall (m :: * -> *) a.
MonadIO m =>
Coordinator -> IO a -> m (Task a)
start (Coordinator Scope
scope) IO a
action = do
  thread <- IO (Thread a) -> m (Thread a)
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Thread a) -> m (Thread a)) -> IO (Thread a) -> m (Thread a)
forall a b. (a -> b) -> a -> b
$ Scope -> IO a -> IO (Thread a)
forall a. Scope -> IO a -> IO (Thread a)
Ki.fork Scope
scope IO a
action
  pure $ Task thread

-- | Provice
startWith :: (MonadIO m) => Coordinator -> Ki.ThreadOptions -> IO a -> m (Task a)
startWith :: forall (m :: * -> *) a.
MonadIO m =>
Coordinator -> ThreadOptions -> IO a -> m (Task a)
startWith (Coordinator Scope
scope) ThreadOptions
to IO a
action = do
  thread <- IO (Thread a) -> m (Thread a)
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Thread a) -> m (Thread a)) -> IO (Thread a) -> m (Thread a)
forall a b. (a -> b) -> a -> b
$ Scope -> ThreadOptions -> IO a -> IO (Thread a)
forall a. Scope -> ThreadOptions -> IO a -> IO (Thread a)
Ki.forkWith Scope
scope ThreadOptions
to IO a
action
  pure $ Task thread

instance Functor Task where
  fmap :: forall a b. (a -> b) -> Task a -> Task b
fmap a -> b
f (Task Thread a
t) = Thread b -> Task b
forall a. Thread a -> Task a
Task ((a -> b) -> Thread a -> Thread b
forall a b. (a -> b) -> Thread a -> Thread b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f Thread a
t)

awaitTask :: (MonadSTM m) => Task a -> m a
awaitTask :: forall (m :: * -> *) a. MonadSTM m => Task a -> m a
awaitTask (Task Thread a
thread) = STM a -> m a
forall a. STM a -> m a
forall (m :: * -> *) a. MonadSTM m => STM a -> m a
liftSTM (STM a -> m a) -> STM a -> m a
forall a b. (a -> b) -> a -> b
$ Thread a -> STM a
forall a. Thread a -> STM a
Ki.await Thread a
thread

awaitAll :: (MonadSTM m) => Coordinator -> m ()
awaitAll :: forall (m :: * -> *). MonadSTM m => Coordinator -> m ()
awaitAll (Coordinator Scope
scope) = STM () -> m ()
forall a. STM a -> m a
forall (m :: * -> *) a. MonadSTM m => STM a -> m a
liftSTM (Scope -> STM ()
Ki.awaitAll Scope
scope)

instance Await (Task a) where
  type Payload (Task a) = a
  await :: forall (m :: * -> *). MonadSTM m => Task a -> m (Payload (Task a))
await = Task a -> m a
Task a -> m (Payload (Task a))
forall (m :: * -> *) a. MonadSTM m => Task a -> m a
awaitTask

instance Await Coordinator where
  type Payload Coordinator = ()
  await :: forall (m :: * -> *).
MonadSTM m =>
Coordinator -> m (Payload Coordinator)
await = Coordinator -> m ()
Coordinator -> m (Payload Coordinator)
forall (m :: * -> *). MonadSTM m => Coordinator -> m ()
awaitAll