Skip to content

Commit

Permalink
Add Data.ByteString.Short.unconsN
Browse files Browse the repository at this point in the history
Fixes #524
  • Loading branch information
hasufell committed Jul 3, 2022
1 parent 2b9416d commit d1896da
Show file tree
Hide file tree
Showing 4 changed files with 66 additions and 0 deletions.
1 change: 1 addition & 0 deletions Data/ByteString/Short.hs
Original file line number Diff line number Diff line change
Expand Up @@ -82,6 +82,7 @@ module Data.ByteString.Short (
last,
tail,
uncons,
unconsN,
head,
init,
unsnoc,
Expand Down
28 changes: 28 additions & 0 deletions Data/ByteString/Short/Internal.hs
Original file line number Diff line number Diff line change
Expand Up @@ -58,6 +58,7 @@ module Data.ByteString.Short.Internal (
last,
tail,
uncons,
unconsN,
head,
init,
unsnoc,
Expand Down Expand Up @@ -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.
--
Expand Down
26 changes: 26 additions & 0 deletions bench/BenchShort.hs
Original file line number Diff line number Diff line change
Expand Up @@ -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
11 changes: 11 additions & 0 deletions tests/Properties/ByteString.hs
Original file line number Diff line number Diff line change
Expand Up @@ -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
]

Expand Down

0 comments on commit d1896da

Please sign in to comment.