branch in conduits

branch.hs

{-# LANGUAGE NoMonomorphismRestriction #-}
import Data.Conduit
import qualified Data.Conduit.List as CL

branch :: Resource m
       => Sink lIn m lOut
       -> Sink rIn m rOut
       -> Sink (Either lIn rIn) m (lOut, rOut)
branch (Sink mlsink) (Sink mrsink) = Sink $ do
    lsink <- mlsink
    rsink <- mrsink
    -- SinkNoData for either sink means we return early
    case (lsink, rsink) of
        (SinkNoData l, SinkNoData r) -> return $ SinkNoData (l, r)
        (SinkNoData l, SinkData _ rclose) -> do
            r <- rclose
            return $ SinkNoData (l, r)
        (SinkData _ lclose, SinkNoData r) -> do
            l <- lclose
            return $ SinkNoData (l, r)
        (SinkData lpush lclose, SinkData rpush rclose) -> do
            let push (Left input) = do
                    res <- lpush input
                    case res of
                        Processing -> return Processing
                        Done lo l -> do
                            let lo' = fmap Left lo
                            r <- rclose
                            return $ Done lo' (l, r)
                push (Right input) = do
                    res <- rpush input
                    case res of
                        Processing -> return Processing
                        Done lo r -> do
                            let lo' = fmap Right lo
                            l <- lclose
                            return $ Done lo' (l, r)
                close = do
                    l <- lclose
                    r <- rclose
                    return (l, r)
            return $ SinkData push close

main :: IO ()
main = runResourceT
    (CL.sourceList [Left (1 :: Int), Right (2 :: Int), Left 3, Right 4]
    $$ branch CL.consume (CL.fold (+) 0)) >>= print