New API to convert directly between lazy and short bytestrings (#619)

* New API to convert directly between Lazy and Short

* fixup! New API to convert directly between Lazy and Short

Matthew Craven review fixes

---------

Co-authored-by: Viktor Dukhovni <[email protected]>

Author: vdukhovni

Changelog.md

@@ -4,9 +4,12 @@
   * [`Data.Data.dataTypeOf` for `StrictByteString` and `LazyByteString` now returns a `DataType` that uses `AlgRep` instead of `NoRep`.](https://github.com/haskell/bytestring/pull/614)
     * This allows utilities like `syb:Data.Generics.Text.gread` to be meaningfully used at these types containing `ByteString`s.
   * [`fromListN` in `instance IsList ByteString` truncates input list if it's longer than the size hint](https://github.com/haskell/bytestring/pull/672)
+
+* API additions and behavior changes:
+  * Data.ByteString.Short now provides `lazyToShort` and `lazyFromShort`.
+
 <!--
 * Bug fixes:
-* API additions and behavior changes:
 * Deprecations:
 * Performance improvements:
 * Miscellaneous:

Data/ByteString/Short.hs

@@ -74,6 +74,8 @@ module Data.ByteString.Short (
     unpack,
     fromShort,
     toShort,
+    lazyFromShort,
+    lazyToShort,
 
     -- * Basic interface
     snoc,

Data/ByteString/Short/Internal.hs

@@ -33,6 +33,8 @@ module Data.ByteString.Short.Internal (
     unpack,
     fromShort,
     toShort,
+    lazyFromShort,
+    lazyToShort,
 
     -- * Basic interface
     snoc,
@@ -242,6 +244,7 @@ import Prelude
   )
 
 import qualified Data.ByteString.Internal.Type as BS
+import qualified Data.ByteString.Lazy.Internal as LBS
 
 import qualified Data.List as List
 import qualified GHC.Exts
@@ -450,6 +453,35 @@ toShortIO (BS fptr len) = do
       stToIO (copyAddrToByteArray ptr mba 0 len)
     ShortByteString <$> stToIO (unsafeFreezeByteArray mba)
 
+-- | A simple wrapper around 'fromShort' that wraps the strict 'ByteString' as
+-- a one-chunk 'LBS.LazyByteString'.
+lazyFromShort :: ShortByteString -> LBS.ByteString
+lazyFromShort = LBS.fromStrict . fromShort
+
+-- | /O(n)/. Convert an 'LBS.LazyByteString' into a 'ShortByteString'.
+--
+-- This makes a copy, so does not retain the input string.  Naturally, best
+-- used only with sufficiently short lazy ByteStrings.  The entire lazy
+-- ByteString is brought into memory before a copy is made.
+--
+lazyToShort :: LBS.ByteString -> ShortByteString
+lazyToShort LBS.Empty = empty
+lazyToShort lbs = unsafeDupablePerformIO $ do
+        mba <- stToIO (newByteArray total)
+        copy mba lbs
+        ShortByteString <$> stToIO (unsafeFreezeByteArray mba)
+  where
+    !total = LBS.foldlChunks (\acc (BS _ l) -> checkedAdd "short.lazyToShort" acc l) 0 lbs
+
+    copy :: MutableByteArray RealWorld -> LBS.ByteString -> IO ()
+    copy mba = go 0
+      where
+        go off (LBS.Chunk (BS fp len) cs) = do
+            BS.unsafeWithForeignPtr fp $ \p ->
+                stToIO $ copyAddrToByteArray p mba off len
+            go (off + len) cs
+        go !_ LBS.Empty = pure ()
+
 -- | /O(n)/. Convert a 'ShortByteString' into a 'ByteString'.
 --
 fromShort :: ShortByteString -> ByteString

tests/Properties.hs

@@ -299,7 +299,7 @@ prop_stimesOverflowScary bs =
 prop_stimesOverflowEmpty = forAll (choose (0, maxBound @Word)) $ \n ->
   stimes n mempty === mempty @P.ByteString
 
-concat32bitOverflow :: (Int -> a) -> ([a] -> a) -> Property
+concat32bitOverflow :: (Int -> a) -> ([a] -> b) -> Property
 concat32bitOverflow replicateLike concatLike = let
   intBits = finiteBitSize @Int 0
   largeBS = concatLike $ replicate (bit 14) $ replicateLike (bit 17)
@@ -315,6 +315,10 @@ prop_32bitOverflow_Lazy_toStrict :: Property
 prop_32bitOverflow_Lazy_toStrict =
   concat32bitOverflow (`P.replicate` 0) (L.toStrict . L.fromChunks)
 
+prop_32bitOverflow_Lazy_toShort :: Property
+prop_32bitOverflow_Lazy_toShort =
+  concat32bitOverflow (`P.replicate` 0) (Short.lazyToShort . L.fromChunks)
+
 prop_32bitOverflow_Short_mconcat :: Property
 prop_32bitOverflow_Short_mconcat =
   concat32bitOverflow makeShort mconcat
@@ -530,6 +534,8 @@ prop_short_pack_unpack xs =
     (Short.unpack . Short.pack) xs == xs
 prop_short_toShort_fromShort bs =
     (Short.fromShort . Short.toShort) bs == bs
+prop_short_lazyToShort_fromShort lbs =
+    (Short.lazyFromShort . Short.lazyToShort) lbs == lbs
 
 prop_short_toShort_unpack bs =
     (Short.unpack . Short.toShort) bs == P.unpack bs
@@ -596,6 +602,7 @@ prop_short_pinned (NonNegative (I# len#)) = runST $ ST $ \s ->
 short_tests =
     [ testProperty "pack/unpack"              prop_short_pack_unpack
     , testProperty "toShort/fromShort"        prop_short_toShort_fromShort
+    , testProperty "lazyToShort/fromShort"    prop_short_lazyToShort_fromShort
     , testProperty "toShort/unpack"           prop_short_toShort_unpack
     , testProperty "pack/fromShort"           prop_short_pack_fromShort
     , testProperty "empty"                    prop_short_empty
@@ -663,6 +670,7 @@ overflow_tests =
     , testProperty "StrictByteString stimes (empty)" prop_stimesOverflowEmpty
     , testProperty "StrictByteString mconcat" prop_32bitOverflow_Strict_mconcat
     , testProperty "LazyByteString toStrict"  prop_32bitOverflow_Lazy_toStrict
+    , testProperty "LazyByteString toShort"   prop_32bitOverflow_Lazy_toShort
     , testProperty "ShortByteString mconcat"  prop_32bitOverflow_Short_mconcat
     ]