{-# LANGUAGE RecursiveDo #-}
{-# OPTIONS_GHC -fplugin UIO #-}
module Test.UIO where

import Data.IORef
import Control.Monad

import UIO

testBoth :: IO ()
testBoth = do
  putStrLn "testUIOFix"
  testUIOFix
  putStrLn "testUIOUnique"
  testUIOUnique
  putStrLn "testUIOReplicate"
  testUIOReplicate
  putStrLn "testUIOBadTimeless"
  testUIOBadTimeless
  putStrLn "testUIOMany"
  testUIOMany
  putStrLn "testUIOCycle"
  testUIOCycle
  putStrLn "testIO"
  testIO
--  putStrLn "testUIOPrintLots"
--  testUIOPrintLots

{-# NOINLINE testUIOUnique #-}
testUIOUnique :: IO ()
testUIOUnique = do
  r <- runUIO2 $ do
    -- The following two lines can be merged by common subexpression elimination (CSE), which is very bad
    r <- timeless $ newIORef 2
    r2 <- timeless $ newIORef 2
    -- Note that the following two lines must be different, otherwise they will *also* be merged by CSE, which will make the test appear to succeed!
    unordered $ atomicModifyIORef' r $ \v -> (v + 5, ())
    unordered $ atomicModifyIORef' r2 $ \v -> (v + 10, ())
    pure r
  7 <- readIORef r
  pure ()

{-# NOINLINE testUIOFix #-}
testUIOFix :: IO ()
testUIOFix = do
  r <- runUIO2 $ mdo
    unordered $ atomicModifyIORef' r $ \v -> (v + 5, ())
    r <- timeless $ newIORef 2
    pure r
  7 <- readIORef r
  pure ()

{-# NOINLINE testUIOBadTimeless #-}
testUIOBadTimeless :: IO ()
testUIOBadTimeless = do
  r <- runUIO2 $ mdo
    timeless $ atomicModifyIORef' r $ \v -> (v + 5, ())
    r <- timeless $ newIORef 2
    pure r
  2 <- readIORef r
  pure ()

{-# NOINLINE testUIOMany #-}
testUIOMany :: IO ()
testUIOMany = do
  r <- runUIO2 $ do
    r <- timeless $ newIORef 0
    unordered $ atomicModifyIORef' r $ \v -> (v + 1, ())
    unordered $ atomicModifyIORef' r $ \v -> (v + 2, ())
    unordered $ atomicModifyIORef' r $ \v -> (v + 3, ())
    unordered $ atomicModifyIORef' r $ \v -> (v + 4, ())
    unordered $ atomicModifyIORef' r $ \v -> (v + 5, ())
    unordered $ atomicModifyIORef' r $ \v -> (v + 6, ())
    unordered $ atomicModifyIORef' r $ \v -> (v + 7, ())
    unordered $ atomicModifyIORef' r $ \v -> (v + 8, ())
    unordered $ atomicModifyIORef' r $ \v -> (v + 9, ())
    unordered $ atomicModifyIORef' r $ \v -> (v + 10, ())
    pure r
  55 <- readIORef r
  pure ()

{-# NOINLINE testUIOReplicate #-}
testUIOReplicate :: IO ()
testUIOReplicate = do
  rs <- runUIO2 $ do
    rs <- replicateM 10 $ timeless $ newIORef 2
    forM_ rs $ \r ->
      unordered $ atomicModifyIORef' r $ \v -> (v + 5, ())
    pure rs
  70 <- fmap sum $ mapM readIORef rs
  pure ()

{-# NOINLINE testUIOPrintLots #-}
testUIOPrintLots :: IO ()
testUIOPrintLots = runUIO2 $ do
  replicateM_ 1000000 $ unordered $ putStrLn "Task"

newtype CycleRef = CycleRef (IORef CycleRef)

{-# NOINLINE testUIOCycle #-}
testUIOCycle :: IO ()
testUIOCycle = runUIO2 $ mdo
  r <- timeless $ newIORef $ CycleRef r
  pure ()

{-# NOINLINE testBoth #-}
testIO :: IO ()
testIO = do
  r <- newIORef 2
  atomicModifyIORef' r $ \v -> (v + 5, ())
  7 <- readIORef r
  pure ()