content: {-# LANGUAGE DeriveAnyClass #-}
  {-# LANGUAGE DeriveGeneric #-}
  
  -- |
  -- Module      : Test.DejaFu.SCT.Internal.Weighted
  -- Copyright   : (c) 2015--2019 Michael Walker
  -- License     : MIT
  -- Maintainer  : Michael Walker <mike@barrucadu.co.uk>
  -- Stability   : experimental
  -- Portability : DeriveAnyClass, DeriveGeneric
  --
  -- Internal types and functions for SCT via weighted random
  -- scheduling.  This module is NOT considered to form part of the
  -- public interface of this library.
  module Test.DejaFu.SCT.Internal.Weighted where
  
  import           Control.DeepSeq      (NFData)
  import           Data.List.NonEmpty   (toList)
  import           Data.Map.Strict      (Map)
  import qualified Data.Map.Strict      as M
  import           GHC.Generics         (Generic)
  import           System.Random        (RandomGen, randomR)
  
  import           Test.DejaFu.Schedule (Scheduler(..))
  import           Test.DejaFu.Types
  
  -------------------------------------------------------------------------------
  -- * Weighted random scheduler
  
  -- | The scheduler state
  data RandSchedState g = RandSchedState
    { schedWeights :: Map ThreadId Int
    -- ^ The thread weights: used in determining which to run.
    , schedLengthBound :: Maybe LengthBound
    -- ^ The optional length bound.
    , schedGen :: g
    -- ^ The random number generator.
    } deriving (Eq, Show, Generic, NFData)
  
  -- | Initial weighted random scheduler state.
  initialRandSchedState :: Maybe LengthBound -> g -> RandSchedState g
  initialRandSchedState = RandSchedState M.empty
  
  -- | Weighted random scheduler: assigns to each new thread a weight,
  -- and makes a weighted random choice out of the runnable threads at
  -- every step.
  randSched :: RandomGen g => (g -> (Int, g)) -> Scheduler (RandSchedState g)
  randSched weightf = Scheduler $ \_ threads _ s ->
    let
      -- Select a thread
      pick idx ((x, f):xs)
        | idx < f = Just x
        | otherwise = pick (idx - f) xs
      pick _ [] = Nothing
      (choice, g'') = randomR (0, sum (map snd enabled) - 1) g'
      enabled = M.toList $ M.filterWithKey (\tid _ -> tid `elem` tids) weights'
  
      -- The weights, with any new threads added.
      (weights', g') = foldr assignWeight (M.empty, schedGen s) tids
      assignWeight tid ~(ws, g0) =
        let (w, g) = maybe (weightf g0) (\w0 -> (w0, g0)) (M.lookup tid (schedWeights s))
        in (M.insert tid w ws, g)
  
      -- The runnable threads.
      tids = map fst (toList threads)
    in case schedLengthBound s of
      Just 0 -> (Nothing, s)
      Just n -> (pick choice enabled, RandSchedState weights' (Just (n - 1)) g'')
      Nothing -> (pick choice enabled, RandSchedState weights' Nothing g'')
  