Fix drop issue

Author: Jaro Reinders

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

@@ -387,15 +387,19 @@ 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' n) t
+drop n (Stream step t) = Stream step' (t, n)
   where
     {-# INLINE_INNER step' #-}
-    step' i s | i > 0 = do
-                          r <- step s
-                          case r of
-                             Yield _ s' -> step' (i - 1) s'
-                             Done       -> return Done
-              | otherwise = step s
+    step' (s, i) | i > 0 = do
+                             r <- step s
+                             case r of
+                               Yield _ s' -> step' (s', i - 1)
+                               Done       -> return Done
+                 | otherwise = liftM (\r ->
+                                 case r of
+                                   Yield x s' -> Yield x (s', i)
+                                   Done -> Done
+                               ) (step s)
 
 
 -- Mapping