From 89a392f10f1e5a55b6bc38296151fd1ceba57b48 Mon Sep 17 00:00:00 2001
From: Torgeir Strand Henriksen <torgeirsh@users.noreply.github.com>
Date: Mon, 16 Apr 2018 21:29:11 -0400
Subject: [PATCH] Interleave effects better in chunksOf

Generalize the type of `chunksOf` to allow the inner and outer
streams of the result to use different underlying monads. This
forces effects to be handled within the inner monad, leading
to an arguably better sort of chunking.

Fixes #45
---
 changelog.md              |  3 +++
 src/Streaming/Internal.hs | 14 ++++++++------
 2 files changed, 11 insertions(+), 6 deletions(-)

diff --git a/changelog.md b/changelog.md
index ca1ae8c..4e63a0d 100644
--- a/changelog.md
+++ b/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.
diff --git a/src/Streaming/Internal.hs b/src/Streaming/Internal.hs
index 25c7491..28b8d3d 100644
--- a/src/Streaming/Internal.hs
+++ b/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.