{-# 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
  -- * 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 :: Comp -> ByteString -> Vector r Word8
fromByteString Comp
comp ByteString
bs = Comp -> Sz Int -> (Int -> Word8) -> Vector r Word8
forall r ix e.
Load r ix e =>
Comp -> Sz ix -> (Int -> e) -> Array r ix e
makeArrayLinear Comp
comp (Int -> Sz Int
forall ix. ix -> Sz ix
SafeSz (ByteString -> Int
S.length ByteString
bs)) (ByteString -> Int -> 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. Index 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'.
(Manifest r 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 :: (Index ix, Source r 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 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 :: 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. 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 :: Array S ix Word8 -> ByteString
castToByteString = (\(ForeignPtr Word8
fp, Int
len) -> ForeignPtr Word8 -> Int -> Int -> ByteString
PS ForeignPtr Word8
fp Int
0 Int
len) ((ForeignPtr Word8, Int) -> ByteString)
-> (Array S ix Word8 -> (ForeignPtr Word8, Int))
-> Array S ix Word8
-> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array S ix Word8 -> (ForeignPtr Word8, Int)
forall ix e. Index ix => Array S ix e -> (ForeignPtr e, Int)
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 Int
offset Int
len) = Comp -> ForeignPtr Word8 -> Int -> Sz Int -> Vector S Word8
forall e.
Storable e =>
Comp -> ForeignPtr e -> Int -> Sz Int -> Array S Int e
unsafeArrayFromForeignPtr Comp
comp ForeignPtr Word8
fp Int
offset (Int -> Sz Int
forall ix. Index ix => ix -> Sz ix
Sz Int
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 :: (e -> Bool) -> Array r ix e -> Maybe ix
findIndex e -> Bool
f Array r ix e
arr = Int -> Maybe ix
go Int
0
  where
    !sz :: Sz ix
sz = Array r ix e -> Sz ix
forall r ix e. Size r => Array r ix e -> Sz ix
size Array r ix e
arr
    !k :: Int
k = Sz ix -> Int
forall ix. Index ix => Sz ix -> Int
totalElem Sz ix
sz
    go :: Int -> Maybe ix
go !Int
i = do
      Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
k)
      if e -> Bool
f (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 Int
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 -> Int -> ix
forall ix. Index ix => Sz ix -> Int -> ix
fromLinearIndex Sz ix
sz Int
i
        else Int -> Maybe ix
go (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
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 :: 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. Size r => Array r ix e -> Sz ix
size Array r ix e
arr
  MArray RealWorld S ix e
marr <- Sz ix -> IO (MArray (PrimState IO) S ix e)
forall ix e (m :: * -> *).
(Index ix, Storable e, PrimMonad m) =>
Sz ix -> m (MArray (PrimState m) S ix e)
unsafeMallocMArray Sz ix
sz
  MArray RealWorld S ix e -> Array r ix e -> IO ()
forall r' ix' e r ix (m :: * -> *).
(Size r', 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
  Comp -> MArray (PrimState IO) S ix e -> IO (Array S ix e)
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 (Array r ix e -> Comp
forall r ix e. Strategy r => 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. Size r => Array r ix e -> Sz ix
size Array S ix e
arr
  MArray RealWorld S ix e
marr <- Sz ix -> IO (MArray (PrimState IO) S ix e)
forall ix e (m :: * -> *).
(Index ix, Storable e, PrimMonad m) =>
Sz ix -> m (MArray (PrimState m) S ix e)
unsafeMallocMArray Sz ix
sz
  Array S ix e
-> Int -> MArray (PrimState IO) S ix e -> Int -> Sz Int -> IO ()
forall r e ix' ix (m :: * -> *).
(Manifest r e, Index ix', Index ix, PrimMonad m) =>
Array r ix' e
-> Int -> MArray (PrimState m) r ix e -> Int -> Sz Int -> m ()
unsafeArrayLinearCopy Array S ix e
arr Int
0 MArray RealWorld S ix e
MArray (PrimState IO) S ix e
marr Int
0 (Int -> Sz Int
forall ix. ix -> Sz ix
SafeSz (Sz ix -> Int
forall ix. Index ix => Sz ix -> Int
totalElem Sz ix
sz))
  Comp -> MArray (PrimState IO) S ix e -> IO (Array S ix e)
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 (Array S ix e -> Comp
forall r ix e. Strategy r => 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 #-}