{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE RankNTypes #-}

-- | Utilities to schedule actions per 'Tick'.
module Test.Util.Schedule (
    Schedule (..)
  , genSchedule
  , joinSchedule
  , lastTick
  , shrinkSchedule
  ) where

import           Data.List (intercalate, unfoldr)
import           Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map
import           Data.Maybe (fromMaybe)
import           Ouroboros.Consensus.Util.Condense (Condense (..))
import           Test.QuickCheck
import           Test.Util.LogicalClock (Tick (..))

-- | A schedule plans actions on certain times.
--
-- TODO Note that a schedule can't express delays between the actions
-- within a single tick. Generating such delays may expose more (most
-- likely concurrency-related) bugs.
newtype Schedule a = Schedule { forall a. Schedule a -> Map Tick [a]
getSchedule :: Map Tick [a] }
  deriving stock (Int -> Schedule a -> ShowS
[Schedule a] -> ShowS
Schedule a -> String
(Int -> Schedule a -> ShowS)
-> (Schedule a -> String)
-> ([Schedule a] -> ShowS)
-> Show (Schedule a)
forall a. Show a => Int -> Schedule a -> ShowS
forall a. Show a => [Schedule a] -> ShowS
forall a. Show a => Schedule a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall a. Show a => Int -> Schedule a -> ShowS
showsPrec :: Int -> Schedule a -> ShowS
$cshow :: forall a. Show a => Schedule a -> String
show :: Schedule a -> String
$cshowList :: forall a. Show a => [Schedule a] -> ShowS
showList :: [Schedule a] -> ShowS
Show, Schedule a -> Schedule a -> Bool
(Schedule a -> Schedule a -> Bool)
-> (Schedule a -> Schedule a -> Bool) -> Eq (Schedule a)
forall a. Eq a => Schedule a -> Schedule a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: forall a. Eq a => Schedule a -> Schedule a -> Bool
== :: Schedule a -> Schedule a -> Bool
$c/= :: forall a. Eq a => Schedule a -> Schedule a -> Bool
/= :: Schedule a -> Schedule a -> Bool
Eq)

instance Condense a => Condense (Schedule a) where
  condense :: Schedule a -> String
condense =
        [String] -> String
unlines
      ([String] -> String)
-> (Schedule a -> [String]) -> Schedule a -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Tick, [a]) -> String) -> [(Tick, [a])] -> [String]
forall a b. (a -> b) -> [a] -> [b]
map ((Tick -> [a] -> String) -> (Tick, [a]) -> String
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry Tick -> [a] -> String
forall {a}. Condense a => Tick -> [a] -> String
showEntry)
      ([(Tick, [a])] -> [String])
-> (Schedule a -> [(Tick, [a])]) -> Schedule a -> [String]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Tick, [a]) -> Bool) -> [(Tick, [a])] -> [(Tick, [a])]
forall a. (a -> Bool) -> [a] -> [a]
filter (Bool -> Bool
not (Bool -> Bool) -> ((Tick, [a]) -> Bool) -> (Tick, [a]) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [a] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null ([a] -> Bool) -> ((Tick, [a]) -> [a]) -> (Tick, [a]) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Tick, [a]) -> [a]
forall a b. (a, b) -> b
snd)
      ([(Tick, [a])] -> [(Tick, [a])])
-> (Schedule a -> [(Tick, [a])]) -> Schedule a -> [(Tick, [a])]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map Tick [a] -> [(Tick, [a])]
forall k a. Map k a -> [(k, a)]
Map.toAscList
      (Map Tick [a] -> [(Tick, [a])])
-> (Schedule a -> Map Tick [a]) -> Schedule a -> [(Tick, [a])]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Schedule a -> Map Tick [a]
forall a. Schedule a -> Map Tick [a]
getSchedule
    where
      showEntry :: Tick -> [a] -> String
showEntry (Tick Word64
tick) [a]
as = Word64 -> String
forall a. Show a => a -> String
show Word64
tick String -> ShowS
forall a. Semigroup a => a -> a -> a
<> String
": " String -> ShowS
forall a. Semigroup a => a -> a -> a
<>
        String -> [String] -> String
forall a. [a] -> [[a]] -> [a]
intercalate String
", " ((a -> String) -> [a] -> [String]
forall a b. (a -> b) -> [a] -> [b]
map a -> String
forall a. Condense a => a -> String
condense [a]
as)

-- | Return the last tick at which an update is planned, if no updates
-- are planned, return 0.
lastTick :: Schedule a -> Tick
lastTick :: forall a. Schedule a -> Tick
lastTick = Tick -> Maybe Tick -> Tick
forall a. a -> Maybe a -> a
fromMaybe (Word64 -> Tick
Tick Word64
0) (Maybe Tick -> Tick)
-> (Schedule a -> Maybe Tick) -> Schedule a -> Tick
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map Tick [a] -> Maybe Tick
forall k v. Map k v -> Maybe k
maxKey (Map Tick [a] -> Maybe Tick)
-> (Schedule a -> Map Tick [a]) -> Schedule a -> Maybe Tick
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Schedule a -> Map Tick [a]
forall a. Schedule a -> Map Tick [a]
getSchedule
  where
    maxKey :: forall k v. Map k v -> Maybe k
    maxKey :: forall k v. Map k v -> Maybe k
maxKey = (((k, v), Map k v) -> k) -> Maybe ((k, v), Map k v) -> Maybe k
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((k, v) -> k
forall a b. (a, b) -> a
fst ((k, v) -> k)
-> (((k, v), Map k v) -> (k, v)) -> ((k, v), Map k v) -> k
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((k, v), Map k v) -> (k, v)
forall a b. (a, b) -> a
fst) (Maybe ((k, v), Map k v) -> Maybe k)
-> (Map k v -> Maybe ((k, v), Map k v)) -> Map k v -> Maybe k
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map k v -> Maybe ((k, v), Map k v)
forall k a. Map k a -> Maybe ((k, a), Map k a)
Map.maxViewWithKey

-- | Spread out elements over a schedule, i.e. schedule a number of
-- elements to be processed on each tick. Most ticks will have no
-- associated elements.
genSchedule :: [a] -> Gen (Schedule a)
genSchedule :: forall a. [a] -> Gen (Schedule a)
genSchedule = (Map Tick [a] -> Schedule a)
-> Gen (Map Tick [a]) -> Gen (Schedule a)
forall a b. (a -> b) -> Gen a -> Gen b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Map Tick [a] -> Schedule a
forall a. Map Tick [a] -> Schedule a
Schedule (Gen (Map Tick [a]) -> Gen (Schedule a))
-> ([a] -> Gen (Map Tick [a])) -> [a] -> Gen (Schedule a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map Tick [a] -> Tick -> [a] -> Gen (Map Tick [a])
forall a. Map Tick [a] -> Tick -> [a] -> Gen (Map Tick [a])
go Map Tick [a]
forall k a. Map k a
Map.empty Tick
1
  where
    go :: Map Tick [a]
       -> Tick
       -> [a]
       -> Gen (Map Tick [a])
    go :: forall a. Map Tick [a] -> Tick -> [a] -> Gen (Map Tick [a])
go !Map Tick [a]
schedule Tick
tick [a]
as
      | [a] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [a]
as = Map Tick [a] -> Gen (Map Tick [a])
forall a. a -> Gen a
forall (m :: * -> *) a. Monad m => a -> m a
return Map Tick [a]
schedule
      | Bool
otherwise    = do
        Int
nbAs <- [(Int, Gen Int)] -> Gen Int
forall a. HasCallStack => [(Int, Gen a)] -> Gen a
frequency [ (Int
2, Int -> Gen Int
forall a. a -> Gen a
forall (m :: * -> *) a. Monad m => a -> m a
return Int
0), (Int
1, (Int, Int) -> Gen Int
forall a. Random a => (a, a) -> Gen a
choose (Int
1, Int
5)) ]
        let ([a]
this, [a]
rest) = Int -> [a] -> ([a], [a])
forall a. Int -> [a] -> ([a], [a])
splitAt Int
nbAs [a]
as
        Map Tick [a] -> Tick -> [a] -> Gen (Map Tick [a])
forall a. Map Tick [a] -> Tick -> [a] -> Gen (Map Tick [a])
go (Tick -> [a] -> Map Tick [a] -> Map Tick [a]
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert Tick
tick [a]
this Map Tick [a]
schedule) (Tick -> Tick
forall a. Enum a => a -> a
succ Tick
tick) [a]
rest

-- | Repeatedly remove the last entry (highest 'Tick')
shrinkSchedule :: Schedule a -> [Schedule a]
shrinkSchedule :: forall a. Schedule a -> [Schedule a]
shrinkSchedule =
      (Map Tick [a] -> Maybe (Schedule a, Map Tick [a]))
-> Map Tick [a] -> [Schedule a]
forall b a. (b -> Maybe (a, b)) -> b -> [a]
unfoldr ((([a], Map Tick [a]) -> (Schedule a, Map Tick [a]))
-> Maybe ([a], Map Tick [a]) -> Maybe (Schedule a, Map Tick [a])
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\([a]
_, Map Tick [a]
m) -> (Map Tick [a] -> Schedule a
forall a. Map Tick [a] -> Schedule a
Schedule Map Tick [a]
m, Map Tick [a]
m)) (Maybe ([a], Map Tick [a]) -> Maybe (Schedule a, Map Tick [a]))
-> (Map Tick [a] -> Maybe ([a], Map Tick [a]))
-> Map Tick [a]
-> Maybe (Schedule a, Map Tick [a])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map Tick [a] -> Maybe ([a], Map Tick [a])
forall k a. Map k a -> Maybe (a, Map k a)
Map.maxView)
    (Map Tick [a] -> [Schedule a])
-> (Schedule a -> Map Tick [a]) -> Schedule a -> [Schedule a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Schedule a -> Map Tick [a]
forall a. Schedule a -> Map Tick [a]
getSchedule

-- | Inverse of 'genSchedule'
joinSchedule :: Schedule a -> [a]
joinSchedule :: forall a. Schedule a -> [a]
joinSchedule = ((Tick, [a]) -> [a]) -> [(Tick, [a])] -> [a]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (Tick, [a]) -> [a]
forall a b. (a, b) -> b
snd ([(Tick, [a])] -> [a])
-> (Schedule a -> [(Tick, [a])]) -> Schedule a -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map Tick [a] -> [(Tick, [a])]
forall k a. Map k a -> [(k, a)]
Map.toAscList (Map Tick [a] -> [(Tick, [a])])
-> (Schedule a -> Map Tick [a]) -> Schedule a -> [(Tick, [a])]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Schedule a -> Map Tick [a]
forall a. Schedule a -> Map Tick [a]
getSchedule