-- |
-- Module      : Test.DejaFu.Schedule
-- Copyright   : (c) 2016--2018 Michael Walker
-- License     : MIT
-- Maintainer  : Michael Walker <mike@barrucadu.co.uk>
-- Stability   : experimental
-- Portability : portable
--
-- Scheduling for concurrent computations.
module Test.DejaFu.Schedule
  ( -- * Scheduling
    Scheduler(..)

  -- ** Preemptive
  , randomSched
  , roundRobinSched

  -- ** Non-preemptive
  , randomSchedNP
  , roundRobinSchedNP

  -- * Utilities
  , makeNonPreemptive
  ) where

import           Data.List.NonEmpty   (NonEmpty(..), toList)
import           System.Random        (RandomGen, randomR)

import           Test.DejaFu.Internal
import           Test.DejaFu.Types

-- | A @Scheduler@ drives the execution of a concurrent program. The
-- parameters it takes are:
--
-- 1. The last thread executed (if this is the first invocation, this
--    is @Nothing@).
--
-- 2. The unblocked threads.
--
-- 3. The concurrency state.
--
-- 4. The scheduler state.
--
-- It returns a thread to execute, or @Nothing@ if execution should
-- abort here, and also a new state.
--
-- @since 2.0.0.0
newtype Scheduler state = Scheduler
  { forall state.
Scheduler state
-> Maybe (ThreadId, ThreadAction)
-> NonEmpty (ThreadId, Lookahead)
-> ConcurrencyState
-> state
-> (Maybe ThreadId, state)
scheduleThread
    :: Maybe (ThreadId, ThreadAction)
    -> NonEmpty (ThreadId, Lookahead)
    -> ConcurrencyState
    -> state
    -> (Maybe ThreadId, state)
  }

-------------------------------------------------------------------------------
-- Preemptive

-- | A simple random scheduler which, at every step, picks a random
-- thread to run.
--
-- @since 0.8.0.0
randomSched :: RandomGen g => Scheduler g
randomSched :: forall g. RandomGen g => Scheduler g
randomSched = (Maybe (ThreadId, ThreadAction)
 -> NonEmpty (ThreadId, Lookahead)
 -> ConcurrencyState
 -> g
 -> (Maybe ThreadId, g))
-> Scheduler g
forall state.
(Maybe (ThreadId, ThreadAction)
 -> NonEmpty (ThreadId, Lookahead)
 -> ConcurrencyState
 -> state
 -> (Maybe ThreadId, state))
-> Scheduler state
Scheduler Maybe (ThreadId, ThreadAction)
-> NonEmpty (ThreadId, Lookahead)
-> ConcurrencyState
-> g
-> (Maybe ThreadId, g)
forall {b} {p} {a} {b} {p}.
RandomGen b =>
p -> NonEmpty (a, b) -> p -> b -> (Maybe a, b)
go where
  go :: p -> NonEmpty (a, b) -> p -> b -> (Maybe a, b)
go p
_ NonEmpty (a, b)
threads p
_ b
g =
    let threads' :: [a]
threads' = ((a, b) -> a) -> [(a, b)] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map (a, b) -> a
forall a b. (a, b) -> a
fst (NonEmpty (a, b) -> [(a, b)]
forall a. NonEmpty a -> [a]
toList NonEmpty (a, b)
threads)
        (Int
choice, b
g') = (Int, Int) -> b -> (Int, b)
forall g. RandomGen g => (Int, Int) -> g -> (Int, g)
forall a g. (Random a, RandomGen g) => (a, a) -> g -> (a, g)
randomR (Int
0, [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
threads' Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) b
g
    in (a -> Maybe a
forall a. a -> Maybe a
Just (a -> Maybe a) -> a -> Maybe a
forall a b. (a -> b) -> a -> b
$ [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
eidx [a]
threads' Int
choice, b
g')

-- | A round-robin scheduler which, at every step, schedules the
-- thread with the next 'ThreadId'.
--
-- @since 0.8.0.0
roundRobinSched :: Scheduler ()
roundRobinSched :: Scheduler ()
roundRobinSched = (Maybe (ThreadId, ThreadAction)
 -> NonEmpty (ThreadId, Lookahead)
 -> ConcurrencyState
 -> ()
 -> (Maybe ThreadId, ()))
-> Scheduler ()
forall state.
(Maybe (ThreadId, ThreadAction)
 -> NonEmpty (ThreadId, Lookahead)
 -> ConcurrencyState
 -> state
 -> (Maybe ThreadId, state))
-> Scheduler state
Scheduler Maybe (ThreadId, ThreadAction)
-> NonEmpty (ThreadId, Lookahead)
-> ConcurrencyState
-> ()
-> (Maybe ThreadId, ())
forall {a} {b} {b} {p} {p}.
Ord a =>
Maybe (a, b) -> NonEmpty (a, b) -> p -> p -> (Maybe a, ())
go where
  go :: Maybe (a, b) -> NonEmpty (a, b) -> p -> p -> (Maybe a, ())
go Maybe (a, b)
Nothing ((a
tid,b
_):|[(a, b)]
_) p
_ p
_ = (a -> Maybe a
forall a. a -> Maybe a
Just a
tid, ())
  go (Just (a
prior, b
_)) NonEmpty (a, b)
threads p
_ p
_ =
    let threads' :: [a]
threads' = ((a, b) -> a) -> [(a, b)] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map (a, b) -> a
forall a b. (a, b) -> a
fst (NonEmpty (a, b) -> [(a, b)]
forall a. NonEmpty a -> [a]
toList NonEmpty (a, b)
threads)
        candidates :: [a]
candidates =
          if a
prior a -> a -> Bool
forall a. Ord a => a -> a -> Bool
>= [a] -> a
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum [a]
threads'
          then [a]
threads'
          else (a -> Bool) -> [a] -> [a]
forall a. (a -> Bool) -> [a] -> [a]
filter (a -> a -> Bool
forall a. Ord a => a -> a -> Bool
>a
prior) [a]
threads'
    in (a -> Maybe a
forall a. a -> Maybe a
Just ([a] -> a
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
minimum [a]
candidates), ())

-------------------------------------------------------------------------------
-- Non-preemptive

-- | A random scheduler which doesn't preempt the running thread. That
-- is, if the previously scheduled thread is not blocked, it is picked
-- again, otherwise schedule randomly.
--
-- @since 0.8.0.0
randomSchedNP :: RandomGen g => Scheduler g
randomSchedNP :: forall g. RandomGen g => Scheduler g
randomSchedNP = Scheduler g -> Scheduler g
forall s. Scheduler s -> Scheduler s
makeNonPreemptive Scheduler g
forall g. RandomGen g => Scheduler g
randomSched

-- | A round-robin scheduler which doesn't preempt the running
-- thread. That is, if the previously scheduled thread is not blocked,
-- it is picked again, otherwise schedule the thread with the next
-- 'ThreadId'.
--
-- @since 0.8.0.0
roundRobinSchedNP :: Scheduler ()
roundRobinSchedNP :: Scheduler ()
roundRobinSchedNP = Scheduler () -> Scheduler ()
forall s. Scheduler s -> Scheduler s
makeNonPreemptive Scheduler ()
roundRobinSched

-------------------------------------------------------------------------------
-- Utilities

-- | Turn a potentially preemptive scheduler into a non-preemptive
-- one.
--
-- @since 0.8.0.0
makeNonPreemptive :: Scheduler s -> Scheduler s
makeNonPreemptive :: forall s. Scheduler s -> Scheduler s
makeNonPreemptive Scheduler s
sched = (Maybe (ThreadId, ThreadAction)
 -> NonEmpty (ThreadId, Lookahead)
 -> ConcurrencyState
 -> s
 -> (Maybe ThreadId, s))
-> Scheduler s
forall state.
(Maybe (ThreadId, ThreadAction)
 -> NonEmpty (ThreadId, Lookahead)
 -> ConcurrencyState
 -> state
 -> (Maybe ThreadId, state))
-> Scheduler state
Scheduler Maybe (ThreadId, ThreadAction)
-> NonEmpty (ThreadId, Lookahead)
-> ConcurrencyState
-> s
-> (Maybe ThreadId, s)
newsched where
  newsched :: Maybe (ThreadId, ThreadAction)
-> NonEmpty (ThreadId, Lookahead)
-> ConcurrencyState
-> s
-> (Maybe ThreadId, s)
newsched p :: Maybe (ThreadId, ThreadAction)
p@(Just (ThreadId
prior, ThreadAction
_)) NonEmpty (ThreadId, Lookahead)
threads ConcurrencyState
cs s
s
    | ThreadId
prior ThreadId -> [ThreadId] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` ((ThreadId, Lookahead) -> ThreadId)
-> [(ThreadId, Lookahead)] -> [ThreadId]
forall a b. (a -> b) -> [a] -> [b]
map (ThreadId, Lookahead) -> ThreadId
forall a b. (a, b) -> a
fst (NonEmpty (ThreadId, Lookahead) -> [(ThreadId, Lookahead)]
forall a. NonEmpty a -> [a]
toList NonEmpty (ThreadId, Lookahead)
threads) = (ThreadId -> Maybe ThreadId
forall a. a -> Maybe a
Just ThreadId
prior, s
s)
    | Bool
otherwise = Scheduler s
-> Maybe (ThreadId, ThreadAction)
-> NonEmpty (ThreadId, Lookahead)
-> ConcurrencyState
-> s
-> (Maybe ThreadId, s)
forall state.
Scheduler state
-> Maybe (ThreadId, ThreadAction)
-> NonEmpty (ThreadId, Lookahead)
-> ConcurrencyState
-> state
-> (Maybe ThreadId, state)
scheduleThread Scheduler s
sched Maybe (ThreadId, ThreadAction)
p NonEmpty (ThreadId, Lookahead)
threads ConcurrencyState
cs s
s
  newsched Maybe (ThreadId, ThreadAction)
Nothing NonEmpty (ThreadId, Lookahead)
threads ConcurrencyState
cs s
s = Scheduler s
-> Maybe (ThreadId, ThreadAction)
-> NonEmpty (ThreadId, Lookahead)
-> ConcurrencyState
-> s
-> (Maybe ThreadId, s)
forall state.
Scheduler state
-> Maybe (ThreadId, ThreadAction)
-> NonEmpty (ThreadId, Lookahead)
-> ConcurrencyState
-> state
-> (Maybe ThreadId, state)
scheduleThread Scheduler s
sched Maybe (ThreadId, ThreadAction)
forall a. Maybe a
Nothing NonEmpty (ThreadId, Lookahead)
threads ConcurrencyState
cs s
s