Add withPrimArrayContents and withMutablePrimArrayContents (#400)

Co-authored-by: konsumlamm <44230978+konsumlamm@users.noreply.github.com>

Data/Primitive/ByteArray.hs

@@ -91,6 +91,8 @@ import System.IO.Unsafe (unsafeDupablePerformIO)
 
 import Data.Array.Byte (ByteArray(..), MutableByteArray(..))
 
+import Data.Primitive.Internal.Operations (mutableByteArrayContentsShim)
+
 -- | Create a new mutable byte array of the specified size in bytes.
 -- The underlying memory is left uninitialized.
 --
@@ -161,15 +163,6 @@ mutableByteArrayContents :: MutableByteArray s -> Ptr Word8
 {-# INLINE mutableByteArrayContents #-}
 mutableByteArrayContents (MutableByteArray arr#) = Ptr (mutableByteArrayContentsShim arr#)
 
-mutableByteArrayContentsShim :: MutableByteArray# s -> Addr#
-{-# INLINE mutableByteArrayContentsShim #-}
-mutableByteArrayContentsShim x =
-#if __GLASGOW_HASKELL__ >= 902
-  mutableByteArrayContents# x
-#else
-  byteArrayContents# (unsafeCoerce# x)
-#endif
-
 -- | A composition of 'mutableByteArrayContents' and 'keepAliveUnlifted'.
 -- The callback function must not return the pointer. The argument byte
 -- array must be /pinned/. See 'byteArrayContents' for an explanation

Data/Primitive/Internal/Operations.hs

@@ -35,6 +35,7 @@ module Data.Primitive.Internal.Operations (
   , keepAliveLiftedLifted#
   , keepAliveUnliftedLifted#
 #endif
+  , mutableByteArrayContentsShim
   , UnliftedType
 ) where
 
@@ -192,3 +193,13 @@ type UnliftedType = TYPE 'PtrRepUnlifted
 #elif __GLASGOW_HASKELL__ < 902
 type UnliftedType = TYPE 'UnliftedRep
 #endif
+
+mutableByteArrayContentsShim :: MutableByteArray# s -> Addr#
+{-# INLINE mutableByteArrayContentsShim #-}
+mutableByteArrayContentsShim x =
+#if __GLASGOW_HASKELL__ >= 902
+  mutableByteArrayContents# x
+#else
+  byteArrayContents# (unsafeCoerce# x)
+#endif
+

Data/Primitive/PrimArray.hs

@@ -63,7 +63,9 @@ module Data.Primitive.PrimArray
   , sizeofMutablePrimArray
   , sizeofPrimArray
   , primArrayContents
+  , withPrimArrayContents
   , mutablePrimArrayContents
+  , withMutablePrimArrayContents
 #if __GLASGOW_HASKELL__ >= 802
   , isPrimArrayPinned
   , isMutablePrimArrayPinned
@@ -133,6 +135,8 @@ import Data.Semigroup
 import qualified GHC.Exts as Exts
 #endif
 
+import Data.Primitive.Internal.Operations (mutableByteArrayContentsShim)
+
 -- | Arrays of unboxed elements. This accepts types like 'Double', 'Char',
 -- 'Int' and 'Word', as well as their fixed-length variants ('Word8',
 -- 'Word16', etc.). Since the elements are unboxed, a 'PrimArray' is strict
@@ -1106,12 +1110,8 @@ primArrayContents (PrimArray arr#) = Ptr (byteArrayContents# arr#)
 -- @since 0.7.1.0
 mutablePrimArrayContents :: MutablePrimArray s a -> Ptr a
 {-# INLINE mutablePrimArrayContents #-}
-mutablePrimArrayContents (MutablePrimArray arr#) = Ptr
-#if __GLASGOW_HASKELL__ >= 902
-  (mutableByteArrayContents# arr#)
-#else
-  (byteArrayContents# (unsafeCoerce# arr#))
-#endif
+mutablePrimArrayContents (MutablePrimArray arr#) =
+  Ptr (mutableByteArrayContentsShim arr#)
 
 -- | Return a newly allocated array with the specified subrange of the
 -- provided array. The provided array should contain the full subrange
@@ -1162,3 +1162,25 @@ unST (GHCST.ST f) = f
 #else /* In older GHCs, runRW# is not available. */
 runPrimArray m = runST $ m >>= unsafeFreezePrimArray
 #endif
+
+-- | A composition of 'primArrayContents' and 'keepAliveUnlifted'.
+-- The callback function must not return the pointer. The argument
+-- array must be /pinned/. See 'primArrayContents' for an explanation
+-- of which primitive arrays are pinned.
+--
+-- Note: This could be implemented with 'keepAlive' instead of
+-- 'keepAliveUnlifted', but 'keepAlive' here would cause GHC to materialize
+-- the wrapper data constructor on the heap.
+withPrimArrayContents :: PrimBase m => PrimArray a -> (Ptr a -> m a) -> m a
+{-# INLINE withPrimArrayContents #-}
+withPrimArrayContents (PrimArray arr#) f =
+  keepAliveUnlifted arr# (f (Ptr (byteArrayContents# arr#)))
+
+-- | A composition of 'mutablePrimArrayContents' and 'keepAliveUnlifted'.
+-- The callback function must not return the pointer. The argument
+-- array must be /pinned/. See 'primArrayContents' for an explanation
+-- of which primitive arrays are pinned.
+withMutablePrimArrayContents :: PrimBase m => MutablePrimArray (PrimState m) a -> (Ptr a -> m a) -> m a
+{-# INLINE withMutablePrimArrayContents #-}
+withMutablePrimArrayContents (MutablePrimArray arr#) f =
+  keepAliveUnlifted arr# (f (Ptr (mutableByteArrayContentsShim arr#)))

changelog.md

@@ -1,5 +1,8 @@
 ## Changes in version 0.9.0.0
 
+  * Add `withByteArrayContents`, `withMutableByteArrayContents`,
+    `withPrimArrayContents`, `withMutablePrimArrayContents`.
+
   * Fix signature of `keepAlive`.
 
   * Remove re-export of `fromList` and `fromListN` from `Data.Primitive.Array`.