module Data.Repa.Stream.Merge
(mergeS)
where
import Data.Vector.Fusion.Stream.Monadic (Stream(..), Step(..))
import Data.Repa.Scalar.Option
import qualified Data.Vector.Fusion.Stream.Size as S
#include "repa-stream.h"
mergeS :: (Monad m, Ord k)
=> (k -> a -> b -> c)
-> (k -> a -> c)
-> (k -> b -> c)
-> Stream m (k, a)
-> Stream m (k, b)
-> Stream m (k, c)
mergeS fBoth fLeft fRight (Stream istepA sA0 _) (Stream istepB sB0 _)
= Stream ostep (sA0, sB0, None2, True, None2, True) S.Unknown
where
ostep (sA, sB, kxA@(Some2 kA xA), hasA
, kxB@(Some2 kB xB), hasB)
= return $ Yield (if | kA == kB -> (kA, fBoth kA xA xB)
| kB < kA -> (kB, fRight kB xB)
| otherwise -> (kA, fLeft kA xA))
(if | kA == kB -> (sA, sB, None2, hasA, None2, hasB)
| kB < kA -> (sA, sB, kxA, hasA, None2, hasB)
| otherwise -> (sA, sB, None2, hasA, kxB, hasB))
ostep (sA, sB, Some2 kA xA, hasA, kxB@None2, hasB@False)
= return $ Yield (kA, fLeft kA xA)
(sA, sB, None2, hasA, kxB, hasB)
ostep (sA, sB, kxA@None2, hasA@False, Some2 kB xB, hasB)
= return $ Yield (kB, fRight kB xB)
(sA, sB, kxA, hasA, None2, hasB)
ostep (sA, sB, kxA@None2, hasA@True, kxB, hasB)
= istepA sA >>= \mA
-> case mA of
Yield (kA, xA) sA'
-> return $ Skip (sA', sB, Some2 kA xA, True, kxB, hasB)
Skip sA'
-> return $ Skip (sA', sB, kxA, hasA, kxB, hasB)
Done
-> return $ Skip (sA, sB, kxA, False, kxB, hasB)
ostep (sA, sB, kxA, hasA, kxB@None2, hasB@True)
= istepB sB >>= \mB
-> case mB of
Yield (kB, xB) sB'
-> return $ Skip (sA, sB', kxA, hasA, Some2 kB xB, True)
Skip sB'
-> return $ Skip (sA, sB', kxA, hasA, kxB, hasB)
Done
-> return $ Skip (sA, sB, kxA, hasA, kxB, False)
ostep (_sA, _sB, None2, False, None2, False)
= return $ Done