{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
module Data.Array.Accelerate.IO.IArray (
fromIArray,
toIArray,
) where
import Data.Array.Accelerate.Array.Sugar
import Data.Array.Accelerate.Type
import Data.Array.IArray ( IArray )
import qualified Data.Array.IArray as IArray
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 UnitTuple () () = ()
go (PairTuple tl tr) (l0, r0) (l,r) = (go tl l0 l, go tr r0 r)
go (SingleTuple (NumScalarType (IntegralNumType TypeInt{}))) i0 i = i0+i
go _ _ _
= error "Data.Array.Accelerate.IO.IArray: error in index offset"
toIArray
:: (IxShapeRepr (EltRepr ix) ~ EltRepr sh, IArray a e, IArray.Ix ix, Shape sh, Elt ix)
=> Array sh e
-> a ix e
toIArray arr = IArray.array bnds [(ix, arr ! toIxShapeRepr ix) | ix <- IArray.range bnds]
where
(lo,hi) = shapeToRange (shape arr)
bnds = (fromIxShapeRepr lo, fromIxShapeRepr hi)
type family IxShapeRepr e where
IxShapeRepr () = ()
IxShapeRepr Int = ((),Int)
IxShapeRepr (t,h) = (IxShapeRepr t, h)
fromIxShapeRepr :: forall ix sh. (IxShapeRepr (EltRepr ix) ~ EltRepr sh, Shape sh, Elt ix) => sh -> ix
fromIxShapeRepr sh = toElt (go (eltType (undefined::ix)) (fromElt sh))
where
go :: forall ix'. TupleType ix' -> IxShapeRepr ix' -> ix'
go UnitTuple () = ()
go (SingleTuple (NumScalarType (IntegralNumType TypeInt{}))) ((),h) = h
go (PairTuple tt _) (t, h) = (go tt t, h)
go _ _ = error "Data.Array.Accelerate.IO.IArray: not a valid IArray.Ix"
toIxShapeRepr :: forall ix sh. (IxShapeRepr (EltRepr ix) ~ EltRepr sh, Shape sh, Elt ix) => ix -> sh
toIxShapeRepr ix = toElt (go (eltType (undefined::ix)) (fromElt ix))
where
go :: forall ix'. TupleType ix' -> ix' -> IxShapeRepr ix'
go UnitTuple () = ()
go (SingleTuple (NumScalarType (IntegralNumType TypeInt{}))) h = ((), h)
go (PairTuple tt _) (t, h) = (go tt t, h)
go _ _ = error "Data.Array.Accelerate.IO.IArray: not a valid IArray.Ix"