module Data.Repa.Array.Internals.Target
( Target (..), TargetI
, empty, singleton
, fromList, fromListInto
, mapMaybeS, mapEitherS
, generateMaybeS, generateEitherS
, unfoldEitherOfLengthIO)
where
import Data.Repa.Array.Generic.Index as A
import Data.Repa.Array.Internals.Bulk as A
import System.IO.Unsafe
import Control.Monad
import qualified Data.Vector.Fusion.Stream.Monadic as S
import Prelude hiding (length)
import qualified Prelude as P
#include "repa-array.h"
class Layout l => Target l a where
data Buffer l a
unsafeNewBuffer :: l -> IO (Buffer l a)
unsafeReadBuffer :: Buffer l a -> Int -> IO a
unsafeWriteBuffer :: Buffer l a -> Int -> a -> IO ()
unsafeGrowBuffer :: Buffer l a -> Int -> IO (Buffer l a)
unsafeSliceBuffer :: Int -> Int -> Buffer l a -> IO (Buffer l a)
unsafeFreezeBuffer :: Buffer l a -> IO (Array l a)
unsafeThawBuffer :: Array l a -> IO (Buffer l a)
touchBuffer :: Buffer l a -> IO ()
bufferLayout :: Buffer l a -> l
type TargetI l a = (Target l a, Index l ~ Int)
empty :: TargetI l a
=> Name l -> Array l a
empty nDst
= unsafePerformIO
$ do let lDst = create nDst 0
buf <- unsafeNewBuffer lDst
unsafeFreezeBuffer buf
singleton
:: TargetI l a
=> Name l -> a -> Array l a
singleton nDst x
= unsafePerformIO
$ do let lDst = create nDst 1
buf <- unsafeNewBuffer lDst
unsafeWriteBuffer buf 0 x
unsafeFreezeBuffer buf
fromList :: TargetI l a
=> Name l -> [a] -> Array l a
fromList nDst xx
= let len = P.length xx
lDst = create nDst len
Just arr = fromListInto lDst xx
in arr
fromListInto :: Target l a
=> l -> [a] -> Maybe (Array l a)
fromListInto lDst xx
= unsafePerformIO
$ do let !len = P.length xx
if len /= size (extent lDst)
then return Nothing
else do
!buf <- unsafeNewBuffer lDst
zipWithM_ (unsafeWriteBuffer buf) [0..] xx
arr <- unsafeFreezeBuffer buf
return $ Just arr
mapMaybeS
:: (BulkI lSrc a, TargetI lDst b)
=> Name lDst
-> (a -> Maybe b)
-> Array lSrc a
-> Maybe (Array lDst b)
mapMaybeS !nDst f arr
= generateMaybeS nDst (length arr) get_maybeS
where
get_maybeS ix
= f (index arr ix)
mapEitherS
:: (BulkI lSrc a, TargetI lDst b)
=> Name lDst
-> (a -> Either err b)
-> Array lSrc a
-> Either err (Array lDst b)
mapEitherS !nDst f arr
= generateEitherS nDst (length arr) get_eitherS
where
get_eitherS ix
= f (index arr ix)
generateMaybeS
:: TargetI l a
=> Name l -> Int -> (Int -> Maybe a)
-> Maybe (Array l a)
generateMaybeS !nDst !len get
= unsafePerformIO
$ do
let lDst = create nDst len
!buf <- unsafeNewBuffer lDst
let fill_generateMaybeS !ix
| ix >= len
= return ix
| otherwise
= case get ix of
Nothing
-> return ix
Just x
-> do unsafeWriteBuffer buf ix $! x
fill_generateMaybeS (ix + 1)
!pos <- fill_generateMaybeS 0
if pos < len
then return Nothing
else fmap Just $! unsafeFreezeBuffer buf
generateEitherS
:: TargetI l a
=> Name l -> Int -> (Int -> Either err a)
-> Either err (Array l a)
generateEitherS !nDst !len get
= unsafePerformIO
$ do
let lDst = create nDst len
!buf <- unsafeNewBuffer lDst
let fill_generateEitherS !ix
| ix >= len
= return Nothing
| otherwise
= case get ix of
Left err
-> return $ Just err
Right x
-> do unsafeWriteBuffer buf ix $! x
fill_generateEitherS (ix + 1)
!mErr <- fill_generateEitherS 0
case mErr of
Just err -> return $ Left err
Nothing -> fmap Right $! unsafeFreezeBuffer buf
unfoldEitherOfLengthIO
:: TargetI l a
=> Name l
-> Int
-> (Int -> acc -> IO (Either err (acc, a)))
-> acc
-> IO (Either err (acc, Array l a))
unfoldEitherOfLengthIO nDst len get acc0
= do
let lDst = create nDst len
!buf <- unsafeNewBuffer lDst
let fill_unfoldEither !sPEC !acc !ix
| ix >= len
= return $ Right acc
| otherwise
= get ix acc >>= \r
-> case r of
Left err
-> return $ Left err
Right (acc', x)
-> do unsafeWriteBuffer buf ix $! x
fill_unfoldEither sPEC acc' (ix + 1)
eErr <- fill_unfoldEither S.SPEC acc0 0
case eErr of
Left err
-> return $ Left err
Right acc
-> do arr <- unsafeFreezeBuffer buf
return $ Right (acc, arr)