module Control.Concurrent.AList (AList(..), filterAList, assocFold, monoid, lenAList, findAList, concatAList) where

import Control.Parallel
import Control.Monad
import Control.Applicative
import Data.Monoid
import Data.Foldable (Foldable, foldMap)
import Data.Traversable (Traversable, traverse, foldMapDefault)

-- | Lists suitable for parallel execution (taken from Hackage's monad-par package). (For converting to regular lists, there is the toList function in Data.Foldable.)
data AList t = Append (AList t) (AList t) | List [t] deriving (Eq, Ord, Show)

instance Monad AList where
	return x = List [x]
	Append ls ls2 >>= f = ((ls >>= f) `par` (ls2 >>= f)) `seq` Append (ls >>= f) (ls2 >>= f)
	List ls >>= f = foldr mplus mzero (map f ls)

instance MonadPlus AList where
	mzero = List []
	mplus m n = (m `par` n) `seq` case (m, n) of
		(List [x], List xs) -> List (x:xs)
		(List [x], Append y z) -> Append (mplus m y) z
		(List [], n) -> n
		(m, List []) -> m
		_ -> Append m n

instance Functor AList where
	fmap f m = m >>= return . f

instance Traversable AList where
	traverse f (Append ls ls2) = Append <$> traverse f ls <*> traverse f ls2
	traverse f (List ls) = List <$> traverse f ls

instance Foldable AList where
	foldMap = foldMapDefault

-- | Filters the AList using a predicate.
filterAList f ls = ls >>= \x -> List $ if f x then [x] else []

noNils (Append m n) = noNils m `mplus` noNils n
noNils ls = ls

assocFold0 f (Append ls ls2) = (assocFold0 f ls `par` assocFold0 f ls2) `seq` f (assocFold0 f ls) (assocFold0 f ls2)
assocFold0 f (List ls) = foldl1 f ls

-- | Folds the AList with a function, that must be associative. This allows parallelism to be introduced.
assocFold f = assocFold0 f . noNils

-- | Combine monoid elements to get a result.
monoid ls = if noNils ls == List [] then
		mempty
	else
		assocFold mappend ls

-- | Length of an AList.
lenAList ls = if noNils ls == List [] then
		0
	else
		assocFold (+) (fmap (const 1) ls)

-- | Find the first element satisfying a predicate.
findAList f = getFirst . monoid . fmap (\x -> First $ if f x then Just x else Nothing)

-- | Concatenate an AList of ALists.
concatAList ls = ls >>= id