{-# LANGUAGE ScopedTypeVariables, FlexibleContexts, Trustworthy #-}
module Control.CUtils.AssociativeFolding (assocFold) where
import Control.CUtils.CPUMultiThreading
import Data.List.Extra
import Data.Array.Unsafe
import Data.Array.IO
import Data.Array
import Control.Monad
import Control.Exception
import System.IO.Unsafe
chunkSize = 10000

assocFold :: forall t. Pool -> (t->t->IO t) -> Int -> (Int -> t) -> IO t
-- | Concurrent evaluation of an associative folding function. The caller is responsible
-- for introducing appropriate strictness for the results in the first argument.
assocFold _ _ n _ | n <= 0 = throwIO$ErrorCall"assocFold: list is empty"
assocFold _ _ n f2 | n == 1 = return(f2 0)
assocFold pool f n f2 = do
	let mx = pred n `div` chunkSize
	let ls2 =[0.. mx]
	let l = succ mx
	-- Using unsafe interleaving with this line seems to help the tasks start faster.

	ar::IOArray Int t <- unsafeInterleaveIO(newArray_(0,mx))
	let ls3 = fmap(f3 ar) ls2
	simpleConc_ pool ls3
	ar2 :: Array Int t <- unsafeFreeze ar
	-- Each recursive call reduces the size of the problem by a factor of 'chunkSize'.
	assocFold pool f l(ar2!)
	where
	f3 ar ii = do
		let x:xs= [chunkSize*ii..pred(min n(chunkSize*ii+chunkSize))]
		let x2 = f2 x
		x3 <- foldM((.f2).f) x2 xs
		writeArray ar ii x3