{-# LANGUAGE BangPatterns #-}
module Data.Array.Repa.Operators.Selection
        (selectP)
where
import Data.Array.Repa.Index
import Data.Array.Repa.Base
import Data.Array.Repa.Eval.Selection
import Data.Array.Repa.Repr.Unboxed             as U
import qualified Data.Vector.Unboxed            as V
import System.IO.Unsafe


-- | Produce an array by applying a predicate to a range of integers.
--   If the predicate matches, then use the second function to generate
--   the element.
--
--   * This is a low-level function helpful for writing filtering
--     operations on arrays.
--
--   * Use the integer as the index into the array you're filtering.
--
selectP :: (Unbox a, Monad m)
        => (Int -> Bool)        -- ^ If the Int matches this predicate,
        -> (Int -> a)           -- ^  ... then pass it to this fn to produce a value
        -> Int                  -- ^ Range between 0 and this maximum.
        -> m (Array U DIM1 a)   -- ^ Array containing produced values.

selectP :: (Int -> Bool) -> (Int -> a) -> Int -> m (Array U DIM1 a)
selectP Int -> Bool
match Int -> a
produce Int
len
 = Array U DIM1 a -> m (Array U DIM1 a)
forall (m :: * -> *) a. Monad m => a -> m a
return
 (Array U DIM1 a -> m (Array U DIM1 a))
-> Array U DIM1 a -> m (Array U DIM1 a)
forall a b. (a -> b) -> a -> b
$ IO (Array U DIM1 a) -> Array U DIM1 a
forall a. IO a -> a
unsafePerformIO
 (IO (Array U DIM1 a) -> Array U DIM1 a)
-> IO (Array U DIM1 a) -> Array U DIM1 a
forall a b. (a -> b) -> a -> b
$ do   (DIM1
sh, Vector a
vec)       <- IO (DIM1, Vector a)
selectIO
        Array U DIM1 a -> IO (Array U DIM1 a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Array U DIM1 a -> IO (Array U DIM1 a))
-> Array U DIM1 a -> IO (Array U DIM1 a)
forall a b. (a -> b) -> a -> b
$ DIM1
sh DIM1 -> Array U DIM1 a -> Array U DIM1 a
`seq` Vector a
vec Vector a -> Array U DIM1 a -> Array U DIM1 a
`seq`
                 DIM1 -> Vector a -> Array U DIM1 a
forall sh e. sh -> Vector e -> Array U sh e
fromUnboxed DIM1
sh Vector a
vec

 where  {-# INLINE selectIO #-}
        selectIO :: IO (DIM1, Vector a)
selectIO
         = do   [IOVector a]
vecs            <- (Int -> Bool) -> (Int -> a) -> Int -> IO [IOVector a]
forall a.
Unbox a =>
(Int -> Bool) -> (Int -> a) -> Int -> IO [IOVector a]
selectChunkedP Int -> Bool
match Int -> a
produce Int
len
                [Vector a]
vecs'           <- (IOVector a -> IO (Vector a)) -> [IOVector a] -> IO [Vector a]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM IOVector a -> IO (Vector a)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
V.unsafeFreeze [IOVector a]
vecs

                -- TODO: avoid copy somehow.
                let result :: Vector a
result      = [Vector a] -> Vector a
forall a. Unbox a => [Vector a] -> Vector a
V.concat [Vector a]
vecs'

                (DIM1, Vector a) -> IO (DIM1, Vector a)
forall (m :: * -> *) a. Monad m => a -> m a
return  (Z
Z Z -> Int -> DIM1
forall tail head. tail -> head -> tail :. head
:. Vector a -> Int
forall a. Unbox a => Vector a -> Int
V.length Vector a
result, Vector a
result)
{-# INLINE [1] selectP #-}