Interleave effects better in chunksOf

changelog.md

@@ -2,6 +2,9 @@
 
     Added `nubOrd` and friends.
 
+    Generalized the type of `chunksOf` and made it handle effects better.
+    (Thanks, Torgeir Strand Henriksen.)
+
 - 0.2.1.0
 
     Adding `Semigroup` instances for GHC 8.4.

src/Streaming/Internal.hs

@@ -722,12 +722,14 @@ takes n = void . splitsAt n
 2
 1
 -}
-chunksOf :: (Monad m, Functor f) => Int -> Stream f m r -> Stream (Stream f m) m r
-chunksOf n0 = loop where
-  loop stream = case stream of
-    Return r  -> Return r
-    Effect m  -> Effect (fmap loop m)
-    Step fs   -> Step (Step (fmap (fmap loop . splitsAt (n0-1)) fs))
+chunksOf :: (Functor f, Monad m, Monad n) => Int -> Stream f m r -> Stream (Stream f m) n r
+chunksOf = go
+  where
+    go !n s = do
+      stage <- yields $ lift . inspect =<< splitsAt n s
+      case stage of
+        Left r -> pure r
+        Right s' -> go n $ wrap s'
 {-# INLINABLE chunksOf #-}        
 
 {- | Make it possible to \'run\' the underlying transformed monad.