{-# LANGUAGE TypeFamilies #-}
import Test.Hspec
import Test.Hspec.QuickCheck
import Test.QuickCheck.Arbitrary
import Test.QuickCheck.Gen
import Data.Mutable.Deque
import Control.Monad (forM_)
import Data.IORef
import Data.Primitive.MutVar (MutVar)
import Control.Monad.Primitive (PrimState)
import Data.Sequence (Seq)
import Data.STRef (STRef)
import Data.Vector (Vector)
import Data.Mutable.URef
import Data.Mutable.SRef
import Data.Mutable.DList

main :: IO ()
main = hspec spec

data RefAction
    = WriteRef Int
    | ModifyRef Int
    | ModifyRef' Int
    | AtomicModifyRef Int
    | AtomicModifyRef' Int
    deriving Show
instance Arbitrary RefAction where
    arbitrary = oneof
        [ fmap WriteRef arbitrary
        , fmap ModifyRef arbitrary
        , fmap ModifyRef' arbitrary
        , fmap AtomicModifyRef arbitrary
        , fmap AtomicModifyRef' arbitrary
        ]

data DequeAction
    = PushFront Int
    | PushBack Int
    | PopFront
    | PopBack
    deriving Show
instance Arbitrary DequeAction where
    arbitrary = oneof $ concat
        [ replicate 25 $ fmap PushFront arbitrary
        , replicate 25 $ fmap PushBack arbitrary
        , [return PopFront, return PopBack]
        ]

spec :: Spec
spec = do
    describe "Deque" $ do
        let test name forceType = prop name $ \actions -> do
                base <- newColl :: IO (IORef [Int])
                tested <- fmap forceType newColl
                forM_ (PopFront : PopBack : actions) $ \action -> do
                    case action of
                        PushFront i -> do
                            pushFront base i
                            pushFront tested i
                        PushBack i -> do
                            pushBack base i
                            pushBack tested i
                        PopFront -> do
                            expected <- popFront base
                            actual <- popFront tested
                            actual `shouldBe` expected
                        PopBack -> do
                            expected <- popBack base
                            actual <- popBack tested
                            actual `shouldBe` expected
                let drain = do
                        expected <- popBack base
                        actual <- popBack tested
                        actual `shouldBe` expected
                        case actual of
                            Just _ -> drain
                            Nothing -> return $! ()
                drain
        test "UDeque" asUDeque
        test "SDeque" asSDeque
        test "DList" asDList
        test "MutVar Seq" (id :: MutVar (PrimState IO) (Seq Int) -> MutVar (PrimState IO) (Seq Int))
        test "STRef Vector" (id :: STRef (PrimState IO) (Vector Int) -> STRef (PrimState IO) (Vector Int))
    describe "Ref" $ do
        let test name forceType atomic atomic' = prop name $ \start actions -> do
                base <- fmap asIORef $ newRef start
                tested <- fmap forceType $ newRef start
                let check = do
                        expected <- readRef base
                        actual <- readRef tested
                        expected `shouldBe` actual
                forM_ actions $ \action -> case action of
                    WriteRef i -> do
                        writeRef base i
                        writeRef tested i
                        check
                    ModifyRef i -> do
                        modifyRef base (+ i)
                        modifyRef tested (+ i)
                        check
                    ModifyRef' i -> do
                        modifyRef' base (subtract i)
                        modifyRef' tested (subtract i)
                        check
                    AtomicModifyRef i -> do
                        let f x = (x + i, ())
                        atomicModifyRef base f
                        atomic tested f
                        check
                    AtomicModifyRef' i -> do
                        atomicModifyRef' base $ \x -> (x - i, ())
                        atomic' tested $ \x -> (x - i, ())
                        check
        test "URef" asURef modifyRefHelper modifyRefHelper'
        test "SRef" asSRef modifyRefHelper modifyRefHelper'
        test "STRef" asSTRef modifyRefHelper modifyRefHelper'
        test "MutVar" asMutVar atomicModifyRef atomicModifyRef'

modifyRefHelper :: (MCState c ~ PrimState IO, RefElement c ~ Int, MutableRef c)
                => c
                -> (Int -> (Int, ()))
                -> IO ()
modifyRefHelper ref f = modifyRef ref $ \i ->
    let (x, y) = f i
     in y `seq` x

modifyRefHelper' :: (MCState c ~ PrimState IO, RefElement c ~ Int, MutableRef c)
                 => c
                 -> (Int -> (Int, ()))
                 -> IO ()
modifyRefHelper' ref f = modifyRef' ref $ \i ->
    let (x, y) = f i
     in y `seq` x