IO operations for ShortByteString

Data/ByteString.hs

@@ -1853,7 +1853,7 @@ mkBigPS _ pss = return $! concat (P.reverse pss)
 -- | Outputs a 'ByteString' to the specified 'Handle'.
 hPut :: Handle -> ByteString -> IO ()
 hPut _ (BS _  0) = return ()
-hPut h (BS ps l) = unsafeWithForeignPtr ps $ \p-> hPutBuf h p l
+hPut h (BS ps l) = unsafeWithForeignPtr ps $ \p -> hPutBuf h p l
 
 -- | Similar to 'hPut' except that it will never block. Instead it returns
 -- any tail that did not get written. This tail may be 'empty' in the case that

Data/ByteString/Short/Internal.hs

@@ -12,6 +12,7 @@
 {-# LANGUAGE TypeFamilies             #-}
 {-# LANGUAGE UnboxedTuples            #-}
 {-# LANGUAGE UnliftedFFITypes         #-}
+{-# LANGUAGE ViewPatterns #-}
 {-# LANGUAGE Unsafe                   #-}
 {-# OPTIONS_HADDOCK not-home #-}
 
@@ -154,6 +155,15 @@ module Data.ByteString.Short.Internal (
     -- ** Using ShortByteStrings as 'Foreign.C.String.CString's
     useAsCString,
     useAsCStringLen,
+
+    -- * Low level I\/O with 'ShortByteString's
+    unsafeHPutOff,
+    hPut,
+    readFile,
+    hGetContents,
+    hGetLine,
+    hGet,
+    hGetSome,
   ) where
 
 import Data.ByteString.Internal
@@ -189,8 +199,6 @@ import Control.DeepSeq
   ( NFData(..) )
 import Control.Exception
   ( assert )
-import Control.Monad
-  ( (>>) )
 import Foreign.C.String
   ( CString
   , CStringLen
@@ -237,7 +245,16 @@ import GHC.Exts
   , indexWord8Array#, indexCharArray#
   , writeWord8Array#
   , unsafeFreezeByteArray#
-  , touch# )
+  , touch#
+#if __GLASGOW_HASKELL__ >= 901
+  , keepAlive#
+#endif
+#if __GLASGOW_HASKELL__ >= 801
+  , getSizeofMutableByteArray#
+#else
+  , sizeofMutableByteArray#
+#endif
+  )
 import GHC.IO
 import GHC.ForeignPtr
   ( ForeignPtr(ForeignPtr)
@@ -263,14 +280,18 @@ import Prelude
   , Maybe(..)
   , not
   , snd
+  , Monad(..)
   )
 
 import qualified Data.ByteString.Internal as BS
-
+import qualified Data.ByteString as BS
 import qualified Data.List as List
 import qualified GHC.Exts
 import qualified Language.Haskell.TH.Lib as TH
 import qualified Language.Haskell.TH.Syntax as TH
+import Foreign.Ptr (plusPtr)
+import qualified System.IO as IO
+import Control.Monad ((=<<))
 
 -- | A compact representation of a 'Word8' vector.
 --
@@ -411,6 +432,37 @@ unsafePackLenLiteral len addr# =
 ------------------------------------------------------------------------
 -- Internal utils
 
+-- | Yields a pinned byte sequence whose contents are identical to those
+-- of the original byte sequence. If the @ByteArray@ backing the argument
+-- was already pinned, this simply aliases the argument and does not perform
+-- any copying.
+pin :: ShortByteString -> ShortByteString
+pin b@(SBS (BA# -> src)) =
+  if isByteArrayPinned src
+    then b
+    else runST $ do
+      let len = sizeofByteArray src
+      dst <- newPinnedByteArray len
+      copyByteArray src 0 dst 0 len
+      (\(BA# res) -> SBS res) <$> unsafeFreezeByteArray dst
+
+
+-- | Invariant: @arr@ must be pinned. This is not checked.
+pinnedWithPtr :: ShortByteString -> (Ptr a -> IO b) -> IO b
+#if __GLASGOW_HASKELL__ >= 901
+pinnedWithPtr (SBS (BA# -> arr)) f = IO $ \s ->
+  case f (byteArrayContents arr) of
+    IO action# -> keepAlive# arr s action#
+#else
+pinnedWithPtr (SBS (BA# -> arr)) f = do
+  x <- f (byteArrayContents arr)
+  touchByteArrayIO arr
+  pure x
+#endif
+
+withPtr :: ShortByteString -> (Ptr a -> IO b) -> IO b
+withPtr arr = pinnedWithPtr $ pin arr
+
 asBA :: ShortByteString -> BA
 asBA (SBS ba#) = BA# ba#
 
@@ -529,6 +581,29 @@ fromShortIO sbs = do
                         (PlainPtr mba#)
     return (BS fp len)
 
+unsafePlainToShort :: ByteString -> ShortByteString
+unsafePlainToShort = unsafeDupablePerformIO . unsafePlainToShortIO
+
+-- | /O(1)/ zero cost conversion between 'ByteString' and 'ShortByteString'.
+-- There are three invariants.
+-- The 'ByteString' must also not be used after the function is called because we use `unsafeFreezeByteArray#' internally.
+-- The 'ByteString' must be created using the @mallocPlain*@ functions.
+-- The 'ByteString' must not be a slice.
+-- That is, the 'ForeignPtr' should point to the start of the pinned 'MutableByteArray#' and
+-- the length should be equal to @sizeofMutableByteArray# marr#@.
+unsafePlainToShortIO :: ByteString -> IO ShortByteString
+unsafePlainToShortIO (BS (ForeignPtr (Ptr -> p) fpc) l) =
+  case fpc of
+    PlainPtr (MBA# -> marr) -> do
+      (BA# arr#) <- stToIO $ unsafeFreezeByteArray marr
+      let baseP = mutableByteArrayContents marr
+      marrSize <- stToIO $ getSizeofMutableByteArray marr
+      if baseP == p && l == marrSize
+        then pure $ SBS arr#
+        else error "Data.ByteString.Short.Internal: cannot be a slice"
+    _ -> error "Data.ByteString.Short.Internal: must be PlainPtr"
+
+
 -- | /O(1)/ Convert a 'Word8' into a 'ShortByteString'
 --
 -- @since 0.11.3.0
@@ -1649,6 +1724,12 @@ newPinnedByteArray (I# len#) =
     ST $ \s -> case newPinnedByteArray# len# s of
                  (# s, mba# #) -> (# s, MBA# mba# #)
 
+sizeofByteArray :: BA -> Int
+sizeofByteArray (BA# ba#) = I# (sizeofByteArray# ba#)
+
+byteArrayContents :: BA -> Ptr a
+byteArrayContents (BA# ba#) = Ptr (byteArrayContents# ba#)
+
 unsafeFreezeByteArray :: MBA s -> ST s BA
 unsafeFreezeByteArray (MBA# mba#) =
     ST $ \s -> case unsafeFreezeByteArray# mba# s of
@@ -1691,6 +1772,40 @@ copyMutableByteArray (MBA# src#) (I# src_off#) (MBA# dst#) (I# dst_off#) (I# len
     ST $ \s -> case copyMutableByteArray# src# src_off# dst# dst_off# len# s of
                  s -> (# s, () #)
 
+-- | Yield a pointer to the array's data. This operation is only safe on
+-- /pinned/ byte arrays allocated by 'newPinnedByteArray' or
+-- 'newAlignedPinnedByteArray'.
+mutableByteArrayContents :: MBA s -> Ptr Word8
+mutableByteArrayContents (MBA# arr#) = Ptr (byteArrayContents# (unsafeCoerce# arr#))
+
+-- | Get the size of a byte array in bytes. Unlike 'sizeofMutableByteArray',
+-- this function ensures sequencing in the presence of resizing.
+getSizeofMutableByteArray :: MBA s -> ST s Int
+#if __GLASGOW_HASKELL__ >= 801
+getSizeofMutableByteArray (MBA# arr#) =
+  ST (\s# -> case getSizeofMutableByteArray# arr# s# of (# s'#, n# #) -> (# s'#, I# n# #))
+#else
+getSizeofMutableByteArray (MBA# arr#) = pure (I# (sizeofMutableByteArray# arr#))
+#endif
+
+-- | Check whether or not the byte array is pinned. Pinned byte arrays cannot
+-- be moved by the garbage collector. It is safe to use 'byteArrayContents' on
+-- such byte arrays.
+--
+-- Caution: This function is only available when compiling with GHC 8.2 or
+-- newer.
+--
+-- @since 0.6.4.0
+isByteArrayPinned :: BA -> Bool
+#if __GLASGOW_HASKELL__ >= 802
+isByteArrayPinned (BA# arr#) = isTrue# (isByteArrayPinned# arr#)
+#else
+isByteArrayPinned _ = False
+#endif
+
+touchByteArrayIO :: BA -> IO ()
+touchByteArrayIO (BA# arr) = IO $ \s -> case touch# arr s of
+  s -> (# s, () #)
 
 ------------------------------------------------------------------------
 -- FFI imports
@@ -1890,3 +2005,29 @@ moduleError :: HasCallStack => String -> String -> a
 moduleError fun msg = error (moduleErrorMsg fun msg)
 {-# NOINLINE moduleError #-}
 
+------------------------------------------------------------------------
+-- IO
+
+-- | Outputs 'ShortByteString' to the specified 'Handle'. This is implemented
+-- with 'IO.hPutBuf'. The offset and length are used because we don't have slices of 'ByteArray' yet.
+-- The function is unsafe because the offset and length is not checked.
+unsafeHPutOff :: IO.Handle -> ShortByteString -> Int -> Int -> IO ()
+unsafeHPutOff handle sbs off len = withPtr sbs $ \p -> IO.hPutBuf handle (p `plusPtr` off) len
+
+hPut :: IO.Handle -> ShortByteString -> IO ()
+hPut h sbs = unsafeHPutOff h sbs 0 (length sbs)
+
+readFile :: FilePath -> IO ShortByteString
+readFile path = unsafePlainToShortIO =<< BS.readFile path 
+
+hGetContents :: IO.Handle -> IO ShortByteString
+hGetContents h = unsafePlainToShortIO =<< BS.hGetContents h
+
+hGetLine :: IO.Handle -> IO ShortByteString
+hGetLine h = unsafePlainToShortIO =<< BS.hGetLine h
+
+hGet :: IO.Handle -> Int -> IO ShortByteString
+hGet h i = unsafePlainToShortIO =<< BS.hGet h i
+
+hGetSome :: IO.Handle -> Int -> IO ShortByteString
+hGetSome h i = unsafePlainToShortIO =<< BS.hGetSome h i