    {-# LANGUAGE BangPatterns, ExplicitForAll, ScopedTypeVariables, PatternGuards #-}
    module Data.Array.Repa.Eval.Selection
            (selectChunkedS, selectChunkedP)
    where
    import Data.Array.Repa.Eval.Gang
    import Data.Array.Repa.Shape
    import Data.Vector.Unboxed                      as V
    import Data.Vector.Unboxed.Mutable              as VM
    import GHC.Base                                 (remInt, quotInt)
    import Prelude                                  as P
    import Control.Monad                            as P
    import Data.IORef

    -- | Select indices matching a predicate.
    --
    --   * This primitive can be useful for writing filtering functions.
    --
    selectChunkedS
            :: Shape sh
            => (sh -> a -> IO ())   -- ^ Update function to write into result.
            -> (sh -> Bool)         -- ^ See if this predicate matches.
            -> (sh -> a)            -- ^  .. and apply fn to the matching index
            -> sh                   -- ^ Extent of indices to apply to predicate.
            -> IO Int               -- ^ Number of elements written to destination array.

    {-# INLINE selectChunkedS #-}
    selectChunkedS fnWrite fnMatch fnProduce !shSize
     = fill 0 0
     where  lenSrc  = size shSize

            fill !nSrc !nDst
             | nSrc >= lenSrc       = return nDst

             | ixSrc        <- fromIndex shSize nSrc
             , fnMatch ixSrc
             = do   fnWrite ixSrc (fnProduce ixSrc)
                    fill (nSrc + 1) (nDst + 1)

             | otherwise
             =      fill (nSrc + 1) nDst

    -- | Select indices matching a predicate, in parallel.
    --
    --   * This primitive can be useful for writing filtering functions.
    --
    --   * The array is split into linear chunks, with one chunk being given to
    --     each thread.
    --
    --   * The number of elements in the result array depends on how many threads
    --     you're running the program with.
    --
    selectChunkedP
            :: forall a
            .  Unbox a
            => (Int -> Bool)        -- ^ See if this predicate matches.
            -> (Int -> a)           --   .. and apply fn to the matching index
            -> Int                  -- Extent of indices to apply to predicate.
            -> IO [IOVector a]      -- Chunks containing array elements.

    {-# INLINE selectChunkedP #-}
    selectChunkedP fnMatch fnProduce !len
     = do
            -- Make IORefs that the threads will write their result chunks to.
            -- We start with a chunk size proportial to the number of threads we have,
            -- but the threads themselves can grow the chunks if they run out of space.
            refs    <- P.replicateM threads
                    $ do    vec     <- VM.new $ len `div` threads
                            newIORef vec

            -- Fire off a thread to fill each chunk.
            gangIO theGang
             $ \thread -> makeChunk (refs !! thread)
                            (splitIx thread)
                            (splitIx (thread + 1) - 1)

            -- Read the result chunks back from the IORefs.
            -- If a thread had to grow a chunk, then these might not be the same ones
            -- we created back in the first step.
            P.mapM readIORef refs

     where  -- See how many threads we have available.
            !threads        = gangSize theGang
            !chunkLen       = len `quotInt` threads
            !chunkLeftover  = len `remInt`  threads

            -- Decide where to split the source array.
            {-# INLINE splitIx #-}
            splitIx thread
             | thread < chunkLeftover = thread * (chunkLen + 1)
             | otherwise              = thread * chunkLen  + chunkLeftover

            -- Fill the given chunk with elements selected from this range of indices.
            makeChunk :: IORef (IOVector a) -> Int -> Int -> IO ()
            makeChunk !ref !ixSrc !ixSrcEnd
             | ixSrc > ixSrcEnd
             = do  vecDst   <- VM.new 0
                   writeIORef ref vecDst

             | otherwise
             = do  vecDst   <- VM.new (len `div` threads)
                   vecDst'  <- fillChunk ixSrc ixSrcEnd vecDst 0 (VM.length vecDst)
                   writeIORef ref vecDst'

            -- The main filling loop.
            fillChunk :: Int -> Int -> IOVector a -> Int -> Int -> IO (IOVector a)
            fillChunk !ixSrc !ixSrcEnd !vecDst !ixDst !ixDstLen
             -- If we've finished selecting elements, then slice the vector down
             -- so it doesn't have any empty space at the end.
             | ixSrc > ixSrcEnd
             =      return  $ VM.slice 0 ixDst vecDst

             -- If we've run out of space in the chunk then grow it some more.
             | ixDst >= ixDstLen
             = do   let ixDstLen'   = (VM.length vecDst + 1) * 2
                    vecDst'         <- VM.grow vecDst ixDstLen'
                    fillChunk ixSrc ixSrcEnd vecDst' ixDst ixDstLen'

             -- We've got a maching element, so add it to the chunk.
             | fnMatch ixSrc
             = do   VM.unsafeWrite vecDst ixDst (fnProduce ixSrc)
                    fillChunk (ixSrc + 1) ixSrcEnd vecDst (ixDst + 1) ixDstLen

             -- The element doesnt match, so keep going.
             | otherwise
             =      fillChunk (ixSrc + 1) ixSrcEnd vecDst ixDst ixDstLen