1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
|
{-# 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 match produce len
= return
$ unsafePerformIO
$ do (sh, vec) <- selectIO
return $ sh `seq` vec `seq`
fromUnboxed sh vec
where {-# INLINE selectIO #-}
selectIO
= do vecs <- selectChunkedP match produce len
vecs' <- mapM V.unsafeFreeze vecs
-- TODO: avoid copy somehow.
let result = V.concat vecs'
return (Z :. V.length result, result)
{-# INLINE [1] selectP #-}
|
