Conduitで遊んでみた



    {-# LANGUAGE RankNTypes, BangPatterns, ViewPatterns #-}
    module Algorithm.Fina.Helper (zipC, teeC, tapC, dupC, sma, ema, fir) where

    import Data.Maybe (isJust, fromJust)
    import Control.Applicative
    import qualified Data.Text as T
    import qualified Data.Conduit as C
    import qualified Data.Conduit.List as C


    {-| 遅延
      NOTE: 先頭にダミーを挟み込んでるだけなので、長さがその分伸びてしまう
    
    >>> C.sourceList [1..10] C.=$= tapC 2 0 C.$$ C.consume
    [0,0,1,2,3,4,5,6,7,8,9,10]
    
    -}
    tapC :: Monad m =>
        Int     -- 遅延サンプル数
        -> o    -- 遅延期間中のダミー出力値
        -> C.Pipe o o m ()
    tapC n o | n < 0  = error "tapC: The n should ge. zero."
                 | otherwise = C.haveMore (C.map id) (return ()) (replicate n o)
    
    
    
    {-| 0次補完でオーバーサンプリング
    
    >>> C.sourceList [1..9] C.=$= dupC 3 C.$$ C.consume
    [1,1,1,2,2,2,3,3,3,4,4,4,5,5,5,6,6,6,7,7,7,8,8,8,9,9,9]
    
    -}
    dupC :: Monad m => 
        Int     -- 倍数
        -> C.Pipe a a m ()
    dupC n | n <= 0 = error "dupC: The n should greater than zero."
           | otherwise = let dc = C.NeedInput (\x -> C.haveMore dc (return ()) (replicate n x)) (return ()) in dc
    
    
    {-| 単純移動平均 -}
    sma :: (Monad m, Real a, Fractional b) =>
        Int                 -- 平均回数
        -> C.Conduit a m b
    sma n = C.conduitState
          (replicate (n - 1) 0)
          (\xs (realToFrac -> y) ->
            return $ let !z = (sum xs + y) / fromIntegral n
                      in C.StateProducing (y:init xs) [z]
          )
          (\_ -> return [])

    {-| 指数平滑移動平均 -}
    ema :: (Monad m, Real a, Fractional b) =>
        Int    -- 平均回数
        -> C.Conduit a m b
    ema n = C.conduitState
              0
              (\y (realToFrac -> x) -> return $ let z = y + alpha * (x - y) in C.StateProducing z [z])
              (\_ -> return [])
      where alpha = 2 / fromIntegral (n + 1)
    
    
    type Queue a = ([a], [a])
    emptyQ = ([], [])
    isEmptyQ ([], []) = True
    isEmptyQ _ = False
    singletonQ x = ([x], [])
    enQ y ([], ys) = (reverse (y:ys), [])
    enQ y (xs, ys) = (xs, y:ys)
    deQ ([], _) = undefined
    deQ (_:[], ys) = (reverse ys, [])
    deQ (_:xs, ys) = (xs, ys)
    headQ ([], _) = undefined
    headQ (x:_, _) = x
    
    {-| 入力を2つのPipeに食わせて、出力をfで束ねて出力するpipe
    
    >>> C.sourceList [0..10] C.=$= (zipC (,) (C.map ((+ 2))) (C.map negate)) C.$$ C.consume
    [(2,0),(3,-1),(4,-2),(5,-3),(6,-4),(7,-5),(8,-6),(9,-7),(10,-8),(11,-9),(12,-10)]
    
    >>> C.sourceList [0..10] C.=$= (zipC (,) (tapC 2 (-1)) (C.map id)) C.$$ C.consume
    [(-1,0),(-1,1),(0,2),(1,3),(2,4),(3,5),(4,6),(5,7),(6,8),(7,9),(8,10)]
    
    >>> C.sourceList [0..10] C.=$= (zipC (,) (dupC 2) (C.map id)) C.$$ C.consume
    [(-1,0),(-1,1),(0,2),(1,3),(2,4),(3,5),(4,6),(5,7),(6,8),(7,9),(8,10)]
    
    --タプルじゃなくてリストにまで拡張したい...そしてconcurrentへ
    joinC :: Monad m => [C.Pipe i o m ()] -> CPipe i [Maybe o] m ()
    parC :: Monad m => [C.Pipe i o m ()] -> CPile [i] [o] m ()
    
    splitC :: C.Pipe i a m () -> (C.Source m a, C.Source m a)
      let (s1,s2) = splitC (c1 =$= c2 =$= c3)
      s1 =$= c4 $$ mysink1
      s2 =$= c5 $$ mysink2
    って方が使いやすいかったかな？
    
    -}
    zipC :: forall a b c i m. (Monad m) => (a -> b -> c) -> C.Pipe i a m () -> C.Pipe i b m () -> C.Pipe i c m ()
    zipC f p1 p2 = zc f emptyQ p1 p2
      where
        -- NOTE: 引数をひっくり返しながら動くので、型シグネチャを書いて多相にしないと駄目
        zc :: Monad m => (a -> b -> c) -> ([i], [i]) -> C.Pipe i a m () -> C.Pipe i b m () -> C.Pipe i c m ()
        -- 出力完成
        zc f b (C.HaveOutput na ca oa) (C.HaveOutput nb cb ob) = C.HaveOutput (zc f b na nb) (ca >> cb) (f oa ob)
        --
        zc f b (C.PipeM mp1 ca) (C.PipeM mp2 cb) = C.PipeM (do { p1 <- mp1; p2 <- mp2; return $ zc f b p1 p2 }) (ca >> cb)
        zc f b (C.PipeM mp clean) p2 = C.PipeM (do { p <- mp; return $ zc f b p p2 }) clean
        zc f b p1 (C.PipeM mp clean) = C.PipeM (do { p <- mp; return $ zc f b p1 p }) clean
        --
        zc f b (C.NeedInput g p1) p2
          | isEmptyQ b  = case p2 of
                            C.NeedInput h p2' -> C.NeedInput (\x -> zc f b (g x) (h x)) (zc f b p1 p2')  -- 需要が合致しているので産地直送するように最適化
                            _ -> C.NeedInput (\x -> zc (flip f) (singletonQ x) p2 (g x)) (zc f b p1 p2) -- NOTE: flipの畳み込みはいつ起こる？
          | otherwise   = zc f (deQ b) (g $ headQ b) p2
        zc f b p1 (C.NeedInput h p2) = C.NeedInput (\x -> zc f (enQ x b) p1 (h x)) (zc f b p1 p2)
        -- FIXME: 片方が閉じたとき、もう片方はどうする？
        zc _ _ (C.Done _ _) _ = C.Done Nothing ()
        zc _ _ _ (C.Done _ _) = C.Done Nothing ()
    
    
    (&&&) :: forall a b i m. (Monad m) => C.Pipe i a m () -> C.Pipe i b m () -> C.Pipe i (a, b) m ()
    (&&&) = zipC (,)
    
    
    {-| tee
    
    >>> C.sourceList [0..5] `teeC` (C.map ((+) 1), C.map negate) C.$$ C.consume
    [(1,0),(2,-1),(3,-2),(4,-3),(5,-4),(6,-5)]
    
     -}
    teeC :: Monad m => C.Pipe i a m () -> (C.Pipe a b m (), C.Pipe a c m ()) -> C.Pipe i (b,c) m ()
    teeC p1 (p2, p3) = p1 C.=$= (p2 &&& p3)

    {-| 3タップのFIRフィルタ
      状態付きがあればtapだのzipだのは要らなかったのだった
    
    >>> C.sourceList [60,60,60,80,40,60,60] C.=$= fir 0.25 0.5 0.25 C.$$ C.consume
    [15.0,45.0,60.0,65.0,65.0,55.0,55.0]
     -}
    fir :: forall a m. (Num a, Monad m) => a -> a -> a -> C.Conduit a m a
    fir a1 a2 a3=
      C.conduitState
        (0, 0)
        (\(z1, z2) x -> return $ C.StateProducing (x, z1) [x * a1 + z1 * a2 + z2 * a3])
        (\_ -> return [])
    
    {-|
     -}
    tsv :: Monad m => C.Conduit T.Text m [T.Text]
    tsv = C.map (T.split ((==) '\t'))
    
    {-| [a]な入力の先頭n個とそれ以外を受け取るConduitに分割する
        Excelで見る事前提なCSVなんかだと、nカラムおきに複数チャンネルが重畳されてたりするから、
        その分割用。
        NOTE: これだと、片方のチャンネルがDoneになった時点で両方とも止まってしまうのでは？
     -}
    multiChannelTSV :: Monad m => Int -> C.Conduit [c] m a -> C.Conduit [c] m b -> C.Conduit [c] m (a,b)
    multiChannelTSV n pa pb = (C.map (take n) C.=$= pa) &&& (C.map (drop n) C.=$= pb)
    
    
    {-| Eitherを使って分岐するPipe
        NOTE: 入力を与えていないのに出力が出てくるようなPipeの場合、両方とも出力を
              持つケースがあるが、その場合は左側から優先して取り出す。
     -}
    splitEitherC :: forall a b c d m. Monad m => C.Conduit a m b -> C.Conduit c m d -> C.Conduit (Either a c) m (Either b d)
    splitEitherC (C.HaveOutput na ca oa) p2 = C.HaveOutput (splitEitherC na p2) ca (Left oa)
    splitEitherC p1 (C.HaveOutput nb cb ob) = C.HaveOutput (splitEitherC p1 nb) cb (Right ob)
    splitEitherC (C.PipeM ma ca) p2 = C.PipeM (do { p1 <- ma; return $ splitEitherC p1 p2 }) ca
    splitEitherC p1 (C.PipeM mb cb) = C.PipeM (do { p2 <- mb; return $ splitEitherC p1 p2 }) cb
    splitEitherC (C.Done _ _) _ = C.Done Nothing ()
    splitEitherC _ (C.Done _ _) = C.Done Nothing ()
    splitEitherC p1@(C.NeedInput f ca) p2@(C.NeedInput g cb)
      = C.NeedInput (\x -> case x of { Left xx -> splitEitherC (f xx) p2; Right xx -> splitEitherC p1 (g xx) }) (splitEitherC ca cb)
    
    
    eitherJoinC :: forall a m. Monad m => C.Conduit (Either a a) m a
    eitherJoinC = C.map (either id id)
    
    {-|
    
    >>> C.sourceList [Left 1,Left 2,Right 3, Left 4, Right 5, Right 6] C.=$= (C.map ((+) 2) ||| C.map ((*) 2)) C.$$ C.consume
    [3,4,6,6,10,12]
    -}
    (|||) :: forall a b c m. Monad m => C.Conduit a m c -> C.Conduit b m c -> C.Conduit (Either a b) m c
    (|||) x y = splitEitherC x y C.=$= eitherJoinC
    
    
    {-|
     -}
    maybeC :: forall a b m. Monad m => C.Conduit a m b -> C.Conduit (Maybe a) m (Maybe b)
    maybeC (C.HaveOutput na ca oa) = C.HaveOutput (maybeC na) ca (Just oa)
    maybeC (C.PipeM ma ca) = C.PipeM (do { p <- ma; return $ maybeC p }) ca
    maybeC (C.Done _ _) = C.Done Nothing ()
    maybeC p@(C.NeedInput f ca) = C.NeedInput (maybe (C.HaveOutput (maybeC p) (return ()) Nothing) $ maybeC . f) (maybeC ca)

    fromJustC :: forall a m. Monad m => C.Conduit (Maybe a) m a
    fromJustC = C.filter isJust C.=$= C.map fromJust

    (?=$=) :: forall a b c d m. Monad m => C.Pipe b (Maybe a) m () -> C.Pipe a c m d -> C.Pipe b c m d
    (?=$=) b c = b C.=$= fromJustC C.=$= c
    
    applyC :: forall a b m. Monad m => C.Conduit (a -> b, a) m b
    applyC = C.map (\(f, x) -> f x) -- もうArrowでやれというか...