Stateful.hs

{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}
module Spec.Stateful where

import Control.Concurrent.STM
import Control.Monad
import Control.Monad.ST
import Data.Proxy
import Data.Typeable
import System.Random (uniformShortByteString)
import System.Random.Stateful hiding (uniformShortByteString)
import Test.SmallCheck.Series
import Test.Tasty
import Test.Tasty.SmallCheck as SC

instance Monad m => Serial m StdGen where
  series = mkStdGen <$> series

instance (Monad m, Serial m g) => Serial m (AtomicGen g) where
  series = AtomicGen <$> series

instance (Monad m, Serial m g) => Serial m (IOGen g) where
  series = IOGen <$> series

instance (Monad m, Serial m g) => Serial m (STGen g) where
  series = STGen <$> series

instance (Monad m, Serial m g) => Serial m (TGen g) where
  series = TGen <$> series

instance (Monad m, Serial m g) => Serial m (StateGen g) where
  series = StateGen <$> series


matchRandomGenSpec ::
     forall f a sg m. (StatefulGen sg m, RandomGen f, Eq f, Show f, Eq a)
  => (forall g n. StatefulGen g n => g -> n a)
  -> (forall g. RandomGen g => g -> (a, g))
  -> (StdGen -> f)
  -> (f -> StdGen)
  -> (f -> (sg -> m a) -> IO (a, f))
  -> Property IO
matchRandomGenSpec genM gen fromStdGen toStdGen runStatefulGen =
  forAll $ \seed -> monadic $ do
    let stdGen = mkStdGen seed
        g = fromStdGen stdGen
        (x1, g1) = gen stdGen
        (x2, g2) = gen g
    (x3, g3) <- runStatefulGen g genM
    pure $ and [x1 == x2, x2 == x3, g1 == toStdGen g2, g1 == toStdGen g3, g2 == g3]

withMutableGenSpec ::
     forall f m. (ThawedGen f m, Eq f, Show f)
  => (forall a. m a -> IO a)
  -> f
  -> Property IO
withMutableGenSpec toIO frozen =
  forAll $ \n -> monadic $ toIO $ do
    let action = uniformListM n
    x@(_, _) :: ([Word], f) <- withMutableGen frozen action
    y@(r, _) <- withMutableGen frozen action
    r' <- withMutableGen_ frozen action
    pure $ x == y && r == r'

overwriteMutableGenSpec ::
     forall f m. (ThawedGen f m, Eq f, Show f)
  => (forall a. m a -> IO a)
  -> f
  -> Property IO
overwriteMutableGenSpec toIO frozen =
  forAll $ \n -> monadic $ toIO $ do
    let action = uniformListM (abs n + 1) -- Non-empty
    ((r1, r2), frozen') :: ((String, String), f) <- withMutableGen frozen $ \mutable -> do
      r1 <- action mutable
      overwriteGen mutable frozen
      r2 <- action mutable
      modifyGen mutable (const ((), frozen))
      pure (r1, r2)
    pure $ r1 == r2 && frozen == frozen'

indepMutableGenSpec ::
     forall f m. (RandomGen f, ThawedGen f m, Eq f, Show f)
  => (forall a. m a -> IO a) -> [f] -> Property IO
indepMutableGenSpec toIO fgs =
  monadic $ toIO $ do
    (fgs ==) <$> (mapM freezeGen =<< mapM thawGen fgs)

immutableFrozenGenSpec ::
     forall f m. (RandomGen f, ThawedGen f m, Eq f, Show f)
  => (forall a. m a -> IO a) -> f -> Property IO
immutableFrozenGenSpec toIO frozen =
  forAll $ \n -> monadic $ toIO $ do
    let action = do
          mg <- thawGen frozen
          (,) <$> uniformWord8 mg <*> freezeGen mg
    x <- action
    xs <- replicateM n action
    pure $ all (x ==) xs

splitMutableGenSpec ::
     forall f m. (SplitGen f, ThawedGen f m, Eq f, Show f)
  => (forall a. m a -> IO a)
  -> f
  -> Property IO
splitMutableGenSpec toIO frozen =
  monadic $ toIO $ do
    (sfg1, fg1) <- withMutableGen frozen splitGenM
    (smg2, fg2) <- withMutableGen frozen splitMutableGenM
    sfg3 <- freezeGen smg2
    pure $ fg1 == fg2 && sfg1 == sfg3

thawedGenSpecFor ::
     forall f m. (SplitGen f, ThawedGen f m, Eq f, Show f, Serial IO f, Typeable f)
  => (forall a. m a -> IO a)
  -> Proxy f
  -> TestTree
thawedGenSpecFor toIO px =
  testGroup
    (showsTypeRep (typeRep px) "")
    [ testProperty "withMutableGen" $
      forAll $ \(f :: f) -> withMutableGenSpec toIO f
    , testProperty "overwriteGen" $
      forAll $ \(f :: f) -> overwriteMutableGenSpec toIO f
    , testProperty "independent mutable generators" $
      forAll $ \(fs :: [f]) -> indepMutableGenSpec toIO fs
    , testProperty "immutable frozen generators" $
      forAll $ \(f :: f) -> immutableFrozenGenSpec toIO f
    , testProperty "splitGen" $
      forAll $ \(f :: f) -> splitMutableGenSpec toIO f
    ]

frozenGenSpecFor ::
     forall f sg m. (RandomGen f, StatefulGen sg m, Eq f, Show f, Typeable f)
  => (StdGen -> f)
  -> (f -> StdGen)
  -> (forall a. f -> (sg -> m a) -> IO (a, f))
  -> TestTree
frozenGenSpecFor fromStdGen toStdGen runStatefulGen =
    testGroup (showsTypeRep (typeRep (Proxy :: Proxy f)) "")
    [ testGroup "matchRandomGenSpec"
      [ testProperty "uniformWord8/genWord8" $
          matchRandomGenSpec uniformWord8 genWord8 fromStdGen toStdGen runStatefulGen
      , testProperty "uniformWord16/genWord16" $
          matchRandomGenSpec uniformWord16 genWord16 fromStdGen toStdGen runStatefulGen
      , testProperty "uniformWord32/genWord32" $
          matchRandomGenSpec uniformWord32 genWord32 fromStdGen toStdGen runStatefulGen
      , testProperty "uniformWord64/genWord64" $
          matchRandomGenSpec uniformWord64 genWord64 fromStdGen toStdGen runStatefulGen
      , testProperty "uniformWord32R/genWord32R" $
        forAll $ \w32 ->
          matchRandomGenSpec (uniformWord32R w32) (genWord32R w32) fromStdGen toStdGen runStatefulGen
      , testProperty "uniformWord64R/genWord64R" $
        forAll $ \w64 ->
          matchRandomGenSpec (uniformWord64R w64) (genWord64R w64) fromStdGen toStdGen runStatefulGen
      , testProperty "uniformShortByteStringM/uniformShortByteString" $
        forAll $ \(NonNegative n') ->
          let n = n' `mod` 100000 -- Ensure it is not too big
          in matchRandomGenSpec
               (uniformShortByteStringM n)
               (uniformShortByteString n)
               fromStdGen
               toStdGen
               runStatefulGen
      , testProperty "uniformByteStringM/uniformByteString" $
        forAll $ \(NonNegative n') ->
          let n = n' `mod` 100000 -- Ensure it is not too big
          in matchRandomGenSpec
               (uniformByteStringM n)
               (uniformByteString n)
               fromStdGen
               toStdGen
               runStatefulGen
      , testProperty "uniformByteArrayM/genByteArray" $
        forAll $ \(NonNegative n', isPinned1 :: Bool, isPinned2 :: Bool) ->
          let n = n' `mod` 100000 -- Ensure it is not too big
          in matchRandomGenSpec
               (uniformByteArrayM isPinned1 n)
               (uniformByteArray isPinned2 n)
               fromStdGen
               toStdGen
               runStatefulGen
      ]
    ]


statefulGenSpec :: TestTree
statefulGenSpec =
  testGroup
    "StatefulGen"
    [ testGroup "ThawedGen"
        [ thawedGenSpecFor id (Proxy :: Proxy (IOGen StdGen))
        , thawedGenSpecFor id (Proxy :: Proxy (AtomicGen StdGen))
        , thawedGenSpecFor stToIO (Proxy :: Proxy (STGen StdGen))
        , thawedGenSpecFor atomically (Proxy :: Proxy (TGen StdGen))
        ]
    , testGroup "FrozenGen"
        [ frozenGenSpecFor StateGen unStateGen runStateGenT
        , frozenGenSpecFor IOGen unIOGen $ \g action -> do
            mg <- newIOGenM (unIOGen g)
            res <- action mg
            g' <- freezeGen mg
            pure (res, g')
        , frozenGenSpecFor AtomicGen unAtomicGen $ \g action -> do
            mg <- newAtomicGenM (unAtomicGen g)
            res <- action mg
            g' <- freezeGen mg
            pure (res, g')
        , frozenGenSpecFor STGen unSTGen $ \g action -> stToIO $ do
            mg <- newSTGenM (unSTGen g)
            res <- action mg
            g' <- freezeGen mg
            pure (res, g')
        , frozenGenSpecFor TGen unTGen $ \g action -> atomically $ do
            mg <- newTGenM (unTGen g)
            res <- action mg
            g' <- freezeGen mg
            pure (res, g')
        ]
    ]