diff --git a/Data/ByteString.hs b/Data/ByteString.hs index 370454892..525b6f357 100644 --- a/Data/ByteString.hs +++ b/Data/ByteString.hs @@ -854,12 +854,14 @@ unfoldrN i f x0 | i < 0 = (empty, Just x0) | otherwise = unsafePerformIO $ createAndTrim' i $ \p -> go p x0 0 where - go !p !x !n - | n == i = return (0, n, Just x) - | otherwise = case f x of - Nothing -> return (0, n, Nothing) - Just (w,x') -> do poke p w - go (p `plusPtr` 1) x' (n+1) + go !p !x !n = go' x n + where + go' !x' !n' + | n' == i = return (0, n', Just x') + | otherwise = case f x' of + Nothing -> return (0, n', Nothing) + Just (w,x'') -> do pokeByteOff p n' w + go' x'' (n'+1) {-# INLINE unfoldrN #-} -- --------------------------------------------------------------------- diff --git a/Data/ByteString/Internal.hs b/Data/ByteString/Internal.hs index f2736058a..8e4da38e2 100644 --- a/Data/ByteString/Internal.hs +++ b/Data/ByteString/Internal.hs @@ -570,6 +570,7 @@ createAndTrim' l f = do else do ps <- create l' $ \p' -> memcpy p' (p `plusPtr` off) l' return (ps, res) +{-# INLINE createAndTrim' #-} -- | Wrapper of 'Foreign.ForeignPtr.mallocForeignPtrBytes' with faster implementation for GHC -- diff --git a/bench/BenchAll.hs b/bench/BenchAll.hs index cab3ba7f4..4234d0473 100644 --- a/bench/BenchAll.hs +++ b/bench/BenchAll.hs @@ -405,6 +405,8 @@ main = do nf (S.foldl' (\acc x -> acc + fromIntegral x) (0 :: Int)) s) foldInputs , bgroup "foldr'" $ map (\s -> bench (show $ S.length s) $ nf (S.foldr' (\x acc -> fromIntegral x + acc) (0 :: Int)) s) foldInputs + , bgroup "unfoldrN" $ map (\s -> bench (show $ S.length s) $ + nf (S.unfoldrN (S.length s) (\a -> Just (a, a + 1))) 0) foldInputs , bgroup "mapAccumL" $ map (\s -> bench (show $ S.length s) $ nf (S.mapAccumL (\acc x -> (acc + fromIntegral x, succ x)) (0 :: Int)) s) foldInputs , bgroup "mapAccumR" $ map (\s -> bench (show $ S.length s) $