Add Data.ByteString.Short.unconsN

Fixes #524

Author: hasufell

Data/ByteString/Short.hs

@@ -82,6 +82,7 @@ module Data.ByteString.Short (
     last,
     tail,
     uncons,
+    unconsN,
     head,
     init,
     unsnoc,

Data/ByteString/Short/Internal.hs

@@ -57,6 +57,7 @@ module Data.ByteString.Short.Internal (
     last,
     tail,
     uncons,
+    unconsN,
     head,
     init,
     unsnoc,
@@ -758,6 +759,33 @@ uncons = \sbs ->
                            t = create nl $ \mba -> copyByteArray (asBA sbs) 1 mba 0 nl
                        in Just (h, t)
 
+-- | /O(n)/ Extract the given number of elements from a ShortByteString and the remainder.
+-- Returns 'Nothing' if the ShortByteString is smaller than the requested number of elements.
+--
+-- >>> unconsN 3 "abcdefg"
+-- Just ([97,98,99],"defg")
+-- >>> unconsN 11 "abcdefg"
+-- Nothing
+-- >>> unconsN 0 "abcdefg"
+-- Nothing
+--
+-- Satisfies the following properties:
+--
+-- > \x i -> unconsN i x == (if length x < i || i < 1 then Nothing else let u = unpack in Just (take i u, pack (drop i u)))
+-- > \x -> unconsN 1 x == fmap (\(x, y) -> ([x], y)) (uncons x)
+-- > \x i -> maybe i (\(xs, _) -> length xs) (unconsN i x) === i
+--
+-- @since 0.11.3.2
+unconsN :: Int -> ShortByteString -> Maybe ([Word8], ShortByteString)
+unconsN n = \sbs ->
+  let l  = length sbs
+      nl = l - n
+      ix = n - 1
+  in if | n < 1 || l <= ix -> Nothing
+        | otherwise -> let h = List.map (indexWord8Array (asBA sbs)) [0..ix]
+                           t = create nl $ \mba -> copyByteArray (asBA sbs) n mba 0 nl
+                       in Just (h, t)
+
 -- | /O(1)/ Extract the first element of a ShortByteString, which must be non-empty.
 -- An exception will be thrown in the case of an empty ShortByteString.
 --

bench/BenchShort.hs

@@ -247,4 +247,31 @@ benchShort = bgroup "ShortByteString"
         , bench "unpack and get last element" $ nf (\x -> last . S.unpack $ x) absurdlong
         , bench "unpack and get first 120 elements" $ nf (\x -> take 120 . S.unpack $ x) absurdlong
         ]
+    , bgroup "ShortByteString unpack/uncons comparison" $
+        [ bench "unpack and look at first 5 elements"   $ nf (unpack5) absurdlong
+        , bench "uncons consecutively 5 times"          $ nf (uncons5) absurdlong
+        , bench "unconsN 5"                             $ nf (unconsN) absurdlong
+        ]
     ]
+
+
+unpack5 :: ShortByteString -> Bool
+unpack5 sbs = case S.unpack sbs of
+                (a:b:c:d:e:_) -> True
+                _             -> False
+
+
+uncons5 :: ShortByteString -> Bool
+uncons5 sbs
+  | Just (a, r1) <- S.uncons sbs
+  , Just (b, r2) <- S.uncons r1
+  , Just (c, r3) <- S.uncons r2
+  , Just (d, r4) <- S.uncons r3
+  , Just (e, xs) <- S.uncons r4 = True
+  | otherwise = False
+
+unconsN :: ShortByteString -> Bool
+unconsN sbs = case S.unconsN 5 sbs of
+                Just ([a, b, c, d, e], xs) -> True
+                Just _ -> error "oops"
+                _      -> False

tests/Properties/ByteString.hs

@@ -655,6 +655,17 @@ tests =
     \s -> fromString s == B.pack (map (fromIntegral . ord :: Char -> Word8) s)
   , testProperty "fromString literal" $
     fromString "\0\1\2\3\4" == B.pack [0,1,2,3,4]
+#endif
+#ifdef BYTESTRING_SHORT
+  , testProperty "unconsN == unpack" $
+    \x i -> B.unconsN i x === (if B.length x < i || i < 1 then Nothing else let u = B.unpack x
+                                                                                l = take i u
+                                                                                r = B.pack (drop i u)
+                                                                            in Just (l, r))
+  , testProperty "unconsN 1 == uncons" $
+    \x -> B.unconsN 1 x === fmap (\(x, y) -> ([x], y)) (B.uncons x)
+  , testProperty "length of items matches input (unconsN)" $
+    \x i -> maybe i (\(xs, _) -> length xs) (B.unconsN i x) === i
 #endif
   ]