module Data.Array.Repa.Repr.Unboxed
        ( U, U.Unbox, Array (..)
        , computeUnboxedS, computeUnboxedP
        , fromListUnboxed
        , fromUnboxed, toUnboxed
        
        , zip,   zip3,   zip4,   zip5,   zip6
        , unzip, unzip3, unzip4, unzip5, unzip6)
where
import Data.Array.Repa.Shape            as R
import Data.Array.Repa.Base             as R
import Data.Array.Repa.Eval             as R
import Data.Array.Repa.Repr.Delayed     as R
import qualified Data.Vector.Unboxed              as U
import qualified Data.Vector.Unboxed.Mutable      as UM
import Control.Monad
import Prelude hiding (zip, zip3, unzip, unzip3)

-- | Unboxed arrays are represented as unboxed vectors.
--
--   The implementation uses @Data.Vector.Unboxed@ which is based on type
--   families and picks an efficient, specialised representation for every
--   element type. In particular, unboxed vectors of pairs are represented
--   as pairs of unboxed vectors.
--   This is the most efficient representation for numerical data.
--
data U

-- | Read elements from an unboxed vector array.
instance U.Unbox a => Source U a where
 data Array U sh a
        = AUnboxed !sh !(U.Vector a)

 linearIndex :: Array U sh a -> Int -> a
linearIndex (AUnboxed _ vec) Int
ix
        = Vector a
vec Vector a -> Int -> a
forall a. Unbox a => Vector a -> Int -> a
U.! Int
ix
 {-# INLINE linearIndex #-}

 unsafeLinearIndex :: Array U sh a -> Int -> a
unsafeLinearIndex (AUnboxed _ vec) Int
ix
        = Vector a
vec Vector a -> Int -> a
forall a. Unbox a => Vector a -> Int -> a
`U.unsafeIndex` Int
ix
 {-# INLINE unsafeLinearIndex #-}

 extent :: Array U sh a -> sh
extent (AUnboxed sh _)
        = sh
sh
 {-# INLINE extent #-}

 deepSeqArray :: Array U sh a -> b -> b
deepSeqArray (AUnboxed sh vec) b
x 
  = sh
sh sh -> b -> b
forall sh a. Shape sh => sh -> a -> a
`deepSeq` Vector a
vec Vector a -> b -> b
`seq` b
x
 {-# INLINE deepSeqArray #-}


deriving instance (Show sh, Show e, U.Unbox e)
        => Show (Array U sh e)

deriving instance (Read sh, Read e, U.Unbox e)
        => Read (Array U sh e)


-- Fill -----------------------------------------------------------------------
-- | Filling of unboxed vector arrays.
instance U.Unbox e => Target U e where
 data MVec U e 
  = UMVec (UM.IOVector e)

 newMVec :: Int -> IO (MVec U e)
newMVec Int
n
  = (IOVector e -> MVec U e) -> IO (IOVector e) -> IO (MVec U e)
forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
liftM IOVector e -> MVec U e
forall e. IOVector e -> MVec U e
UMVec (Int -> IO (MVector (PrimState IO) e)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
UM.new Int
n)
 {-# INLINE newMVec #-}

 unsafeWriteMVec :: MVec U e -> Int -> e -> IO ()
unsafeWriteMVec (UMVec v) Int
ix
  = MVector (PrimState IO) e -> Int -> e -> IO ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
UM.unsafeWrite IOVector e
MVector (PrimState IO) e
v Int
ix
 {-# INLINE unsafeWriteMVec #-}

 unsafeFreezeMVec :: sh -> MVec U e -> IO (Array U sh e)
unsafeFreezeMVec sh
sh (UMVec mvec)     
  = do  Vector e
vec     <- MVector (PrimState IO) e -> IO (Vector e)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
U.unsafeFreeze IOVector e
MVector (PrimState IO) e
mvec
        Array U sh e -> IO (Array U sh e)
forall (m :: * -> *) a. Monad m => a -> m a
return  (Array U sh e -> IO (Array U sh e))
-> Array U sh e -> IO (Array U sh e)
forall a b. (a -> b) -> a -> b
$  sh -> Vector e -> Array U sh e
forall sh a. sh -> Vector a -> Array U sh a
AUnboxed sh
sh Vector e
vec
 {-# INLINE unsafeFreezeMVec #-}

 deepSeqMVec :: MVec U e -> a -> a
deepSeqMVec (UMVec vec) a
x
  = IOVector e
vec IOVector e -> a -> a
`seq` a
x
 {-# INLINE deepSeqMVec #-}

 touchMVec :: MVec U e -> IO ()
touchMVec MVec U e
_ 
  = () -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
 {-# INLINE touchMVec #-}


-- Conversions ----------------------------------------------------------------
-- | Sequential computation of array elements..
--
--   * This is an alias for `computeS` with a more specific type.
--
computeUnboxedS
        :: (Load r1 sh e, U.Unbox e)
        => Array r1 sh e -> Array U sh e
computeUnboxedS :: Array r1 sh e -> Array U sh e
computeUnboxedS = Array r1 sh e -> Array U sh e
forall r1 sh e r2.
(Load r1 sh e, Target r2 e) =>
Array r1 sh e -> Array r2 sh e
computeS
{-# INLINE computeUnboxedS #-}


-- | Parallel computation of array elements.
--
--   * This is an alias for `computeP` with a more specific type.
--
computeUnboxedP
        :: (Load r1 sh e, Monad m, U.Unbox e)
        => Array r1 sh e -> m (Array U sh e)
computeUnboxedP :: Array r1 sh e -> m (Array U sh e)
computeUnboxedP = Array r1 sh e -> m (Array U sh e)
forall r1 sh e r2 (m :: * -> *).
(Load r1 sh e, Target r2 e, Source r2 e, Monad m) =>
Array r1 sh e -> m (Array r2 sh e)
computeP
{-# INLINE computeUnboxedP #-}


-- | O(n). Convert a list to an unboxed vector array.
-- 
--   * This is an alias for `fromList` with a more specific type.
--
fromListUnboxed
        :: (Shape sh, U.Unbox a)
        => sh -> [a] -> Array U sh a
fromListUnboxed :: sh -> [a] -> Array U sh a
fromListUnboxed = sh -> [a] -> Array U sh a
forall sh r e. (Shape sh, Target r e) => sh -> [e] -> Array r sh e
R.fromList
{-# INLINE fromListUnboxed #-}


-- | O(1). Wrap an unboxed vector as an array.
fromUnboxed :: sh -> U.Vector e -> Array U sh e
fromUnboxed :: sh -> Vector e -> Array U sh e
fromUnboxed sh
sh Vector e
vec
        = sh -> Vector e -> Array U sh e
forall sh a. sh -> Vector a -> Array U sh a
AUnboxed sh
sh Vector e
vec
{-# INLINE fromUnboxed #-}


-- | O(1). Unpack an unboxed vector from an array.
toUnboxed :: Array U sh e -> U.Vector e
toUnboxed :: Array U sh e -> Vector e
toUnboxed (AUnboxed _ vec)
        = Vector e
vec
{-# INLINE toUnboxed #-}


-- Zip ------------------------------------------------------------------------
-- | O(1). Zip some unboxed arrays.
--         The shapes must be identical else `error`.
zip     :: (Shape sh, U.Unbox a, U.Unbox b)
        => Array U sh a -> Array U sh b
        -> Array U sh (a, b)
zip :: Array U sh a -> Array U sh b -> Array U sh (a, b)
zip (AUnboxed sh1 vec1) (AUnboxed sh2 vec2)
 | sh
sh1 sh -> sh -> Bool
forall a. Eq a => a -> a -> Bool
/= sh
sh2   = String -> Array U sh (a, b)
forall a. HasCallStack => String -> a
error String
"Repa: zip array shapes not identical"
 | Bool
otherwise    = sh -> Vector (a, b) -> Array U sh (a, b)
forall sh a. sh -> Vector a -> Array U sh a
AUnboxed sh
sh1 (Vector a -> Vector b -> Vector (a, b)
forall a b.
(Unbox a, Unbox b) =>
Vector a -> Vector b -> Vector (a, b)
U.zip Vector a
vec1 Vector b
vec2)
{-# INLINE zip #-}


-- | O(1). Zip some unboxed arrays.
--         The shapes must be identical else `error`.
zip3    :: (Shape sh, U.Unbox a, U.Unbox b, U.Unbox c)
        => Array U sh a -> Array U sh b -> Array U sh c
        -> Array U sh (a, b, c)
zip3 :: Array U sh a
-> Array U sh b -> Array U sh c -> Array U sh (a, b, c)
zip3 (AUnboxed sh1 vec1) (AUnboxed sh2 vec2) (AUnboxed sh3 vec3)
 | sh
sh1 sh -> sh -> Bool
forall a. Eq a => a -> a -> Bool
/= sh
sh2 Bool -> Bool -> Bool
|| sh
sh1 sh -> sh -> Bool
forall a. Eq a => a -> a -> Bool
/= sh
sh3
 = String -> Array U sh (a, b, c)
forall a. HasCallStack => String -> a
error String
"Repa: zip array shapes not identical"
 | Bool
otherwise    = sh -> Vector (a, b, c) -> Array U sh (a, b, c)
forall sh a. sh -> Vector a -> Array U sh a
AUnboxed sh
sh1 (Vector a -> Vector b -> Vector c -> Vector (a, b, c)
forall a b c.
(Unbox a, Unbox b, Unbox c) =>
Vector a -> Vector b -> Vector c -> Vector (a, b, c)
U.zip3 Vector a
vec1 Vector b
vec2 Vector c
vec3)
{-# INLINE zip3 #-}


-- | O(1). Zip some unboxed arrays.
--         The shapes must be identical else `error`.
zip4    :: (Shape sh, U.Unbox a, U.Unbox b, U.Unbox c, U.Unbox d)
        => Array U sh a -> Array U sh b -> Array U sh c -> Array U sh d
        -> Array U sh (a, b, c, d)
zip4 :: Array U sh a
-> Array U sh b
-> Array U sh c
-> Array U sh d
-> Array U sh (a, b, c, d)
zip4 (AUnboxed sh1 vec1) (AUnboxed sh2 vec2) (AUnboxed sh3 vec3) (AUnboxed sh4 vec4)
 | sh
sh1 sh -> sh -> Bool
forall a. Eq a => a -> a -> Bool
/= sh
sh2 Bool -> Bool -> Bool
|| sh
sh1 sh -> sh -> Bool
forall a. Eq a => a -> a -> Bool
/= sh
sh3 Bool -> Bool -> Bool
|| sh
sh1 sh -> sh -> Bool
forall a. Eq a => a -> a -> Bool
/= sh
sh4
 = String -> Array U sh (a, b, c, d)
forall a. HasCallStack => String -> a
error String
"Repa: zip array shapes not identical"
 | Bool
otherwise    = sh -> Vector (a, b, c, d) -> Array U sh (a, b, c, d)
forall sh a. sh -> Vector a -> Array U sh a
AUnboxed sh
sh1 (Vector a -> Vector b -> Vector c -> Vector d -> Vector (a, b, c, d)
forall a b c d.
(Unbox a, Unbox b, Unbox c, Unbox d) =>
Vector a -> Vector b -> Vector c -> Vector d -> Vector (a, b, c, d)
U.zip4 Vector a
vec1 Vector b
vec2 Vector c
vec3 Vector d
vec4)
{-# INLINE zip4 #-}


-- | O(1). Zip some unboxed arrays.
--         The shapes must be identical else `error`.
zip5    :: (Shape sh, U.Unbox a, U.Unbox b, U.Unbox c, U.Unbox d, U.Unbox e)
        => Array U sh a -> Array U sh b -> Array U sh c -> Array U sh d -> Array U sh e
        -> Array U sh (a, b, c, d, e)
zip5 :: Array U sh a
-> Array U sh b
-> Array U sh c
-> Array U sh d
-> Array U sh e
-> Array U sh (a, b, c, d, e)
zip5 (AUnboxed sh1 vec1) (AUnboxed sh2 vec2) (AUnboxed sh3 vec3) (AUnboxed sh4 vec4) (AUnboxed sh5 vec5)
 | sh
sh1 sh -> sh -> Bool
forall a. Eq a => a -> a -> Bool
/= sh
sh2 Bool -> Bool -> Bool
|| sh
sh1 sh -> sh -> Bool
forall a. Eq a => a -> a -> Bool
/= sh
sh3 Bool -> Bool -> Bool
|| sh
sh1 sh -> sh -> Bool
forall a. Eq a => a -> a -> Bool
/= sh
sh4 Bool -> Bool -> Bool
|| sh
sh1 sh -> sh -> Bool
forall a. Eq a => a -> a -> Bool
/= sh
sh5
 = String -> Array U sh (a, b, c, d, e)
forall a. HasCallStack => String -> a
error String
"Repa: zip array shapes not identical"
 | Bool
otherwise    = sh -> Vector (a, b, c, d, e) -> Array U sh (a, b, c, d, e)
forall sh a. sh -> Vector a -> Array U sh a
AUnboxed sh
sh1 (Vector a
-> Vector b
-> Vector c
-> Vector d
-> Vector e
-> Vector (a, b, c, d, e)
forall a b c d e.
(Unbox a, Unbox b, Unbox c, Unbox d, Unbox e) =>
Vector a
-> Vector b
-> Vector c
-> Vector d
-> Vector e
-> Vector (a, b, c, d, e)
U.zip5 Vector a
vec1 Vector b
vec2 Vector c
vec3 Vector d
vec4 Vector e
vec5)
{-# INLINE zip5 #-}


-- | O(1). Zip some unboxed arrays.
--         The shapes must be identical else `error`.
zip6    :: (Shape sh, U.Unbox a, U.Unbox b, U.Unbox c, U.Unbox d, U.Unbox e, U.Unbox f)
        => Array U sh a -> Array U sh b -> Array U sh c -> Array U sh d -> Array U sh e -> Array U sh f
        -> Array U sh (a, b, c, d, e, f)
zip6 :: Array U sh a
-> Array U sh b
-> Array U sh c
-> Array U sh d
-> Array U sh e
-> Array U sh f
-> Array U sh (a, b, c, d, e, f)
zip6 (AUnboxed sh1 vec1) (AUnboxed sh2 vec2) (AUnboxed sh3 vec3) (AUnboxed sh4 vec4) (AUnboxed sh5 vec5) (AUnboxed sh6 vec6)
 | sh
sh1 sh -> sh -> Bool
forall a. Eq a => a -> a -> Bool
/= sh
sh2 Bool -> Bool -> Bool
|| sh
sh1 sh -> sh -> Bool
forall a. Eq a => a -> a -> Bool
/= sh
sh3 Bool -> Bool -> Bool
|| sh
sh1 sh -> sh -> Bool
forall a. Eq a => a -> a -> Bool
/= sh
sh4 Bool -> Bool -> Bool
|| sh
sh1 sh -> sh -> Bool
forall a. Eq a => a -> a -> Bool
/= sh
sh5 Bool -> Bool -> Bool
|| sh
sh1 sh -> sh -> Bool
forall a. Eq a => a -> a -> Bool
/= sh
sh6
 = String -> Array U sh (a, b, c, d, e, f)
forall a. HasCallStack => String -> a
error String
"Repa: zip array shapes not identical"
 | Bool
otherwise    = sh -> Vector (a, b, c, d, e, f) -> Array U sh (a, b, c, d, e, f)
forall sh a. sh -> Vector a -> Array U sh a
AUnboxed sh
sh1 (Vector a
-> Vector b
-> Vector c
-> Vector d
-> Vector e
-> Vector f
-> Vector (a, b, c, d, e, f)
forall a b c d e f.
(Unbox a, Unbox b, Unbox c, Unbox d, Unbox e, Unbox f) =>
Vector a
-> Vector b
-> Vector c
-> Vector d
-> Vector e
-> Vector f
-> Vector (a, b, c, d, e, f)
U.zip6 Vector a
vec1 Vector b
vec2 Vector c
vec3 Vector d
vec4 Vector e
vec5 Vector f
vec6)
{-# INLINE zip6 #-}
 

-- Unzip ----------------------------------------------------------------------
-- | O(1). Unzip an unboxed array.
unzip   :: (U.Unbox a, U.Unbox b)
        => Array U sh (a, b)
        -> (Array U sh a, Array U sh b)
unzip :: Array U sh (a, b) -> (Array U sh a, Array U sh b)
unzip (AUnboxed sh vec)
 = let  (Vector a
as, Vector b
bs)        = Vector (a, b) -> (Vector a, Vector b)
forall a b.
(Unbox a, Unbox b) =>
Vector (a, b) -> (Vector a, Vector b)
U.unzip Vector (a, b)
vec
   in   (sh -> Vector a -> Array U sh a
forall sh a. sh -> Vector a -> Array U sh a
AUnboxed sh
sh Vector a
as, sh -> Vector b -> Array U sh b
forall sh a. sh -> Vector a -> Array U sh a
AUnboxed sh
sh Vector b
bs)
{-# INLINE unzip #-}


-- | O(1). Unzip an unboxed array.
unzip3   :: (U.Unbox a, U.Unbox b, U.Unbox c)
        => Array U sh (a, b, c)
        -> (Array U sh a, Array U sh b, Array U sh c)
unzip3 :: Array U sh (a, b, c) -> (Array U sh a, Array U sh b, Array U sh c)
unzip3 (AUnboxed sh vec)
 = let  (Vector a
as, Vector b
bs, Vector c
cs) = Vector (a, b, c) -> (Vector a, Vector b, Vector c)
forall a b c.
(Unbox a, Unbox b, Unbox c) =>
Vector (a, b, c) -> (Vector a, Vector b, Vector c)
U.unzip3 Vector (a, b, c)
vec
   in   (sh -> Vector a -> Array U sh a
forall sh a. sh -> Vector a -> Array U sh a
AUnboxed sh
sh Vector a
as, sh -> Vector b -> Array U sh b
forall sh a. sh -> Vector a -> Array U sh a
AUnboxed sh
sh Vector b
bs, sh -> Vector c -> Array U sh c
forall sh a. sh -> Vector a -> Array U sh a
AUnboxed sh
sh Vector c
cs)
{-# INLINE unzip3 #-}


-- | O(1). Unzip an unboxed array.
unzip4   :: (U.Unbox a, U.Unbox b, U.Unbox c, U.Unbox d)
        => Array U sh (a, b, c, d)
        -> (Array U sh a, Array U sh b, Array U sh c, Array U sh d)
unzip4 :: Array U sh (a, b, c, d)
-> (Array U sh a, Array U sh b, Array U sh c, Array U sh d)
unzip4 (AUnboxed sh vec)
 = let  (Vector a
as, Vector b
bs, Vector c
cs, Vector d
ds) = Vector (a, b, c, d) -> (Vector a, Vector b, Vector c, Vector d)
forall a b c d.
(Unbox a, Unbox b, Unbox c, Unbox d) =>
Vector (a, b, c, d) -> (Vector a, Vector b, Vector c, Vector d)
U.unzip4 Vector (a, b, c, d)
vec
   in   (sh -> Vector a -> Array U sh a
forall sh a. sh -> Vector a -> Array U sh a
AUnboxed sh
sh Vector a
as, sh -> Vector b -> Array U sh b
forall sh a. sh -> Vector a -> Array U sh a
AUnboxed sh
sh Vector b
bs, sh -> Vector c -> Array U sh c
forall sh a. sh -> Vector a -> Array U sh a
AUnboxed sh
sh Vector c
cs, sh -> Vector d -> Array U sh d
forall sh a. sh -> Vector a -> Array U sh a
AUnboxed sh
sh Vector d
ds)
{-# INLINE unzip4 #-}


-- | O(1). Unzip an unboxed array.
unzip5   :: (U.Unbox a, U.Unbox b, U.Unbox c, U.Unbox d, U.Unbox e)
        => Array U sh (a, b, c, d, e)
        -> (Array U sh a, Array U sh b, Array U sh c, Array U sh d, Array U sh e)
unzip5 :: Array U sh (a, b, c, d, e)
-> (Array U sh a, Array U sh b, Array U sh c, Array U sh d,
    Array U sh e)
unzip5 (AUnboxed sh vec)
 = let  (Vector a
as, Vector b
bs, Vector c
cs, Vector d
ds, Vector e
es) = Vector (a, b, c, d, e)
-> (Vector a, Vector b, Vector c, Vector d, Vector e)
forall a b c d e.
(Unbox a, Unbox b, Unbox c, Unbox d, Unbox e) =>
Vector (a, b, c, d, e)
-> (Vector a, Vector b, Vector c, Vector d, Vector e)
U.unzip5 Vector (a, b, c, d, e)
vec
   in   (sh -> Vector a -> Array U sh a
forall sh a. sh -> Vector a -> Array U sh a
AUnboxed sh
sh Vector a
as, sh -> Vector b -> Array U sh b
forall sh a. sh -> Vector a -> Array U sh a
AUnboxed sh
sh Vector b
bs, sh -> Vector c -> Array U sh c
forall sh a. sh -> Vector a -> Array U sh a
AUnboxed sh
sh Vector c
cs, sh -> Vector d -> Array U sh d
forall sh a. sh -> Vector a -> Array U sh a
AUnboxed sh
sh Vector d
ds, sh -> Vector e -> Array U sh e
forall sh a. sh -> Vector a -> Array U sh a
AUnboxed sh
sh Vector e
es)
{-# INLINE unzip5 #-}


-- | O(1). Unzip an unboxed array.
unzip6  :: (U.Unbox a, U.Unbox b, U.Unbox c, U.Unbox d, U.Unbox e, U.Unbox f)
        => Array U sh (a, b, c, d, e, f)
        -> (Array U sh a, Array U sh b, Array U sh c, Array U sh d, Array U sh e, Array U sh f)
unzip6 :: Array U sh (a, b, c, d, e, f)
-> (Array U sh a, Array U sh b, Array U sh c, Array U sh d,
    Array U sh e, Array U sh f)
unzip6 (AUnboxed sh vec)
 = let  (Vector a
as, Vector b
bs, Vector c
cs, Vector d
ds, Vector e
es, Vector f
fs) = Vector (a, b, c, d, e, f)
-> (Vector a, Vector b, Vector c, Vector d, Vector e, Vector f)
forall a b c d e f.
(Unbox a, Unbox b, Unbox c, Unbox d, Unbox e, Unbox f) =>
Vector (a, b, c, d, e, f)
-> (Vector a, Vector b, Vector c, Vector d, Vector e, Vector f)
U.unzip6 Vector (a, b, c, d, e, f)
vec
   in   (sh -> Vector a -> Array U sh a
forall sh a. sh -> Vector a -> Array U sh a
AUnboxed sh
sh Vector a
as, sh -> Vector b -> Array U sh b
forall sh a. sh -> Vector a -> Array U sh a
AUnboxed sh
sh Vector b
bs, sh -> Vector c -> Array U sh c
forall sh a. sh -> Vector a -> Array U sh a
AUnboxed sh
sh Vector c
cs, sh -> Vector d -> Array U sh d
forall sh a. sh -> Vector a -> Array U sh a
AUnboxed sh
sh Vector d
ds, sh -> Vector e -> Array U sh e
forall sh a. sh -> Vector a -> Array U sh a
AUnboxed sh
sh Vector e
es, sh -> Vector f -> Array U sh f
forall sh a. sh -> Vector a -> Array U sh a
AUnboxed sh
sh Vector f
fs)
{-# INLINE unzip6 #-}