module Data.Array.Repa.Repr.ForeignPtr
        ( F, Array (..)
        , fromForeignPtr, toForeignPtr
        , computeIntoS,   computeIntoP)
where
import Data.Array.Repa.Shape
import Data.Array.Repa.Base
import Data.Array.Repa.Eval.Load
import Data.Array.Repa.Eval.Target
import Data.Array.Repa.Repr.Delayed
import Foreign.Storable
import Foreign.ForeignPtr
import Foreign.Marshal.Alloc
import System.IO.Unsafe
import qualified Foreign.ForeignPtr.Unsafe      as Unsafe

-- | Arrays represented as foreign buffers in the C heap.
data F

-- | Read elements from a foreign buffer.
instance Storable a => Source F a where
 data Array F sh a
        = AForeignPtr !sh !Int !(ForeignPtr a)

 linearIndex :: Array F sh a -> Int -> a
linearIndex (AForeignPtr _ len fptr) Int
ix
  | Int
ix Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
len  
        = IO a -> a
forall a. IO a -> a
unsafePerformIO 
        (IO a -> a) -> IO a -> a
forall a b. (a -> b) -> a -> b
$ ForeignPtr a -> (Ptr a -> IO a) -> IO a
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr a
fptr
        ((Ptr a -> IO a) -> IO a) -> (Ptr a -> IO a) -> IO a
forall a b. (a -> b) -> a -> b
$ \Ptr a
ptr -> Ptr a -> Int -> IO a
forall a. Storable a => Ptr a -> Int -> IO a
peekElemOff Ptr a
ptr Int
ix
  
  | Bool
otherwise
  = [Char] -> a
forall a. HasCallStack => [Char] -> a
error [Char]
"Repa: foreign array index out of bounds"
 {-# INLINE linearIndex #-}
 
 unsafeLinearIndex :: Array F sh a -> Int -> a
unsafeLinearIndex (AForeignPtr _ _ fptr) Int
ix
        = IO a -> a
forall a. IO a -> a
unsafePerformIO
        (IO a -> a) -> IO a -> a
forall a b. (a -> b) -> a -> b
$ ForeignPtr a -> (Ptr a -> IO a) -> IO a
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr a
fptr 
        ((Ptr a -> IO a) -> IO a) -> (Ptr a -> IO a) -> IO a
forall a b. (a -> b) -> a -> b
$ \Ptr a
ptr -> Ptr a -> Int -> IO a
forall a. Storable a => Ptr a -> Int -> IO a
peekElemOff Ptr a
ptr Int
ix
 {-# INLINE unsafeLinearIndex #-}

 extent :: Array F sh a -> sh
extent (AForeignPtr sh _ _)
        = sh
sh
 {-# INLINE extent #-}

 deepSeqArray :: Array F sh a -> b -> b
deepSeqArray (AForeignPtr sh len fptr) b
x 
  = sh
sh sh -> b -> b
forall sh a. Shape sh => sh -> a -> a
`deepSeq` Int
len Int -> b -> b
`seq` ForeignPtr a
fptr ForeignPtr a -> b -> b
`seq` b
x
 {-# INLINE deepSeqArray #-}
 

-- Load -----------------------------------------------------------------------
-- | Filling foreign buffers.
instance Storable e => Target F e where
 data MVec F e 
  = FPVec !Int !(ForeignPtr e)

 newMVec :: Int -> IO (MVec F e)
newMVec Int
n
  = do  let (e
proxy :: e) = e
forall a. HasCallStack => a
undefined
        Ptr e
ptr              <- Int -> IO (Ptr e)
forall a. Int -> IO (Ptr a)
mallocBytes (e -> Int
forall a. Storable a => a -> Int
sizeOf e
proxy Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
n)
        e
_                <- Ptr e -> IO e
forall a. Storable a => Ptr a -> IO a
peek Ptr e
ptr  IO e -> IO e -> IO e
forall a. a -> a -> a
`asTypeOf` e -> IO e
forall (m :: * -> *) a. Monad m => a -> m a
return e
proxy
        
        ForeignPtr e
fptr             <- FinalizerPtr e -> Ptr e -> IO (ForeignPtr e)
forall a. FinalizerPtr a -> Ptr a -> IO (ForeignPtr a)
newForeignPtr FinalizerPtr e
forall a. FinalizerPtr a
finalizerFree Ptr e
ptr
        MVec F e -> IO (MVec F e)
forall (m :: * -> *) a. Monad m => a -> m a
return           (MVec F e -> IO (MVec F e)) -> MVec F e -> IO (MVec F e)
forall a b. (a -> b) -> a -> b
$ Int -> ForeignPtr e -> MVec F e
forall e. Int -> ForeignPtr e -> MVec F e
FPVec Int
n ForeignPtr e
fptr
 {-# INLINE newMVec #-}

 -- CAREFUL: Unwrapping the foreignPtr like this means we need to be careful
 -- to touch it after the last use, otherwise the finaliser might run too early.
 unsafeWriteMVec :: MVec F e -> Int -> e -> IO ()
unsafeWriteMVec (FPVec _ fptr) !Int
ix !e
x
  = Ptr e -> Int -> e -> IO ()
forall a. Storable a => Ptr a -> Int -> a -> IO ()
pokeElemOff (ForeignPtr e -> Ptr e
forall a. ForeignPtr a -> Ptr a
Unsafe.unsafeForeignPtrToPtr ForeignPtr e
fptr) Int
ix e
x
 {-# INLINE unsafeWriteMVec #-}

 unsafeFreezeMVec :: sh -> MVec F e -> IO (Array F sh e)
unsafeFreezeMVec !sh
sh (FPVec len fptr)
  =     Array F sh e -> IO (Array F sh e)
forall (m :: * -> *) a. Monad m => a -> m a
return  (Array F sh e -> IO (Array F sh e))
-> Array F sh e -> IO (Array F sh e)
forall a b. (a -> b) -> a -> b
$ sh -> Int -> ForeignPtr e -> Array F sh e
forall sh a. sh -> Int -> ForeignPtr a -> Array F sh a
AForeignPtr sh
sh Int
len ForeignPtr e
fptr
 {-# INLINE unsafeFreezeMVec #-}

 deepSeqMVec :: MVec F e -> a -> a
deepSeqMVec !(FPVec _ fptr) a
x
  = ForeignPtr e -> Ptr e
forall a. ForeignPtr a -> Ptr a
Unsafe.unsafeForeignPtrToPtr ForeignPtr e
fptr Ptr e -> a -> a
`seq` a
x
 {-# INLINE deepSeqMVec #-}

 touchMVec :: MVec F e -> IO ()
touchMVec (FPVec _ fptr)
  = ForeignPtr e -> IO ()
forall a. ForeignPtr a -> IO ()
touchForeignPtr ForeignPtr e
fptr
 {-# INLINE touchMVec #-}


-- Conversions ----------------------------------------------------------------
-- | O(1). Wrap a `ForeignPtr` as an array.
fromForeignPtr
        :: Shape sh
        => sh -> ForeignPtr e -> Array F sh e
fromForeignPtr :: sh -> ForeignPtr e -> Array F sh e
fromForeignPtr !sh
sh !ForeignPtr e
fptr
        = sh -> Int -> ForeignPtr e -> Array F sh e
forall sh a. sh -> Int -> ForeignPtr a -> Array F sh a
AForeignPtr sh
sh (sh -> Int
forall sh. Shape sh => sh -> Int
size sh
sh) ForeignPtr e
fptr
{-# INLINE fromForeignPtr #-}


-- | O(1). Unpack a `ForeignPtr` from an array.
toForeignPtr :: Array F sh e -> ForeignPtr e
toForeignPtr :: Array F sh e -> ForeignPtr e
toForeignPtr (AForeignPtr _ _ fptr)
        = ForeignPtr e
fptr
{-# INLINE toForeignPtr #-}


-- | Compute an array sequentially and write the elements into a foreign
--   buffer without intermediate copying. If you want to copy a
--   pre-existing manifest array to a foreign buffer then `delay` it first.
computeIntoS
        :: (Load r1 sh e, Storable e)
        => ForeignPtr e -> Array r1 sh e -> IO ()
computeIntoS :: ForeignPtr e -> Array r1 sh e -> IO ()
computeIntoS !ForeignPtr e
fptr !Array r1 sh e
arr
 = Array r1 sh e -> MVec F e -> IO ()
forall r1 sh e r2.
(Load r1 sh e, Target r2 e) =>
Array r1 sh e -> MVec r2 e -> IO ()
loadS Array r1 sh e
arr (Int -> ForeignPtr e -> MVec F e
forall e. Int -> ForeignPtr e -> MVec F e
FPVec Int
0 ForeignPtr e
fptr)
{-# INLINE computeIntoS #-}


-- | Compute an array in parallel and write the elements into a foreign
--   buffer without intermediate copying. If you want to copy a
--   pre-existing manifest array to a foreign buffer then `delay` it first.
computeIntoP
        :: (Load r1 sh e, Storable e)
        => ForeignPtr e -> Array r1 sh e -> IO ()
computeIntoP :: ForeignPtr e -> Array r1 sh e -> IO ()
computeIntoP !ForeignPtr e
fptr !Array r1 sh e
arr
 = Array r1 sh e -> MVec F e -> IO ()
forall r1 sh e r2.
(Load r1 sh e, Target r2 e) =>
Array r1 sh e -> MVec r2 e -> IO ()
loadP Array r1 sh e
arr (Int -> ForeignPtr e -> MVec F e
forall e. Int -> ForeignPtr e -> MVec F e
FPVec Int
0 ForeignPtr e
fptr)
{-# INLINE computeIntoP #-}