module CfnFlip.Conduit
  ( takeBalancedC
  , module Conduit
  , module Data.Conduit.List
  ) where

import CfnFlip.Prelude

import Conduit
import Data.Conduit.List (sourceList)

-- | Take until an ending element, including any reopened-ended in between
--
-- >>> :{
-- runIdentity
--   $ runConduit
--   $ yieldMany "this is (a thing) here) and more"
--   .| takeBalancedC (== '(') (==')')
--   .| sinkList
-- :}
-- "this is (a thing) here)"
--
-- Note that imbalance will terminate early,
--
-- >>> :{
-- runIdentity
--   $ runConduit
--   $ yieldMany "this is, a) unexpected but b) the best we can do)"
--   .| takeBalancedC (== '(') (==')')
--   .| sinkList
-- :}
-- "this is, a)"
--
-- Or just run to the end
--
-- >>> :{
-- runIdentity
--   $ runConduit
--   $ yieldMany "this is (pretty unlikely"
--   .| takeBalancedC (== '(') (==')')
--   .| sinkList
-- :}
-- "this is (pretty unlikely"
--
takeBalancedC :: Monad m => (a -> Bool) -> (a -> Bool) -> ConduitT a a m ()
takeBalancedC :: (a -> Bool) -> (a -> Bool) -> ConduitT a a m ()
takeBalancedC a -> Bool
reopens a -> Bool
closes = Int -> ConduitT a a m ()
go (Int
0 :: Int)
 where
  go :: Int -> ConduitT a a m ()
go Int
balance = do
    Maybe a
me <- ConduitT a a m (Maybe a)
forall (m :: * -> *) i. Monad m => Consumer i m (Maybe i)
await

    Maybe a -> (a -> ConduitT a a m ()) -> ConduitT a a m ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ Maybe a
me ((a -> ConduitT a a m ()) -> ConduitT a a m ())
-> (a -> ConduitT a a m ()) -> ConduitT a a m ()
forall a b. (a -> b) -> a -> b
$ \a
a -> do
      let
        loop :: ConduitT a a m ()
loop
          | a -> Bool
closes a
a = Bool -> ConduitT a a m () -> ConduitT a a m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (Int
balance Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0) (ConduitT a a m () -> ConduitT a a m ())
-> ConduitT a a m () -> ConduitT a a m ()
forall a b. (a -> b) -> a -> b
$ Int -> ConduitT a a m ()
go (Int -> ConduitT a a m ()) -> Int -> ConduitT a a m ()
forall a b. (a -> b) -> a -> b
$ Int
balance Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
          | a -> Bool
reopens a
a = Int -> ConduitT a a m ()
go (Int -> ConduitT a a m ()) -> Int -> ConduitT a a m ()
forall a b. (a -> b) -> a -> b
$ Int
balance Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
          | Bool
otherwise = Int -> ConduitT a a m ()
go Int
balance

      a -> ConduitT a a m ()
forall (m :: * -> *) o i. Monad m => o -> ConduitT i o m ()
yield a
a ConduitT a a m () -> ConduitT a a m () -> ConduitT a a m ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> ConduitT a a m ()
loop