Skip to content

Commit 51a99e4

Browse files
author
Jaro Reinders
committed
Individual streaming optimizations
1 parent 25793d3 commit 51a99e4

File tree

1 file changed

+112
-116
lines changed

1 file changed

+112
-116
lines changed

‎vector-stream/src/Data/Stream/Monadic.hs

Lines changed: 112 additions & 116 deletions
Original file line numberDiff line numberDiff line change
@@ -241,20 +241,19 @@ infixr 5 ++
241241
Stream stepa ta ++ Stream stepb tb = Stream step (Left ta)
242242
where
243243
{-# INLINE_INNER step #-}
244-
step s0 =
245-
let
246-
-- go is a join point
247-
go (Left sa) = do
248-
r <- stepa sa
249-
case r of
250-
Yield x sa' -> return $ Yield x (Left sa')
251-
Done -> go (Right tb)
252-
go (Right sb) = do
253-
r <- stepb sb
254-
case r of
255-
Yield x sb' -> return $ Yield x (Right sb')
256-
Done -> return $ Done
257-
in go s0
244+
step (Left sa) = do
245+
r <- stepa sa
246+
case r of
247+
Yield x sa' -> return $ Yield x (Left sa')
248+
Done -> step' tb
249+
step (Right sb) = step' sb
250+
251+
{-# INLINE_INNER step' #-}
252+
step' s = do
253+
r <- stepb s
254+
case r of
255+
Yield x s' -> return $ Yield x (Right s')
256+
Done -> return $ Done
258257

259258
-- Accessing elements
260259
-- ------------------
@@ -340,43 +339,40 @@ init :: (HasCallStack, Monad m) => Stream m a -> Stream m a
340339
init (Stream step t) = Stream step' (Nothing, t)
341340
where
342341
{-# INLINE_INNER step' #-}
343-
step' s0 =
344-
let
345-
-- go is a join point
346-
go (Nothing, s) = do
347-
r <- step s
348-
case r of
349-
Yield x s' -> go (Just x, s')
350-
Done -> return (error emptyStream)
342+
step' (Nothing, s) = do
343+
r <- step s
344+
case r of
345+
Yield x s' -> step'' x s'
346+
Done -> return (error emptyStream)
351347

352-
go (Just x, s) = liftM (\r ->
353-
case r of
354-
Yield y s' -> Yield x (Just y, s')
355-
Done -> Done
356-
) (step s)
357-
in go s0
348+
step' (Just x, s) = step'' x s
349+
350+
{-# INLINE_INNER step'' #-}
351+
step'' x s = liftM (\r ->
352+
case r of
353+
Yield y s' -> Yield x (Just y, s')
354+
Done -> Done
355+
) (step s)
358356

359357
-- | All but the first element
360358
tail :: (HasCallStack, Monad m) => Stream m a -> Stream m a
361359
{-# INLINE_FUSED tail #-}
362360
tail (Stream step t) = Stream step' (Left t)
363361
where
364362
{-# INLINE_INNER step' #-}
365-
step' s0 =
366-
let
367-
-- go is a join point
368-
go (Left s) = do
369-
r <- step s
370-
case r of
371-
Yield _ s' -> go (Right s')
372-
Done -> return (error emptyStream)
363+
step' (Left s) = do
364+
r <- step s
365+
case r of
366+
Yield _ s' -> step'' s'
367+
Done -> return (error emptyStream)
368+
step' (Right s) = step'' s
373369

374-
go (Right s) = liftM (\r ->
375-
case r of
376-
Yield x s' -> Yield x (Right s')
377-
Done -> Done
378-
) (step s)
379-
in go s0
370+
{-# INLINE_INNER step'' #-}
371+
step'' s = liftM (\r ->
372+
case r of
373+
Yield x s' -> Yield x (Right s')
374+
Done -> Done
375+
) (step s)
380376

381377
-- | The first @n@ elements
382378
take :: Monad m => Int -> Stream m a -> Stream m a
@@ -394,25 +390,28 @@ take n (Stream step t) = n `seq` Stream step' (t, 0)
394390
-- | All but the first @n@ elements
395391
drop :: Monad m => Int -> Stream m a -> Stream m a
396392
{-# INLINE_FUSED drop #-}
397-
drop n (Stream step t) = Stream step' (t, Just n)
393+
drop n (Stream step t) = Stream step' (t, n)
398394
where
399395
{-# INLINE_INNER step' #-}
400-
step' s0 =
401-
let
402-
-- go is a join point
403-
go (s, Just i) | i > 0 = do
404-
r <- step s
405-
case r of
406-
Yield _ s' -> go (s', Just (i-1))
407-
Done -> return Done
408-
| otherwise = go (s, Nothing)
409-
410-
go (s, Nothing) = liftM (\r ->
411-
case r of
412-
Yield x s' -> Yield x (s', Nothing)
413-
Done -> Done
414-
) (step s)
415-
in go s0
396+
step' (s, i) | i > 0 = go s i
397+
step' (s, _) = step'' s
398+
399+
-- go is a recursive join point
400+
{-# INLINABLE go #-}
401+
go s i | i > 0 = do
402+
r <- step s
403+
case r of
404+
Yield _ s' -> go s' (i-1)
405+
Done -> return Done
406+
| otherwise = step'' s
407+
408+
409+
{-# INLINE_INNER step'' #-}
410+
step'' s = liftM (\r ->
411+
case r of
412+
Yield x s' -> Yield x (s', 0)
413+
Done -> Done
414+
) (step s)
416415

417416
-- Mapping
418417
-- -------
@@ -510,24 +509,22 @@ zipWithM :: Monad m => (a -> b -> m c) -> Stream m a -> Stream m b -> Stream m c
510509
zipWithM f (Stream stepa ta) (Stream stepb tb) = Stream step (ta, tb, Nothing)
511510
where
512511
{-# INLINE_INNER step #-}
513-
step s0 =
514-
let
515-
-- go is a join point
516-
go (sa, sb, Nothing) = do
517-
r <- stepa sa
518-
case r of
519-
Yield x sa' -> go (sa', sb, Just x)
520-
Done -> return Done
521-
522-
go (sa, sb, Just x) = do
523-
r <- stepb sb
524-
case r of
525-
Yield y sb' ->
526-
do
527-
z <- f x y
528-
return $ Yield z (sa, sb', Nothing)
529-
Done -> return Done
530-
in go s0
512+
step (sa, sb, Nothing) = do
513+
r <- stepa sa
514+
case r of
515+
Yield x sa' -> step' sa' sb x
516+
Done -> return Done
517+
step (sa, sb, Just x) = step' sa sb x
518+
519+
{-# INLINE_INNER step' #-}
520+
step' sa sb x = do
521+
r <- stepb sb
522+
case r of
523+
Yield y sb' ->
524+
do
525+
z <- f x y
526+
return $ Yield z (sa, sb', Nothing)
527+
Done -> return Done
531528

532529
zipWithM_ :: Monad m => (a -> b -> m c) -> Stream m a -> Stream m b -> m ()
533530
{-# INLINE zipWithM_ #-}
@@ -540,27 +537,27 @@ zipWith3M f (Stream stepa ta)
540537
(Stream stepc tc) = Stream step (ta, tb, tc, Nothing)
541538
where
542539
{-# INLINE_INNER step #-}
543-
step s0 =
544-
let
545-
-- go is a join point
546-
go (sa, sb, sc, Nothing) = do
547-
r <- stepa sa
548-
case r of
549-
Yield x sa' -> go (sa', sb, sc, Just (x, Nothing))
550-
Done -> return Done
551-
552-
go (sa, sb, sc, Just (x, Nothing)) = do
553-
r <- stepb sb
554-
case r of
555-
Yield y sb' -> go (sa, sb', sc, Just (x, Just y))
556-
Done -> return Done
557-
558-
go (sa, sb, sc, Just (x, Just y)) = do
559-
r <- stepc sc
560-
case r of
561-
Yield z sc' -> f x y z >>= (\res -> return $ Yield res (sa, sb, sc', Nothing))
562-
Done -> return $ Done
563-
in go s0
540+
step (sa, sb, sc, Nothing) = do
541+
r <- stepa sa
542+
case r of
543+
Yield x sa' -> step' sa' sb sc x
544+
Done -> return Done
545+
step (sa, sb, sc, Just (x, Nothing)) = step' sa sb sc x
546+
step (sa, sb, sc, Just (x, Just y)) = step'' sa sb sc x y
547+
548+
{-# INLINE_INNER step' #-}
549+
step' sa sb sc x = do
550+
r <- stepb sb
551+
case r of
552+
Yield y sb' -> step'' sa sb' sc x y
553+
Done -> return Done
554+
555+
{-# INLINE_INNER step'' #-}
556+
step'' sa sb sc x y = do
557+
r <- stepc sc
558+
case r of
559+
Yield z sc' -> f x y z >>= (\res -> return $ Yield res (sa, sb, sc', Nothing))
560+
Done -> return $ Done
564561

565562
zipWith4M :: Monad m => (a -> b -> c -> d -> m e)
566563
-> Stream m a -> Stream m b -> Stream m c -> Stream m d
@@ -702,14 +699,13 @@ mapMaybe f (Stream step t) = Stream step' t
702699
{-# INLINE_INNER step' #-}
703700
step' s0 =
704701
let
705-
-- go is a join point
702+
-- go is a recursive join point
706703
go s = do
707704
r <- step s
708705
case r of
709-
Yield x s' -> do
710-
case f x of
711-
Nothing -> go s'
712-
Just b' -> return $ Yield b' s'
706+
Yield x s' -> case f x of
707+
Nothing -> go s'
708+
Just b' -> return $ Yield b' s'
713709
Done -> return $ Done
714710
in go s0
715711

@@ -763,19 +759,19 @@ uniq :: (Eq a, Monad m) => Stream m a -> Stream m a
763759
uniq (Stream step st) = Stream step' (Nothing,st)
764760
where
765761
{-# INLINE_INNER step' #-}
766-
step' s0 =
767-
let
768-
-- go is a join point
769-
go (Nothing, s) = do r <- step s
770-
case r of
771-
Yield x s' -> return $ Yield x (Just x , s')
772-
Done -> return Done
773-
go (Just x0, s) = do r <- step s
774-
case r of
775-
Yield x s' | x == x0 -> go (Just x0, s')
776-
| otherwise -> return $ Yield x (Just x , s')
777-
Done -> return Done
778-
in go s0
762+
step' (Nothing, s) = do r <- step s
763+
case r of
764+
Yield x s' -> return $ Yield x (Just x , s')
765+
Done -> return Done
766+
step' (Just x, s) = go x s
767+
768+
-- go is a recursive join point
769+
{-# INLINABLE go #-}
770+
go x0 s = do r <- step s
771+
case r of
772+
Yield x s' | x == x0 -> go x0 s'
773+
| otherwise -> return $ Yield x (Just x , s')
774+
Done -> return Done
779775

780776
-- | Longest prefix of elements that satisfy the predicate
781777
takeWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a

0 commit comments

Comments
 (0)