{-# 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-2022
-- License     : BSD3
-- Maintainer  : Alexey Kuleshevich <lehins@yandex.ru>
-- Stability   : experimental
-- Portability : non-portable
--
module Data.Massiv.Array.Manifest
  ( -- * Manifest
    Manifest
  -- ** Generate
  , generateArray
  , generateArrayLinear
  , generateArrayS
  , generateArrayLinearS
  , generateSplitSeedArray
  -- ** Stateful worker threads
  , generateArrayWS
  , generateArrayLinearWS
  -- ** Unfold
  , unfoldrPrimM_
  , iunfoldrPrimM_
  , unfoldrPrimM
  , iunfoldrPrimM
  , unfoldlPrimM_
  , iunfoldlPrimM_
  , unfoldlPrimM
  , iunfoldlPrimM
  -- ** Mapping
  , forPrimM
  , forPrimM_
  , iforPrimM
  , iforPrimM_
  , iforLinearPrimM
  , iforLinearPrimM_
  , for2PrimM_
  , ifor2PrimM_
  -- * 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(n)/ - 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 ::
     Load r Ix1 Word8
  => Comp -- ^ Computation strategy
  -> ByteString -- ^ Strict ByteString to use as a source.
  -> Vector r Word8
fromByteString :: forall r. Load r Ix1 Word8 => Comp -> ByteString -> Vector r Word8
fromByteString Comp
comp ByteString
bs = forall r ix e.
Load r ix e =>
Comp -> Sz ix -> (Ix1 -> e) -> Array r ix e
makeArrayLinear Comp
comp (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 :: forall r ix. Load r ix Word8 => Array r ix Word8 -> ByteString
toByteString = forall ix. Index ix => Array S ix Word8 -> ByteString
castToByteString forall b c a. (b -> c) -> (a -> b) -> a -> c
.
#if __GLASGOW_HASKELL__ >= 820
  forall r ix e r'.
(Manifest r e, Load r' ix e) =>
Array r' ix e -> Array r ix e
convert
  {- For ghc-8.0 `convert` results in "internal error: ARR_WORDS object entered!" -}
#else
  compute
#endif
{-# INLINE toByteString #-}

-- | /O(n)/ - Conversion of array monoidally into a ByteString `Builder`.
--
-- @since 0.2.1
toBuilder :: (Index ix, Source r e) => (e -> Builder) -> Array r ix e -> Builder
toBuilder :: forall ix r e.
(Index ix, Source r e) =>
(e -> Builder) -> Array r ix e -> Builder
toBuilder = forall ix r e m.
(Index ix, Source r 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 :: Index ix => Array S ix Word8 -> Builder
castToBuilder :: forall ix. Index ix => Array S ix Word8 -> Builder
castToBuilder = ByteString -> Builder
byteString forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix. Index 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 :: Index ix => Array S ix Word8 -> ByteString
castToByteString :: forall ix. Index ix => Array S ix Word8 -> ByteString
castToByteString = (\(ForeignPtr Word8
fp, Ix1
len) -> ForeignPtr Word8 -> Ix1 -> Ix1 -> ByteString
PS ForeignPtr Word8
fp Ix1
0 Ix1
len) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix e. Index ix => 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 -> Vector S Word8
castFromByteString :: Comp -> ByteString -> Vector S Word8
castFromByteString Comp
comp (PS ForeignPtr Word8
fp Ix1
offset Ix1
len) = forall e.
Storable e =>
Comp -> ForeignPtr e -> Ix1 -> Sz1 -> Array S Ix1 e
unsafeArrayFromForeignPtr Comp
comp ForeignPtr Word8
fp Ix1
offset (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 :: (Index ix, Manifest r e) => (e -> Bool) -> Array r ix e -> Maybe ix
findIndex :: forall ix r e.
(Index ix, Manifest r e) =>
(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 = forall r ix e. Size r => Array r ix e -> Sz ix
size Array r ix e
arr
    !k :: Ix1
k = forall ix. Index ix => Sz ix -> Ix1
totalElem Sz ix
sz
    go :: Ix1 -> Maybe ix
go !Ix1
i = do
      forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Ix1
i forall a. Ord a => a -> a -> Bool
< Ix1
k)
      if e -> Bool
f (forall r e ix. (Source r e, Index ix) => Array r ix e -> Ix1 -> e
unsafeLinearIndex Array r ix e
arr Ix1
i)
        then forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall ix. Index ix => Sz ix -> Ix1 -> ix
fromLinearIndex Sz ix
sz Ix1
i
        else Ix1 -> Maybe ix
go (Ix1
i 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. (Size r, Load r ix e, Storable e) => Array r ix e -> IO (Array S ix e)
mallocCompute :: forall r ix e.
(Size r, Load r ix e, Storable e) =>
Array r ix e -> IO (Array S ix e)
mallocCompute Array r ix e
arr = do
  let sz :: Sz ix
sz = forall r ix e. Size r => Array r ix e -> Sz ix
size Array r ix e
arr
  MArray RealWorld S ix e
marr <- forall ix e (m :: * -> *).
(Index ix, Storable e, PrimMonad m) =>
Sz ix -> m (MArray (PrimState m) S ix e)
unsafeMallocMArray Sz ix
sz
  forall r' ix' e r ix (m :: * -> *).
(Load r' ix' e, Manifest r e, Index ix, MonadIO m) =>
MArray RealWorld r ix e -> Array r' ix' e -> m ()
computeInto MArray RealWorld S ix e
marr Array r ix e
arr
  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 (forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r ix e
arr) MArray RealWorld 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 :: forall ix e.
(Index ix, Storable e) =>
Array S ix e -> IO (Array S ix e)
mallocCopy Array S ix e
arr = do
  let sz :: Sz ix
sz = forall r ix e. Size r => Array r ix e -> Sz ix
size Array S ix e
arr
  MArray RealWorld S ix e
marr <- forall ix e (m :: * -> *).
(Index ix, Storable e, PrimMonad m) =>
Sz ix -> m (MArray (PrimState m) S ix e)
unsafeMallocMArray Sz ix
sz
  forall r e ix' ix (m :: * -> *).
(Manifest r e, Index ix', Index ix, PrimMonad m) =>
Array r ix' e
-> Ix1 -> MArray (PrimState m) r ix e -> Ix1 -> Sz1 -> m ()
unsafeArrayLinearCopy Array S ix e
arr Ix1
0 MArray RealWorld S ix e
marr Ix1
0 (forall ix. ix -> Sz ix
SafeSz (forall ix. Index ix => Sz ix -> Ix1
totalElem Sz ix
sz))
  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 (forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array S ix e
arr) MArray RealWorld S ix e
marr
{-# INLINE mallocCopy #-}