Conversion.hs

{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE ConstraintKinds #-}
module Data.Array.Accelerate.CArray.Conversion where

import qualified Data.Array.Accelerate.IO.Foreign.ForeignPtr as AIO
import qualified Data.Array.Accelerate as A
import Data.Array.Accelerate.Sugar.Elt (EltR)
import Data.Array.Accelerate (Array, DIM1, DIM2, Z(Z), (:.)((:.)))

import qualified Data.Array.CArray as CArray
import Data.Array.CArray (CArray, unsafeForeignPtrToCArray, rangeSize)
import Data.Ix (Ix)
import Foreign.Storable.Complex ()

import Data.Array.Accelerate.Data.Complex (Complex)

import System.IO.Unsafe (unsafePerformIO)
import Foreign.ForeignPtr (ForeignPtr, castForeignPtr)
import Foreign.Storable (Storable)


type RealElement a =
      (A.Elt a, Storable a, AIO.ForeignPtrs (EltR a) ~ ForeignPtr a)

type ComplexElement a ac =
      (Storable a, RealFloat a, Complex a ~ ac, A.Elt ac,
       AIO.ForeignPtrs (EltR ac) ~ ForeignPtr a)

accFromCArrayReal1D ::
   (RealElement a) =>
   CArray Int a -> Array DIM1 a
accFromCArrayReal1D arr =
   let (len, fptr) = CArray.toForeignPtr arr
   in  AIO.fromForeignPtrs (Z :. len) fptr

accFromCArrayReal2D ::
   (RealElement a) =>
   CArray (Int, Int) a -> Array DIM2 a
accFromCArrayReal2D arr =
   AIO.fromForeignPtrs
      (accDimsFromArrayBounds2D $ CArray.bounds arr)
      (snd $ CArray.toForeignPtr arr)


accFromCArrayComplex1D ::
   (ComplexElement a ac) =>
   CArray Int (Complex a) -> Array DIM1 (Complex a)
accFromCArrayComplex1D arr =
   accFromCArrayComplex (accDimsFromArrayBounds1D $ CArray.bounds arr) arr

accFromCArrayComplex2D ::
   (ComplexElement a ac) =>
   CArray (Int, Int) (Complex a) -> Array DIM2 (Complex a)
accFromCArrayComplex2D arr =
   accFromCArrayComplex (accDimsFromArrayBounds2D $ CArray.bounds arr) arr

accFromCArrayComplex ::
   (Ix i, A.Shape sh, ComplexElement a ac) =>
   sh -> CArray i ac -> Array sh ac
accFromCArrayComplex sh arr =
   AIO.fromForeignPtrs sh $
   castForeignPtrFromComplex $ snd $ CArray.toForeignPtr arr

castForeignPtrFromComplex :: ForeignPtr (Complex a) -> ForeignPtr a
castForeignPtrFromComplex = castForeignPtr



accDimsFromArrayBounds1D :: (CArray.Ix i) => (i, i) -> DIM1
accDimsFromArrayBounds1D bnds = Z :. rangeSize bnds

accDimsFromArrayBounds2D ::
   (CArray.Ix i, CArray.Ix j) => ((i, j), (i, j)) -> DIM2
accDimsFromArrayBounds2D ((li,lj), (ri,rj)) =
   Z :. rangeSize (li,ri) :. rangeSize (lj,rj)


cArrayFromAccReal1D ::
   (RealElement a) =>
   Array DIM1 a -> CArray Int a
cArrayFromAccReal1D arr =
   let (Z:.len) = A.arrayShape arr
   in  unsafePerformIO $
       unsafeForeignPtrToCArray (AIO.toForeignPtrs arr) (0,len-1)

cArrayFromAccReal2D ::
   (RealElement a) =>
   Array DIM2 a -> CArray (Int, Int) a
cArrayFromAccReal2D arr =
   unsafePerformIO $
   unsafeForeignPtrToCArray (AIO.toForeignPtrs arr)
      (arrayBounds2DFromAccDims $ A.arrayShape arr)

cArrayFromAccComplex1D ::
   (ComplexElement a ac) =>
   Array DIM1 ac -> CArray Int ac
cArrayFromAccComplex1D arr =
   cArrayFromAccComplex
      (arrayBounds1DFromAccDims $ A.arrayShape arr) arr

cArrayFromAccComplex2D ::
   (ComplexElement a ac) =>
   Array DIM2 ac -> CArray (Int,Int) ac
cArrayFromAccComplex2D arr =
   cArrayFromAccComplex
      (arrayBounds2DFromAccDims $ A.arrayShape arr) arr

cArrayFromAccComplex ::
   (Ix i, A.Shape sh, ComplexElement a ac) =>
   (i,i) -> Array sh ac -> CArray i ac
cArrayFromAccComplex bnds arr =
   unsafePerformIO $
   unsafeForeignPtrToCArray
      (castForeignPtrToComplex $ AIO.toForeignPtrs arr) bnds

castForeignPtrToComplex :: ForeignPtr a -> ForeignPtr (Complex a)
castForeignPtrToComplex = castForeignPtr


withCArrayComplex ::
   (Ix i, A.Shape sh, ComplexElement a ac) =>
   (sh -> (i,i)) ->
   (CArray i ac -> CArray i ac) ->
   Array sh ac -> Array sh ac
withCArrayComplex bndsFromSh f arr =
   let sh = A.arrayShape arr
   in  accFromCArrayComplex sh $ f $
       cArrayFromAccComplex (bndsFromSh sh) arr


arrayBounds1DFromAccDims :: DIM1 -> (Int, Int)
arrayBounds1DFromAccDims (Z :. i) = (0, i-1)

arrayBounds2DFromAccDims :: DIM2 -> ((Int, Int), (Int, Int))
arrayBounds2DFromAccDims (Z :. i :. j) = ((0,0), (i-1,j-1))