content: {-# LANGUAGE MagicHash #-}
  {-# LANGUAGE UnboxedTuples #-}
  
  module Main ( main ) where
  
  import Control.Concurrent
  import Control.Concurrent.MVar
  import Control.Monad (when)
  import Foreign.Storable
  import GHC.Exts
  import GHC.IO
  
  -- | Iterations per worker.
  iters :: Int
  iters = 1000000
  
  main :: IO ()
  main = do
      fetchAddSubTest
      fetchAndTest
      fetchNandTest
      fetchOrTest
      fetchXorTest
      casTest
      readWriteTest
  
  -- | Test fetchAddIntArray# by having two threads concurrenctly
  -- increment a counter and then checking the sum at the end.
  fetchAddSubTest :: IO ()
  fetchAddSubTest = do
      tot <- race 0
          (\ mba -> work fetchAddIntArray mba iters 2)
          (\ mba -> work fetchSubIntArray mba iters 1)
      assertEq 1000000 tot "fetchAddSubTest"
    where
      work :: (MByteArray -> Int -> Int -> IO ()) -> MByteArray -> Int -> Int
           -> IO ()
      work op mba 0 val = return ()
      work op mba n val = op mba 0 val >> work op mba (n-1) val
  
  -- | Test fetchXorIntArray# by having two threads concurrenctly XORing
  -- and then checking the result at the end. Works since XOR is
  -- commutative.
  --
  -- Covers the code paths for AND, NAND, and OR as well.
  fetchXorTest :: IO ()
  fetchXorTest = do
      res <- race n0
          (\ mba -> work mba iters t1pat)
          (\ mba -> work mba iters t2pat)
      assertEq expected res "fetchXorTest"
    where
      work :: MByteArray -> Int -> Int -> IO ()
      work mba 0 val = return ()
      work mba n val = fetchXorIntArray mba 0 val >> work mba (n-1) val
  
      -- Initial value is a large prime and the two patterns are 1010...
      -- and 0101...
      (n0, t1pat, t2pat)
          | sizeOf (undefined :: Int) == 8 =
              (0x00000000ffffffff, 0x5555555555555555, 0x9999999999999999)
          | otherwise = (0x0000ffff, 0x55555555, 0x99999999)
      expected
          | sizeOf (undefined :: Int) == 8 = 4294967295
          | otherwise = 65535
  
  -- The tests for AND, NAND, and OR are trivial for two reasons:
  --
  --  * The code path is already well exercised by 'fetchXorTest'.
  --
  --  * It's harder to test these operations, as a long sequence of them
  --    convert to a single value but we'd like to write a test in the
  --    style of 'fetchXorTest' that applies the operation repeatedly,
  --    to make it likely that any race conditions are detected.
  --
  -- Right now we only test that they return the correct value for a
  -- single op on each thread.
  
  -- | Test an associative operation.
  fetchOpTest :: (MByteArray -> Int -> Int -> IO ())
              -> Int -> String -> IO ()
  fetchOpTest op expected name = do
      res <- race n0
          (\ mba -> work mba t1pat)
          (\ mba -> work mba t2pat)
      assertEq expected res name
    where
      work :: MByteArray -> Int -> IO ()
      work mba val = op mba 0 val
  
  -- | Initial value and operation arguments for race test.
  --
  -- Initial value is a large prime and the two patterns are 1010...
  -- and 0101...
  n0, t1pat, t2pat :: Int
  (n0, t1pat, t2pat)
      | sizeOf (undefined :: Int) == 8 =
          (0x00000000ffffffff, 0x5555555555555555, 0x9999999999999999)
      | otherwise = (0x0000ffff, 0x55555555, 0x99999999)
  
  fetchAndTest :: IO ()
  fetchAndTest = fetchOpTest fetchAndIntArray expected "fetchAndTest"
    where expected
              | sizeOf (undefined :: Int) == 8 = 286331153
              | otherwise = 4369
  
  -- | Test NAND without any race, as NAND isn't associative.
  fetchNandTest :: IO ()
  fetchNandTest = do
      mba <- newByteArray (sizeOf (undefined :: Int))
      writeIntArray mba 0 n0
      fetchNandIntArray mba 0 t1pat
      fetchNandIntArray mba 0 t2pat
      res <- readIntArray mba 0
      assertEq expected res "fetchNandTest"
    where expected
              | sizeOf (undefined :: Int) == 8 = 7378697629770151799
              | otherwise = -2576976009
  
  fetchOrTest :: IO ()
  fetchOrTest = fetchOpTest fetchOrIntArray expected "fetchOrTest"
    where expected
              | sizeOf (undefined :: Int) == 8 = 15987178197787607039
              | otherwise = -572653569
  
  -- | Test casIntArray# by using it to emulate fetchAddIntArray# and
  -- then having two threads concurrenctly increment a counter,
  -- checking the sum at the end.
  casTest :: IO ()
  casTest = do
      tot <- race 0
          (\ mba -> work mba iters 1)
          (\ mba -> work mba iters 2)
      assertEq 3000000 tot "casTest"
    where
      work :: MByteArray -> Int -> Int -> IO ()
      work mba 0 val = return ()
      work mba n val = add mba 0 val >> work mba (n-1) val
  
      -- Fetch-and-add implemented using CAS.
      add :: MByteArray -> Int -> Int -> IO ()
      add mba ix n = do
          old <- readIntArray mba ix
          old' <- casIntArray mba ix old (old + n)
          when (old /= old') $ add mba ix n
  
  -- | Tests atomic reads and writes by making sure that one thread sees
  -- updates that are done on another. This test isn't very good at the
  -- moment, as this might work even without atomic ops, but at least it
  -- exercises the code.
  readWriteTest :: IO ()
  readWriteTest = do
      mba <- newByteArray (sizeOf (undefined :: Int))
      writeIntArray mba 0 0
      latch <- newEmptyMVar
      done <- newEmptyMVar
      forkIO $ do
          takeMVar latch
          n <- atomicReadIntArray mba 0
          assertEq 1 n "readWriteTest"
          putMVar done ()
      atomicWriteIntArray mba 0 1
      putMVar latch ()
      takeMVar done
  
  -- | Create two threads that mutate the byte array passed to them
  -- concurrently. The array is one word large.
  race :: Int                    -- ^ Initial value of array element
       -> (MByteArray -> IO ())  -- ^ Thread 1 action
       -> (MByteArray -> IO ())  -- ^ Thread 2 action
       -> IO Int                 -- ^ Final value of array element
  race n0 thread1 thread2 = do
      done1 <- newEmptyMVar
      done2 <- newEmptyMVar
      mba <- newByteArray (sizeOf (undefined :: Int))
      writeIntArray mba 0 n0
      forkIO $ thread1 mba >> putMVar done1 ()
      forkIO $ thread2 mba >> putMVar done2 ()
      mapM_ takeMVar [done1, done2]
      readIntArray mba 0
  
  ------------------------------------------------------------------------
  -- Test helper
  
  assertEq :: (Eq a, Show a) => a -> a -> String -> IO ()
  assertEq expected actual name
      | expected == actual = putStrLn $ name ++ ": OK"
      | otherwise = do
          putStrLn $ name ++ ": FAIL"
          putStrLn $ "Expected: " ++ show expected
          putStrLn $ "  Actual: " ++ show actual
  
  ------------------------------------------------------------------------
  -- Wrappers around MutableByteArray#
  
  data MByteArray = MBA (MutableByteArray# RealWorld)
  
  fetchAddIntArray :: MByteArray -> Int -> Int -> IO ()
  fetchAddIntArray (MBA mba#) (I# ix#) (I# n#) = IO $ \ s# ->
      case fetchAddIntArray# mba# ix# n# s# of
          (# s2#, _ #) -> (# s2#, () #)
  
  fetchSubIntArray :: MByteArray -> Int -> Int -> IO ()
  fetchSubIntArray (MBA mba#) (I# ix#) (I# n#) = IO $ \ s# ->
      case fetchSubIntArray# mba# ix# n# s# of
          (# s2#, _ #) -> (# s2#, () #)
  
  fetchAndIntArray :: MByteArray -> Int -> Int -> IO ()
  fetchAndIntArray (MBA mba#) (I# ix#) (I# n#) = IO $ \ s# ->
      case fetchAndIntArray# mba# ix# n# s# of
          (# s2#, _ #) -> (# s2#, () #)
  
  fetchNandIntArray :: MByteArray -> Int -> Int -> IO ()
  fetchNandIntArray (MBA mba#) (I# ix#) (I# n#) = IO $ \ s# ->
      case fetchNandIntArray# mba# ix# n# s# of
          (# s2#, _ #) -> (# s2#, () #)
  
  fetchOrIntArray :: MByteArray -> Int -> Int -> IO ()
  fetchOrIntArray (MBA mba#) (I# ix#) (I# n#) = IO $ \ s# ->
      case fetchOrIntArray# mba# ix# n# s# of
          (# s2#, _ #) -> (# s2#, () #)
  
  fetchXorIntArray :: MByteArray -> Int -> Int -> IO ()
  fetchXorIntArray (MBA mba#) (I# ix#) (I# n#) = IO $ \ s# ->
      case fetchXorIntArray# mba# ix# n# s# of
          (# s2#, _ #) -> (# s2#, () #)
  
  newByteArray :: Int -> IO MByteArray
  newByteArray (I# n#) = IO $ \ s# ->
      case newByteArray# n# s# of
          (# s2#, mba# #) -> (# s2#, MBA mba# #)
  
  writeIntArray :: MByteArray -> Int -> Int -> IO ()
  writeIntArray (MBA mba#) (I# ix#) (I# n#) = IO $ \ s# ->
      case writeIntArray# mba# ix# n# s# of
          s2# -> (# s2#, () #)
  
  readIntArray :: MByteArray -> Int -> IO Int
  readIntArray (MBA mba#) (I# ix#) = IO $ \ s# ->
      case readIntArray# mba# ix# s# of
          (# s2#, n# #) -> (# s2#, I# n# #)
  
  atomicWriteIntArray :: MByteArray -> Int -> Int -> IO ()
  atomicWriteIntArray (MBA mba#) (I# ix#) (I# n#) = IO $ \ s# ->
      case atomicWriteIntArray# mba# ix# n# s# of
          s2# -> (# s2#, () #)
  
  atomicReadIntArray :: MByteArray -> Int -> IO Int
  atomicReadIntArray (MBA mba#) (I# ix#) = IO $ \ s# ->
      case atomicReadIntArray# mba# ix# s# of
          (# s2#, n# #) -> (# s2#, I# n# #)
  
  casIntArray :: MByteArray -> Int -> Int -> Int -> IO Int
  casIntArray (MBA mba#) (I# ix#) (I# old#) (I# new#) = IO $ \ s# ->
      case casIntArray# mba# ix# old# new# s# of
          (# s2#, old2# #) -> (# s2#, I# old2# #)
  