{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
module Data.Massiv.Array.Manifest.Vector (
fromVectorM,
fromVector',
castFromVector,
toVector,
castToVector,
ARepr,
VRepr,
) where
import Control.Monad (guard, join, msum)
import Data.Kind
import Data.Massiv.Array.Manifest.Boxed
import Data.Massiv.Array.Manifest.Internal
import Data.Massiv.Array.Manifest.Primitive
import Data.Massiv.Array.Manifest.Storable
import Data.Massiv.Array.Manifest.Unboxed
import Data.Massiv.Core.Common
import Data.Maybe (fromMaybe)
import Data.Typeable
import qualified Data.Vector as VB
import qualified Data.Vector.Generic as VG
import qualified Data.Vector.Primitive as VP
import qualified Data.Vector.Storable as VS
import qualified Data.Vector.Unboxed as VU
type family ARepr (v :: Type -> Type) :: Type where
ARepr VU.Vector = U
ARepr VS.Vector = S
ARepr VP.Vector = P
ARepr VB.Vector = BL
type family VRepr r :: Type -> Type where
VRepr U = VU.Vector
VRepr S = VS.Vector
VRepr P = VP.Vector
VRepr B = VB.Vector
VRepr BN = VB.Vector
VRepr BL = VB.Vector
castFromVector
:: forall v r ix e
. (VG.Vector v e, Typeable v, Index ix, ARepr v ~ r)
=> Comp
-> Sz ix
-> v e
-> Maybe (Array r ix e)
castFromVector :: forall (v :: * -> *) r ix e.
(Vector v e, Typeable v, Index ix, ARepr v ~ r) =>
Comp -> Sz ix -> v e -> Maybe (Array r ix e)
castFromVector Comp
comp Sz ix
sz v e
vector = do
Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Sz ix -> Int
forall ix. Index ix => Sz ix -> Int
totalElem Sz ix
sz Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== v e -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
VG.length v e
vector)
[Maybe (Array r ix e)] -> Maybe (Array r ix e)
forall (t :: * -> *) (m :: * -> *) a.
(Foldable t, MonadPlus m) =>
t (m a) -> m a
msum
[ do
v :~: Vector
Refl <- Maybe (v :~: Vector)
forall {k} (a :: k) (b :: k).
(Typeable a, Typeable b) =>
Maybe (a :~: b)
eqT :: Maybe (v :~: VU.Vector)
Vector e
uVector <- Maybe (Maybe (Vector e)) -> Maybe (Vector e)
forall (m :: * -> *) a. Monad m => m (m a) -> m a
join (Maybe (Maybe (Vector e)) -> Maybe (Vector e))
-> Maybe (Maybe (Vector e)) -> Maybe (Vector e)
forall a b. (a -> b) -> a -> b
$ Maybe (v e) -> Maybe (Maybe (Vector e))
forall {k1} {k2} (c :: k1 -> *) (t :: k2 -> k1) (t' :: k2 -> k1)
(a :: k2).
(Typeable t, Typeable t') =>
c (t a) -> Maybe (c (t' a))
gcast1 (v e -> Maybe (v e)
forall a. a -> Maybe a
Just v e
vector)
Array r ix e -> Maybe (Array r ix e)
forall a. a -> Maybe a
forall (m :: * -> *) a. Monad m => a -> m a
return (Array r ix e -> Maybe (Array r ix e))
-> Array r ix e -> Maybe (Array r ix e)
forall a b. (a -> b) -> a -> b
$ UArray{uComp :: Comp
uComp = Comp
comp, uSize :: Sz ix
uSize = Sz ix
sz, uData :: Vector e
uData = Vector e
uVector}
, do
v :~: Vector
Refl <- Maybe (v :~: Vector)
forall {k} (a :: k) (b :: k).
(Typeable a, Typeable b) =>
Maybe (a :~: b)
eqT :: Maybe (v :~: VS.Vector)
Vector e
sVector <- Maybe (Maybe (Vector e)) -> Maybe (Vector e)
forall (m :: * -> *) a. Monad m => m (m a) -> m a
join (Maybe (Maybe (Vector e)) -> Maybe (Vector e))
-> Maybe (Maybe (Vector e)) -> Maybe (Vector e)
forall a b. (a -> b) -> a -> b
$ Maybe (v e) -> Maybe (Maybe (Vector e))
forall {k1} {k2} (c :: k1 -> *) (t :: k2 -> k1) (t' :: k2 -> k1)
(a :: k2).
(Typeable t, Typeable t') =>
c (t a) -> Maybe (c (t' a))
gcast1 (v e -> Maybe (v e)
forall a. a -> Maybe a
Just v e
vector)
Array r ix e -> Maybe (Array r ix e)
forall a. a -> Maybe a
forall (m :: * -> *) a. Monad m => a -> m a
return (Array r ix e -> Maybe (Array r ix e))
-> Array r ix e -> Maybe (Array r ix e)
forall a b. (a -> b) -> a -> b
$ Sz ix -> Array r Int e -> Array r ix e
forall r ix ix' e.
(Size r, Index ix, Index ix') =>
Sz ix' -> Array r ix e -> Array r ix' e
forall ix ix' e.
(Index ix, Index ix') =>
Sz ix' -> Array r ix e -> Array r ix' e
unsafeResize Sz ix
sz (Array r Int e -> Array r ix e) -> Array r Int e -> Array r ix e
forall a b. (a -> b) -> a -> b
$ Comp -> Vector e -> Vector S e
forall e. Comp -> Vector e -> Vector S e
fromStorableVector Comp
comp Vector e
sVector
, do
v :~: Vector
Refl <- Maybe (v :~: Vector)
forall {k} (a :: k) (b :: k).
(Typeable a, Typeable b) =>
Maybe (a :~: b)
eqT :: Maybe (v :~: VP.Vector)
VP.Vector Int
o Int
_ ByteArray
ba <- Maybe (Maybe (Vector e)) -> Maybe (Vector e)
forall (m :: * -> *) a. Monad m => m (m a) -> m a
join (Maybe (Maybe (Vector e)) -> Maybe (Vector e))
-> Maybe (Maybe (Vector e)) -> Maybe (Vector e)
forall a b. (a -> b) -> a -> b
$ Maybe (v e) -> Maybe (Maybe (Vector e))
forall {k1} {k2} (c :: k1 -> *) (t :: k2 -> k1) (t' :: k2 -> k1)
(a :: k2).
(Typeable t, Typeable t') =>
c (t a) -> Maybe (c (t' a))
gcast1 (v e -> Maybe (v e)
forall a. a -> Maybe a
Just v e
vector)
Array r ix e -> Maybe (Array r ix e)
forall a. a -> Maybe a
forall (m :: * -> *) a. Monad m => a -> m a
return (Array r ix e -> Maybe (Array r ix e))
-> Array r ix e -> Maybe (Array r ix e)
forall a b. (a -> b) -> a -> b
$ PArray{pComp :: Comp
pComp = Comp
comp, pSize :: Sz ix
pSize = Sz ix
sz, pOffset :: Int
pOffset = Int
o, pData :: ByteArray
pData = ByteArray
ba}
, do
v :~: Vector
Refl <- Maybe (v :~: Vector)
forall {k} (a :: k) (b :: k).
(Typeable a, Typeable b) =>
Maybe (a :~: b)
eqT :: Maybe (v :~: VB.Vector)
Vector e
bVector <- Maybe (Maybe (Vector e)) -> Maybe (Vector e)
forall (m :: * -> *) a. Monad m => m (m a) -> m a
join (Maybe (Maybe (Vector e)) -> Maybe (Vector e))
-> Maybe (Maybe (Vector e)) -> Maybe (Vector e)
forall a b. (a -> b) -> a -> b
$ Maybe (v e) -> Maybe (Maybe (Vector e))
forall {k1} {k2} (c :: k1 -> *) (t :: k2 -> k1) (t' :: k2 -> k1)
(a :: k2).
(Typeable t, Typeable t') =>
c (t a) -> Maybe (c (t' a))
gcast1 (v e -> Maybe (v e)
forall a. a -> Maybe a
Just v e
vector)
Array r ix e -> Maybe (Array r ix e)
forall a. a -> Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Array r ix e -> Maybe (Array r ix e))
-> Array r ix e -> Maybe (Array r ix e)
forall a b. (a -> b) -> a -> b
$ Sz ix -> Array r Int e -> Array r ix e
forall r ix ix' e.
(Size r, Index ix, Index ix') =>
Sz ix' -> Array r ix e -> Array r ix' e
forall ix ix' e.
(Index ix, Index ix') =>
Sz ix' -> Array r ix e -> Array r ix' e
unsafeResize Sz ix
sz (Array r Int e -> Array r ix e) -> Array r Int e -> Array r ix e
forall a b. (a -> b) -> a -> b
$ Comp -> Array r Int e -> Array r Int e
forall r ix e. Strategy r => Comp -> Array r ix e -> Array r ix e
forall ix e. Comp -> Array r ix e -> Array r ix e
setComp Comp
comp (Array r Int e -> Array r Int e) -> Array r Int e -> Array r Int e
forall a b. (a -> b) -> a -> b
$ Vector e -> Vector BL e
forall a. Vector a -> Vector BL a
fromBoxedVector Vector e
bVector
]
{-# NOINLINE castFromVector #-}
fromVectorM
:: (MonadThrow m, Typeable v, VG.Vector v a, Manifest r a, Load (ARepr v) ix a, Load r ix a)
=> Comp
-> Sz ix
-> v a
-> m (Array r ix a)
fromVectorM :: forall (m :: * -> *) (v :: * -> *) a r ix.
(MonadThrow m, Typeable v, Vector v a, Manifest r a,
Load (ARepr v) ix a, Load r ix a) =>
Comp -> Sz ix -> v a -> m (Array r ix a)
fromVectorM Comp
comp Sz ix
sz v a
v =
case Comp -> Sz ix -> v a -> Maybe (Array (ARepr v) ix a)
forall (v :: * -> *) r ix e.
(Vector v e, Typeable v, Index ix, ARepr v ~ r) =>
Comp -> Sz ix -> v e -> Maybe (Array r ix e)
castFromVector Comp
comp Sz ix
sz v a
v of
Just Array (ARepr v) ix a
arr -> Array r ix a -> m (Array r ix a)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Array r ix a -> m (Array r ix a))
-> Array r ix a -> m (Array r ix a)
forall a b. (a -> b) -> a -> b
$ Array (ARepr v) ix a -> Array r ix a
forall r ix e r'.
(Manifest r e, Load r' ix e) =>
Array r' ix e -> Array r ix e
convert Array (ARepr v) ix a
arr
Maybe (Array (ARepr v) ix a)
Nothing -> do
Sz ix -> Sz Int -> m ()
forall (m :: * -> *) ix ix'.
(MonadThrow m, Index ix, Index ix') =>
Sz ix -> Sz ix' -> m ()
guardNumberOfElements Sz ix
sz (Int -> Sz Int
forall ix. Index ix => ix -> Sz ix
Sz (v a -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
VG.length v a
v))
Array r ix a -> m (Array r ix a)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Comp -> Sz ix -> (Int -> a) -> Array r ix a
forall r ix e.
Load r ix e =>
Comp -> Sz ix -> (Int -> e) -> Array r ix e
makeArrayLinear Comp
comp Sz ix
sz (v a -> Int -> a
forall (v :: * -> *) a. Vector v a => v a -> Int -> a
VG.unsafeIndex v a
v))
{-# NOINLINE fromVectorM #-}
fromVector'
:: (HasCallStack, Typeable v, VG.Vector v a, Load (ARepr v) ix a, Load r ix a, Manifest r a)
=> Comp
-> Sz ix
-> v a
-> Array r ix a
fromVector' :: forall (v :: * -> *) a ix r.
(HasCallStack, Typeable v, Vector v a, Load (ARepr v) ix a,
Load r ix a, Manifest r a) =>
Comp -> Sz ix -> v a -> Array r ix a
fromVector' Comp
comp Sz ix
sz = Either SomeException (Array r ix a) -> Array r ix a
forall a. HasCallStack => Either SomeException a -> a
throwEither (Either SomeException (Array r ix a) -> Array r ix a)
-> (v a -> Either SomeException (Array r ix a))
-> v a
-> Array r ix a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Comp -> Sz ix -> v a -> Either SomeException (Array r ix a)
forall (m :: * -> *) (v :: * -> *) a r ix.
(MonadThrow m, Typeable v, Vector v a, Manifest r a,
Load (ARepr v) ix a, Load r ix a) =>
Comp -> Sz ix -> v a -> m (Array r ix a)
fromVectorM Comp
comp Sz ix
sz
{-# INLINE fromVector' #-}
castToVector
:: forall v r ix e
. (Manifest r e, Index ix, VRepr r ~ v)
=> Array r ix e
-> Maybe (v e)
castToVector :: forall (v :: * -> *) r ix e.
(Manifest r e, Index ix, VRepr r ~ v) =>
Array r ix e -> Maybe (v e)
castToVector Array r ix e
arr =
[Maybe (v e)] -> Maybe (v e)
forall (t :: * -> *) (m :: * -> *) a.
(Foldable t, MonadPlus m) =>
t (m a) -> m a
msum
[ do
r :~: U
Refl <- Maybe (r :~: U)
forall {k} (a :: k) (b :: k).
(Typeable a, Typeable b) =>
Maybe (a :~: b)
eqT :: Maybe (r :~: U)
Array U ix e
uArr <- Array r ix e -> Maybe (Array U ix e)
forall r ix e r'.
(Typeable r, Typeable r') =>
Array r' ix e -> Maybe (Array r ix e)
gcastArr Array r ix e
arr
v e -> Maybe (v e)
forall a. a -> Maybe a
forall (m :: * -> *) a. Monad m => a -> m a
return (v e -> Maybe (v e)) -> v e -> Maybe (v e)
forall a b. (a -> b) -> a -> b
$ Array U ix e -> Vector e
forall ix e. Array U ix e -> Vector e
uData Array U ix e
uArr
, do
r :~: S
Refl <- Maybe (r :~: S)
forall {k} (a :: k) (b :: k).
(Typeable a, Typeable b) =>
Maybe (a :~: b)
eqT :: Maybe (r :~: S)
Array S ix e
sArr <- Array r ix e -> Maybe (Array S ix e)
forall r ix e r'.
(Typeable r, Typeable r') =>
Array r' ix e -> Maybe (Array r ix e)
gcastArr Array r ix e
arr
v e -> Maybe (v e)
forall a. a -> Maybe a
forall (m :: * -> *) a. Monad m => a -> m a
return (v e -> Maybe (v e)) -> v e -> Maybe (v e)
forall a b. (a -> b) -> a -> b
$ Array S ix e -> Vector e
forall ix e. Index ix => Array S ix e -> Vector e
toStorableVector Array S ix e
sArr
, do
r :~: P
Refl <- Maybe (r :~: P)
forall {k} (a :: k) (b :: k).
(Typeable a, Typeable b) =>
Maybe (a :~: b)
eqT :: Maybe (r :~: P)
Array P ix e
pArr <- Array r ix e -> Maybe (Array P ix e)
forall r ix e r'.
(Typeable r, Typeable r') =>
Array r' ix e -> Maybe (Array r ix e)
gcastArr Array r ix e
arr
v e -> Maybe (v e)
forall a. a -> Maybe a
forall (m :: * -> *) a. Monad m => a -> m a
return (v e -> Maybe (v e)) -> v e -> Maybe (v e)
forall a b. (a -> b) -> a -> b
$ Int -> Int -> ByteArray -> Vector e
forall a. Int -> Int -> ByteArray -> Vector a
VP.Vector (Array P ix e -> Int
forall ix e. Array P ix e -> Int
pOffset Array P ix e
pArr) (Sz ix -> Int
forall ix. Index ix => Sz ix -> Int
totalElem (Array r ix e -> Sz ix
forall r ix e. Size r => Array r ix e -> Sz ix
forall ix e. Array r ix e -> Sz ix
size Array r ix e
arr)) (ByteArray -> Vector e) -> ByteArray -> Vector e
forall a b. (a -> b) -> a -> b
$ Array P ix e -> ByteArray
forall ix e. Array P ix e -> ByteArray
pData Array P ix e
pArr
, do
r :~: B
Refl <- Maybe (r :~: B)
forall {k} (a :: k) (b :: k).
(Typeable a, Typeable b) =>
Maybe (a :~: b)
eqT :: Maybe (r :~: B)
Array B ix e
bArr <- Array r ix e -> Maybe (Array B ix e)
forall r ix e r'.
(Typeable r, Typeable r') =>
Array r' ix e -> Maybe (Array r ix e)
gcastArr Array r ix e
arr
v e -> Maybe (v e)
forall a. a -> Maybe a
forall (m :: * -> *) a. Monad m => a -> m a
return (v e -> Maybe (v e)) -> v e -> Maybe (v e)
forall a b. (a -> b) -> a -> b
$ Array BL ix e -> Vector e
forall ix a. Index ix => Array BL ix a -> Vector a
toBoxedVector (Array BL ix e -> Vector e) -> Array BL ix e -> Vector e
forall a b. (a -> b) -> a -> b
$ Array B ix e -> Array BL ix e
forall ix e. Array B ix e -> Array BL ix e
toLazyArray Array B ix e
bArr
, do
r :~: BN
Refl <- Maybe (r :~: BN)
forall {k} (a :: k) (b :: k).
(Typeable a, Typeable b) =>
Maybe (a :~: b)
eqT :: Maybe (r :~: BN)
Array BN ix e
bArr <- Array r ix e -> Maybe (Array BN ix e)
forall r ix e r'.
(Typeable r, Typeable r') =>
Array r' ix e -> Maybe (Array r ix e)
gcastArr Array r ix e
arr
v e -> Maybe (v e)
forall a. a -> Maybe a
forall (m :: * -> *) a. Monad m => a -> m a
return (v e -> Maybe (v e)) -> v e -> Maybe (v e)
forall a b. (a -> b) -> a -> b
$ Array BL ix e -> Vector e
forall ix a. Index ix => Array BL ix a -> Vector a
toBoxedVector (Array BL ix e -> Vector e) -> Array BL ix e -> Vector e
forall a b. (a -> b) -> a -> b
$ Array B ix e -> Array BL ix e
forall ix e. Array B ix e -> Array BL ix e
toLazyArray (Array B ix e -> Array BL ix e) -> Array B ix e -> Array BL ix e
forall a b. (a -> b) -> a -> b
$ Array BN ix e -> Array B ix e
forall ix e. Array BN ix e -> Array B ix e
unwrapNormalForm Array BN ix e
bArr
, do
r :~: BL
Refl <- Maybe (r :~: BL)
forall {k} (a :: k) (b :: k).
(Typeable a, Typeable b) =>
Maybe (a :~: b)
eqT :: Maybe (r :~: BL)
Array BL ix e
bArr <- Array r ix e -> Maybe (Array BL ix e)
forall r ix e r'.
(Typeable r, Typeable r') =>
Array r' ix e -> Maybe (Array r ix e)
gcastArr Array r ix e
arr
v e -> Maybe (v e)
forall a. a -> Maybe a
forall (m :: * -> *) a. Monad m => a -> m a
return (v e -> Maybe (v e)) -> v e -> Maybe (v e)
forall a b. (a -> b) -> a -> b
$ Array BL ix e -> Vector e
forall ix a. Index ix => Array BL ix a -> Vector a
toBoxedVector Array BL ix e
bArr
]
{-# NOINLINE castToVector #-}
toVector
:: forall r ix e v
. ( Manifest r e
, Load r ix e
, Manifest (ARepr v) e
, VG.Vector v e
, VRepr (ARepr v) ~ v
)
=> Array r ix e
-> v e
toVector :: forall r ix e (v :: * -> *).
(Manifest r e, Load r ix e, Manifest (ARepr v) e, Vector v e,
VRepr (ARepr v) ~ v) =>
Array r ix e -> v e
toVector Array r ix e
arr =
v e -> Maybe (v e) -> v e
forall a. a -> Maybe a -> a
fromMaybe
(Int -> (Int -> e) -> v e
forall (v :: * -> *) a. Vector v a => Int -> (Int -> a) -> v a
VG.generate (Sz ix -> Int
forall ix. Index ix => Sz ix -> Int
totalElem (Array r ix e -> Sz ix
forall r ix e. Size r => Array r ix e -> Sz ix
forall ix e. Array r ix e -> Sz ix
size Array r ix e
arr)) (Array r ix e -> Int -> e
forall ix. Index ix => Array r ix e -> Int -> e
forall r e ix. (Source r e, Index ix) => Array r ix e -> Int -> e
unsafeLinearIndex Array r ix e
arr))
(Array (ARepr v) ix e -> Maybe (v e)
forall (v :: * -> *) r ix e.
(Manifest r e, Index ix, VRepr r ~ v) =>
Array r ix e -> Maybe (v e)
castToVector (Array r ix e -> Array (ARepr v) ix e
forall r ix e r'.
(Manifest r e, Load r' ix e) =>
Array r' ix e -> Array r ix e
convert Array r ix e
arr :: Array (ARepr v) ix e))
{-# NOINLINE toVector #-}