----------------------------------------------------------------------------
-- |
-- Module      :  TestUtils
-- Copyright   :  (c) Sergey Vinokurov 2022
-- License     :  Apache-2.0 (see LICENSE)
-- Maintainer  :  serg.foo@gmail.com
----------------------------------------------------------------------------

{-# LANGUAGE BangPatterns   #-}
{-# LANGUAGE DeriveGeneric  #-}
{-# LANGUAGE NamedFieldPuns #-}

module TestUtils
  ( Delay(..)
  , sleep
  , Iterations(..)
  , callN
  , Thread(..)
  , runThread
  , Threads(..)
  , spawnAndCall
  ) where

import Control.Concurrent
import Control.Concurrent.Async
import Control.Monad
import Control.Monad.IO.Class
import Data.Foldable
import Data.List.NonEmpty (NonEmpty(..))
import GHC.Generics (Generic)
import Test.QuickCheck

-- In microseconds
newtype Delay = Delay { Delay -> Int
unDelay :: Int }
  deriving (Delay -> Delay -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Delay -> Delay -> Bool
$c/= :: Delay -> Delay -> Bool
== :: Delay -> Delay -> Bool
$c== :: Delay -> Delay -> Bool
Eq, Int -> Delay -> ShowS
[Delay] -> ShowS
Delay -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Delay] -> ShowS
$cshowList :: [Delay] -> ShowS
show :: Delay -> String
$cshow :: Delay -> String
showsPrec :: Int -> Delay -> ShowS
$cshowsPrec :: Int -> Delay -> ShowS
Show)

sleep :: MonadIO m => Delay -> m ()
sleep :: forall (m :: * -> *). MonadIO m => Delay -> m ()
sleep (Delay Int
n) = case Int
n of
  Int
0 -> forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
  Int
k -> forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ Int -> IO ()
threadDelay Int
k

instance Arbitrary Delay where
  arbitrary :: Gen Delay
arbitrary = Int -> Delay
Delay forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Int, Int) -> Gen Int
chooseInt (Int
0, Int
10)
  shrink :: Delay -> [Delay]
shrink = forall a b. (a -> b) -> [a] -> [b]
map Int -> Delay
Delay forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. (a -> Bool) -> [a] -> [a]
filter (\Int
x -> Int
0 forall a. Ord a => a -> a -> Bool
<= Int
x Bool -> Bool -> Bool
&& Int
x forall a. Ord a => a -> a -> Bool
<= Int
25) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Arbitrary a => a -> [a]
shrink forall b c a. (b -> c) -> (a -> b) -> a -> c
. Delay -> Int
unDelay

newtype Iterations = Iterations { Iterations -> Int
unIterations :: Int }
  deriving (Iterations -> Iterations -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Iterations -> Iterations -> Bool
$c/= :: Iterations -> Iterations -> Bool
== :: Iterations -> Iterations -> Bool
$c== :: Iterations -> Iterations -> Bool
Eq, Int -> Iterations -> ShowS
[Iterations] -> ShowS
Iterations -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Iterations] -> ShowS
$cshowList :: [Iterations] -> ShowS
show :: Iterations -> String
$cshow :: Iterations -> String
showsPrec :: Int -> Iterations -> ShowS
$cshowsPrec :: Int -> Iterations -> ShowS
Show)

instance Arbitrary Iterations where
  arbitrary :: Gen Iterations
arbitrary = Int -> Iterations
Iterations forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Int, Int) -> Gen Int
chooseInt (Int
0, Int
50)
  shrink :: Iterations -> [Iterations]
shrink = forall a b. (a -> b) -> [a] -> [b]
map Int -> Iterations
Iterations forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. (a -> Bool) -> [a] -> [a]
filter (\Int
x -> Int
0 forall a. Ord a => a -> a -> Bool
<= Int
x Bool -> Bool -> Bool
&& Int
x forall a. Ord a => a -> a -> Bool
<= Int
50) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Arbitrary a => a -> [a]
shrink forall b c a. (b -> c) -> (a -> b) -> a -> c
. Iterations -> Int
unIterations

callN :: Applicative m => Iterations -> m a -> m ()
callN :: forall (m :: * -> *) a. Applicative m => Iterations -> m a -> m ()
callN (Iterations !Int
n) m a
action = forall {t}. (Ord t, Num t) => t -> m ()
go Int
n
  where
    go :: t -> m ()
go !t
k =
      if t
k forall a. Ord a => a -> a -> Bool
> t
0
      then m a
action forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> t -> m ()
go (t
k forall a. Num a => a -> a -> a
- t
1)
      else forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

data Thread = Thread
  { Thread -> Delay
tDelay      :: Delay
  , Thread -> Int
tIncrement  :: Int
  , Thread -> Iterations
tIterations :: Iterations
  } deriving (Thread -> Thread -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Thread -> Thread -> Bool
$c/= :: Thread -> Thread -> Bool
== :: Thread -> Thread -> Bool
$c== :: Thread -> Thread -> Bool
Eq, Int -> Thread -> ShowS
[Thread] -> ShowS
Thread -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Thread] -> ShowS
$cshowList :: [Thread] -> ShowS
show :: Thread -> String
$cshow :: Thread -> String
showsPrec :: Int -> Thread -> ShowS
$cshowsPrec :: Int -> Thread -> ShowS
Show, forall x. Rep Thread x -> Thread
forall x. Thread -> Rep Thread x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cto :: forall x. Rep Thread x -> Thread
$cfrom :: forall x. Thread -> Rep Thread x
Generic)

instance Arbitrary Thread where
  arbitrary :: Gen Thread
arbitrary = Delay -> Int -> Iterations -> Thread
Thread forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. Arbitrary a => Gen a
arbitrary forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (Int, Int) -> Gen Int
chooseInt (-Int
1000, Int
1000) forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall a. Arbitrary a => Gen a
arbitrary
  shrink :: Thread -> [Thread]
shrink = forall a. (a -> Bool) -> [a] -> [a]
filter ((forall a. Ord a => a -> a -> Bool
<= Int
1000) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Num a => a -> a
abs forall b c a. (b -> c) -> (a -> b) -> a -> c
. Thread -> Int
tIncrement) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a.
(Generic a, RecursivelyShrink (Rep a), GSubterms (Rep a) a) =>
a -> [a]
genericShrink

runThread :: MonadIO m => Thread -> (Delay -> m a) -> (Int -> m b) -> m ()
runThread :: forall (m :: * -> *) a b.
MonadIO m =>
Thread -> (Delay -> m a) -> (Int -> m b) -> m ()
runThread Thread{Delay
tDelay :: Delay
tDelay :: Thread -> Delay
tDelay, Int
tIncrement :: Int
tIncrement :: Thread -> Int
tIncrement, Iterations
tIterations :: Iterations
tIterations :: Thread -> Iterations
tIterations} Delay -> m a
doSleep Int -> m b
f =
  forall (m :: * -> *) a. Applicative m => Iterations -> m a -> m ()
callN Iterations
tIterations (Int -> m b
f Int
tIncrement forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> Delay -> m a
doSleep Delay
tDelay)

newtype Threads = Threads { Threads -> NonEmpty Thread
unThreads :: NonEmpty Thread }
  deriving (Threads -> Threads -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Threads -> Threads -> Bool
$c/= :: Threads -> Threads -> Bool
== :: Threads -> Threads -> Bool
$c== :: Threads -> Threads -> Bool
Eq, Int -> Threads -> ShowS
[Threads] -> ShowS
Threads -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Threads] -> ShowS
$cshowList :: [Threads] -> ShowS
show :: Threads -> String
$cshow :: Threads -> String
showsPrec :: Int -> Threads -> ShowS
$cshowsPrec :: Int -> Threads -> ShowS
Show)

instance Arbitrary Threads where
  arbitrary :: Gen Threads
arbitrary = do
    Int
n <- (Int, Int) -> Gen Int
chooseInt (Int
0, Int
31)
    NonEmpty Thread -> Threads
Threads forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (forall a. a -> [a] -> NonEmpty a
(:|) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. Arbitrary a => Gen a
arbitrary forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
n forall a. Arbitrary a => Gen a
arbitrary)
  shrink :: Threads -> [Threads]
shrink = forall a b. (a -> b) -> [a] -> [b]
map NonEmpty Thread -> Threads
Threads forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a.
(Generic a, RecursivelyShrink (Rep a), GSubterms (Rep a) a) =>
a -> [a]
genericShrink forall b c a. (b -> c) -> (a -> b) -> a -> c
. Threads -> NonEmpty Thread
unThreads

spawnAndCall :: Traversable f => f b -> IO a -> (a -> b -> IO ()) -> IO a
spawnAndCall :: forall (f :: * -> *) b a.
Traversable f =>
f b -> IO a -> (a -> b -> IO ()) -> IO a
spawnAndCall f b
threads IO a
mkRes a -> b -> IO ()
action = do
  a
res <- IO a
mkRes
  forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
(a -> f b) -> t a -> f ()
traverse_ forall a. Async a -> IO a
wait forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (forall a. IO a -> IO (Async a)
async forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> b -> IO ()
action a
res) f b
threads
  forall (f :: * -> *) a. Applicative f => a -> f a
pure a
res