{-# LANGUAGE ScopedTypeVariables #-}

-- |
-- Module      :  ELynx.Tools.Concurrent
-- Description :  Tools for concurrent random calculations
-- Copyright   :  (c) Dominik Schrempf 2021
-- License     :  GPL-3.0-or-later
--
-- Maintainer  :  dominik.schrempf@gmail.com
-- Stability   :  unstable
-- Portability :  portable
--
-- Creation date: Tue May  7 10:33:24 2019.
module ELynx.Tools.Concurrent
  ( -- * MWC
    splitGen,

    -- * Concurrent calculations
    getChunks,
  )
where

import Control.Monad
import Control.Monad.Primitive
import qualified Data.Vector as V
import Data.Word
import System.Random.MWC

-- | Split a generator.
splitGen :: PrimMonad m => Int -> Gen (PrimState m) -> m [Gen (PrimState m)]
splitGen :: Int -> Gen (PrimState m) -> m [Gen (PrimState m)]
splitGen Int
n Gen (PrimState m)
gen
  | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0 = [Gen (PrimState m)] -> m [Gen (PrimState m)]
forall (m :: * -> *) a. Monad m => a -> m a
return []
  | Bool
otherwise = do
    [Vector Word32]
seeds :: [V.Vector Word32] <- Int -> m (Vector Word32) -> m [Vector Word32]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) (m (Vector Word32) -> m [Vector Word32])
-> m (Vector Word32) -> m [Vector Word32]
forall a b. (a -> b) -> a -> b
$ Gen (PrimState m) -> Int -> m (Vector Word32)
forall (m :: * -> *) g a (v :: * -> *).
(PrimMonad m, StatefulGen g m, Uniform a, Vector v a) =>
g -> Int -> m (v a)
uniformVector Gen (PrimState m)
gen Int
256
    ([Gen (PrimState m)] -> [Gen (PrimState m)])
-> m [Gen (PrimState m)] -> m [Gen (PrimState m)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Gen (PrimState m)
gen Gen (PrimState m) -> [Gen (PrimState m)] -> [Gen (PrimState m)]
forall a. a -> [a] -> [a]
:) ((Vector Word32 -> m (Gen (PrimState m)))
-> [Vector Word32] -> m [Gen (PrimState m)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Vector Word32 -> m (Gen (PrimState m))
forall (m :: * -> *) (v :: * -> *).
(PrimMonad m, Vector v Word32) =>
v Word32 -> m (Gen (PrimState m))
initialize [Vector Word32]
seeds)

-- | For a given number of capabilities and number of calculations, get chunk
-- sizes. The chunk sizes will be as evenly distributed as possible and sum up
-- to the number of calculations.
getChunks :: Int -> Int -> [Int]
getChunks :: Int -> Int -> [Int]
getChunks Int
c Int
n = [Int]
ns
  where
    n' :: Int
n' = Int
n Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
c
    r :: Int
r = Int
n Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
c
    ns :: [Int]
ns = Int -> Int -> [Int]
forall a. Int -> a -> [a]
replicate Int
r (Int
n' Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) [Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
++ Int -> Int -> [Int]
forall a. Int -> a -> [a]
replicate (Int
c Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
r) Int
n'