Skip to content

Commit cfaf0d4

Browse files
author
Jaro Reinders
committed
Explicit loopification and more INLINE_INNER
1 parent 114fde6 commit cfaf0d4

File tree

2 files changed

+142
-103
lines changed

2 files changed

+142
-103
lines changed

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

Lines changed: 141 additions & 103 deletions
Original file line numberDiff line numberDiff line change
@@ -247,6 +247,7 @@ Stream stepa ta ++ Stream stepb tb = Stream step (Left ta)
247247
Yield x sa' -> return $ Yield x (Left sa')
248248
Done -> step' tb
249249
step (Right sb) = step' sb
250+
{-# INLINE_INNER step' #-}
250251
step' sb = do
251252
r <- stepb sb
252253
case r of
@@ -390,16 +391,20 @@ drop :: Monad m => Int -> Stream m a -> Stream m a
390391
drop n (Stream step t) = Stream step' (t, n)
391392
where
392393
{-# INLINE_INNER step' #-}
393-
step' (s, i) | i > 0 = do
394-
r <- step s
395-
case r of
396-
Yield _ s' -> step' (s', i - 1)
397-
Done -> return Done
398-
| otherwise = liftM (\r ->
399-
case r of
394+
step' s0 =
395+
let
396+
-- go is a join point
397+
go (s, i) | i > 0 = do
398+
r <- step s
399+
case r of
400+
Yield _ s' -> go (s', i - 1)
401+
Done -> return Done
402+
| otherwise = liftM (\r ->
403+
case r of
400404
Yield x s' -> Yield x (s', i)
401405
Done -> Done
402-
) (step s)
406+
) (step s)
407+
in go s0
403408

404409

405410
-- Mapping
@@ -526,23 +531,27 @@ zipWith3M f (Stream stepa ta)
526531
(Stream stepc tc) = Stream step (ta, tb, tc, Nothing)
527532
where
528533
{-# INLINE_INNER step #-}
529-
step (sa, sb, sc, Nothing) = do
530-
r <- stepa sa
531-
case r of
532-
Yield x sa' -> step (sa', sb, sc, Just (x, Nothing))
533-
Done -> return Done
534-
535-
step (sa, sb, sc, Just (x, Nothing)) = do
536-
r <- stepb sb
537-
case r of
538-
Yield y sb' -> step (sa, sb', sc, Just (x, Just y))
539-
Done -> return Done
540-
541-
step (sa, sb, sc, Just (x, Just y)) = do
542-
r <- stepc sc
543-
case r of
544-
Yield z sc' -> f x y z >>= (\res -> return $ Yield res (sa, sb, sc', Nothing))
545-
Done -> return $ Done
534+
step s0 =
535+
let
536+
-- go is a join point
537+
go (sa, sb, sc, Nothing) = do
538+
r <- stepa sa
539+
case r of
540+
Yield x sa' -> go (sa', sb, sc, Just (x, Nothing))
541+
Done -> return Done
542+
543+
go (sa, sb, sc, Just (x, Nothing)) = do
544+
r <- stepb sb
545+
case r of
546+
Yield y sb' -> go (sa, sb', sc, Just (x, Just y))
547+
Done -> return Done
548+
549+
go (sa, sb, sc, Just (x, Just y)) = do
550+
r <- stepc sc
551+
case r of
552+
Yield z sc' -> f x y z >>= (\res -> return $ Yield res (sa, sb, sc', Nothing))
553+
Done -> return $ Done
554+
in go s0
546555

547556
zipWith4M :: Monad m => (a -> b -> c -> d -> m e)
548557
-> Stream m a -> Stream m b -> Stream m c -> Stream m d
@@ -682,14 +691,18 @@ mapMaybe :: Monad m => (a -> Maybe b) -> Stream m a -> Stream m b
682691
mapMaybe f (Stream step t) = Stream step' t
683692
where
684693
{-# INLINE_INNER step' #-}
685-
step' s = do
686-
r <- step s
687-
case r of
688-
Yield x s' -> do
689-
case f x of
690-
Nothing -> step' s'
691-
Just b' -> return $ Yield b' s'
692-
Done -> return $ Done
694+
step' s0 =
695+
let
696+
-- go is a join point
697+
go s = do
698+
r <- step s
699+
case r of
700+
Yield x s' -> do
701+
case f x of
702+
Nothing -> go s'
703+
Just b' -> return $ Yield b' s'
704+
Done -> return $ Done
705+
in go s0
693706

694707
catMaybes :: Monad m => Stream m (Maybe a) -> Stream m a
695708
catMaybes = mapMaybe id
@@ -700,14 +713,18 @@ filterM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
700713
filterM f (Stream step t) = Stream step' t
701714
where
702715
{-# INLINE_INNER step' #-}
703-
step' s = do
704-
r <- step s
705-
case r of
706-
Yield x s' -> do
707-
b <- f x
708-
if b then return $ Yield x s'
709-
else step' s'
710-
Done -> return $ Done
716+
step' s0 =
717+
let
718+
-- go is a join point
719+
go s = do
720+
r <- step s
721+
case r of
722+
Yield x s' -> do
723+
b <- f x
724+
if b then return $ Yield x s'
725+
else go s'
726+
Done -> return $ Done
727+
in go s0
711728

712729
-- | Apply monadic function to each element and drop all Nothings
713730
--
@@ -717,31 +734,39 @@ mapMaybeM :: Monad m => (a -> m (Maybe b)) -> Stream m a -> Stream m b
717734
mapMaybeM f (Stream step t) = Stream step' t
718735
where
719736
{-# INLINE_INNER step' #-}
720-
step' s = do
721-
r <- step s
722-
case r of
723-
Yield x s' -> do
724-
fx <- f x
725-
case fx of
726-
Nothing -> step' s'
727-
Just b -> return $ Yield b s'
728-
Done -> return $ Done
737+
step' s0 =
738+
let
739+
-- go is a join point
740+
go s = do
741+
r <- step s
742+
case r of
743+
Yield x s' -> do
744+
fx <- f x
745+
case fx of
746+
Nothing -> go s'
747+
Just b -> return $ Yield b s'
748+
Done -> return $ Done
749+
in go s0
729750

730751
-- | Drop repeated adjacent elements.
731752
uniq :: (Eq a, Monad m) => Stream m a -> Stream m a
732753
{-# INLINE_FUSED uniq #-}
733754
uniq (Stream step st) = Stream step' (Nothing,st)
734755
where
735756
{-# INLINE_INNER step' #-}
736-
step' (Nothing, s) = do r <- step s
737-
case r of
738-
Yield x s' -> return $ Yield x (Just x , s')
739-
Done -> return Done
740-
step' (Just x0, s) = do r <- step s
741-
case r of
742-
Yield x s' | x == x0 -> step' (Just x0, s')
743-
| otherwise -> return $ Yield x (Just x , s')
744-
Done -> return Done
757+
step' s0 =
758+
let
759+
-- go is a join point
760+
go (Nothing, s) = do r <- step s
761+
case r of
762+
Yield x s' -> return $ Yield x (Just x , s')
763+
Done -> return Done
764+
go (Just x0, s) = do r <- step s
765+
case r of
766+
Yield x s' | x == x0 -> go (Just x0, s')
767+
| otherwise -> return $ Yield x (Just x , s')
768+
Done -> return Done
769+
in go s0
745770

746771
-- | Longest prefix of elements that satisfy the predicate
747772
takeWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
@@ -778,24 +803,28 @@ dropWhileM f (Stream step t) = Stream step' (DropWhile_Drop t)
778803
-- declarations would be nice!
779804

780805
{-# INLINE_INNER step' #-}
781-
step' (DropWhile_Drop s)
782-
= do
783-
r <- step s
784-
case r of
785-
Yield x s' -> do
786-
b <- f x
787-
if b then step' (DropWhile_Drop s')
788-
else step' (DropWhile_Yield x s')
789-
Done -> return $ Done
790-
791-
step' (DropWhile_Yield x s) = return $ Yield x (DropWhile_Next s)
792-
793-
step' (DropWhile_Next s)
794-
= do
795-
r <- step s
796-
case r of
797-
Yield x s' -> step' (DropWhile_Yield x s')
798-
Done -> return Done
806+
step' s0 =
807+
let
808+
-- go is a join point
809+
go (DropWhile_Drop s)
810+
= do
811+
r <- step s
812+
case r of
813+
Yield x s' -> do
814+
b <- f x
815+
if b then go (DropWhile_Drop s')
816+
else go (DropWhile_Yield x s')
817+
Done -> return $ Done
818+
819+
go (DropWhile_Yield x s) = return $ Yield x (DropWhile_Next s)
820+
821+
go (DropWhile_Next s)
822+
= do
823+
r <- step s
824+
case r of
825+
Yield x s' -> go (DropWhile_Yield x s')
826+
Done -> return Done
827+
in go s0
799828

800829
-- Searching
801830
-- ---------
@@ -1031,41 +1060,50 @@ concatMap f = concatMapM (return . f)
10311060

10321061
concatMapM :: Monad m => (a -> m (Stream m b)) -> Stream m a -> Stream m b
10331062
{-# INLINE_FUSED concatMapM #-}
1034-
concatMapM f (Stream step t) = Stream concatMap_go (Left t)
1035-
where
1036-
concatMap_go (Left s) = do
1037-
r <- step s
1038-
case r of
1039-
Yield a s' -> do
1040-
b_stream <- f a
1041-
concatMap_go (Right (b_stream, s'))
1042-
Done -> return Done
1043-
concatMap_go (Right (Stream inner_step inner_s, s)) = do
1044-
r <- inner_step inner_s
1045-
case r of
1046-
Yield b inner_s' -> return $ Yield b (Right (Stream inner_step inner_s', s))
1047-
Done -> concatMap_go (Left s)
1063+
concatMapM f (Stream step t) = Stream step' (Left t)
1064+
where
1065+
{-# INLINE_INNER step' #-}
1066+
step' s0 =
1067+
let
1068+
-- go is a join point
1069+
go (Left s) = do
1070+
r <- step s
1071+
case r of
1072+
Yield a s' -> do
1073+
b_stream <- f a
1074+
go (Right (b_stream, s'))
1075+
Done -> return Done
1076+
go (Right (Stream inner_step inner_s, s)) = do
1077+
r <- inner_step inner_s
1078+
case r of
1079+
Yield b inner_s' -> return $ Yield b (Right (Stream inner_step inner_s', s))
1080+
Done -> go (Left s)
1081+
in go s0
10481082

10491083
-- | Create a 'Stream' of values from a 'Stream' of streamable things
10501084
flatten :: Monad m => (a -> m s) -> (s -> m (Step s b)) -> Stream m a -> Stream m b
10511085
{-# INLINE_FUSED flatten #-}
10521086
flatten mk istep (Stream ostep u) = Stream step (Left u)
10531087
where
10541088
{-# INLINE_INNER step #-}
1055-
step (Left t) = do
1056-
r <- ostep t
1057-
case r of
1058-
Yield a t' -> do
1059-
s <- mk a
1060-
s `seq` step (Right (s,t'))
1061-
Done -> return $ Done
1089+
step s0 =
1090+
let
1091+
-- go is a join point
1092+
go (Left t) = do
1093+
r <- ostep t
1094+
case r of
1095+
Yield a t' -> do
1096+
s <- mk a
1097+
s `seq` go (Right (s,t'))
1098+
Done -> return $ Done
10621099

10631100

1064-
step (Right (s,t)) = do
1065-
r <- istep s
1066-
case r of
1067-
Yield x s' -> return $ Yield x (Right (s',t))
1068-
Done -> step (Left t)
1101+
go (Right (s,t)) = do
1102+
r <- istep s
1103+
case r of
1104+
Yield x s' -> return $ Yield x (Right (s',t))
1105+
Done -> go (Left t)
1106+
in go s0
10691107

10701108
-- Unfolding
10711109
-- ---------

‎vector/src/Data/Vector/Fusion/Bundle/Monadic.hs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -143,6 +143,7 @@ fromStream :: Monad m => Stream m a -> Size -> Bundle m v a
143143
{-# INLINE fromStream #-}
144144
fromStream (Stream step t) sz = Bundle (Stream step t) (Stream step' t) Nothing sz
145145
where
146+
{-# INLINE_INNER step' #-}
146147
step' s = do r <- step s
147148
return $ fmap (\x -> Chunk 1 (\v -> stToPrim $ M.basicUnsafeWrite v 0 x)) r
148149

0 commit comments

Comments
 (0)