Small streaming optimizations

Author: Jaro Reinders

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

@@ -245,12 +245,13 @@ Stream stepa ta ++ Stream stepb tb = Stream step (Left ta)
                         r <- stepa sa
                         case r of
                           Yield x sa' -> return $ Yield x (Left  sa')
-                          Done        -> step (Right tb)
-    step (Right sb) = do
-                        r <- stepb sb
-                        case r of
-                          Yield x sb' -> return $ Yield x (Right sb')
-                          Done        -> return $ Done
+                          Done        -> step' tb
+    step (Right sb) = step' sb
+    step' sb = do
+                 r <- stepb sb
+                 case r of
+                   Yield x sb' -> return $ Yield x (Right sb')
+                   Done        -> return $ Done
 
 -- Accessing elements
 -- ------------------
@@ -339,14 +340,16 @@ init (Stream step t) = Stream step' (Nothing, t)
     step' (Nothing, s) = do
                            r <- step s
                            case r of
-                             Yield x s' -> step' (Just x,  s')
+                             Yield x s' -> step'' x  s'
                              Done       -> return (error emptyStream)
 
-    step' (Just x,  s) = liftM (\r ->
-                           case r of
-                             Yield y s' -> Yield x (Just y, s')
-                             Done       -> Done
-                         ) (step s)
+    step' (Just x,  s) = step'' x s
+    {-# INLINE_INNER step'' #-}
+    step'' x s = liftM (\r ->
+                   case r of
+                     Yield y s' -> Yield x (Just y, s')
+                     Done       -> Done
+                 ) (step s)
 
 -- | All but the first element
 tail :: (HasCallStack, Monad m) => Stream m a -> Stream m a
@@ -357,14 +360,16 @@ tail (Stream step t) = Stream step' (Left t)
     step' (Left  s) = do
                         r <- step s
                         case r of
-                          Yield _ s' -> step' (Right s')
+                          Yield _ s' -> step'' s'
                           Done       -> return (error emptyStream)
 
-    step' (Right s) = liftM (\r ->
-                        case r of
-                          Yield x s' -> Yield x (Right s')
-                          Done       -> Done
-                      ) (step s)
+    step' (Right s) = step'' s
+    {-# INLINE_INNER step'' #-}
+    step'' s = liftM (\r ->
+                 case r of
+                   Yield x s' -> Yield x (Right s')
+                   Done       -> Done
+               ) (step s)
 
 -- | The first @n@ elements
 take :: Monad m => Int -> Stream m a -> Stream m a
@@ -382,21 +387,16 @@ take n (Stream step t) = n `seq` Stream step' (t, 0)
 -- | All but the first @n@ elements
 drop :: Monad m => Int -> Stream m a -> Stream m a
 {-# INLINE_FUSED drop #-}
-drop n (Stream step t) = Stream step' (t, Just n)
+drop n (Stream step t) = Stream (step' n) t
   where
     {-# INLINE_INNER step' #-}
-    step' (s, Just i) | i > 0 = do
-                                  r <- step s
-                                  case r of
-                                     Yield _ s' -> step' (s', Just (i-1))
-                                     Done       -> return Done
-                      | otherwise = step' (s, Nothing)
-
-    step' (s, Nothing) = liftM (\r ->
-                           case r of
-                             Yield x s' -> Yield x (s', Nothing)
-                             Done       -> Done
-                           ) (step s)
+    step' i s | i > 0 = do
+                          r <- step s
+                          case r of
+                             Yield _ s' -> step' (i - 1) s'
+                             Done       -> return Done
+              | otherwise = step s
+
 
 -- Mapping
 -- -------
@@ -497,17 +497,19 @@ zipWithM f (Stream stepa ta) (Stream stepb tb) = Stream step (ta, tb, Nothing)
     step (sa, sb, Nothing) = do 
                                r <- stepa sa
                                case r of
-                                 Yield x sa' -> step (sa', sb, Just x)
+                                 Yield x sa' -> step' sa' sb x
                                  Done        -> return Done
 
-    step (sa, sb, Just x)  = do
-                               r <- stepb sb
-                               case r of
-                                 Yield y sb' ->
-                                   do
-                                     z <- f x y
-                                     return $ Yield z (sa, sb', Nothing)
-                                 Done        -> return Done
+    step (sa, sb, Just x)  = step' sa sb x
+    {-# INLINE_INNER step' #-}
+    step' sa sb x = do
+                      r <- stepb sb
+                      case r of
+                        Yield y sb' ->
+                          do
+                            z <- f x y
+                            return $ Yield z (sa, sb', Nothing)
+                        Done        -> return Done
 
 zipWithM_ :: Monad m => (a -> b -> m c) -> Stream m a -> Stream m b -> m ()
 {-# INLINE zipWithM_ #-}