Use unsafeWithForeignPtr in Data.ByteString.Unsafe (haskell#401)

* Use unsafeWithForeignPtr in Data.ByteString.Unsafe

* Review remaining cases of withForeignPtr

* Review suggestions

Data/ByteString.hs

@@ -1146,8 +1146,7 @@ split w (BS x l) = loop 0
                              w (fromIntegral (l-n))
             in if q == nullPtr
                 then [BS (plusForeignPtr x n) (l-n)]
-                else let i = accursedUnutterablePerformIO $ withForeignPtr x $ \p ->
-                               return (q `minusPtr` p)
+                else let i = q `minusPtr` unsafeForeignPtrToPtr x
                       in BS (plusForeignPtr x n) (i-n) : loop (i+1)
 
 {-# INLINE split #-}
@@ -1288,7 +1287,7 @@ count w (BS x m) = accursedUnutterablePerformIO $ unsafeWithForeignPtr x $ \p ->
 -- returns the index of the first element in the ByteString
 -- satisfying the predicate.
 findIndex :: (Word8 -> Bool) -> ByteString -> Maybe Int
-findIndex k (BS x l) = accursedUnutterablePerformIO $ withForeignPtr x g
+findIndex k (BS x l) = accursedUnutterablePerformIO $ unsafeWithForeignPtr x g
   where
     g !ptr = go 0
       where
@@ -1305,7 +1304,7 @@ findIndex k (BS x l) = accursedUnutterablePerformIO $ withForeignPtr x g
 --
 -- @since 0.10.12.0
 findIndexEnd :: (Word8 -> Bool) -> ByteString -> Maybe Int
-findIndexEnd k (BS x l) = accursedUnutterablePerformIO $ withForeignPtr x g
+findIndexEnd k (BS x l) = accursedUnutterablePerformIO $ unsafeWithForeignPtr x g
   where
     g !ptr = go (l-1)
       where
@@ -1358,7 +1357,7 @@ filter k = \ps@(BS x l) ->
   if null ps
     then ps
     else
-      unsafePerformIO $ createAndTrim l $ \pOut -> withForeignPtr x $ \pIn -> do
+      unsafePerformIO $ createAndTrim l $ \pOut -> unsafeWithForeignPtr x $ \pIn -> do
         let
           go' pf pt = go pf pt
             where
@@ -1417,7 +1416,7 @@ find f p = case findIndex f p of
 partition :: (Word8 -> Bool) -> ByteString -> (ByteString, ByteString)
 partition f s = unsafeDupablePerformIO $
     do fp' <- mallocByteString len
-       withForeignPtr fp' $ \p ->
+       unsafeWithForeignPtr fp' $ \p ->
            do let end = p `plusPtr` (len - 1)
               mid <- sep 0 p end
               rev mid end
@@ -1456,7 +1455,7 @@ isPrefixOf (BS x1 l1) (BS x2 l2)
     | l1 == 0   = True
     | l2 < l1   = False
     | otherwise = accursedUnutterablePerformIO $ unsafeWithForeignPtr x1 $ \p1 ->
-        withForeignPtr x2 $ \p2 -> do
+        unsafeWithForeignPtr x2 $ \p2 -> do
             i <- memcmp p1 p2 (fromIntegral l1)
             return $! i == 0
 
@@ -1484,7 +1483,7 @@ isSuffixOf (BS x1 l1) (BS x2 l2)
     | l1 == 0   = True
     | l2 < l1   = False
     | otherwise = accursedUnutterablePerformIO $ unsafeWithForeignPtr x1 $ \p1 ->
-        withForeignPtr x2 $ \p2 -> do
+        unsafeWithForeignPtr x2 $ \p2 -> do
             i <- memcmp p1 (p2 `plusPtr` (l2 - l1)) (fromIntegral l1)
             return $! i == 0
 
@@ -1608,8 +1607,8 @@ zipWith f ps qs = case uncons ps of
 -- @since 0.11.1.0
 packZipWith :: (Word8 -> Word8 -> Word8) -> ByteString -> ByteString -> ByteString
 packZipWith f (BS fp l) (BS fq m) = unsafeDupablePerformIO $
-    withForeignPtr fp $ \a ->
-    withForeignPtr fq $ \b ->
+    unsafeWithForeignPtr fp $ \a ->
+    unsafeWithForeignPtr fq $ \b ->
     create len $ go a b
   where
     go p1 p2 = zipWith_ 0
@@ -1653,7 +1652,7 @@ tails p | null p    = [empty]
 sort :: ByteString -> ByteString
 sort (BS input l)
   -- qsort outperforms counting sort for small arrays
-  | l <= 20 = unsafeCreate l $ \ptr -> withForeignPtr input $ \inp -> do
+  | l <= 20 = unsafeCreate l $ \ptr -> unsafeWithForeignPtr input $ \inp -> do
     memcpy ptr inp (fromIntegral l)
     c_sort ptr (fromIntegral l)
   | otherwise = unsafeCreate l $ \p -> allocaArray 256 $ \arr -> do
@@ -1688,11 +1687,12 @@ sort (BS input l)
 -- subcomputation finishes.
 useAsCString :: ByteString -> (CString -> IO a) -> IO a
 useAsCString (BS fp l) action =
- allocaBytes (l+1) $ \buf ->
-   withForeignPtr fp $ \p -> do
-     memcpy buf p (fromIntegral l)
-     pokeByteOff buf l (0::Word8)
-     action (castPtr buf)
+  allocaBytes (l+1) $ \buf ->
+    -- Cannot use unsafeWithForeignPtr, because action can diverge
+    withForeignPtr fp $ \p -> do
+      memcpy buf p (fromIntegral l)
+      pokeByteOff buf l (0::Word8)
+      action (castPtr buf)
 
 -- | /O(n) construction/ Use a @ByteString@ with a function requiring a @CStringLen@.
 -- As for @useAsCString@ this function makes a copy of the original @ByteString@.

Data/ByteString/Internal.hs

@@ -299,7 +299,7 @@ instance TH.Lift ByteString where
 -- of the string if no element is found, rather than Nothing.
 findIndexOrLength :: (Word8 -> Bool) -> ByteString -> Int
 findIndexOrLength k (BS x l) =
-    accursedUnutterablePerformIO $ withForeignPtr x g
+    accursedUnutterablePerformIO $ unsafeWithForeignPtr x g
   where
     g ptr = go 0
       where
@@ -565,19 +565,21 @@ unsafeCreateUptoN' l f = unsafeDupablePerformIO (createUptoN' l f)
 
 -- | Create ByteString of size @l@ and use action @f@ to fill its contents.
 create :: Int -> (Ptr Word8 -> IO ()) -> IO ByteString
-create l f = do
+create l action = do
     fp <- mallocByteString l
-    withForeignPtr fp $ \p -> f p
+    -- Cannot use unsafeWithForeignPtr, because action can diverge
+    withForeignPtr fp $ \p -> action p
     return $! BS fp l
 {-# INLINE create #-}
 
 -- | Given a maximum size @l@ and an action @f@ that fills the 'ByteString'
 -- starting at the given 'Ptr' and returns the actual utilized length,
 -- @`createUpToN'` l f@ returns the filled 'ByteString'.
 createUptoN :: Int -> (Ptr Word8 -> IO Int) -> IO ByteString
-createUptoN l f = do
+createUptoN l action = do
     fp <- mallocByteString l
-    l' <- withForeignPtr fp $ \p -> f p
+    -- Cannot use unsafeWithForeignPtr, because action can diverge
+    l' <- withForeignPtr fp $ \p -> action p
     assert (l' <= l) $ return $! BS fp l'
 {-# INLINE createUptoN #-}
 
@@ -586,9 +588,10 @@ createUptoN l f = do
 --
 -- @since 0.10.12.0
 createUptoN' :: Int -> (Ptr Word8 -> IO (Int, a)) -> IO (ByteString, a)
-createUptoN' l f = do
+createUptoN' l action = do
     fp <- mallocByteString l
-    (l', res) <- withForeignPtr fp $ \p -> f p
+    -- Cannot use unsafeWithForeignPtr, because action can diverge
+    (l', res) <- withForeignPtr fp $ \p -> action p
     assert (l' <= l) $ return (BS fp l', res)
 {-# INLINE createUptoN' #-}
 
@@ -601,20 +604,22 @@ createUptoN' l f = do
 -- ByteString functions, using Haskell or C functions to fill the space.
 --
 createAndTrim :: Int -> (Ptr Word8 -> IO Int) -> IO ByteString
-createAndTrim l f = do
+createAndTrim l action = do
     fp <- mallocByteString l
+    -- Cannot use unsafeWithForeignPtr, because action can diverge
     withForeignPtr fp $ \p -> do
-        l' <- f p
+        l' <- action p
         if assert (l' <= l) $ l' >= l
             then return $! BS fp l
             else create l' $ \p' -> memcpy p' p l'
 {-# INLINE createAndTrim #-}
 
 createAndTrim' :: Int -> (Ptr Word8 -> IO (Int, Int, a)) -> IO (ByteString, a)
-createAndTrim' l f = do
+createAndTrim' l action = do
     fp <- mallocByteString l
+    -- Cannot use unsafeWithForeignPtr, because action can diverge
     withForeignPtr fp $ \p -> do
-        (off, l', res) <- f p
+        (off, l', res) <- action p
         if assert (l' <= l) $ l' >= l
             then return (BS fp l, res)
             else do ps <- create l' $ \p' ->

Data/ByteString/Lazy/Char8.hs

@@ -227,7 +227,6 @@ import Data.Int (Int64)
 import Data.Word
 import qualified Data.List as List
 import Foreign.Ptr (Ptr, plusPtr)
-import Foreign.ForeignPtr (withForeignPtr)
 import Foreign.Storable (peek)
 
 import Prelude hiding
@@ -868,7 +867,7 @@ readInt bs = case L.uncons bs of
         -- the provided digits (end of input or non-digit encountered).
         accumWord acc (BI.BS fp len) =
             BI.accursedUnutterablePerformIO $
-                withForeignPtr fp $ \ptr -> do
+                BI.unsafeWithForeignPtr fp $ \ptr -> do
                     let end = ptr `plusPtr` len
                     x@(!_, !_, !_) <- if positive
                         then digits intmaxQuot10 intmaxRem10 end ptr 0 acc

Data/ByteString/Lazy/Internal.hs

@@ -52,7 +52,6 @@ import Prelude hiding (concat)
 import qualified Data.ByteString.Internal as S
 
 import Data.Word (Word8)
-import Foreign.ForeignPtr (withForeignPtr)
 import Foreign.Ptr (plusPtr)
 import Foreign.Storable (Storable(sizeOf))
 
@@ -318,7 +317,7 @@ toStrict = \cs -> goLen0 cs cs
     goCopy Empty                    !_   = return ()
     goCopy (Chunk (S.BS _  0  ) cs) !ptr = goCopy cs ptr
     goCopy (Chunk (S.BS fp len) cs) !ptr =
-      withForeignPtr fp $ \p -> do
+      S.unsafeWithForeignPtr fp $ \p -> do
         S.memcpy ptr p len
         goCopy cs (ptr `plusPtr` len)
 -- See the comment on Data.ByteString.Internal.concat for some background on

Data/ByteString/Unsafe.hs

@@ -73,7 +73,7 @@ import GHC.Ptr                  (Ptr(..))
 -- to provide a proof that the ByteString is non-empty.
 unsafeHead :: ByteString -> Word8
 unsafeHead (BS x l) = assert (l > 0) $
-    accursedUnutterablePerformIO $ withForeignPtr x $ \p -> peek p
+    accursedUnutterablePerformIO $ unsafeWithForeignPtr x $ \p -> peek p
 {-# INLINE unsafeHead #-}
 
 -- | A variety of 'tail' for non-empty ByteStrings. 'unsafeTail' omits the
@@ -95,7 +95,7 @@ unsafeInit (BS ps l) = assert (l > 0) $ BS ps (l-1)
 -- provide a separate proof that the ByteString is non-empty.
 unsafeLast :: ByteString -> Word8
 unsafeLast (BS x l) = assert (l > 0) $
-    accursedUnutterablePerformIO $ withForeignPtr x $ \p -> peekByteOff p (l-1)
+    accursedUnutterablePerformIO $ unsafeWithForeignPtr x $ \p -> peekByteOff p (l-1)
 {-# INLINE unsafeLast #-}
 
 -- | Unsafe 'ByteString' index (subscript) operator, starting from 0, returning a 'Word8'
@@ -104,7 +104,7 @@ unsafeLast (BS x l) = assert (l > 0) $
 -- other way.
 unsafeIndex :: ByteString -> Int -> Word8
 unsafeIndex (BS x l) i = assert (i >= 0 && i < l) $
-    accursedUnutterablePerformIO $ withForeignPtr x $ \p -> peekByteOff p i
+    accursedUnutterablePerformIO $ unsafeWithForeignPtr x $ \p -> peekByteOff p i
 {-# INLINE unsafeIndex #-}
 
 -- | A variety of 'take' which omits the checks on @n@ so there is an
@@ -262,7 +262,8 @@ unsafePackMallocCStringLen (cstr, len) = do
 -- after this.
 --
 unsafeUseAsCString :: ByteString -> (CString -> IO a) -> IO a
-unsafeUseAsCString (BS ps _) ac = withForeignPtr ps $ \p -> ac (castPtr p)
+unsafeUseAsCString (BS ps _) action = withForeignPtr ps $ \p -> action (castPtr p)
+-- Cannot use unsafeWithForeignPtr, because action can diverge
 
 -- | /O(1) construction/ Use a 'ByteString' with a function requiring a
 -- 'CStringLen'.
@@ -281,4 +282,5 @@ unsafeUseAsCString (BS ps _) ac = withForeignPtr ps $ \p -> ac (castPtr p)
 --
 -- If 'Data.ByteString.empty' is given, it will pass @('Foreign.Ptr.nullPtr', 0)@.
 unsafeUseAsCStringLen :: ByteString -> (CStringLen -> IO a) -> IO a
-unsafeUseAsCStringLen (BS ps l) f = withForeignPtr ps $ \p -> f (castPtr p,l)
+unsafeUseAsCStringLen (BS ps l) action = withForeignPtr ps $ \p -> action (castPtr p, l)
+-- Cannot use unsafeWithForeignPtr, because action can diverge