@@ -166,6 +166,7 @@ module Streamly.Internal.Data.Parser.ParserD
166166where
167167
168168import Control.Exception (assert, Exception)
169+import Control.Monad (when)
169170import Control.Monad.Catch (MonadCatch, MonadThrow(..))
170171import Fusion.Plugin.Types (Fuse(..))
171172import Streamly.Internal.Data.Fold.Type (Fold(..))
@@ -352,6 +353,10 @@ either parser = Parser step initial extract
352353-- Taking elements
353354-------------------------------------------------------------------------------
354355
356+-- Required to fuse "take" with "many" in "chunksOf", for ghc-9.x
357+{-# ANN type Tuple'Fused Fuse #-}
358+data Tuple'Fused a b = Tuple'Fused !a !b deriving Show
359+
355360-- | See 'Streamly.Internal.Data.Parser.takeBetween'.
356361--
357362-- /Pre-release/
@@ -360,61 +365,69 @@ either parser = Parser step initial extract
360365takeBetween :: MonadCatch m => Int -> Int -> Fold m a b -> Parser m a b
361366takeBetween low high (Fold fstep finitial fextract) =
362367
363− Parser step initial extract
368+ Parser step initial (extract streamErr)
364369
365370 where
366371
367− initial = do
368− res <- finitial
369− return $ case res of
370− FL.Partial s -> IPartial $ Tuple' 0 s
372+ streamErr i =
373+ "takeBetween: Expecting alteast " ++ show low
374+ ++ " elements, got " ++ show i
375+
376+ invalidRange =
377+ "takeBetween: lower bound - " ++ show low
378+ ++ " is greater than higher bound - " ++ show high
379+
380+ foldErr :: Int -> String
381+ foldErr i =
382+ "takeBetween: the collecting fold terminated after"
383+ ++ " consuming" ++ show i ++ " elements"
384+ ++ " minimum" ++ show low ++ " elements needed"
385+
386+ -- Exactly the same as snext except different constructors, we can possibly
387+ -- deduplicate the two.
388+ {-# INLINE inext #-}
389+ inext i res =
390+ let i1 = i + 1
391+ in case res of
392+ FL.Partial s -> do
393+ let s1 = Tuple'Fused i1 s
394+ if i1 < high
395+ -- XXX ideally this should be a Continue instead
396+ then return $ IPartial s1
397+ else IDone <$> extract foldErr s1
371398 FL.Done b ->
372− if low <= 0
373− then IDone b
374− else IError
375− $ "takeBetween: the collecting fold terminated without"
376− ++ " consuming any elements"
377− ++ " minimum" ++ show low ++ " elements needed"
399+ return
400+ $ if i1 >= low
401+ then IDone b
402+ else IError (foldErr i1)
378403
379− step (Tuple' i s) a
380− | low > high =
381− throwM
382− $ ParseError
383− $ "takeBetween: lower bound - " ++ show low
384− ++ " is greater than higher bound - " ++ show high
385− | high <= 0 = Done 1 <$> fextract s
386− | i1 < low = do
387− res <- fstep s a
388− return
389− $ case res of
390− FL.Partial s1 -> Continue 0 $ Tuple' i1 s1
391− FL.Done _ ->
392− Error
393− $ "takeBetween: the collecting fold terminated after"
394− ++ " consuming" ++ show i1 ++ " elements"
395− ++ " minimum" ++ show low ++ " elements needed"
396− | otherwise = do
397− res <- fstep s a
398− case res of
399− FL.Partial s1 ->
400− if i1 >= high
401− then Done 0 <$> fextract s1
402− else return $ Partial 0 $ Tuple' i1 s1
403− FL.Done b -> return $ Done 0 b
404−
405− where
404+ initial = do
405+ when (low >= 0 && high >= 0 && low > high)
406+ $ throwM $ ParseError invalidRange
407+
408+ finitial >>= inext (-1)
409+
410+ -- Keep the impl same as inext
411+ {-# INLINE snext #-}
412+ snext i res =
413+ let i1 = i + 1
414+ in case res of
415+ FL.Partial s -> do
416+ let s1 = Tuple'Fused i1 s
417+ if i1 < high
418+ then return $ Continue 0 s1
419+ else Done 0 <$> extract foldErr s1
420+ FL.Done b ->
421+ return
422+ $ if i1 >= low
423+ then Done 0 b
424+ else Error (foldErr i1)
406425
407− i1 = i + 1
426+ step (Tuple'Fused i s) a = fstep s a >>= snext i
408427
409− extract (Tuple' i s)
428+ extract f (Tuple'Fused i s)
410429 | i >= low && i <= high = fextract s
411− | otherwise = throwM $ ParseError err
412−
413− where
414−
415− err =
416− "takeBetween: Expecting alteast " ++ show low
417− ++ " elements, got " ++ show i
430+ | otherwise = throwM $ ParseError (f i)
418431
419432-- | See 'Streamly.Internal.Data.Parser.takeEQ'.
420433--