Conversion.hs

{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
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.Array.Sugar (EltRepr)
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,)
import Foreign.Storable (Storable)


accFromCArrayReal1D ::
   (A.Elt a, Storable a, AIO.ForeignPtrs (EltRepr a) ~ ForeignPtr a) =>
   CArray Int a -> Array DIM1 a
accFromCArrayReal1D arr =
   let (len, fptr) = CArray.toForeignPtr arr
   in  AIO.fromForeignPtrs (Z :. len) fptr

accFromCArrayReal2D ::
   (A.Elt a, Storable a, AIO.ForeignPtrs (EltRepr a) ~ ForeignPtr a) =>
   CArray (Int, Int) a -> Array DIM2 a
accFromCArrayReal2D arr =
   AIO.fromForeignPtrs
      (accDimsFromArrayBounds2D $ CArray.bounds arr)
      (snd $ CArray.toForeignPtr arr)


accFromCArrayComplex1D ::
   (Storable a, RealFloat a, Complex a ~ ac, A.Elt ac,
    AIO.ForeignPtrs (EltRepr ac) ~ ForeignPtr ac) =>
   CArray Int (Complex a) -> Array DIM1 (Complex a)
accFromCArrayComplex1D arr =
   accFromCArrayComplex (accDimsFromArrayBounds1D $ CArray.bounds arr) arr

accFromCArrayComplex2D ::
   (Storable a, RealFloat a, Complex a ~ ac, A.Elt ac,
    AIO.ForeignPtrs (EltRepr ac) ~ ForeignPtr ac) =>
   CArray (Int, Int) (Complex a) -> Array DIM2 (Complex a)
accFromCArrayComplex2D arr =
   accFromCArrayComplex (accDimsFromArrayBounds2D $ CArray.bounds arr) arr

accFromCArrayComplex ::
   (Ix i, A.Shape sh, Storable a, RealFloat a, Complex a ~ ac, A.Elt ac,
    AIO.ForeignPtrs (EltRepr ac) ~ ForeignPtr ac) =>
   sh -> CArray i ac -> Array sh ac
accFromCArrayComplex sh arr =
   AIO.fromForeignPtrs sh (snd $ CArray.toForeignPtr arr)



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 ::
   (A.Elt a, Storable a, AIO.ForeignPtrs (EltRepr a) ~ ForeignPtr 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 ::
   (A.Elt a, Storable a, AIO.ForeignPtrs (EltRepr a) ~ ForeignPtr a) =>
   Array DIM2 a -> CArray (Int, Int) a
cArrayFromAccReal2D arr =
   unsafePerformIO $
   unsafeForeignPtrToCArray (AIO.toForeignPtrs arr)
      (arrayBounds2DFromAccDims $ A.arrayShape arr)

cArrayFromAccComplex1D ::
   (Storable a, RealFloat a, Complex a ~ ac, A.Elt ac,
    AIO.ForeignPtrs (EltRepr ac) ~ ForeignPtr ac) =>
   Array DIM1 ac -> CArray Int ac
cArrayFromAccComplex1D arr =
   cArrayFromAccComplex
      (arrayBounds1DFromAccDims $ A.arrayShape arr) arr

cArrayFromAccComplex2D ::
   (Storable a, RealFloat a, Complex a ~ ac, A.Elt ac,
    AIO.ForeignPtrs (EltRepr ac) ~ ForeignPtr ac) =>
   Array DIM2 ac -> CArray (Int,Int) ac
cArrayFromAccComplex2D arr =
   cArrayFromAccComplex
      (arrayBounds2DFromAccDims $ A.arrayShape arr) arr

cArrayFromAccComplex ::
   (Ix i, A.Shape sh,
    Storable a, RealFloat a, Complex a ~ ac, A.Elt ac,
    AIO.ForeignPtrs (EltRepr ac) ~ ForeignPtr ac) =>
   (i,i) -> Array sh ac -> CArray i ac
cArrayFromAccComplex bnds arr =
   unsafePerformIO $
   unsafeForeignPtrToCArray (AIO.toForeignPtrs arr) bnds


withCArrayComplex ::
   (Ix i, A.Shape sh,
    Storable a, RealFloat a, Complex a ~ ac, A.Elt ac,
    AIO.ForeignPtrs (EltRepr ac) ~ ForeignPtr ac) =>
   (sh -> (i,i)) ->
   (CArray i (Complex a) -> CArray i (Complex a)) ->
   Array sh (Complex a) -> Array sh (Complex a)
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))