Chunk better

Control/Parallel/Strategies.hs

@@ -1,5 +1,6 @@
 {-# LANGUAGE BangPatterns, CPP, MagicHash, UnboxedTuples #-}
 {-# LANGUAGE GeneralizedNewtypeDeriving #-}
+{-# LANGUAGE RecursiveDo #-}
 -----------------------------------------------------------------------------
 -- |
 -- Module      :  Control.Parallel.Strategies
@@ -592,20 +593,35 @@ parListNth n strat = evalListNth n (rparWith strat)
 -- | Divides a list into chunks, and applies the strategy
 -- @'evalList' strat@ to each chunk in parallel.
 --
--- It is expected that this function will be replaced by a more
--- generic clustering infrastructure in the future.
---
 -- If the chunk size is 1 or less, 'parListChunk' is equivalent to
 -- 'parList'
 --
+-- This function may be replaced by a more
+-- generic clustering infrastructure in the future.
 parListChunk :: Int -> Strategy a -> Strategy [a]
-parListChunk n strat xs
-  | n <= 1    = parList strat xs
-  | otherwise = concat `fmap` parList (evalList strat) (chunk n xs)
-
-chunk :: Int -> [a] -> [[a]]
-chunk _ [] = []
-chunk n xs = as : chunk n bs where (as,bs) = splitAt n xs
+parListChunk n strat
+  | n <= 1 = parList strat
+  | otherwise = go
+  where
+    go [] = pure []
+    go as = mdo
+      -- Calculate the first chunk in parallel, passing it the result
+      -- of calculating the rest
+      bs <- rpar $ runEval $ evalChunk strat more n as
+
+      -- Calculate the rest
+      more <- go (drop n as)
+      return bs
+
+-- | @evalChunk strat end n as@ uses @strat@ to evaluate the first @n@
+-- elements of @as@ (ignoring the rest) and appends @end@ to the result.
+evalChunk :: Strategy a -> [a] -> Int -> Strategy [a]
+evalChunk strat = \end ->
+  let
+    go !_n [] = pure end
+    go 0 _ = pure end
+    go n (a:as) = (:) <$> strat a <*> go (n - 1) as
+  in go
 
 -- --------------------------------------------------------------------------
 -- Convenience