{-# LANGUAGE CPP #-}
#include "fusion-phases.h"

-- | Parallel combinators for unlifted arrays. 
module Data.Array.Parallel.Unlifted.Parallel.Combinators 
        ( mapUP
        , filterUP
        , packUP
        , combineUP, combine2UP
        , zipWithUP
        , foldUP, foldlUP, fold1UP, foldl1UP
        , scanUP)
where
import Data.Array.Parallel.Base
import Data.Array.Parallel.Unlifted.Distributed
import Data.Array.Parallel.Unlifted.Parallel.UPSel
import Data.Array.Parallel.Unlifted.Sequential.Vector as Seq

here :: String -> String
here s = "Data.Array.Parallel.Unlifted.Parallel.Combinators." Prelude.++ s


-- | Apply a worker to all elements of an array.
mapUP :: (Unbox a, Unbox b) => (a -> b) -> Vector a -> Vector b
mapUP f xs 
        = splitJoinD theGang (mapD theGang (Seq.map f)) xs
{-# INLINE_UP mapUP #-}


-- | Keep elements that match the given predicate.
filterUP :: Unbox a => (a -> Bool) -> Vector a -> Vector a
filterUP f
        = joinD  theGang unbalanced
        . mapD   theGang (Seq.filter f)
        . splitD theGang unbalanced
{-# INLINE_UP filterUP #-}


-- | Take elements of an array where a flag value is true, and pack them into
--   the result. 
-- 
--   * The souce and flag arrays must have the same length, but this is not checked.
--
packUP :: Unbox e => Vector e -> Vector Bool -> Vector e
packUP xs flags 
        = Seq.fsts . filterUP snd $ Seq.zip xs flags
{-# INLINE_UP packUP #-}


-- | Combine two vectors based on a selector. 
--   If the selector is true then take the element from the first vector, 
--   otherwise take it from the second.
--
--   * The data vectors must have enough elements to satisfy the flag vector, 
--     but this is not checked.
--  
combineUP :: Unbox a => Vector Bool -> Vector a -> Vector a -> Vector a
combineUP flags xs ys 
        = checkEq (here "combineUP")
                ("tags length /= sum of args length")
                (Seq.length flags) (Seq.length xs + Seq.length ys)
        $ combine2UP tags (mkUPSelRep2 tags) xs ys
        where tags = Seq.map (fromBool . not) flags
{-# INLINE combineUP #-}


-- | Combine two vectors based on a selector. 
--
--   * The data vectors must have enough elements to satisfy the selector,
--     but this is not checked.
--
combine2UP :: Unbox a => Vector Tag -> UPSelRep2 -> Vector a -> Vector a -> Vector a
combine2UP tags rep !xs !ys 
        = checkEq (here "combine2UP")
                ("tags length /= sum of args length")
                (Seq.length tags) (Seq.length xs + Seq.length ys)
        $ joinD    theGang balanced
        $ zipWithD theGang go rep
        $ splitD   theGang balanced tags
        where   go ((i,j), (m,n)) ts 
                 = Seq.combine2ByTag ts 
                        (Seq.slice (here "combine2UP") xs i m)
                        (Seq.slice (here "combine2UP") ys j n)
{-# INLINE_UP combine2UP #-}


-- | Apply a worker function to correponding elements of two arrays.
zipWithUP :: (Unbox a, Unbox b, Unbox c) 
          => (a -> b -> c) -> Vector a -> Vector b -> Vector c
zipWithUP f xs ys
        = splitJoinD theGang 
                (mapD theGang (Seq.map (uncurry f))) 
                (Seq.zip xs ys)
{-# INLINE_UP zipWithUP #-}


-- | Undirected fold.
--   Note that this function has more constraints on its parameters than the
--   standard fold function from the Haskell Prelude.
--
--   * The worker function must be associative.
--
--   * The provided starting element must be neutral with respect to the worker.
--     For example 0 is neutral wrt (+) and 1 is neutral wrt (*).
--
--   We need these constraints so that we can partition the fold across 
--   several threads. Each thread folds a chunk of the input vector, 
--   then we fold together all the results in the main thread.
--
foldUP  :: (Unbox a, DT a) => (a -> a -> a) -> a -> Vector a -> a
foldUP f !z xs
        = foldD theGang f
                (mapD   theGang (Seq.fold f z)
                (splitD theGang unbalanced xs))
{-# INLINE_UP foldUP #-}


-- | Left fold over an array. 
--
--   * If the vector is empty then this returns the provided neural element.
--
--   * The worker function must be associative.
--
--   * The provided starting element must be neutral with respect to the worker,
--     see `foldUP` for discussion.
--
foldlUP :: (DT a, Unbox a) => (a -> a -> a) -> a -> Vector a -> a
foldlUP f z arr 
  | Seq.null arr = z
  | otherwise    = foldl1UP f arr
{-# INLINE_UP foldlUP #-}


-- | Alias for `foldl1UP`
fold1UP :: (DT a, Unbox a) => (a -> a -> a) -> Vector a -> a
fold1UP = foldl1UP
{-# INLINE_UP fold1UP #-}


-- | Left fold over an array, using the first element of the vector as the
--   neural element.
--
--   * If the vector contains no elements then you'll get a bounds-check error.
--
--   * The worker function must be associative.
--
--   * The provided starting element must be neutral with respect to the worker,
--     see `foldUP` for discussion.
--
foldl1UP :: (DT a, Unbox a) => (a -> a -> a) -> Vector a -> a
foldl1UP f arr 
        = (maybe z (f z)
        . foldD  theGang combine'
        . mapD   theGang (Seq.foldl1Maybe f)
        . splitD theGang unbalanced) arr
        where
                z = Seq.index (here "fold1UP") arr 0
                combine' (Just x) (Just y) = Just (f x y)
                combine' (Just x) Nothing  = Just x
                combine' Nothing  (Just y) = Just y
                combine' Nothing  Nothing  = Nothing
{-# INLINE_UP foldl1UP #-}


-- | Prefix scan. Similar to fold, but produce an array of the intermediate states.
--
--   * The worker function must be associative.
-- 
--   * The provided starting element must be neutral with respect to the worker,
--     see `foldUP` for discussion.
--
scanUP :: (DT a, Unbox a) => (a -> a -> a) -> a -> Vector a -> Vector a
scanUP f z 
        = splitJoinD theGang go
        where   go xs = let (ds,zs) = unzipD $ mapD theGang (Seq.scanRes f z) xs
                            zs'     = fst (scanD theGang f z zs)
                        in  zipWithD theGang (Seq.map . f) zs' ds
{-# INLINE_UP scanUP #-}