{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
-- |
-- Module      : Data.Massiv.Array.Manifest.Vector
-- Copyright   : (c) Alexey Kuleshevich 2018-2021
-- License     : BSD3
-- Maintainer  : Alexey Kuleshevich <lehins@yandex.ru>
-- Stability   : experimental
-- Portability : non-portable
--
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

-- | Match vector type to array representation
type family ARepr (v :: Type -> Type) :: Type where
  ARepr VU.Vector = U
  ARepr VS.Vector = S
  ARepr VP.Vector = P
  ARepr VB.Vector = BL

-- | Match array representation to a vector type
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


-- | /O(1)/ - conversion from vector to an array with a corresponding representation. Will
-- return `Nothing` if there is a size mismatch or if some non-standard vector type is
-- supplied. Is suppplied is the boxed `Data.Vector.Vector` then it's all elements will be
-- evaluated toWHNF, therefore complexity will be /O(n)/
castFromVector :: forall v r ix e. (VG.Vector v e, Typeable v, Index ix, ARepr v ~ r)
               => Comp
               -> Sz ix -- ^ Size of the result Array
               -> v e -- ^ Source Vector
               -> Maybe (Array r ix e)
castFromVector :: 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 U ix e -> Maybe (Array U ix e)
forall (m :: * -> *) a. Monad m => a -> m a
return (Array U ix e -> Maybe (Array U ix e))
-> Array U ix e -> Maybe (Array U ix e)
forall a b. (a -> b) -> a -> b
$ UArray :: forall ix e. Comp -> Sz ix -> Vector e -> Array U ix e
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 S ix e -> Maybe (Array S ix e)
forall (m :: * -> *) a. Monad m => a -> m a
return (Array S ix e -> Maybe (Array S ix e))
-> Array S ix e -> Maybe (Array S ix e)
forall a b. (a -> b) -> a -> b
$ Sz ix -> Array S Int e -> Array S ix e
forall r ix ix' e.
(Size r, Index ix, Index ix') =>
Sz ix' -> Array r ix e -> Array r ix' e
unsafeResize Sz ix
sz (Array S Int e -> Array S ix e) -> Array S Int e -> Array S ix e
forall a b. (a -> b) -> a -> b
$ Comp -> Vector e -> Array S Int 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 P ix e -> Maybe (Array P ix e)
forall (m :: * -> *) a. Monad m => a -> m a
return (Array P ix e -> Maybe (Array P ix e))
-> Array P ix e -> Maybe (Array P ix e)
forall a b. (a -> b) -> a -> b
$ PArray :: forall ix e. Comp -> Sz ix -> Int -> ByteArray -> Array P ix e
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 BL ix e -> Maybe (Array BL ix e)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Array BL ix e -> Maybe (Array BL ix e))
-> Array BL ix e -> Maybe (Array BL ix e)
forall a b. (a -> b) -> a -> b
$ Sz ix -> Array BL Int e -> Array BL ix e
forall r ix ix' e.
(Size r, Index ix, Index ix') =>
Sz ix' -> Array r ix e -> Array r ix' e
unsafeResize Sz ix
sz (Array BL Int e -> Array BL ix e)
-> Array BL Int e -> Array BL ix e
forall a b. (a -> b) -> a -> b
$ Comp -> Array BL Int e -> Array BL Int e
forall r ix e. Strategy r => Comp -> Array r ix e -> Array r ix e
setComp Comp
comp (Array BL Int e -> Array BL Int e)
-> Array BL Int e -> Array BL Int e
forall a b. (a -> b) -> a -> b
$ Vector e -> Array BL Int e
forall a. Vector a -> Vector BL a
fromBoxedVector Vector e
bVector
    ]
{-# NOINLINE castFromVector #-}


-- | In case when resulting array representation matches the one of vector's it
-- will do a /O(1)/ - conversion using `castFromVector`, otherwise Vector elements
-- will be copied into a new array. Will throw an error if length of resulting
-- array doesn't match the source vector length.
--
-- @since 0.3.0
fromVectorM ::
     (MonadThrow m, Typeable v, VG.Vector v a, Manifest r a, Load (ARepr v) ix a, Load r ix a)
  => Comp
  -> Sz ix -- ^ Resulting size of the array
  -> v a -- ^ Source Vector
  -> m (Array r ix a)
fromVectorM :: 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 (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 (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 #-}


-- | Just like `fromVectorM`, but will throw an exception on a mismatched size.
--
-- @since 0.3.0
fromVector' ::
     (HasCallStack, Typeable v, VG.Vector v a, Load (ARepr v) ix a, Load r ix a, Manifest r a)
  => Comp
  -> Sz ix -- ^ Resulting size of the array
  -> v a -- ^ Source Vector
  -> Array r ix a
fromVector' :: 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' #-}

-- | /O(1)/ - conversion from `Mutable` array to a corresponding vector. Will
-- return `Nothing` only if source array representation was not one of `B`, `N`,
-- `P`, `S` or `U`.
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 -> 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
         Vector e -> Maybe (Vector e)
forall (m :: * -> *) a. Monad m => a -> m a
return (Vector e -> Maybe (Vector e)) -> Vector e -> Maybe (Vector 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
         Vector e -> Maybe (Vector e)
forall (m :: * -> *) a. Monad m => a -> m a
return (Vector e -> Maybe (Vector e)) -> Vector e -> Maybe (Vector 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
         Vector e -> Maybe (Vector e)
forall (m :: * -> *) a. Monad m => a -> m a
return (Vector e -> Maybe (Vector e)) -> Vector e -> Maybe (Vector 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
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
         Vector e -> Maybe (Vector e)
forall (m :: * -> *) a. Monad m => a -> m a
return (Vector e -> Maybe (Vector e)) -> Vector e -> Maybe (Vector 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
         Vector e -> Maybe (Vector e)
forall (m :: * -> *) a. Monad m => a -> m a
return (Vector e -> Maybe (Vector e)) -> Vector e -> Maybe (Vector 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
         Vector e -> Maybe (Vector e)
forall (m :: * -> *) a. Monad m => a -> m a
return (Vector e -> Maybe (Vector e)) -> Vector e -> Maybe (Vector 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 #-}


-- | Convert an array into a vector. Will perform a cast if resulting vector is
-- of compatible representation, otherwise memory copy will occur.
--
-- ==== __Examples__
--
-- In this example a `S`torable Array is created and then casted into a Storable
-- `VS.Vector` in costant time:
--
-- >>> import Data.Massiv.Array as A
-- >>> import Data.Massiv.Array.Manifest.Vector (toVector)
-- >>> import qualified Data.Vector.Storable as VS
-- >>> toVector (makeArrayR S Par (Sz2 5 6) (\(i :. j) -> i + j)) :: VS.Vector Int
-- [0,1,2,3,4,5,1,2,3,4,5,6,2,3,4,5,6,7,3,4,5,6,7,8,4,5,6,7,8,9]
--
-- While in this example `S`torable Array will first be converted into `U`nboxed
-- representation in `Par`allel and only after that will be coverted into Unboxed
-- `VU.Vector` in constant time.
--
-- >>> import qualified Data.Vector.Unboxed as VU
-- >>> toVector (makeArrayR S Par (Sz2 5 6) (\(i :. j) -> i + j)) :: VU.Vector Int
-- [0,1,2,3,4,5,1,2,3,4,5,6,2,3,4,5,6,7,3,4,5,6,7,8,4,5,6,7,8,9]
--
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 :: 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
size Array r ix e
arr)) (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 #-}