combine two Sinks

combine.hs

import Data.Conduit
import Data.Conduit.Internal
import Data.Void (Void, absurd)
import qualified Data.Conduit.List as CL
import Control.Monad (zipWithM_)
import Control.Monad.Trans.Class (lift)

combine :: Monad m
        => Sink i1 m r1
        -> Sink i2 m r2
        -> Sink (Either i1 i2) m (r1, r2)
combine s1' s2' =
    go (injectLeftovers s1') (injectLeftovers s2')
  where
    go :: Monad m
       => Pipe Void i1 Void () m r1
       -> Pipe Void i2 Void () m r2
       -> Pipe l (Either i1 i2) Void () m (r1, r2)
    go (HaveOutput _ _ o) _ = absurd o
    go _ (HaveOutput _ _ o) = absurd o
    go (Leftover _ l) _ = absurd l
    go _ (Leftover _ l) = absurd l
    go (Done r1) s2 = do
        r2 <- lift $ runPipe $ return () >+> s2
        return (r1, r2)
    go s1 (Done r2) = do
        r1 <- lift $ runPipe $ return () >+> s1
        return (r1, r2)
    go (PipeM m1) s2 = do
        s1 <- lift m1
        go s1 s2
    go s1 (PipeM m2) = do
        s2 <- lift m2
        go s1 s2
    go (NeedInput p1 c1) (NeedInput p2 c2) =
        NeedInput p c
      where
        p (Left i1) = go (p1 i1) (NeedInput p2 c2)
        p (Right i2) = go (NeedInput p1 c1) (p2 i2)
        c () = go (c1 ()) (c2 ())

source :: Monad m => Source m (Either Char Int)
source =
    zipWithM_ go "hello world" [1..]
  where
    go x y = yield (Left x) >> yield (Right y)

sink :: Monad m => Sink (Either Char Int) m (String, Int)
sink = combine CL.consume (CL.fold (+) 0)

main :: IO ()
main = (source $$ sink) >>= print