{-# LANGUAGE RecordWildCards #-}

import Control.Concurrent
import Control.Concurrent.Async
import Control.Concurrent.FVar
import Control.Monad
import Control.Monad.IO.Class
import Data.IORef
import Test.Hspec

data Blob a = Blob { x :: a, y :: a, z :: a } deriving (Eq, Show)

main :: IO ()
main = hspec $
  describe "FVar Tests" $ do
    it "Tests FVar fast read only access" $ do
      fvar <- newFVar $ Blob 0 0 0
      a    <- async . readFVar fvar $ return . x
      b    <- async . readFVar fvar $ return . y
      c    <- async . readFVar fvar $ return . x
      blob <- Blob <$> wait a <*> wait b <*> wait c
      blob `shouldBe` Blob 0 0 0
    it "Tests FVar write waits for reads to finish" $ do
      fvar <- newFVar $ Blob 0 0 0
      a    <- async . readFVar fvar $ \Blob {..} -> threadDelay 1000000 >> return x
      b    <- async . readFVar fvar $ \Blob {..} -> threadDelay 1000000 >> return y
      c    <- async . readFVar fvar $ \Blob {..} -> threadDelay 1000000 >> return z
      threadDelay 100000
      blob2 <- modifyFVar fvar $ \blob@Blob {..} -> do
                let blob' = blob { y = y + 1 }
                return (blob', blob')
      blob1 <- Blob <$> wait a <*> wait b <*> wait c
      blob1 `shouldBe` Blob 0 0 0
      blob2 `shouldBe` Blob 0 1 0
    it "Tests FVar reads waits for write to finish" $ do
      fvar <- newFVar $ Blob 0 0 0
      a  <- async . modifyFVar fvar $ \blob@Blob {..} ->
              threadDelay 1000000 >> return (blob { x = x + 1}, ())
      threadDelay 100000
      b <- async . readFVar fvar $ \Blob {..} -> return x
      c <- async . readFVar fvar $ \Blob {..} -> return y
      d <- async . readFVar fvar $ \Blob {..} -> return z
      wait a
      blob <- Blob <$> wait b <*> wait c <*> wait d
      blob `shouldBe` Blob 1 0 0
    it "Tests multiple writers" $ do
      fvar <- newFVar $ Blob 0 0 0
      a <- async . modifyFVar fvar $ \blob@Blob {..} ->
            threadDelay 1000000 >> return (blob { x = x + 1}, ())
      b <- async . modifyFVar fvar $ \blob@Blob {..} ->
            threadDelay 1000000 >> return (blob { x = x + 1}, ())
      wait a
      wait b
      x' <- readFVar fvar $ return . x
      y' <- readFVar fvar $ return . y
      z' <- readFVar fvar $ return . z
      x' `shouldBe` 2
    it "Tests interspersed reads with writes" $ do
      fvar <- newFVar $ Blob 0 0 0
      _    <- async . modifyFVar fvar $ \blob@Blob {..} ->
               threadDelay 1000000 >> return (blob { x = x + 1}, ())
      a    <- async . readFVar fvar $ return . x
      c    <- async . readFVar fvar $ return . z
      _    <- async . modifyFVar fvar $ \blob@Blob {..} -> do
               threadDelay 1000000
               return (blob { z = z + 1}, ())
      _    <- wait a
      _    <- wait c
      a'   <- async . readFVar fvar $ return . x
      b'   <- async . readFVar fvar $ return . y
      c'   <- async . readFVar fvar $ return . x
      blob <- Blob <$> wait a' <*> wait b' <*> wait c'
      blob `shouldBe` Blob 1 0 1
    it "Tests that fast reads are reenabled" $ do
      fvar <- newFVar $ Blob 0 0 0
      a <- async . modifyFVar fvar $ \blob@Blob {..} ->
            threadDelay 1000000 >> return (blob { x = x + 1}, ())
      threadDelay 100000
      b <- async . readFVar fvar $ return . x
      c <- async . readFVar fvar $ return . z
      d <- async . modifyFVar fvar $ \blob@Blob {..} -> do
            threadDelay 1000000
            return (blob { z = z + 1}, ())
      _ <- wait a
      _ <- wait b
      _ <- wait c
      _ <- wait d
      a'    <- async . readFVar fvar $ return . x
      b'    <- async . readFVar fvar $ return . y
      c'    <- async . readFVar fvar $ return . x
      blob <- Blob <$> wait a' <*> wait b' <*> wait c'
      blob `shouldBe` Blob 1 0 1
    it "Does an acid test" $ do
      let n = 100
          n' = (n * (n+1)) `div` 2
          actions = do
                      i <- newIORef 0
                      j <- newIORef 0
                      k <- newIORef 0
                      c <- newIORef 0
                      fvar <- newFVar $ Blob i j k
                      let write blob@Blob {..} = do
                            modifyIORef' i $ (+ 1)
                            x' <- readIORef i
                            modifyIORef' j $ (+ x')
                            x'' <- readIORef i
                            modifyIORef' k $ (+ x'')
                            return (blob, ())
                      writers <- async $ do
                          as <- replicateM n . async $ modifyFVar fvar write
                          mapM_ wait as
                      readers <- async $ do
                          as <- replicateM n . async $ readFVar fvar return
                          mapM_ wait as
                      wait writers
                      wait readers
                      blob <- readFVar fvar $ \Blob{..} ->
                                Blob <$> readIORef i <*> readIORef j <*> readIORef k
                      blob `shouldBe` Blob n n' n'
      replicateM_ 1000 actions