{-# OPTIONS_HADDOCK hide, not-home #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE ScopedTypeVariables #-}
-- |
-- Module      : Control.Scheduler.Computation
-- Copyright   : (c) Alexey Kuleshevich 2018-2019
-- License     : BSD3
-- Maintainer  : Alexey Kuleshevich <lehins@yandex.ru>
-- Stability   : experimental
-- Portability : non-portable
--
module Control.Scheduler.Computation
  ( Comp(.., Par, Par'), getCompWorkers
  ) where

import Control.Concurrent (getNumCapabilities)
import Control.DeepSeq (NFData(..), deepseq)
import Control.Monad.IO.Class
#if !MIN_VERSION_base(4,11,0)
import Data.Semigroup
#endif
import Data.IORef
import Data.Word
import System.IO.Unsafe (unsafePerformIO)

-- | Computation strategy to use when scheduling work.
data Comp
  = Seq -- ^ Sequential computation
  | ParOn ![Int]
  -- ^ Schedule workers to run on specific capabilities. Specifying an empty list @`ParOn` []@ or
  -- using `Par` will result in utilization of all available capabilities.
  | ParN {-# UNPACK #-} !Word16
  -- ^ Specify the number of workers that will be handling all the jobs. Difference from `ParOn` is
  -- that workers can jump between cores. Using @`ParN` 0@ will result in using all available
  -- capabilities.
  deriving Comp -> Comp -> Bool
(Comp -> Comp -> Bool) -> (Comp -> Comp -> Bool) -> Eq Comp
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Comp -> Comp -> Bool
$c/= :: Comp -> Comp -> Bool
== :: Comp -> Comp -> Bool
$c== :: Comp -> Comp -> Bool
Eq

-- | Parallel computation using all available cores. Same as @`ParOn` []@
--
-- @since 1.0.0
pattern Par :: Comp
pattern $bPar :: Comp
$mPar :: forall r. Comp -> (Void# -> r) -> (Void# -> r) -> r
Par <- ParOn [] where
        Par =  [Int] -> Comp
ParOn []

-- | Parallel computation using all available cores. Same as @`ParN` 0@
--
-- @since 1.1.0
pattern Par' :: Comp
pattern $bPar' :: Comp
$mPar' :: forall r. Comp -> (Void# -> r) -> (Void# -> r) -> r
Par' <- ParN 0 where
        Par' =  Word16 -> Comp
ParN Word16
0

instance Show Comp where
  show :: Comp -> String
show Comp
Seq        = String
"Seq"
  show Comp
Par        = String
"Par"
  show (ParOn [Int]
ws) = String
"ParOn " String -> ShowS
forall a. [a] -> [a] -> [a]
++ [Int] -> String
forall a. Show a => a -> String
show [Int]
ws
  show (ParN Word16
n)   = String
"ParN " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Word16 -> String
forall a. Show a => a -> String
show Word16
n
  showsPrec :: Int -> Comp -> ShowS
showsPrec Int
_ Comp
Seq  = (String
"Seq" String -> ShowS
forall a. [a] -> [a] -> [a]
++)
  showsPrec Int
_ Comp
Par  = (String
"Par" String -> ShowS
forall a. [a] -> [a] -> [a]
++)
  showsPrec Int
0 Comp
comp = (Comp -> String
forall a. Show a => a -> String
show Comp
comp String -> ShowS
forall a. [a] -> [a] -> [a]
++)
  showsPrec Int
_ Comp
comp = ((String
"(" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Comp -> String
forall a. Show a => a -> String
show Comp
comp String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
")") String -> ShowS
forall a. [a] -> [a] -> [a]
++)

instance NFData Comp where
  rnf :: Comp -> ()
rnf Comp
comp =
    case Comp
comp of
      Comp
Seq        -> ()
      ParOn [Int]
wIds -> [Int]
wIds [Int] -> () -> ()
forall a b. NFData a => a -> b -> b
`deepseq` ()
      ParN Word16
n     -> Word16
n Word16 -> () -> ()
forall a b. NFData a => a -> b -> b
`deepseq` ()
  {-# INLINE rnf #-}

instance Monoid Comp where
  mempty :: Comp
mempty = Comp
Seq
  {-# INLINE mempty #-}
  mappend :: Comp -> Comp -> Comp
mappend = Comp -> Comp -> Comp
joinComp
  {-# INLINE mappend #-}

instance Semigroup Comp where
  <> :: Comp -> Comp -> Comp
(<>) = Comp -> Comp -> Comp
joinComp
  {-# INLINE (<>) #-}

joinComp :: Comp -> Comp -> Comp
joinComp :: Comp -> Comp -> Comp
joinComp Comp
x Comp
y =
  case Comp
x of
    Comp
Seq -> Comp
y
    Comp
Par -> Comp
Par
    Comp
Par' -> Comp
Par'
    ParOn [Int]
xs ->
      case Comp
y of
        Comp
Par      -> Comp
Par
        Comp
Par'     -> Comp
Par'
        ParOn [Int]
ys -> [Int] -> Comp
ParOn ([Int]
xs [Int] -> [Int] -> [Int]
forall a. Semigroup a => a -> a -> a
<> [Int]
ys)
        Comp
_        -> Comp
x
    ParN Word16
n1 ->
      case Comp
y of
        Comp
Seq     -> Comp
x
        Comp
Par     -> Comp
Par
        ParOn [Int]
_ -> Comp
y
        Comp
Par'    -> Comp
y
        ParN Word16
n2 -> Word16 -> Comp
ParN (Word16 -> Word16 -> Word16
forall a. Ord a => a -> a -> a
max Word16
n1 Word16
n2)
{-# NOINLINE joinComp #-}

numCapsRef :: IORef Int
numCapsRef :: IORef Int
numCapsRef = IO (IORef Int) -> IORef Int
forall a. IO a -> a
unsafePerformIO (IO (IORef Int) -> IORef Int) -> IO (IORef Int) -> IORef Int
forall a b. (a -> b) -> a -> b
$ do
  Int
caps <- IO Int
getNumCapabilities
  Int -> IO (IORef Int)
forall a. a -> IO (IORef a)
newIORef Int
caps
{-# NOINLINE numCapsRef #-}

-- | Figure out how many workers will this computation strategy create.
--
-- /Note/ - If at any point during program execution global number of capabilities gets
-- changed with `Control.Concurrent.setNumCapabilities`, it will have no affect on this
-- function, unless it hasn't yet been called with `Par` or `Par'` arguments.
--
-- @since 1.1.0
getCompWorkers :: MonadIO m => Comp -> m Int
getCompWorkers :: Comp -> m Int
getCompWorkers =
  \case
    Comp
Seq -> Int -> m Int
forall (m :: * -> *) a. Monad m => a -> m a
return Int
1
    Comp
Par -> IO Int -> m Int
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IORef Int -> IO Int
forall a. IORef a -> IO a
readIORef IORef Int
numCapsRef)
    ParOn [Int]
ws -> Int -> m Int
forall (m :: * -> *) a. Monad m => a -> m a
return (Int -> m Int) -> Int -> m Int
forall a b. (a -> b) -> a -> b
$ [Int] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
ws
    Comp
Par' -> IO Int -> m Int
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IORef Int -> IO Int
forall a. IORef a -> IO a
readIORef IORef Int
numCapsRef)
    ParN Word16
n -> Int -> m Int
forall (m :: * -> *) a. Monad m => a -> m a
return (Int -> m Int) -> Int -> m Int
forall a b. (a -> b) -> a -> b
$ Word16 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word16
n