{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
-- |
-- Module      : Data.Massiv.Array.Manifest
-- Copyright   : (c) Alexey Kuleshevich 2018-2021
-- License     : BSD3
-- Maintainer  : Alexey Kuleshevich <lehins@yandex.ru>
-- Stability   : experimental
-- Portability : non-portable
--
module Data.Massiv.Array.Manifest
  ( -- * Manifest
    Manifest
  , toManifest
  , M
  -- * Boxed
  , B(..)
  , BL(..)
  , BN(..)
  , N
  , pattern N
  , Uninitialized(..)
  -- ** Access
  , findIndex
  -- ** Conversion
  -- $boxed_conversion_note
  , toLazyArray
  , evalLazyArray
  , forceLazyArray
  , unwrapNormalForm
  , evalNormalForm
  -- *** Primitive Boxed Array
  , unwrapLazyArray
  , wrapLazyArray
  , unwrapArray
  , evalArray
  , unwrapMutableArray
  , unwrapMutableLazyArray
  , evalMutableArray
  , unwrapNormalFormArray
  , evalNormalFormArray
  , unwrapNormalFormMutableArray
  , evalNormalFormMutableArray
  -- *** Boxed Vector
  , toBoxedVector
  , toBoxedMVector
  , fromBoxedVector
  , fromBoxedMVector
  , evalBoxedVector
  , evalBoxedMVector
  -- * Primitive
  , P(..)
  , Prim
  -- ** Conversion
  -- *** Primitive ByteArray
  , toByteArray
  , toByteArrayM
  , unwrapByteArray
  , unwrapByteArrayOffset
  , fromByteArray
  , fromByteArrayM
  , fromByteArrayOffsetM
  , toMutableByteArray
  , unwrapMutableByteArray
  , unwrapMutableByteArrayOffset
  , fromMutableByteArray
  , fromMutableByteArrayM
  , fromMutableByteArrayOffsetM
  -- *** Primitive Vector
  , toPrimitiveVector
  , toPrimitiveMVector
  , fromPrimitiveVector
  , fromPrimitiveMVector
  -- * Storable
  , S(..)
  , Storable
  , mallocCompute
  , mallocCopy
  -- ** Conversion
  -- *** Storable Vector
  , toStorableVector
  , toStorableMVector
  , fromStorableVector
  , fromStorableMVector
  -- *** Direct Pointer Access
  , withPtr
  -- * Unboxed
  , U(..)
  , Unbox
  -- ** Conversion
  -- *** Unboxed Vector
  , toUnboxedVector
  , toUnboxedMVector
  , fromUnboxedVector
  , fromUnboxedMVector
  -- * ByteString Conversion
  , fromByteString
  , castFromByteString
  , toByteString
  , castToByteString
  , toBuilder
  , castToBuilder
  ) where

import Control.Monad
import Data.ByteString as S hiding (findIndex)
import Data.ByteString.Builder
import Data.ByteString.Internal
import Data.ByteString.Unsafe as SU
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.Array.Mutable
import Data.Massiv.Array.Ops.Fold
import Data.Massiv.Core.Common
import Data.Word (Word8)

-- | /O(1)/ - Convert a strict ByteString into a manifest array. Will return `Nothing` if length
-- doesn't match the total number of elements of new array.
--
-- @since 0.2.1
fromByteString ::
     Comp -- ^ Computation strategy
  -> ByteString -- ^ Strict ByteString to use as a source.
  -> Array M Ix1 Word8
fromByteString :: Comp -> ByteString -> Array M Ix1 Word8
fromByteString Comp
comp ByteString
bs = Comp -> Sz Ix1 -> (Ix1 -> Word8) -> Array M Ix1 Word8
forall ix e. Comp -> Sz ix -> (Ix1 -> e) -> Array M ix e
MArray Comp
comp (Ix1 -> Sz Ix1
forall ix. ix -> Sz ix
SafeSz (ByteString -> Ix1
S.length ByteString
bs)) (ByteString -> Ix1 -> Word8
SU.unsafeIndex ByteString
bs)
{-# INLINE fromByteString #-}

-- | /O(n)/ - Convert any source array into a strict `ByteString`. In case when the source array is
-- actually storable, no memory copy will occur.
--
-- @since 0.2.1
toByteString ::
     Load r ix Word8
  => Array r ix Word8 -- ^ Source array
  -> ByteString
toByteString :: Array r ix Word8 -> ByteString
toByteString = Array S ix Word8 -> ByteString
forall ix. Array S ix Word8 -> ByteString
castToByteString (Array S ix Word8 -> ByteString)
-> (Array r ix Word8 -> Array S ix Word8)
-> Array r ix Word8
-> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
.
#if __GLASGOW_HASKELL__ >= 820
  convert
  {- For ghc-8.0 `covert` results in "internal error: ARR_WORDS object entered!" -}
#else
  Array r ix Word8 -> Array S ix Word8
forall r ix e r'.
(Mutable r ix e, Load r' ix e) =>
Array r' ix e -> Array r ix e
compute
#endif
{-# INLINE toByteString #-}

-- | /O(n)/ - Conversion of array monoidally into a ByteString `Builder`.
--
-- @since 0.2.1
toBuilder :: Source r ix e => (e -> Builder) -> Array r ix e -> Builder
toBuilder :: (e -> Builder) -> Array r ix e -> Builder
toBuilder = (e -> Builder) -> Array r ix e -> Builder
forall r ix e m.
(Source r ix e, Monoid m) =>
(e -> m) -> Array r ix e -> m
foldMono
{-# INLINE toBuilder #-}

-- | /O(1)/ - Cast a storable array of `Word8` to ByteString `Builder`.
--
-- @since 0.5.0
castToBuilder :: Array S ix Word8 -> Builder
castToBuilder :: Array S ix Word8 -> Builder
castToBuilder = ByteString -> Builder
byteString (ByteString -> Builder)
-> (Array S ix Word8 -> ByteString) -> Array S ix Word8 -> Builder
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array S ix Word8 -> ByteString
forall ix. Array S ix Word8 -> ByteString
castToByteString
{-# INLINE castToBuilder #-}

-- | /O(1)/ - Cast a `S`torable array into a strict `ByteString`
--
-- @since 0.3.0
castToByteString :: Array S ix Word8 -> ByteString
castToByteString :: Array S ix Word8 -> ByteString
castToByteString = (\(ForeignPtr Word8
fp, Ix1
len) -> ForeignPtr Word8 -> Ix1 -> Ix1 -> ByteString
PS ForeignPtr Word8
fp Ix1
0 Ix1
len) ((ForeignPtr Word8, Ix1) -> ByteString)
-> (Array S ix Word8 -> (ForeignPtr Word8, Ix1))
-> Array S ix Word8
-> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array S ix Word8 -> (ForeignPtr Word8, Ix1)
forall e ix. Storable e => Array S ix e -> (ForeignPtr e, Ix1)
unsafeArrayToForeignPtr
{-# INLINE castToByteString #-}

-- | /O(1)/ - Cast a strict `ByteString` into a `S`torable array
--
-- @since 0.3.0
castFromByteString :: Comp -> ByteString -> Array S Ix1 Word8
castFromByteString :: Comp -> ByteString -> Array S Ix1 Word8
castFromByteString Comp
comp (PS ForeignPtr Word8
fp Ix1
offset Ix1
len) = Comp -> ForeignPtr Word8 -> Ix1 -> Sz Ix1 -> Array S Ix1 Word8
forall e.
Storable e =>
Comp -> ForeignPtr e -> Ix1 -> Sz Ix1 -> Array S Ix1 e
unsafeArrayFromForeignPtr Comp
comp ForeignPtr Word8
fp Ix1
offset (Ix1 -> Sz Ix1
forall ix. Index ix => ix -> Sz ix
Sz Ix1
len)
{-# INLINE castFromByteString #-}

-- $boxed_conversion_note
--
-- Important part of all conversions in this section is that the actual boxed
-- `Data.Primitive.Array.Array`, which holds the pointers to values isn't copied around, it is always
-- kept as the same array. Conversion to Massiv boxed array will undergo evaluation during which
-- computation strategies will be respected.



-- | /O(n)/ - Perform a row-major search starting at @0@ for an element. Returns the index
-- of the first occurance of an element or `Nothing` if a predicate could not be satisifed
-- after it was applyied to all elements of the array.
--
-- @since 0.5.5
findIndex :: Manifest r ix e => (e -> Bool) -> Array r ix e -> Maybe ix
findIndex :: (e -> Bool) -> Array r ix e -> Maybe ix
findIndex e -> Bool
f Array r ix e
arr = Ix1 -> Maybe ix
go Ix1
0
  where
    !sz :: Sz ix
sz = Array r ix e -> Sz ix
forall r ix e. Load r ix e => Array r ix e -> Sz ix
size Array r ix e
arr
    !k :: Ix1
k = Sz ix -> Ix1
forall ix. Index ix => Sz ix -> Ix1
totalElem Sz ix
sz
    go :: Ix1 -> Maybe ix
go !Ix1
i = do
      Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Ix1
i Ix1 -> Ix1 -> Bool
forall a. Ord a => a -> a -> Bool
< Ix1
k)
      if e -> Bool
f (Array r ix e -> Ix1 -> e
forall r ix e. Source r ix e => Array r ix e -> Ix1 -> e
unsafeLinearIndex Array r ix e
arr Ix1
i)
        then ix -> Maybe ix
forall a. a -> Maybe a
Just (ix -> Maybe ix) -> ix -> Maybe ix
forall a b. (a -> b) -> a -> b
$ Sz ix -> Ix1 -> ix
forall ix. Index ix => Sz ix -> Ix1 -> ix
fromLinearIndex Sz ix
sz Ix1
i
        else Ix1 -> Maybe ix
go (Ix1
i Ix1 -> Ix1 -> Ix1
forall a. Num a => a -> a -> a
+ Ix1
1)
{-# INLINE findIndex #-}


-- | Very similar to @`computeAs` `S`@ except load the source array into memory allocated
-- with @malloc@ on C heap. It can potentially be useful when iteroperating with some C
-- programs.
--
-- @since 0.5.9
mallocCompute :: forall r ix e. (Source r ix e, Storable e) => Array r ix e -> IO (Array S ix e)
mallocCompute :: Array r ix e -> IO (Array S ix e)
mallocCompute Array r ix e
arr = do
  let sz :: Sz ix
sz = Array r ix e -> Sz ix
forall r ix e. Load r ix e => Array r ix e -> Sz ix
size Array r ix e
arr
  MArray RealWorld S ix e
marr <- Sz ix -> IO (MArray RealWorld S ix e)
forall ix e (m :: * -> *).
(Index ix, Storable e, MonadIO m) =>
Sz ix -> m (MArray RealWorld S ix e)
unsafeMallocMArray Sz ix
sz
  MArray RealWorld S ix e -> Array r ix e -> IO ()
forall r' ix' e r ix (m :: * -> *).
(Load r' ix' e, Mutable r ix e, MonadIO m) =>
MArray RealWorld r ix e -> Array r' ix' e -> m ()
computeInto MArray RealWorld S ix e
marr Array r ix e
arr
  Comp -> MArray (PrimState IO) S ix e -> IO (Array S ix e)
forall r ix e (m :: * -> *).
(Mutable r ix e, PrimMonad m) =>
Comp -> MArray (PrimState m) r ix e -> m (Array r ix e)
unsafeFreeze (Array r ix e -> Comp
forall r ix e. Load r ix e => Array r ix e -> Comp
getComp Array r ix e
arr) MArray RealWorld S ix e
MArray (PrimState IO) S ix e
marr
{-# INLINE mallocCompute #-}

-- | Allocate memory on C heap with @malloc@ and copy the source array over.
--
-- @since 0.5.9
mallocCopy :: forall ix e. (Index ix, Storable e) => Array S ix e -> IO (Array S ix e)
mallocCopy :: Array S ix e -> IO (Array S ix e)
mallocCopy Array S ix e
arr = do
  let sz :: Sz ix
sz = Array S ix e -> Sz ix
forall r ix e. Load r ix e => Array r ix e -> Sz ix
size Array S ix e
arr
  MArray RealWorld S ix e
marr <- Sz ix -> IO (MArray RealWorld S ix e)
forall ix e (m :: * -> *).
(Index ix, Storable e, MonadIO m) =>
Sz ix -> m (MArray RealWorld S ix e)
unsafeMallocMArray Sz ix
sz
  Array S ix e
-> Ix1 -> MArray (PrimState IO) S ix e -> Ix1 -> Sz Ix1 -> IO ()
forall r ix e ix' (m :: * -> *).
(Mutable r ix e, Mutable r ix' e, PrimMonad m) =>
Array r ix' e
-> Ix1 -> MArray (PrimState m) r ix e -> Ix1 -> Sz Ix1 -> m ()
unsafeArrayLinearCopy Array S ix e
arr Ix1
0 MArray RealWorld S ix e
MArray (PrimState IO) S ix e
marr Ix1
0 (Ix1 -> Sz Ix1
forall ix. ix -> Sz ix
SafeSz (Sz ix -> Ix1
forall ix. Index ix => Sz ix -> Ix1
totalElem Sz ix
sz))
  Comp -> MArray (PrimState IO) S ix e -> IO (Array S ix e)
forall r ix e (m :: * -> *).
(Mutable r ix e, PrimMonad m) =>
Comp -> MArray (PrimState m) r ix e -> m (Array r ix e)
unsafeFreeze (Array S ix e -> Comp
forall r ix e. Load r ix e => Array r ix e -> Comp
getComp Array S ix e
arr) MArray RealWorld S ix e
MArray (PrimState IO) S ix e
marr
{-# INLINE mallocCopy #-}