Conversion.hs

{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE ConstraintKinds #-}
-- ToDo: extract to a separate package
module Data.Array.Accelerate.CArray.Conversion where

import qualified Data.Array.Accelerate.IO.Foreign.ForeignPtr as AIO
import qualified Data.Array.Accelerate.Sugar.Shape as Shape
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.Comfort.Storable.Private as CArray
import qualified Data.Array.Comfort.Shape as CShape
import Foreign.Storable.Complex ()

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

import Foreign.ForeignPtr (ForeignPtr, castForeignPtr)
import Foreign.Storable (Storable)


type CArray = CArray.Array
type ShapeInt = CShape.ZeroBased Int

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 ShapeInt a -> Array DIM1 a
accFromCArrayReal1D arr =
   AIO.fromForeignPtrs
      (accDimsFromArrayBounds1D $ CArray.shape arr)
      (CArray.buffer arr)

accFromCArrayReal2D ::
   (RealElement a) =>
   CArray (ShapeInt, ShapeInt) a -> Array DIM2 a
accFromCArrayReal2D arr =
   AIO.fromForeignPtrs
      (accDimsFromArrayBounds2D $ CArray.shape arr)
      (CArray.buffer arr)


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

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

accFromCArrayComplex ::
   (A.Shape sh, ComplexElement a ac) =>
   sh -> CArray csh ac -> Array sh ac
accFromCArrayComplex sh arr =
   AIO.fromForeignPtrs sh $ castForeignPtrFromComplex $ CArray.buffer arr

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



accDimsFromArrayBounds1D :: (CShape.C sh) => sh -> DIM1
accDimsFromArrayBounds1D sh = Z :. CShape.size sh

accDimsFromArrayBounds2D ::
   (CShape.C sh0, CShape.C sh1) => (sh0, sh1) -> DIM2
accDimsFromArrayBounds2D (sh0, sh1) =
   Z :. CShape.size sh0 :. CShape.size sh1


cArrayFromAccReal1D ::
   (RealElement a) =>
   Array DIM1 a -> CArray ShapeInt a
cArrayFromAccReal1D arr =
   CArray.Array
      (arrayBounds1DFromAccDims $ A.arrayShape arr)
      (AIO.toForeignPtrs arr)

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

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

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

cArrayFromAccComplex ::
   (CShape.C csh, A.Shape sh, ComplexElement a ac) =>
   csh -> Array sh ac -> CArray csh ac
cArrayFromAccComplex bnds arr =
   CArray.Array bnds (castForeignPtrToComplex $ AIO.toForeignPtrs arr)

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


withCArrayComplex ::
   (CShape.C csh, A.Shape sh, ComplexElement a ac) =>
   (sh -> csh) ->
   (CArray csh ac -> CArray csh 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 -> ShapeInt
arrayBounds1DFromAccDims (Z :. i) = CShape.ZeroBased i

arrayBounds2DFromAccDims :: DIM2 -> (ShapeInt, ShapeInt)
arrayBounds2DFromAccDims (Z :. i :. j) =
   (CShape.ZeroBased i, CShape.ZeroBased j)


newtype Cube sh = Cube sh

instance (A.Shape sh) => CShape.C (Cube sh) where
   size (Cube sh) = Shape.size sh