{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UndecidableInstances #-}

-- |
-- Module      : Data.Massiv.Array.Manifest.Unboxed
-- Copyright   : (c) Alexey Kuleshevich 2018-2022
-- License     : BSD3
-- Maintainer  : Alexey Kuleshevich <lehins@yandex.ru>
-- Stability   : experimental
-- Portability : non-portable
module Data.Massiv.Array.Manifest.Unboxed (
  U (..),
  Unbox,
  Array (..),
  MArray (..),
  toUnboxedVector,
  toUnboxedMVector,
  fromUnboxedVector,
  fromUnboxedMVector,
) where

import Control.DeepSeq (NFData (..), deepseq)
import Control.Monad.Primitive (stToPrim)
import Data.Massiv.Array.Delayed.Pull (D, compareArrays, eqArrays)
import Data.Massiv.Array.Manifest.Internal
import Data.Massiv.Array.Manifest.List as A
import Data.Massiv.Array.Mutable
import Data.Massiv.Core.Common
import Data.Massiv.Core.List
import Data.Massiv.Core.Operations
import Data.Massiv.Vector.Stream as S (isteps, steps)
import qualified Data.Vector.Generic.Mutable as VGM
import Data.Vector.Unboxed (Unbox)
import qualified Data.Vector.Unboxed as VU
import qualified Data.Vector.Unboxed.Mutable as MVU
import GHC.Exts as GHC (IsList (..))
import System.IO.Unsafe (unsafePerformIO)
import Prelude hiding (mapM)

-- | Representation for `Unbox`ed elements
data U = U deriving (Int -> U -> ShowS
[U] -> ShowS
U -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [U] -> ShowS
$cshowList :: [U] -> ShowS
show :: U -> String
$cshow :: U -> String
showsPrec :: Int -> U -> ShowS
$cshowsPrec :: Int -> U -> ShowS
Show)

data instance Array U ix e = UArray
  { forall ix e. Array U ix e -> Comp
uComp :: !Comp
  , forall ix e. Array U ix e -> Sz ix
uSize :: !(Sz ix)
  , forall ix e. Array U ix e -> Vector e
uData :: !(VU.Vector e)
  }
data instance MArray s U ix e = MUArray !(Sz ix) !(VU.MVector s e)

instance (Ragged L ix e, Show e, Unbox e) => Show (Array U ix e) where
  showsPrec :: Int -> Array U ix e -> ShowS
showsPrec = forall r r' ix e.
(Ragged L ix e, Load r ix e, Load r' ix e, Source r' e, Show e) =>
(Array r ix e -> Array r' ix e) -> Int -> Array r ix e -> ShowS
showsArrayPrec forall a. a -> a
id
  showList :: [Array U ix e] -> ShowS
showList = forall arr. Show arr => [arr] -> ShowS
showArrayList

instance NFData ix => NFData (Array U ix e) where
  rnf :: Array U ix e -> ()
rnf (UArray Comp
c Sz ix
sz Vector e
v) = Comp
c forall a b. NFData a => a -> b -> b
`deepseq` Sz ix
sz forall a b. NFData a => a -> b -> b
`deepseq` Vector e
v forall a b. NFData a => a -> b -> b
`deepseq` ()
  {-# INLINE rnf #-}

instance NFData ix => NFData (MArray s U ix e) where
  rnf :: MArray s U ix e -> ()
rnf (MUArray Sz ix
sz MVector s e
mv) = Sz ix
sz forall a b. NFData a => a -> b -> b
`deepseq` MVector s e
mv forall a b. NFData a => a -> b -> b
`deepseq` ()
  {-# INLINE rnf #-}

instance Strategy U where
  getComp :: forall ix e. Array U ix e -> Comp
getComp = forall ix e. Array U ix e -> Comp
uComp
  {-# INLINE getComp #-}
  setComp :: forall ix e. Comp -> Array U ix e -> Array U ix e
setComp Comp
c Array U ix e
arr = Array U ix e
arr{uComp :: Comp
uComp = Comp
c}
  {-# INLINE setComp #-}
  repr :: U
repr = U
U

instance (Unbox e, Eq e, Index ix) => Eq (Array U ix e) where
  == :: Array U ix e -> Array U ix e -> Bool
(==) = forall ix r1 e1 r2 e2.
(Index ix, Source r1 e1, Source r2 e2) =>
(e1 -> e2 -> Bool) -> Array r1 ix e1 -> Array r2 ix e2 -> Bool
eqArrays forall a. Eq a => a -> a -> Bool
(==)
  {-# INLINE (==) #-}

instance (Unbox e, Ord e, Index ix) => Ord (Array U ix e) where
  compare :: Array U ix e -> Array U ix e -> Ordering
compare = forall ix r1 e1 r2 e2.
(Index ix, Source r1 e1, Source r2 e2) =>
(e1 -> e2 -> Ordering)
-> Array r1 ix e1 -> Array r2 ix e2 -> Ordering
compareArrays forall a. Ord a => a -> a -> Ordering
compare
  {-# INLINE compare #-}

instance Unbox e => Source U e where
  unsafeLinearIndex :: forall ix. Index ix => Array U ix e -> Int -> e
unsafeLinearIndex (UArray Comp
_ Sz ix
_ Vector e
v) =
    forall a ix e.
String -> (a -> Sz ix) -> (a -> ix -> e) -> a -> ix -> e
indexAssert String
"U.unsafeLinearIndex" (forall ix. ix -> Sz ix
SafeSz forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Unbox a => Vector a -> Int
VU.length) forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector e
v
  {-# INLINE unsafeLinearIndex #-}

  unsafeOuterSlice :: forall ix.
(Index ix, Index (Lower ix)) =>
Array U ix e -> Sz (Lower ix) -> Int -> Array U (Lower ix) e
unsafeOuterSlice (UArray Comp
c Sz ix
_ Vector e
v) Sz (Lower ix)
szL Int
i =
    let k :: Int
k = forall ix. Index ix => Sz ix -> Int
totalElem Sz (Lower ix)
szL
     in forall ix e. Comp -> Sz ix -> Vector e -> Array U ix e
UArray Comp
c Sz (Lower ix)
szL forall a b. (a -> b) -> a -> b
$ forall a. Unbox a => Int -> Int -> Vector a -> Vector a
VU.unsafeSlice (Int
i forall a. Num a => a -> a -> a
* Int
k) Int
k Vector e
v
  {-# INLINE unsafeOuterSlice #-}

  unsafeLinearSlice :: forall ix. Index ix => Int -> Sz1 -> Array U ix e -> Array U Int e
unsafeLinearSlice Int
i Sz1
k (UArray Comp
c Sz ix
_ Vector e
v) = forall ix e. Comp -> Sz ix -> Vector e -> Array U ix e
UArray Comp
c Sz1
k forall a b. (a -> b) -> a -> b
$ forall a. Unbox a => Int -> Int -> Vector a -> Vector a
VU.unsafeSlice Int
i (forall ix. Sz ix -> ix
unSz Sz1
k) Vector e
v
  {-# INLINE unsafeLinearSlice #-}

instance Index ix => Shape U ix where
  maxLinearSize :: forall e. Array U ix e -> Maybe Sz1
maxLinearSize = forall a. a -> Maybe a
Just forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix. ix -> Sz ix
SafeSz forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix r e. (Index ix, Size r) => Array r ix e -> Int
elemsCount
  {-# INLINE maxLinearSize #-}

instance Size U where
  size :: forall ix e. Array U ix e -> Sz ix
size = forall ix e. Array U ix e -> Sz ix
uSize
  {-# INLINE size #-}
  unsafeResize :: forall ix ix' e.
(Index ix, Index ix') =>
Sz ix' -> Array U ix e -> Array U ix' e
unsafeResize !Sz ix'
sz !Array U ix e
arr = Array U ix e
arr{uSize :: Sz ix'
uSize = Sz ix'
sz}
  {-# INLINE unsafeResize #-}

instance (Unbox e, Index ix) => Load U ix e where
  makeArray :: Comp -> Sz ix -> (ix -> e) -> Array U ix e
makeArray Comp
comp Sz ix
sz ix -> e
f = forall r ix e r'.
(Manifest r e, Load r' ix e) =>
Array r' ix e -> Array r ix e
compute (forall r ix e.
Load r ix e =>
Comp -> Sz ix -> (ix -> e) -> Array r ix e
makeArray Comp
comp Sz ix
sz ix -> e
f :: Array D ix e)
  {-# INLINE makeArray #-}

  makeArrayLinear :: Comp -> Sz ix -> (Int -> e) -> Array U ix e
makeArrayLinear !Comp
comp !Sz ix
sz Int -> e
f = forall a. IO a -> a
unsafePerformIO forall a b. (a -> b) -> a -> b
$ forall r ix e (m :: * -> *).
(MonadUnliftIO m, Manifest r e, Index ix) =>
Comp -> Sz ix -> (Int -> m e) -> m (Array r ix e)
generateArrayLinear Comp
comp Sz ix
sz (forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> e
f)
  {-# INLINE makeArrayLinear #-}

  replicate :: Comp -> Sz ix -> e -> Array U ix e
replicate Comp
comp !Sz ix
sz !e
e = forall a. (forall s. ST s a) -> a
runST (forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
Sz ix -> e -> m (MArray (PrimState m) r ix e)
newMArray Sz ix
sz e
e forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
Comp -> MArray (PrimState m) r ix e -> m (Array r ix e)
unsafeFreeze Comp
comp)
  {-# INLINE replicate #-}

  iterArrayLinearST_ :: forall s.
Scheduler s () -> Array U ix e -> (Int -> e -> ST s ()) -> ST s ()
iterArrayLinearST_ !Scheduler s ()
scheduler !Array U ix e
arr =
    forall s (m :: * -> *) b.
MonadPrimBase s m =>
Scheduler s () -> Int -> (Int -> b) -> (Int -> b -> m ()) -> m ()
splitLinearlyWith_ Scheduler s ()
scheduler (forall ix r e. (Index ix, Size r) => Array r ix e -> Int
elemsCount Array U ix e
arr) (forall r e ix. (Source r e, Index ix) => Array r ix e -> Int -> e
unsafeLinearIndex Array U ix e
arr)
  {-# INLINE iterArrayLinearST_ #-}

instance (Unbox e, Index ix) => StrideLoad U ix e

instance Unbox e => Manifest U e where
  unsafeLinearIndexM :: forall ix. Index ix => Array U ix e -> Int -> e
unsafeLinearIndexM (UArray Comp
_ Sz ix
_ Vector e
v) =
    forall a ix e.
String -> (a -> Sz ix) -> (a -> ix -> e) -> a -> ix -> e
indexAssert String
"S.unsafeLinearIndexM" (forall ix. ix -> Sz ix
SafeSz forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Unbox a => Vector a -> Int
VU.length) forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector e
v
  {-# INLINE unsafeLinearIndexM #-}

  sizeOfMArray :: forall ix s. Index ix => MArray s U ix e -> Sz ix
sizeOfMArray (MUArray Sz ix
sz MVector s e
_) = Sz ix
sz
  {-# INLINE sizeOfMArray #-}

  unsafeResizeMArray :: forall ix' ix s.
(Index ix', Index ix) =>
Sz ix' -> MArray s U ix e -> MArray s U ix' e
unsafeResizeMArray Sz ix'
sz (MUArray Sz ix
_ MVector s e
mv) = forall s ix e. Sz ix -> MVector s e -> MArray s U ix e
MUArray Sz ix'
sz MVector s e
mv
  {-# INLINE unsafeResizeMArray #-}

  unsafeLinearSliceMArray :: forall ix s.
Index ix =>
Int -> Sz1 -> MArray s U ix e -> MVector s U e
unsafeLinearSliceMArray Int
i Sz1
k (MUArray Sz ix
_ MVector s e
mv) = forall s ix e. Sz ix -> MVector s e -> MArray s U ix e
MUArray Sz1
k forall a b. (a -> b) -> a -> b
$ forall a s. Unbox a => Int -> Int -> MVector s a -> MVector s a
MVU.unsafeSlice Int
i (forall ix. Sz ix -> ix
unSz Sz1
k) MVector s e
mv
  {-# INLINE unsafeLinearSliceMArray #-}

  unsafeThaw :: forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
Array U ix e -> m (MArray (PrimState m) U ix e)
unsafeThaw (UArray Comp
_ Sz ix
sz Vector e
v) = forall s ix e. Sz ix -> MVector s e -> MArray s U ix e
MUArray Sz ix
sz forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
Vector a -> m (MVector (PrimState m) a)
VU.unsafeThaw Vector e
v
  {-# INLINE unsafeThaw #-}

  unsafeFreeze :: forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
Comp -> MArray (PrimState m) U ix e -> m (Array U ix e)
unsafeFreeze Comp
comp (MUArray Sz ix
sz MVector (PrimState m) e
v) = forall ix e. Comp -> Sz ix -> Vector e -> Array U ix e
UArray Comp
comp Sz ix
sz forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze MVector (PrimState m) e
v
  {-# INLINE unsafeFreeze #-}

  unsafeNew :: forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
Sz ix -> m (MArray (PrimState m) U ix e)
unsafeNew Sz ix
sz = forall s ix e. Sz ix -> MVector s e -> MArray s U ix e
MUArray Sz ix
sz forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
MVU.unsafeNew (forall ix. Index ix => Sz ix -> Int
totalElem Sz ix
sz)
  {-# INLINE unsafeNew #-}

  initialize :: forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) U ix e -> m ()
initialize (MUArray Sz ix
_ MVector (PrimState m) e
marr) = forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (forall (v :: * -> * -> *) a s. MVector v a => v s a -> ST s ()
VGM.basicInitialize MVector (PrimState m) e
marr)
  {-# INLINE initialize #-}

  unsafeLinearCopy :: forall ix' ix (m :: * -> *).
(Index ix', Index ix, PrimMonad m) =>
MArray (PrimState m) U ix' e
-> Int -> MArray (PrimState m) U ix e -> Int -> Sz1 -> m ()
unsafeLinearCopy (MUArray Sz ix'
_ MVector (PrimState m) e
mvFrom) Int
iFrom (MUArray Sz ix
_ MVector (PrimState m) e
mvTo) Int
iTo (Sz Int
k) =
    forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> MVector (PrimState m) a -> m ()
MVU.unsafeCopy (forall a s. Unbox a => Int -> Int -> MVector s a -> MVector s a
MVU.unsafeSlice Int
iTo Int
k MVector (PrimState m) e
mvTo) (forall a s. Unbox a => Int -> Int -> MVector s a -> MVector s a
MVU.unsafeSlice Int
iFrom Int
k MVector (PrimState m) e
mvFrom)
  {-# INLINE unsafeLinearCopy #-}

  unsafeLinearRead :: forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) U ix e -> Int -> m e
unsafeLinearRead (MUArray Sz ix
_ MVector (PrimState m) e
mv) =
    forall a ix e.
String -> (a -> Sz ix) -> (a -> ix -> e) -> a -> ix -> e
indexAssert String
"U.unsafeLinearRead" (forall ix. Index ix => ix -> Sz ix
Sz forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a s. Unbox a => MVector s a -> Int
MVU.length) forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
MVU.unsafeRead MVector (PrimState m) e
mv
  {-# INLINE unsafeLinearRead #-}

  unsafeLinearWrite :: forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) U ix e -> Int -> e -> m ()
unsafeLinearWrite (MUArray Sz ix
_ MVector (PrimState m) e
mv) =
    forall a ix e.
String -> (a -> Sz ix) -> (a -> ix -> e) -> a -> ix -> e
indexAssert String
"U.unsafeLinearWrite" (forall ix. Index ix => ix -> Sz ix
Sz forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a s. Unbox a => MVector s a -> Int
MVU.length) forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
MVU.unsafeWrite MVector (PrimState m) e
mv
  {-# INLINE unsafeLinearWrite #-}

  unsafeLinearGrow :: forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) U ix e
-> Sz ix -> m (MArray (PrimState m) U ix e)
unsafeLinearGrow (MUArray Sz ix
_ MVector (PrimState m) e
mv) Sz ix
sz = forall s ix e. Sz ix -> MVector s e -> MArray s U ix e
MUArray Sz ix
sz forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m (MVector (PrimState m) a)
MVU.unsafeGrow MVector (PrimState m) e
mv (forall ix. Index ix => Sz ix -> Int
totalElem Sz ix
sz)
  {-# INLINE unsafeLinearGrow #-}

instance (Index ix, Unbox e) => Stream U ix e where
  toStream :: Array U ix e -> Steps Id e
toStream = forall r ix e (m :: * -> *).
(Monad m, Index ix, Source r e) =>
Array r ix e -> Steps m e
S.steps
  {-# INLINE toStream #-}
  toStreamIx :: Array U ix e -> Steps Id (ix, e)
toStreamIx = forall r ix e (m :: * -> *).
(Monad m, Index ix, Source r e) =>
Array r ix e -> Steps m (ix, e)
S.isteps
  {-# INLINE toStreamIx #-}

instance (Unbox e, IsList (Array L ix e), Ragged L ix e) => IsList (Array U ix e) where
  type Item (Array U ix e) = Item (Array L ix e)
  fromList :: [Item (Array U ix e)] -> Array U ix e
fromList = forall r ix e.
(HasCallStack, Ragged L ix e, Manifest r e) =>
Comp -> [ListItem ix e] -> Array r ix e
A.fromLists' Comp
Seq
  {-# INLINE fromList #-}
  toList :: Array U ix e -> [Item (Array U ix e)]
toList = forall l. IsList l => l -> [Item l]
GHC.toList forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix e r.
(Ragged L ix e, Shape r ix, Source r e) =>
Array r ix e -> Array L ix e
toListArray
  {-# INLINE toList #-}

instance (VU.Unbox e, Num e) => FoldNumeric U e where
  unsafeDotProduct :: forall ix. Index ix => Array U ix e -> Array U ix e -> e
unsafeDotProduct = forall e ix r.
(Num e, Index ix, Source r e) =>
Array r ix e -> Array r ix e -> e
defaultUnsafeDotProduct
  {-# INLINE unsafeDotProduct #-}
  powerSumArray :: forall ix. Index ix => Array U ix e -> Int -> e
powerSumArray = forall ix r e.
(Index ix, Source r e, Num e) =>
Array r ix e -> Int -> e
defaultPowerSumArray
  {-# INLINE powerSumArray #-}
  foldArray :: forall ix. Index ix => (e -> e -> e) -> e -> Array U ix e -> e
foldArray = forall ix r e.
(Index ix, Source r e) =>
(e -> e -> e) -> e -> Array r ix e -> e
defaultFoldArray
  {-# INLINE foldArray #-}

instance (VU.Unbox e, Num e) => Numeric U e where
  unsafeLiftArray :: forall ix. Index ix => (e -> e) -> Array U ix e -> Array U ix e
unsafeLiftArray = forall r ix e.
(Load r ix e, Source r e) =>
(e -> e) -> Array r ix e -> Array r ix e
defaultUnsafeLiftArray
  {-# INLINE unsafeLiftArray #-}
  unsafeLiftArray2 :: forall ix.
Index ix =>
(e -> e -> e) -> Array U ix e -> Array U ix e -> Array U ix e
unsafeLiftArray2 = forall r ix e.
(Load r ix e, Source r e) =>
(e -> e -> e) -> Array r ix e -> Array r ix e -> Array r ix e
defaultUnsafeLiftArray2
  {-# INLINE unsafeLiftArray2 #-}

-- | /O(1)/ - Unwrap unboxed array and pull out the underlying unboxed vector.
--
-- @since 0.2.1
toUnboxedVector :: Array U ix e -> VU.Vector e
toUnboxedVector :: forall ix e. Array U ix e -> Vector e
toUnboxedVector = forall ix e. Array U ix e -> Vector e
uData
{-# INLINE toUnboxedVector #-}

-- | /O(1)/ - Unwrap unboxed mutable array and pull out the underlying unboxed mutable vector.
--
-- @since 0.2.1
toUnboxedMVector :: MArray s U ix e -> VU.MVector s e
toUnboxedMVector :: forall s ix e. MArray s U ix e -> MVector s e
toUnboxedMVector (MUArray Sz ix
_ MVector s e
mv) = MVector s e
mv
{-# INLINE toUnboxedMVector #-}

-- | /O(1)/ - Wrap an unboxed vector and produce an unboxed flat array.
--
-- @since 0.6.0
fromUnboxedVector :: VU.Unbox e => Comp -> VU.Vector e -> Vector U e
fromUnboxedVector :: forall e. Unbox e => Comp -> Vector e -> Vector U e
fromUnboxedVector Comp
comp Vector e
v = forall ix e. Comp -> Sz ix -> Vector e -> Array U ix e
UArray Comp
comp (forall ix. ix -> Sz ix
SafeSz (forall a. Unbox a => Vector a -> Int
VU.length Vector e
v)) Vector e
v
{-# INLINE fromUnboxedVector #-}

-- | /O(1)/ - Wrap an unboxed mutable vector and produce a mutable unboxed flat array.
--
-- @since 0.5.0
fromUnboxedMVector :: Unbox e => VU.MVector s e -> MVector s U e
fromUnboxedMVector :: forall e s. Unbox e => MVector s e -> MVector s U e
fromUnboxedMVector MVector s e
mv = forall s ix e. Sz ix -> MVector s e -> MArray s U ix e
MUArray (forall ix. ix -> Sz ix
SafeSz (forall a s. Unbox a => MVector s a -> Int
MVU.length MVector s e
mv)) MVector s e
mv
{-# INLINE fromUnboxedMVector #-}