@@ -764,6 +764,9 @@ takeBetween low high (Fold fstep finitial fextract) =
764764-- * Fails - if the stream or the collecting fold ends before it can collect
765765-- exactly @n@ elements.
766766--
767+-- >>> Stream.parse (Parser.takeEQ 2 Fold.toList) $ Stream.fromList [1,0,1]
768+-- Right [1,0]
769+--
767770-- >>> Stream.parse (Parser.takeEQ 4 Fold.toList) $ Stream.fromList [1,0,1]
768771-- Left (ParseError "takeEQ: Expecting exactly 4 elements, input terminated on 3")
769772--
@@ -773,51 +776,46 @@ takeEQ n (Fold fstep finitial fextract) = Parser step initial extract
773776
774777 where
775778
776− cnt = max n 0
777−
778779 initial = do
779780 res <- finitial
780− return $ case res of
781− FL.Partial s -> IPartial $ Tuple' 0 s
782− FL.Done b ->
783− if cnt == 0
784− then IDone b
785− else IError
786− $ "takeEQ: Expecting exactly " ++ show cnt
781+ case res of
782+ FL.Partial s ->
783+ if n > 0
784+ then return $ IPartial $ Tuple'Fused 1 s
785+ else fmap IDone (fextract s)
786+ FL.Done b -> return $
787+ if n > 0
788+ then IError
789+ $ "takeEQ: Expecting exactly " ++ show n
787790 ++ " elements, fold terminated without"
788791 ++ " consuming any elements"
792+ else IDone b
789793
790− step (Tuple' i r) a
791− | i1 < cnt = do
792− res <- fstep r a
794+ step (Tuple'Fused i1 r) a = do
795+ res <- fstep r a
796+ if n > i1
797+ then
793798 return
794799 $ case res of
795− FL.Partial s -> Continue 0 $ Tuple' i1 s
800+ FL.Partial s -> Continue 0 $ Tuple'Fused (i1 + 1) s
796801 FL.Done _ ->
797802 Error
798− $ "takeEQ: Expecting exactly " ++ show cnt
803+ $ "takeEQ: Expecting exactly " ++ show n
799804 ++ " elements, fold terminated on " ++ show i1
800− | i1 == cnt = do
801− res <- fstep r a
805+ else
806+ -- assert (n == i1)
802807 Done 0
803808 <$> case res of
804809 FL.Partial s -> fextract s
805810 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
809−
810− where
811811
812− i1 = i + 1
813−
814− extract (Tuple' i r)
815− | i == 0 && cnt == 0 = fmap (Done 0) $ fextract r
816− | otherwise =
817− return
818− $ Error
819− $ "takeEQ: Expecting exactly " ++ show cnt
820− ++ " elements, input terminated on " ++ show i
812+ extract (Tuple'Fused i _) =
813+ -- Using the count "i" in the message below causes large performance
814+ -- regression unless we use Fuse annotation on Tuple.
815+ return
816+ $ Error
817+ $ "takeEQ: Expecting exactly " ++ show n
818+ ++ " elements, input terminated on " ++ show (i - 1)
821819
822820-- | Take at least @n@ input elements, but can collect more.
823821--