From d1896da5e04799aa244b0833e41fb5ad1719aa18 Mon Sep 17 00:00:00 2001
From: Julian Ospald <hasufell@posteo.de>
Date: Sun, 3 Jul 2022 15:47:49 +0200
Subject: [PATCH] Add Data.ByteString.Short.unconsN

Fixes #524
---
 Data/ByteString/Short.hs          |  1 +
 Data/ByteString/Short/Internal.hs | 28 ++++++++++++++++++++++++++++
 bench/BenchShort.hs               | 26 ++++++++++++++++++++++++++
 tests/Properties/ByteString.hs    | 11 +++++++++++
 4 files changed, 66 insertions(+)

diff --git a/Data/ByteString/Short.hs b/Data/ByteString/Short.hs
index e1339e970..8ca0c0c67 100644
--- a/Data/ByteString/Short.hs
+++ b/Data/ByteString/Short.hs
@@ -82,6 +82,7 @@ module Data.ByteString.Short (
     last,
     tail,
     uncons,
+    unconsN,
     head,
     init,
     unsnoc,
diff --git a/Data/ByteString/Short/Internal.hs b/Data/ByteString/Short/Internal.hs
index cc18e3a31..93f262260 100644
--- a/Data/ByteString/Short/Internal.hs
+++ b/Data/ByteString/Short/Internal.hs
@@ -58,6 +58,7 @@ module Data.ByteString.Short.Internal (
     last,
     tail,
     uncons,
+    unconsN,
     head,
     init,
     unsnoc,
@@ -743,6 +744,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.
 --
diff --git a/bench/BenchShort.hs b/bench/BenchShort.hs
index f6c37662e..8447ac20a 100644
--- a/bench/BenchShort.hs
+++ b/bench/BenchShort.hs
@@ -231,5 +231,31 @@ benchShort = bgroup "ShortByteString"
         , bench "FindIndex/inlined"       $ nf (S.findIndex      (== nl)) absurdlong
         , bench "FindIndex/non-inlined"   $ nf (S.findIndex   (nilEq nl)) 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
diff --git a/tests/Properties/ByteString.hs b/tests/Properties/ByteString.hs
index 757059701..1f7599452 100644
--- a/tests/Properties/ByteString.hs
+++ b/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
   ]