Storable.hs

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

-- |
-- Module      : Data.Massiv.Array.Manifest.Storable
-- Copyright   : (c) Alexey Kuleshevich 2018-2022
-- License     : BSD3
-- Maintainer  : Alexey Kuleshevich <lehins@yandex.ru>
-- Stability   : experimental
-- Portability : non-portable
module Data.Massiv.Array.Manifest.Storable (
  S (..),
  Array (..),
  MArray (..),
  Storable,
  toStorableVector,
  toStorableMVector,
  fromStorableVector,
  fromStorableMVector,
  withPtr,
  unsafeWithPtr,
  unsafeMallocMArray,
  unsafeArrayToForeignPtr,
  unsafeMArrayToForeignPtr,
  unsafeArrayFromForeignPtr,
  unsafeArrayFromForeignPtr0,
  unsafeMArrayFromForeignPtr,
  unsafeMArrayFromForeignPtr0,
) where

import Control.DeepSeq (NFData (..), deepseq)
import Control.Exception
import Control.Monad
import Control.Monad.IO.Unlift
import Control.Monad.Primitive
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 Data.Primitive.ByteArray
import Data.Primitive.Ptr (setPtr)
import Data.Word
import Foreign.ForeignPtr
import Foreign.Marshal.Alloc
import Foreign.Marshal.Array (advancePtr, copyArray)
import Foreign.Ptr
import Foreign.Storable
import GHC.Exts as GHC
import GHC.ForeignPtr
import System.IO.Unsafe (unsafePerformIO)
import Unsafe.Coerce
import Prelude hiding (mapM)

import qualified Data.Vector.Generic.Mutable as MVG
import qualified Data.Vector.Storable as VS
import qualified Data.Vector.Storable.Mutable as MVS

-- | Representation for `Storable` elements
data S = S deriving (Show)

data instance Array S ix e = SArray
  { sComp :: !Comp
  , sSize :: !(Sz ix)
  , sData :: {-# UNPACK #-} !(ForeignPtr e)
  }

data instance MArray s S ix e = MSArray !(Sz ix) {-# UNPACK #-} !(ForeignPtr e)

instance (Ragged L ix e, Show e, Storable e) => Show (Array S ix e) where
  showsPrec = showsArrayPrec id
  showList = showArrayList

instance NFData ix => NFData (Array S ix e) where
  rnf (SArray c sz _v) = c `deepseq` sz `deepseq` ()
  {-# INLINE rnf #-}

instance NFData ix => NFData (MArray s S ix e) where
  rnf (MSArray sz _mv) = sz `deepseq` ()
  {-# INLINE rnf #-}

instance (Storable e, Eq e, Index ix) => Eq (Array S ix e) where
  (==) = eqArrays (==)
  {-# INLINE (==) #-}

instance (Storable e, Ord e, Index ix) => Ord (Array S ix e) where
  compare = compareArrays compare
  {-# INLINE compare #-}

instance Strategy S where
  getComp = sComp
  {-# INLINE getComp #-}
  setComp c arr = arr{sComp = c}
  {-# INLINE setComp #-}
  repr = S

plusFp :: ForeignPtr a -> Int -> ForeignPtr b
plusFp (ForeignPtr addr c) (I# d) = ForeignPtr (plusAddr# addr d) c

advanceForeignPtr :: forall e. Storable e => ForeignPtr e -> Int -> ForeignPtr e
advanceForeignPtr fp i = plusFp fp (i * sizeOf (undefined :: e))
{-# INLINE advanceForeignPtr #-}

indexForeignPtr :: Storable e => ForeignPtr e -> Int -> e
indexForeignPtr fp i = unsafeInlineIO $ unsafeWithForeignPtr fp $ \p -> peekElemOff p i
{-# INLINE indexForeignPtr #-}

instance Storable e => Source S e where
  unsafeLinearIndex (SArray _ _sz fp) =
    indexAssert "S.unsafeLinearIndex" (const (toLinearSz _sz)) indexForeignPtr fp
  {-# INLINE unsafeLinearIndex #-}

  unsafeOuterSlice (SArray c _ fp) szL i =
    let k = totalElem szL
     in SArray c szL $ advanceForeignPtr fp (i * k)
  {-# INLINE unsafeOuterSlice #-}

  unsafeLinearSlice i k (SArray c _ fp) =
    SArray c k $ advanceForeignPtr fp i
  {-# INLINE unsafeLinearSlice #-}

instance Index ix => Shape S ix where
  maxLinearSize = Just . SafeSz . elemsCount
  {-# INLINE maxLinearSize #-}

instance Size S where
  size = sSize
  {-# INLINE size #-}
  unsafeResize !sz !arr = arr{sSize = sz}
  {-# INLINE unsafeResize #-}

instance Storable e => Manifest S e where
  unsafeLinearIndexM (SArray _ _sz fp) =
    indexAssert "S.unsafeLinearIndex" (const (toLinearSz _sz)) indexForeignPtr fp
  {-# INLINE unsafeLinearIndexM #-}

  sizeOfMArray (MSArray sz _) = sz
  {-# INLINE sizeOfMArray #-}

  unsafeResizeMArray sz (MSArray _ fp) = MSArray sz fp
  {-# INLINE unsafeResizeMArray #-}

  unsafeLinearSliceMArray i k (MSArray _ fp) = MSArray k $ advanceForeignPtr fp i
  {-# INLINE unsafeLinearSliceMArray #-}

  unsafeThaw (SArray _ sz fp) = pure $ MSArray sz fp
  {-# INLINE unsafeThaw #-}

  unsafeFreeze comp (MSArray sz v) = pure $ SArray comp sz v
  {-# INLINE unsafeFreeze #-}

  unsafeNew sz = do
    let !n = totalElem sz
        dummy = undefined :: e
        !eSize = sizeOf dummy
    when (n > (maxBound :: Int) `div` eSize) $ error $ "Array size is too big: " ++ show sz
    unsafeIOToPrim $ do
      fp <- mallocPlainForeignPtrAlignedBytes (n * sizeOf dummy) (alignment dummy)
      pure $ MSArray sz fp
  {-# INLINE unsafeNew #-}

  initialize (MSArray sz fp) =
    unsafeIOToPrim $
      unsafeWithForeignPtr fp $ \p ->
        setPtr (castPtr p) (totalElem sz * sizeOf (undefined :: e)) (0 :: Word8)
  {-# INLINE initialize #-}

  unsafeLinearRead (MSArray _sz fp) o =
    unsafeIOToPrim $
      indexAssert
        "S.unsafeLinearRead"
        (const (toLinearSz _sz))
        (\_ _ -> unsafeWithForeignPtr fp (`peekElemOff` o))
        fp
        o
  {-# INLINE unsafeLinearRead #-}

  unsafeLinearWrite (MSArray _sz fp) o e =
    unsafeIOToPrim $
      indexAssert
        "S.unsafeLinearWrite"
        (const (toLinearSz _sz))
        (\_ _ -> unsafeWithForeignPtr fp (\p -> pokeElemOff p o e))
        fp
        o
  {-# INLINE unsafeLinearWrite #-}

  unsafeLinearSet (MSArray _ fp) i k e =
    stToPrim (MVG.basicSet (MVS.unsafeFromForeignPtr0 (advanceForeignPtr fp i) (unSz k)) e)
  {-# INLINE unsafeLinearSet #-}

  unsafeLinearCopy (MSArray _ fpFrom) iFrom (MSArray _ fpTo) iTo (Sz k) = do
    unsafePrimToPrim $
      withForeignPtr fpFrom $ \ptrFrom ->
        withForeignPtr fpTo $ \ptrTo -> do
          let ptrFrom' = advancePtr ptrFrom iFrom
              ptrTo' = advancePtr ptrTo iTo
          copyArray ptrTo' ptrFrom' k
  {-# INLINE unsafeLinearCopy #-}

  unsafeArrayLinearCopy arrFrom iFrom marrTo iTo sz = do
    marrFrom <- unsafeThaw arrFrom
    unsafeLinearCopy marrFrom iFrom marrTo iTo sz
  {-# INLINE unsafeArrayLinearCopy #-}

  unsafeLinearShrink marr@(MSArray _ fp@(ForeignPtr _ fpc)) sz = do
    let shrinkMBA :: MutableByteArray RealWorld -> IO ()
        shrinkMBA mba = shrinkMutableByteArray mba (totalElem sz * sizeOf (undefined :: e))
        {-# INLINE shrinkMBA #-}
    case fpc of
      MallocPtr mba# _ -> do
        unsafePrimToPrim $ shrinkMBA (MutableByteArray mba#)
        pure $ MSArray sz fp
      PlainPtr mba# -> do
        unsafePrimToPrim $ shrinkMBA (MutableByteArray mba#)
        pure $ MSArray sz fp
      _ -> unsafeDefaultLinearShrink marr sz
  {-# INLINE unsafeLinearShrink #-}

instance (Index ix, Storable e) => Load S ix e where
  makeArray comp sz f = compute (makeArray comp sz f :: Array D ix e)
  {-# INLINE makeArray #-}

  makeArrayLinear !comp !sz f = unsafePerformIO $ generateArrayLinear comp sz (pure . f)
  {-# INLINE makeArrayLinear #-}

  replicate comp !sz !e = runST (newMArray sz e >>= unsafeFreeze comp)
  {-# INLINE replicate #-}

  iterArrayLinearST_ !scheduler !arr =
    splitLinearlyWith_ scheduler (elemsCount arr) (unsafeLinearIndex arr)
  {-# INLINE iterArrayLinearST_ #-}

instance (Index ix, Storable e) => StrideLoad S ix e

instance (Index ix, Storable e) => Stream S ix e where
  toStream = S.steps
  {-# INLINE toStream #-}
  toStreamIx = S.isteps
  {-# INLINE toStreamIx #-}

instance (Storable e, Num e) => FoldNumeric S e where
  unsafeDotProduct = defaultUnsafeDotProduct
  {-# INLINE unsafeDotProduct #-}
  powerSumArray = defaultPowerSumArray
  {-# INLINE powerSumArray #-}
  foldArray = defaultFoldArray
  {-# INLINE foldArray #-}

instance (Storable e, Num e) => Numeric S e where
  unsafeLiftArray = defaultUnsafeLiftArray
  {-# INLINE unsafeLiftArray #-}
  unsafeLiftArray2 = defaultUnsafeLiftArray2
  {-# INLINE unsafeLiftArray2 #-}

instance (Storable e, Floating e) => NumericFloat S e

instance (Storable e, IsList (Array L ix e), Ragged L ix e) => IsList (Array S ix e) where
  type Item (Array S ix e) = Item (Array L ix e)
  fromList = A.fromLists' Seq
  {-# INLINE fromList #-}
  toList = GHC.toList . toListArray
  {-# INLINE toList #-}

-- | A pointer to the beginning of the storable array. It is unsafe since, if mutated, it can break
-- referential transparency.
--
-- @since 0.1.3
unsafeWithPtr :: MonadUnliftIO m => Array S ix e -> (Ptr e -> m b) -> m b
unsafeWithPtr arr f = withRunInIO $ \run -> unsafeWithForeignPtr (sData arr) (run . f)
{-# INLINE unsafeWithPtr #-}

-- | A pointer to the beginning of the mutable array.
--
-- @since 0.1.3
withPtr :: MonadUnliftIO m => MArray RealWorld S ix e -> (Ptr e -> m b) -> m b
withPtr (MSArray _ fp) f = withRunInIO $ \run -> unsafeWithForeignPtr fp (run . f)
{-# INLINE withPtr #-}

-- | /O(1)/ - Unwrap storable array and pull out the underlying storable vector.
--
-- @since 0.2.1
toStorableVector :: Index ix => Array S ix e -> VS.Vector e
toStorableVector arr =
  -- this hack is needed to workaround the redundant Storable constraint
  -- see haskell/vector#394
  unsafeCoerce $
    VS.unsafeFromForeignPtr0 (castForeignPtr (sData arr) :: ForeignPtr Word) (totalElem (sSize arr))
{-# INLINE toStorableVector #-}

-- | /O(1)/ - Unwrap storable mutable array and pull out the underlying storable mutable vector.
--
-- @since 0.2.1
toStorableMVector :: Index ix => MArray s S ix e -> VS.MVector s e
toStorableMVector (MSArray sz fp) = MVS.MVector (totalElem sz) fp
{-# INLINE toStorableMVector #-}

-- | /O(1)/ - Cast a storable vector to a storable array.
--
-- @since 0.5.0
fromStorableVector :: Comp -> VS.Vector e -> Vector S e
fromStorableVector comp v =
  -- unasfeCoerce hack below is needed to workaround the redundant Storable
  -- constraint see haskell/vector#394
  case VS.unsafeToForeignPtr0 (unsafeCoerce v :: VS.Vector Word) of
    (fp, k) -> SArray{sComp = comp, sSize = SafeSz k, sData = castForeignPtr fp}
{-# INLINE fromStorableVector #-}

-- | /O(1)/ - Cast a mutable storable vector to a mutable storable array.
--
-- @since 0.5.0
fromStorableMVector :: MVS.MVector s e -> MVector s S e
fromStorableMVector (MVS.MVector n fp) = MSArray (SafeSz n) fp
{-# INLINE fromStorableMVector #-}

-- | /O(1)/ - Yield the underlying `ForeignPtr` together with its length.
--
-- @since 0.3.0
unsafeArrayToForeignPtr :: Index ix => Array S ix e -> (ForeignPtr e, Int)
unsafeArrayToForeignPtr (SArray _ sz fp) = (fp, totalElem sz)
{-# INLINE unsafeArrayToForeignPtr #-}

-- | /O(1)/ - Yield the underlying `ForeignPtr` together with its length.
--
-- @since 0.3.0
unsafeMArrayToForeignPtr :: Index ix => MArray s S ix e -> (ForeignPtr e, Int)
unsafeMArrayToForeignPtr (MSArray sz fp) = (fp, totalElem sz)
{-# INLINE unsafeMArrayToForeignPtr #-}

-- | /O(1)/ - Wrap a `ForeignPtr` and it's size into a pure storable array.
--
-- @since 0.3.0
unsafeArrayFromForeignPtr0 :: Comp -> ForeignPtr e -> Sz1 -> Vector S e
unsafeArrayFromForeignPtr0 comp fp sz = SArray{sComp = comp, sSize = sz, sData = fp}
{-# INLINE unsafeArrayFromForeignPtr0 #-}

-- | /O(1)/ - Wrap a `ForeignPtr`, an offset and it's size into a pure storable array.
--
-- @since 0.3.0
unsafeArrayFromForeignPtr :: Storable e => Comp -> ForeignPtr e -> Int -> Sz1 -> Array S Ix1 e
unsafeArrayFromForeignPtr comp ptr offset sz =
  SArray{sComp = comp, sSize = sz, sData = advanceForeignPtr ptr offset}
{-# INLINE unsafeArrayFromForeignPtr #-}

-- | /O(1)/ - Wrap a `ForeignPtr` and it's size into a mutable storable array. It is still safe to
-- modify the pointer, unless the array gets frozen prior to modification.
--
-- @since 0.3.0
unsafeMArrayFromForeignPtr0 :: ForeignPtr e -> Sz1 -> MArray s S Ix1 e
unsafeMArrayFromForeignPtr0 fp sz = MSArray sz fp
{-# INLINE unsafeMArrayFromForeignPtr0 #-}

-- | /O(1)/ - Wrap a `ForeignPtr`, an offset and it's size into a mutable storable array. It is
-- still safe to modify the pointer, unless the array gets frozen prior to modification.
--
-- @since 0.3.0
unsafeMArrayFromForeignPtr :: Storable e => ForeignPtr e -> Int -> Sz1 -> MArray s S Ix1 e
unsafeMArrayFromForeignPtr fp offset sz = MSArray sz (advanceForeignPtr fp offset)
{-# INLINE unsafeMArrayFromForeignPtr #-}

-- | Allocate memory using @malloc@ on C heap, instead of on Haskell heap. Memory is left
-- uninitialized
--
-- @since 0.5.9
unsafeMallocMArray
  :: forall ix e m
   . (Index ix, Storable e, PrimMonad m)
  => Sz ix
  -> m (MArray (PrimState m) S ix e)
unsafeMallocMArray sz = unsafePrimToPrim $ do
  let n = totalElem sz
  foreignPtr <- mask_ $ do
    ptr <- mallocBytes (sizeOf (undefined :: e) * n)
    newForeignPtr finalizerFree ptr
  pure $ MSArray sz foreignPtr
{-# INLINE unsafeMallocMArray #-}

#if !MIN_VERSION_base(4,15,0)
-- | A compatibility wrapper for 'GHC.ForeignPtr.unsafeWithForeignPtr' provided
-- by GHC 9.0.1 and later.
--
-- Only to be used when the continuation is known not to
-- unconditionally diverge lest unsoundness can result.
unsafeWithForeignPtr :: ForeignPtr a -> (Ptr a -> IO b) -> IO b
unsafeWithForeignPtr = withForeignPtr
#endif