Ticket #3028: par1.hs

module Par1 (
             Slot, 
             CachedSlots, 
             newCachedSlots,
             getSlot,
             waitSlot,
             doubles,
             numSlots,
             numRandom,
             startingGenerator,
             startThread,
             totalNeeded,
            )
where

import Control.Concurrent
import Control.Concurrent.STM
import Control.Monad
import Data.Array.Unboxed
import System.Random
import Data.IORef

data Slot = Slot { doubles :: UArray Int Double }

data CachedSlots = CachedSlots {
  numSlots          :: Int,
  numRandom         :: Int,
  firstReady        :: TVar Int,
  filledSlots       :: TVar Int,
  slots             :: Array Int (TVar Slot),
  order             :: TVar Int,
  startingGenerator :: StdGen,
  nextStart         :: MVar StdGen,
  currentGenerator  :: IORef StdGen,
  totalNeeded       :: TVar Int,
  alreadyGenerated  :: IORef Int
}

-- Generates a given number of random doubles using provided IORef on generator
-- 
genArray :: (RandomGen t) => Int -> IORef t -> IO (UArray Int Double)
genArray len generator = 
               do 
                  listValues <- forM [1..len] (\_ -> genAction)
                  return (listArray (0, len - 1) listValues)
               where 
                  genAction ::  IO (Double)
                  genAction = do
                    g <- readIORef (generator)
                    (number, g') <- return (randomR (0.0, 1.0) g)
                    writeIORef generator g'
                    return number
                  
genSlot numRandom currentRef n = do
    arr <- genArray numRandom currentRef
    return (Slot { doubles = arr})

-- Creates an instance of CachedSlots
-- 
newCachedSlots :: Int -> Int -> StdGen -> IO CachedSlots
newCachedSlots numSlots numRandom argGen = 
  do
    firstReady       <- newTVarIO (0 :: Int)
    filledSlots      <- newTVarIO (numSlots)
    currentRef       <- newIORef argGen
    nextStart        <- newMVar argGen
    listSlots        <- forM [1..numSlots] (genSlot numRandom currentRef)
    listTSlots       <- mapM newTVarIO listSlots 
    order            <- newTVarIO (0)
    totalNeeded      <- newTVarIO (0)
    alreadyGenerated <- newIORef (numSlots)
    
    return (CachedSlots {
      numSlots = numSlots,
      numRandom = numRandom,
      firstReady = firstReady,
      filledSlots = filledSlots,
      
      slots = listArray (0, numSlots - 1) listTSlots,
      order = order,
      startingGenerator = argGen,
      currentGenerator  = currentRef,
      nextStart = nextStart,
      totalNeeded = totalNeeded,
      alreadyGenerated = alreadyGenerated
      })


startThread :: CachedSlots -> IO ThreadId
startThread cs = forkIO (randomThreadS cs)
 
randomThreadS ::  CachedSlots -> IO ()
randomThreadS ca = 
  do 
   takeMVar (nextStart ca) >>= writeIORef randomRef 
   readIORef (alreadyGenerated ca) >>= addSlots
   readIORef randomRef >>= putMVar (nextStart ca) 
    where
      nSlots = numSlots ca 
      randomRef = currentGenerator ca 
      addSlots threadGenerated = 
       do
        precomputed <- genSlot (numRandom ca) randomRef threadGenerated
        fContinue <- atomically (updateCA precomputed)
        if fContinue 
          then addSlots (threadGenerated + 1) 
          else do 
            writeIORef (alreadyGenerated ca) threadGenerated
            return ()
       where 
         updateCA precomputed = do
           first   <- readTVar (firstReady ca)
           nFilled <- readTVar (filledSlots ca)
           needed  <- readTVar (totalNeeded ca)
           if threadGenerated >= needed then return False
             else if (nFilled == nSlots) 
              then retry
              else do 
                    slotPosition <- return ((first + nFilled) `mod` nSlots)
                    writeTVar ((slots ca)! slotPosition) precomputed
                    writeTVar (filledSlots ca) (nFilled + 1)
                    return True

getSlot :: CachedSlots -> STM (Int, Slot)
getSlot cs = do 
    nFilled <- readTVar (filledSlots cs)
    if nFilled == 0 then retry
       else do  
             first <- readTVar (firstReady cs)
             nextFirst <- return ( (first + 1) `mod` (numSlots cs))
             writeTVar (firstReady cs) nextFirst
             slot <- readTVar ((slots cs)! first)
             idx <- readTVar (order cs) 
             writeTVar (order cs) (idx + 1)
             writeTVar (filledSlots cs) (nFilled - 1)
             return ((idx, slot))

waitSlot :: CachedSlots -> Int -> STM (Int, Slot)
waitSlot cs n = do
       currentOrder <- readTVar (order cs)
       if (currentOrder < n) 
          then retry
          else getSlot cs