Chunk better (#45)

treeowl · konsumlamm · web-flow · commit be2398bc3c13 · 2025-05-30T21:38:33.000+02:00
`parListChunk` previously split a list up into chunks, applied
the given strategy to each chunk, and then put them all together
again. This led to two extra copies of the list.

We get very little benefit from actually splitting the list, because the
parallel computations need to traverse their part anyway; we can instead just
hand off the whole list and let them count out their chunk.  We count each
chunk twice, but that shouldn't cost enough to matter.

Now that `Eval` has a `MonadFix` instance, we can avoid actually having
to put together lists at the end; instead, we pass each parallel
computation the (as-yet-uncomputed) result of calculating the rest
of the list.

---------

Co-authored-by: konsumlamm &lt;44230978+konsumlamm@users.noreply.github.com&gt;
diff --git a/Control/Parallel/Strategies.hs b/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
