@@ -764,6 +764,9 @@ takeBetween low high (Fold fstep finitial fextract) =
764-- * Fails - if the stream or the collecting fold ends before it can collect
764-- * Fails - if the stream or the collecting fold ends before it can collect
765-- exactly @n@ elements.
765-- exactly @n@ elements.
767+-- >>> Stream.parse (Parser.takeEQ 2 Fold.toList) $ Stream.fromList [1,0,1]
767-- >>> Stream.parse (Parser.takeEQ 4 Fold.toList) $ Stream.fromList [1,0,1]
770-- >>> Stream.parse (Parser.takeEQ 4 Fold.toList) $ Stream.fromList [1,0,1]
768-- Left (ParseError "takeEQ: Expecting exactly 4 elements, input terminated on 3")
771-- Left (ParseError "takeEQ: Expecting exactly 4 elements, input terminated on 3")
@@ -773,51 +776,46 @@ takeEQ n (Fold fstep finitial fextract) = Parser step initial extract
778 initial = do
779 initial = do
779 res <- finitial
780 res <- finitial
780− return $ case res of
781+ case res of
781− FL.Partial s -> IPartial $ Tuple' 0 s
782+ FL.Partial s ->
782− FL.Done b ->
783+ if n > 0
783− if cnt == 0
784+ then return $ IPartial $ Tuple'Fused 1 s
784− then IDone b
785+ else fmap IDone (fextract s)
786+ FL.Done b -> return $
785− else IError
788+ then IError
786− $ "takeEQ: Expecting exactly " ++ show cnt
789+ $ "takeEQ: Expecting exactly " ++ show n
787 ++ " elements, fold terminated without"
790 ++ " elements, fold terminated without"
788 ++ " consuming any elements"
791 ++ " consuming any elements"
790− step (Tuple' i r) a
794+ step (Tuple'Fused i1 r) a = do
792− res <- fstep r a
795+ res <- fstep r a
794 $ case res of
799 $ case res of
795− FL.Partial s -> Continue 0 $ Tuple' i1 s
800+ FL.Partial s -> Continue 0 $ Tuple'Fused (i1 + 1) s
796 FL.Done _ ->
801 FL.Done _ ->
798− $ "takeEQ: Expecting exactly " ++ show cnt
803+ $ "takeEQ: Expecting exactly " ++ show n
799 ++ " elements, fold terminated on " ++ show i1
804 ++ " elements, fold terminated on " ++ show i1
800− | i1 == cnt = do
805+ else
801− res <- fstep r a
806+ -- assert (n == i1)
803 <$> case res of
808 <$> case res of
804 FL.Partial s -> fextract s
809 FL.Partial s -> fextract s
805 FL.Done b -> return b
810 FL.Done b -> return b
806− -- XXX we should not reach here when initial returns Step type
807− -- reachable only when n == 0
808− | otherwise = Done 1 <$> fextract r
814− extract (Tuple' i r)
812+ extract (Tuple'Fused i _) =
815− | i == 0 && cnt == 0 = fmap (Done 0) $ fextract r
813+ -- Using the count "i" in the message below causes large performance
816− | otherwise =
814+ -- regression unless we use Fuse annotation on Tuple.
819− $ "takeEQ: Expecting exactly " ++ show cnt
817+ $ "takeEQ: Expecting exactly " ++ show n
820− ++ " elements, input terminated on " ++ show i
818+ ++ " elements, input terminated on " ++ show (i - 1)
822-- | Take at least @n@ input elements, but can collect more.
820-- | Take at least @n@ input elements, but can collect more.