{-|
Module      : Data.Conduit.Algorithms.Utils
Copyright   : 2013-2021 Luis Pedro Coelho
License     : MIT
Maintainer  : luis@luispedro.org

A few miscellaneous conduit utils
-}
module Data.Conduit.Algorithms.Utils
    ( awaitJust
    , enumerateC
    , groupC
    , dispatchC
    , dispatchC_
    ) where

import qualified Data.Conduit as C
import qualified Data.Conduit.List as CL
import           Data.Conduit ((.|))
import           Data.Maybe (maybe)
import           Control.Monad (unless, void)

-- | Act on the next input (do nothing if no input). @awaitJust f@ is equivalent to
--
--
-- @ do
--      next <- C.await
--      case next of
--          Just val -> f val
--          Nothing -> return ()
-- @
--
-- This is a simple utility adapted from
-- http://neilmitchell.blogspot.de/2015/07/thoughts-on-conduits.html
awaitJust :: Monad m => (a -> C.ConduitT a b m ()) -> C.ConduitT a b m ()
awaitJust :: (a -> ConduitT a b m ()) -> ConduitT a b m ()
awaitJust a -> ConduitT a b m ()
f = ConduitT a b m (Maybe a)
forall (m :: * -> *) i. Monad m => Consumer i m (Maybe i)
C.await ConduitT a b m (Maybe a)
-> (Maybe a -> ConduitT a b m ()) -> ConduitT a b m ()
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= ConduitT a b m ()
-> (a -> ConduitT a b m ()) -> Maybe a -> ConduitT a b m ()
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (() -> ConduitT a b m ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()) a -> ConduitT a b m ()
f
{-# INLINE awaitJust #-}

-- | Conduit analogue to Python's enumerate function
enumerateC :: Monad m => C.ConduitT a (Int, a) m ()
enumerateC :: ConduitT a (Int, a) m ()
enumerateC = Int -> ConduitT a (Int, a) m ()
forall (m :: * -> *) t b.
(Monad m, Num t) =>
t -> ConduitT b (t, b) m ()
enumerateC' Int
0
    where
        enumerateC' :: t -> ConduitT b (t, b) m ()
enumerateC' !t
i = (b -> ConduitT b (t, b) m ()) -> ConduitT b (t, b) m ()
forall (m :: * -> *) a b.
Monad m =>
(a -> ConduitT a b m ()) -> ConduitT a b m ()
awaitJust ((b -> ConduitT b (t, b) m ()) -> ConduitT b (t, b) m ())
-> (b -> ConduitT b (t, b) m ()) -> ConduitT b (t, b) m ()
forall a b. (a -> b) -> a -> b
$ \b
v -> do
                                        (t, b) -> ConduitT b (t, b) m ()
forall (m :: * -> *) o i. Monad m => o -> ConduitT i o m ()
C.yield (t
i, b
v)
                                        t -> ConduitT b (t, b) m ()
enumerateC' (t
i t -> t -> t
forall a. Num a => a -> a -> a
+ t
1)
{-# INLINE enumerateC #-}

-- | This function is deprecated; use 'Data.Conduit.List.chunksOf'
--
-- groupC yields the input as groups of 'n' elements. If the input is not a
-- multiple of 'n', the last element will be incomplete
--
-- Example:
--
-- @
--      CC.yieldMany [0..10] .| groupC 3 .| CC.consumeList
-- @
--
-- results in @[ [0,1,2], [3,4,5], [6,7,8], [9, 10] ]@
--
groupC :: (Monad m) => Int -> C.ConduitT a [a] m ()
groupC :: Int -> ConduitT a [a] m ()
groupC Int
n = Int -> [a] -> ConduitT a [a] m ()
forall (m :: * -> *) a.
Monad m =>
Int -> [a] -> ConduitT a [a] m ()
loop Int
n []
    where
        loop :: Int -> [a] -> ConduitT a [a] m ()
loop Int
0 [a]
ps = [a] -> ConduitT a [a] m ()
forall (m :: * -> *) o i. Monad m => o -> ConduitT i o m ()
C.yield ([a] -> [a]
forall a. [a] -> [a]
reverse [a]
ps) ConduitT a [a] m () -> ConduitT a [a] m () -> ConduitT a [a] m ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Int -> [a] -> ConduitT a [a] m ()
loop Int
n []
        loop Int
c [a]
ps = ConduitT a [a] m (Maybe a)
forall (m :: * -> *) i. Monad m => Consumer i m (Maybe i)
C.await ConduitT a [a] m (Maybe a)
-> (Maybe a -> ConduitT a [a] m ()) -> ConduitT a [a] m ()
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
            Maybe a
Nothing -> Bool -> ConduitT a [a] m () -> ConduitT a [a] m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless ([a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [a]
ps) (ConduitT a [a] m () -> ConduitT a [a] m ())
-> ConduitT a [a] m () -> ConduitT a [a] m ()
forall a b. (a -> b) -> a -> b
$ [a] -> ConduitT a [a] m ()
forall (m :: * -> *) o i. Monad m => o -> ConduitT i o m ()
C.yield ([a] -> [a]
forall a. [a] -> [a]
reverse [a]
ps)
            Just a
p -> Int -> [a] -> ConduitT a [a] m ()
loop (Int
cInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) (a
pa -> [a] -> [a]
forall a. a -> [a] -> [a]
:[a]
ps)
{-# WARNING groupC "This function is deprecated; use 'Data.Conduit.List.chunksOf'" #-}

-- | dispatchC dispatches indexed input to the respective sink
-- 
-- Example:
--
-- @
-- 	let input = [(0, "one")
-- 	            ,(1, "two")
-- 	            ,(0, "three")
-- 	            ]
-- 	    CC.yieldMany input .| dispatches [sink1, sink2]
-- @
--
-- Then 'sink1' will receive "one" and "three", while 'sink2' will receive "two"
--
-- Out of bounds indices are clipped to the 0..n-1 range (where 'n' is 'length sinks')
dispatchC :: Monad m => [C.ConduitT a C.Void m r] -> C.ConduitT (Int, a) C.Void m [r]
dispatchC :: [ConduitT a Void m r] -> ConduitT (Int, a) Void m [r]
dispatchC [ConduitT a Void m r]
sinks = [Sink (Int, a) m r] -> ConduitT (Int, a) Void m [r]
forall (f :: * -> *) (m :: * -> *) i r.
(Traversable f, Monad m) =>
f (Sink i m r) -> Sink i m (f r)
C.sequenceSinks [Int -> ConduitT (Int, a) a m ()
forall (m :: * -> *) b. Monad m => Int -> ConduitT (Int, b) b m ()
select Int
i ConduitT (Int, a) a m ()
-> ConduitT a Void m r -> Sink (Int, a) m r
forall (m :: * -> *) a b c r.
Monad m =>
ConduitM a b m () -> ConduitM b c m r -> ConduitM a c m r
.| ConduitT a Void m r
s | (Int
i,ConduitT a Void m r
s) <- [Int] -> [ConduitT a Void m r] -> [(Int, ConduitT a Void m r)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0..] [ConduitT a Void m r]
sinks]
    where
        n :: Int
n = [ConduitT a Void m r] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [ConduitT a Void m r]
sinks
        select :: Int -> ConduitT (Int, b) b m ()
select Int
i = ((Int, b) -> Maybe b) -> ConduitT (Int, b) b m ()
forall (m :: * -> *) a b.
Monad m =>
(a -> Maybe b) -> ConduitT a b m ()
CL.mapMaybe (((Int, b) -> Maybe b) -> ConduitT (Int, b) b m ())
-> ((Int, b) -> Maybe b) -> ConduitT (Int, b) b m ()
forall a b. (a -> b) -> a -> b
$ \(Int
j,b
val) ->
            if Int
j Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
i Bool -> Bool -> Bool
|| (Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1 Bool -> Bool -> Bool
&& Int
j Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
n) Bool -> Bool -> Bool
|| (Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 Bool -> Bool -> Bool
&& Int
j Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0)
                then b -> Maybe b
forall a. a -> Maybe a
Just b
val
                else Maybe b
forall a. Maybe a
Nothing

-- | Version of 'dispatchC' that returns ()
dispatchC_ :: Monad m => [C.ConduitT a C.Void m ()] -> C.ConduitT (Int, a) C.Void m ()
dispatchC_ :: [ConduitT a Void m ()] -> ConduitT (Int, a) Void m ()
dispatchC_ [ConduitT a Void m ()]
sinks = ConduitT (Int, a) Void m [()] -> ConduitT (Int, a) Void m ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (ConduitT (Int, a) Void m [()] -> ConduitT (Int, a) Void m ())
-> ConduitT (Int, a) Void m [()] -> ConduitT (Int, a) Void m ()
forall a b. (a -> b) -> a -> b
$ [ConduitT a Void m ()] -> ConduitT (Int, a) Void m [()]
forall (m :: * -> *) a r.
Monad m =>
[ConduitT a Void m r] -> ConduitT (Int, a) Void m [r]
dispatchC [ConduitT a Void m ()]
sinks