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 8639a3e36..780c3f896 100644 --- a/Data/ByteString/Short/Internal.hs +++ b/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. -- diff --git a/bench/BenchShort.hs b/bench/BenchShort.hs index 17609b4b1..45ada06e9 100644 --- a/bench/BenchShort.hs +++ b/bench/BenchShort.hs @@ -1,4 +1,5 @@ {-# LANGUAGE BangPatterns #-} +{-# LANGUAGE ViewPatterns #-} {-# LANGUAGE PackageImports #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE MagicHash #-} @@ -247,4 +248,37 @@ 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 + , bench "unconsNViaSplit 5" $ nf (unconsNViaSplit) 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 + +unconsNViaSplit :: ShortByteString -> Bool +unconsNViaSplit sbs = case S.splitAt 5 sbs of + (S.unpack -> [a, b, c, d, e], xs) -> True + _ -> 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 ]