{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE TypeFamilies #-}
module Data.Array.Accelerate.IO.Data.Array.IArray (
IxShapeRepr,
fromIArray,
toIArray,
) where
import Data.Array.Accelerate.Array.Sugar
import Data.Array.Accelerate.Type
import Data.Array.Accelerate.Error
import Data.Array.Accelerate.IO.Data.Array.Internal
import Data.Array.IArray ( IArray )
import qualified Data.Array.IArray as IArray
{-# INLINE fromIArray #-}
fromIArray
:: (IxShapeRepr (EltRepr ix) ~ EltRepr sh, IArray a e, IArray.Ix ix, Shape sh, Elt ix, Elt e)
=> a ix e
-> Array sh e
fromIArray iarr = fromFunction sh (\ix -> iarr IArray.! fromIxShapeRepr (offset lo' ix))
where
(lo,hi) = IArray.bounds iarr
lo' = toIxShapeRepr lo
hi' = toIxShapeRepr hi
sh = rangeToShape (lo', hi')
offset :: forall sh. Shape sh => sh -> sh -> sh
offset ix0 ix = toElt $ go (eltType (undefined::sh)) (fromElt ix0) (fromElt ix)
where
go :: TupleType ix -> ix -> ix -> ix
go TypeRunit () () = ()
go (TypeRpair tl tr) (l0, r0) (l,r) = (go tl l0 l, go tr r0 r)
go (TypeRscalar (SingleScalarType (NumSingleType (IntegralNumType TypeInt{})))) i0 i = i0+i
go _ _ _ =
$internalError "fromIArray" "error in index offset"
{-# INLINE toIArray #-}
toIArray
:: forall ix sh a e. (IxShapeRepr (EltRepr ix) ~ EltRepr sh, IArray a e, IArray.Ix ix, Shape sh, Elt ix)
=> Maybe ix
-> Array sh e
-> a ix e
toIArray mix0 arr = IArray.array bnds0 [(offset ix, arr ! toIxShapeRepr ix) | ix <- IArray.range bnds]
where
(u,v) = shapeToRange (shape arr)
bnds@(lo,hi) = (fromIxShapeRepr u, fromIxShapeRepr v)
bnds0 = (offset lo, offset hi)
offset :: ix -> ix
offset ix =
case mix0 of
Nothing -> ix
Just ix0 -> offset' ix0 ix
offset' :: ix -> ix -> ix
offset' ix0 ix
= fromIxShapeRepr
. (toElt :: EltRepr sh -> sh)
$ go (eltType (undefined::sh)) (fromElt (toIxShapeRepr ix0 :: sh)) (fromElt (toIxShapeRepr ix :: sh))
where
go :: TupleType sh' -> sh' -> sh' -> sh'
go TypeRunit () () = ()
go (TypeRpair tl tr) (l0,r0) (l,r) = (go tl l0 l, go tr r0 r)
go (TypeRscalar (SingleScalarType (NumSingleType (IntegralNumType TypeInt{})))) i0 i = i0+i
go _ _ _ =
$internalError "toIArray" "error in index offset"