--- > unfoldrN 10 (\x -> Just (x, chr (ord x + 1))) '0' == "0123456789"
---
--- The following equation connects the depth-limited unfoldr to the List unfoldr:
---
--- > unfoldrN n == take n $ List.unfoldr
-unfoldrN :: Int -> (a -> Maybe (Word8, a)) -> a -> ByteString
-unfoldrN i f w
- | i <= 0 = empty
- | otherwise = inlinePerformIO $ generate i $ \p -> go p w 0
- where
- STRICT3(go)
- go q c n | n == i = return n -- stop if we reach `i'
- | otherwise = case f c of
- Nothing -> return n
- Just (a,new_c) -> do
- poke q a
- go (q `plusPtr` 1) new_c (n+1)
+-- > unfoldr (\x -> if x <= 5 then Just (x, x + 1) else Nothing) 0
+-- > == pack [0, 1, 2, 3, 4, 5]
+--
+unfoldr :: (a -> Maybe (Word8, a)) -> a -> ByteString
+unfoldr f = concat . unfoldChunk 32 64
+ where unfoldChunk n n' x =
+ case unfoldrN n f x of
+ (s, Nothing) -> s : []
+ (s, Just x') -> s : unfoldChunk n' (n+n') x'
+
+-- | /O(n)/ Like 'unfoldr', 'unfoldrN' builds a ByteString from a seed
+-- value. However, the length of the result is limited by the first
+-- argument to 'unfoldrN'. This function is more efficient than 'unfoldr'
+-- when the maximum length of the result is known.
+--
+-- The following equation relates 'unfoldrN' and 'unfoldr':
+--
+-- > unfoldrN n f s == take n (unfoldr f s)
+--
+unfoldrN :: Int -> (a -> Maybe (Word8, a)) -> a -> (ByteString, Maybe a)
+unfoldrN i f x0
+ | i < 0 = (empty, Just x0)
+ | otherwise = inlinePerformIO $ do
+ fp <- mallocByteString i
+ withForeignPtr fp (\p -> go fp p x0 0)
+ where STRICT4(go)
+ go fp p x n =
+ case f x of
+ Nothing -> let s = copy (PS fp 0 n)
+ in s `seq` return (s, Nothing)
+ Just (w,x')
+ | n == i -> return (PS fp 0 i, Just x)
+ | otherwise -> do poke p w
+ go fp (p `plusPtr` 1) x' (n+1)