    {-# 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 match :: Int -> Bool
    match produce :: Int -> a
    produce len :: 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   (sh :: DIM1
    sh, vec :: 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
    forall a b. a -> b -> b
    `seq` Vector a
    vec Vector a -> Array U DIM1 a -> Array U DIM1 a
    forall a b. a -> b -> b
    `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 #-}