Fix a bug that passes around an unpinned array as a foreign pointer

src/Streamly/External/ByteString.hs

@@ -49,12 +49,13 @@ import GHC.ForeignPtr (plusForeignPtr)
 #if MIN_VERSION_streamly_core(0,2,0)
 import Streamly.Internal.Data.Array (Array(..))
 import Streamly.Internal.Data.MutByteArray (MutByteArray(..))
-import qualified Streamly.Internal.Data.MutByteArray as MutBA (nil)
+import qualified Streamly.Internal.Data.Array as Array
+import qualified Streamly.Internal.Data.MutByteArray as MutBA
 import qualified Streamly.Internal.Data.Stream as StreamD (Step(Yield))
 #else
 import Streamly.Internal.Data.Array.Type (Array(..))
 import Streamly.Internal.Data.Unboxed (MutableByteArray(..))
-import qualified Streamly.Internal.Data.Unboxed as MutBA (nil)
+import qualified Streamly.Internal.Data.Unboxed as MutBA
 import qualified Streamly.Internal.Data.Stream.StreamD as StreamD (Step(Yield))
 #endif
 
@@ -66,6 +67,12 @@ import Prelude hiding (read)
 #define MUT_BYTE_ARRAY MutableByteArray
 #endif
 
+#if MIN_VERSION_streamly_core(0,2,2)
+#define NIL MutBA.empty
+#else
+#define NIL MutBA.nil
+#endif
+
 #if MIN_VERSION_bytestring(0,11,0)
 #define CONSTRUCTOR(a, b, c) BS a c
 #define WHEN_0_10_12(x)
@@ -74,6 +81,28 @@ import Prelude hiding (read)
 #define WHEN_0_10_12(x) x
 #endif
 
+
+{-# INLINE ensurePinned #-}
+ensurePinned :: Array a -> IO (Array a)
+{-# INLINE pinnedCreateOf #-}
+pinnedCreateOf :: MonadIO m => Int -> Fold m Word8 (Array Word8)
+{-# INLINE pinnedCreate #-}
+pinnedCreate :: MonadIO m => Fold m Word8 (Array Word8)
+
+#if MIN_VERSION_streamly_core(0,2,2)
+ensurePinned = Array.pin
+pinnedCreateOf = Array.pinnedCreateOf
+pinnedCreate = Array.pinnedCreate
+#elif MIN_VERSION_streamly_core(0,2,0)
+ensurePinned = Array.pin
+pinnedCreateOf = Array.pinnedWriteN
+pinnedCreate = Array.pinnedWrite
+#else
+ensurePinned = pure
+pinnedCreateOf = Array.writeN
+pinnedCreate = Array.write
+#endif
+
 {-# INLINE mutableByteArrayContents# #-}
 mutableByteArrayContents# :: MutableByteArray# RealWorld -> Addr#
 mutableByteArrayContents# marr# = byteArrayContents# (unsafeCoerce# marr#)
@@ -92,7 +121,7 @@ makeForeignPtr (MUT_BYTE_ARRAY marr#) (I# off#) =
 {-# INLINE toArray #-}
 toArray :: ByteString -> Array Word8
 toArray (CONSTRUCTOR((ForeignPtr addr# _), _, _))
-    | Ptr addr# == nullPtr = Array MutBA.nil 0 0
+    | Ptr addr# == nullPtr = Array NIL 0 0
 toArray (CONSTRUCTOR((ForeignPtr addr# (PlainPtr marr#)), off0, len)) =
     let off = I# (addr# `minusAddr#` mutableByteArrayContents# marr#)
                   WHEN_0_10_12(+ off0)
@@ -108,18 +137,24 @@ toArray (CONSTRUCTOR(fptr, off, len)) =
         Unfold.mkUnfoldrM
             (\ptr -> flip StreamD.Yield (ptr `plusPtr` 1) <$> peek ptr)
 
--- | Convert an array of 'Word8' to a 'ByteString'. This function unwraps the
--- 'Array' and wraps it with 'ByteString' constructors and hence the operation
--- is performed in constant time.
+-- | Convert an array of 'Word8' to a 'ByteString'.
+--
+-- Please ensure that the array is pinned when using this function.
+
+-- If the array is pinned, the operation is performed in constant time. Whereas
+-- for an unpinned array a copy is involved to pin it.
+--
 {-# INLINE fromArray #-}
 fromArray :: Array Word8 -> ByteString
-fromArray (Array {..})
+fromArray arr
     | aLen == 0 = mempty
-    | otherwise = CONSTRUCTOR((makeForeignPtr arrContents arrStart), 0, aLen)
+    | otherwise = unsafeInlineIO $ do
+        Array{..} <- ensurePinned arr
+        pure $ CONSTRUCTOR((makeForeignPtr arrContents arrStart), 0, aLen)
 
     where
 
-    aLen = arrEnd - arrStart
+    aLen = arrEnd arr - arrStart arr
 
 -- | Unfold a strict ByteString to a stream of Word8.
 {-# INLINE reader #-}
@@ -129,12 +164,12 @@ reader = lmap toArray Array.reader
 -- | Fold a stream of Word8 to a strict ByteString of given size in bytes.
 {-# INLINE writeN #-}
 writeN :: MonadIO m => Int -> Fold m Word8 ByteString
-writeN i = fromArray <$> Array.writeN i
+writeN i = fromArray <$> pinnedCreateOf i
 
 -- | Fold a stream of Word8 to a strict ByteString of appropriate size.
 {-# INLINE write #-}
 write :: MonadIO m => Fold m Word8 ByteString
-write = fromArray <$> Array.write
+write = fromArray <$> pinnedCreate
 
 --------------------------------------------------------------------------------
 -- Deprecated