{-# LANGUAGE LambdaCase #-}
module Data.BAByNF.Util.Stream where
import Data.Bifunctor qualified as Bifunctor
import Data.Maybe (isJust, isNothing)
import Control.Applicative (liftA2)
import Control.Monad (when)
import Prelude hiding (take, drop, takeWhile, dropWhile)
import Data.Kind (Type)

newtype Stream e a = Stream
    { forall e a. Stream e a -> [e] -> ([e], a)
runStream :: [e] -> ([e], a)
    }

runStream_ :: Stream e a -> [e] -> a
runStream_ :: forall e a. Stream e a -> [e] -> a
runStream_ Stream e a
stream = ([e], a) -> a
forall a b. (a, b) -> b
snd (([e], a) -> a) -> ([e] -> ([e], a)) -> [e] -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Stream e a -> [e] -> ([e], a)
forall e a. Stream e a -> [e] -> ([e], a)
runStream Stream e a
stream
instance Functor (Stream e) where
  fmap :: forall a b. (a -> b) -> Stream e a -> Stream e b
fmap a -> b
func Stream e a
stream = Stream { runStream :: [e] -> ([e], b)
runStream =  (a -> b) -> ([e], a) -> ([e], b)
forall b c a. (b -> c) -> (a, b) -> (a, c)
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
Bifunctor.second a -> b
func (([e], a) -> ([e], b)) -> ([e] -> ([e], a)) -> [e] -> ([e], b)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Stream e a -> [e] -> ([e], a)
forall e a. Stream e a -> [e] -> ([e], a)
runStream Stream e a
stream}
instance Applicative (Stream e) where
  pure :: forall a. a -> Stream e a
pure a
a = Stream { runStream :: [e] -> ([e], a)
runStream = (, a
a) }
  liftA2 :: forall a b c.
(a -> b -> c) -> Stream e a -> Stream e b -> Stream e c
liftA2 a -> b -> c
func Stream e a
s1 Stream e b
s2 = Stream { runStream :: [e] -> ([e], c)
runStream = (\([e]
es1, a
a) -> (b -> c) -> ([e], b) -> ([e], c)
forall b c a. (b -> c) -> (a, b) -> (a, c)
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
Bifunctor.second (a -> b -> c
func a
a) (([e], b) -> ([e], c)) -> ([e] -> ([e], b)) -> [e] -> ([e], c)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Stream e b -> [e] -> ([e], b)
forall e a. Stream e a -> [e] -> ([e], a)
runStream Stream e b
s2 ([e] -> ([e], c)) -> [e] -> ([e], c)
forall a b. (a -> b) -> a -> b
$ [e]
es1) (([e], a) -> ([e], c)) -> ([e] -> ([e], a)) -> [e] -> ([e], c)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Stream e a -> [e] -> ([e], a)
forall e a. Stream e a -> [e] -> ([e], a)
runStream Stream e a
s1 }
instance Monad (Stream e) where
  >>= :: forall a b. Stream e a -> (a -> Stream e b) -> Stream e b
(>>=) Stream e a
s a -> Stream e b
sf = Stream { runStream :: [e] -> ([e], b)
runStream = (\([e]
es, a
a) -> Stream e b -> [e] -> ([e], b)
forall e a. Stream e a -> [e] -> ([e], a)
runStream (a -> Stream e b
sf a
a) [e]
es ) (([e], a) -> ([e], b)) -> ([e] -> ([e], a)) -> [e] -> ([e], b)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Stream e a -> [e] -> ([e], a)
forall e a. Stream e a -> [e] -> ([e], a)
runStream Stream e a
s }


hasNext :: Stream e Bool
hasNext :: forall e. Stream e Bool
hasNext = Stream { runStream :: [e] -> ([e], Bool)
runStream = \[e]
es -> case [e]
es of [] -> ([e]
es, Bool
False); [e]
_ -> ([e]
es, Bool
True) }
take :: Stream e (Maybe e)
take :: forall e. Stream e (Maybe e)
take = Stream { runStream :: [e] -> ([e], Maybe e)
runStream = \[e]
es -> case [e]
es of [] -> ([e]
es, Maybe e
forall a. Maybe a
Nothing); e
x:[e]
xs -> ([e]
xs, e -> Maybe e
forall a. a -> Maybe a
Just e
x)}
drop :: Stream e ()
drop :: forall e. Stream e ()
drop = Stream { runStream :: [e] -> ([e], ())
runStream = \[e]
es -> case [e]
es of [] -> ([e]
es, ()); e
_:[e]
xs -> ([e]
xs, ())}
peek :: Stream e (Maybe e)
peek :: forall e. Stream e (Maybe e)
peek = Stream { runStream :: [e] -> ([e], Maybe e)
runStream = \[e]
es -> case [e]
es of [] -> ([e]
es, Maybe e
forall a. Maybe a
Nothing); e
x:[e]
_ -> ([e]
es, e -> Maybe e
forall a. a -> Maybe a
Just e
x)}
takeIf :: (e -> Bool) -> Stream e (Maybe e)
takeIf :: forall e. (e -> Bool) -> Stream e (Maybe e)
takeIf e -> Bool
cond = do
    Maybe e
opt <- Stream e (Maybe e)
forall e. Stream e (Maybe e)
peek
    case Maybe e
opt of
        Maybe e
Nothing -> Maybe e -> Stream e (Maybe e)
forall a. a -> Stream e a
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe e
forall a. Maybe a
Nothing
        Just e
x -> if e -> Bool
cond e
x then Stream e ()
forall e. Stream e ()
drop Stream e () -> Stream e (Maybe e) -> Stream e (Maybe e)
forall a b. Stream e a -> Stream e b -> Stream e b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Maybe e -> Stream e (Maybe e)
forall a. a -> Stream e a
forall (m :: * -> *) a. Monad m => a -> m a
return (e -> Maybe e
forall a. a -> Maybe a
Just e
x) else Maybe e -> Stream e (Maybe e)
forall a. a -> Stream e a
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe e
forall a. Maybe a
Nothing
dropIf :: (e -> Bool) -> Stream e ()
dropIf :: forall e. (e -> Bool) -> Stream e ()
dropIf e -> Bool
cond = do
    Maybe e
opt <- Stream e (Maybe e)
forall e. Stream e (Maybe e)
peek
    case Maybe e
opt of
        Maybe e
Nothing -> () -> Stream e ()
forall a. a -> Stream e a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
        Just e
x -> Bool -> Stream e () -> Stream e ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (e -> Bool
cond e
x) Stream e ()
forall e. Stream e ()
drop
takeWhile :: (e -> Bool) -> Stream e [e]
takeWhile :: forall e. (e -> Bool) -> Stream e [e]
takeWhile e -> Bool
cond = do
    Maybe e
opt <- (e -> Bool) -> Stream e (Maybe e)
forall e. (e -> Bool) -> Stream e (Maybe e)
takeIf e -> Bool
cond
    case Maybe e
opt of
        Maybe e
Nothing -> [e] -> Stream e [e]
forall a. a -> Stream e a
forall (m :: * -> *) a. Monad m => a -> m a
return []
        Just e
x -> do
            [e]
xs <- (e -> Bool) -> Stream e [e]
forall e. (e -> Bool) -> Stream e [e]
takeWhile e -> Bool
cond
            [e] -> Stream e [e]
forall a. a -> Stream e a
forall (m :: * -> *) a. Monad m => a -> m a
return (e
xe -> [e] -> [e]
forall a. a -> [a] -> [a]
:[e]
xs)
dropWhile :: (e -> Bool) -> Stream e ()
dropWhile :: forall e. (e -> Bool) -> Stream e ()
dropWhile e -> Bool
cond = do
    Maybe e
opt <- Stream e (Maybe e)
forall e. Stream e (Maybe e)
peek
    case Maybe e
opt of
        Maybe e
Nothing -> () -> Stream e ()
forall a. a -> Stream e a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
        Just e
x -> Bool -> Stream e () -> Stream e ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (e -> Bool
cond e
x) (Stream e () -> Stream e ()) -> Stream e () -> Stream e ()
forall a b. (a -> b) -> a -> b
$ Stream e ()
forall e. Stream e ()
drop Stream e () -> Stream e () -> Stream e ()
forall a b. Stream e a -> Stream e b -> Stream e b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> (e -> Bool) -> Stream e ()
forall e. (e -> Bool) -> Stream e ()
dropWhile e -> Bool
cond

find :: (e -> Bool) -> Stream e (Maybe e)
find :: forall e. (e -> Bool) -> Stream e (Maybe e)
find e -> Bool
cond = (e -> Bool) -> Stream e ()
forall e. (e -> Bool) -> Stream e ()
dropWhile (Bool -> Bool
not (Bool -> Bool) -> (e -> Bool) -> e -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. e -> Bool
cond) Stream e () -> Stream e (Maybe e) -> Stream e (Maybe e)
forall a b. Stream e a -> Stream e b -> Stream e b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> (e -> Bool) -> Stream e (Maybe e)
forall e. (e -> Bool) -> Stream e (Maybe e)
takeIf e -> Bool
cond

findSeq :: (e -> Bool) -> Stream e [e]
findSeq :: forall e. (e -> Bool) -> Stream e [e]
findSeq e -> Bool
cond = (e -> Bool) -> Stream e ()
forall e. (e -> Bool) -> Stream e ()
dropWhile (Bool -> Bool
not (Bool -> Bool) -> (e -> Bool) -> e -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. e -> Bool
cond) Stream e () -> Stream e [e] -> Stream e [e]
forall a b. Stream e a -> Stream e b -> Stream e b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> (e -> Bool) -> Stream e [e]
forall e. (e -> Bool) -> Stream e [e]
takeWhile e -> Bool
cond

takeIfMap :: (e -> Maybe a) -> Stream e (Maybe a)
takeIfMap :: forall e a. (e -> Maybe a) -> Stream e (Maybe a)
takeIfMap e -> Maybe a
func = do
    Maybe a
opt <- Stream e (Maybe e)
forall e. Stream e (Maybe e)
peek Stream e (Maybe e)
-> (Maybe e -> Stream e (Maybe a)) -> Stream e (Maybe a)
forall a b. Stream e a -> (a -> Stream e b) -> Stream e b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (\Maybe e
maybeE -> Maybe a -> Stream e (Maybe a)
forall a. a -> Stream e a
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe a -> Stream e (Maybe a)) -> Maybe a -> Stream e (Maybe a)
forall a b. (a -> b) -> a -> b
$ Maybe e
maybeE Maybe e -> (e -> Maybe a) -> Maybe a
forall a b. Maybe a -> (a -> Maybe b) -> Maybe b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= e -> Maybe a
func)
    Bool -> Stream e () -> Stream e ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Maybe a -> Bool
forall a. Maybe a -> Bool
isJust Maybe a
opt) Stream e ()
forall e. Stream e ()
drop Stream e () -> Stream e (Maybe a) -> Stream e (Maybe a)
forall a b. Stream e a -> Stream e b -> Stream e b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Maybe a -> Stream e (Maybe a)
forall a. a -> Stream e a
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe a
opt

takeWhileMap :: (e -> Maybe a) -> Stream e [a]
takeWhileMap :: forall e a. (e -> Maybe a) -> Stream e [a]
takeWhileMap e -> Maybe a
func = do
    Maybe a
opt <- (e -> Maybe a) -> Stream e (Maybe a)
forall e a. (e -> Maybe a) -> Stream e (Maybe a)
takeIfMap e -> Maybe a
func
    case Maybe a
opt of
        Maybe a
Nothing -> [a] -> Stream e [a]
forall a. a -> Stream e a
forall (m :: * -> *) a. Monad m => a -> m a
return []
        Just a
e -> (e -> Maybe a) -> Stream e [a]
forall e a. (e -> Maybe a) -> Stream e [a]
takeWhileMap e -> Maybe a
func Stream e [a] -> ([a] -> Stream e [a]) -> Stream e [a]
forall a b. Stream e a -> (a -> Stream e b) -> Stream e b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (\[a]
es -> [a] -> Stream e [a]
forall a. a -> Stream e a
forall (m :: * -> *) a. Monad m => a -> m a
return ([a] -> Stream e [a]) -> [a] -> Stream e [a]
forall a b. (a -> b) -> a -> b
$ a
e a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a]
es)

findMap :: (e -> Maybe a) -> Stream e (Maybe a)
findMap :: forall e a. (e -> Maybe a) -> Stream e (Maybe a)
findMap e -> Maybe a
func = (e -> Bool) -> Stream e ()
forall e. (e -> Bool) -> Stream e ()
dropWhile (Maybe a -> Bool
forall a. Maybe a -> Bool
isNothing (Maybe a -> Bool) -> (e -> Maybe a) -> e -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. e -> Maybe a
func) Stream e () -> Stream e (Maybe a) -> Stream e (Maybe a)
forall a b. Stream e a -> Stream e b -> Stream e b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> (e -> Maybe a) -> Stream e (Maybe a)
forall e a. (e -> Maybe a) -> Stream e (Maybe a)
takeIfMap e -> Maybe a
func

findSeqMap :: (e -> Maybe a) -> Stream e [a]
findSeqMap :: forall e a. (e -> Maybe a) -> Stream e [a]
findSeqMap e -> Maybe a
func = (e -> Bool) -> Stream e ()
forall e. (e -> Bool) -> Stream e ()
dropWhile (Maybe a -> Bool
forall a. Maybe a -> Bool
isNothing (Maybe a -> Bool) -> (e -> Maybe a) -> e -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. e -> Maybe a
func) Stream e () -> Stream e [a] -> Stream e [a]
forall a b. Stream e a -> Stream e b -> Stream e b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> (e -> Maybe a) -> Stream e [a]
forall e a. (e -> Maybe a) -> Stream e [a]
takeWhileMap e -> Maybe a
func

either :: Stream e (Either l a) -> (a -> Stream e b) -> Stream e (Either l b)
either :: forall e l a b.
Stream e (Either l a) -> (a -> Stream e b) -> Stream e (Either l b)
either Stream e (Either l a)
stream a -> Stream e b
action = Stream e (Either l a)
stream Stream e (Either l a)
-> (Either l a -> Stream e (Either l b)) -> Stream e (Either l b)
forall a b. Stream e a -> (a -> Stream e b) -> Stream e b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case Left l
l -> Either l b -> Stream e (Either l b)
forall a. a -> Stream e a
forall (m :: * -> *) a. Monad m => a -> m a
return (l -> Either l b
forall a b. a -> Either a b
Left l
l); Right a
a -> (b -> Either l b) -> Stream e b -> Stream e (Either l b)
forall a b. (a -> b) -> Stream e a -> Stream e b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap b -> Either l b
forall a b. b -> Either a b
Right (a -> Stream e b
action a
a)


either' :: Stream e (Either l a) -> (a -> Stream e (Either l b)) -> Stream e (Either l b)
either' :: forall e l a b.
Stream e (Either l a)
-> (a -> Stream e (Either l b)) -> Stream e (Either l b)
either' Stream e (Either l a)
stream a -> Stream e (Either l b)
action = Stream e (Either l a)
stream Stream e (Either l a)
-> (Either l a -> Stream e (Either l b)) -> Stream e (Either l b)
forall a b. Stream e a -> (a -> Stream e b) -> Stream e b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case Left l
l -> Either l b -> Stream e (Either l b)
forall a. a -> Stream e a
forall (m :: * -> *) a. Monad m => a -> m a
return (l -> Either l b
forall a b. a -> Either a b
Left l
l); Right a
a -> a -> Stream e (Either l b)
action a
a


class Propagate (p :: Type -> Type)  where
    propagate :: (Monad m) => m (p a) -> (a -> m (p b)) -> m (p b)

instance Propagate Maybe where
    propagate :: forall (m :: * -> *) a b.
Monad m =>
m (Maybe a) -> (a -> m (Maybe b)) -> m (Maybe b)
propagate m (Maybe a)
m1 a -> m (Maybe b)
m2 = m (Maybe a)
m1 m (Maybe a) -> (Maybe a -> m (Maybe b)) -> m (Maybe b)
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case Maybe a
Nothing -> Maybe b -> m (Maybe b)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe b
forall a. Maybe a
Nothing; Just a
a -> a -> m (Maybe b)
m2 a
a

instance Propagate (Either a) where
    propagate :: forall (m :: * -> *) a b.
Monad m =>
m (Either a a) -> (a -> m (Either a b)) -> m (Either a b)
propagate m (Either a a)
m1 a -> m (Either a b)
m2 = m (Either a a)
m1 m (Either a a) -> (Either a a -> m (Either a b)) -> m (Either a b)
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case Left a
e -> Either a b -> m (Either a b)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (a -> Either a b
forall a b. a -> Either a b
Left a
e); Right a
a -> a -> m (Either a b)
m2 a
a

propagate' :: (Propagate p, Monad m) => m (p a) -> m (p b) -> m (p b)
propagate' :: forall (p :: * -> *) (m :: * -> *) a b.
(Propagate p, Monad m) =>
m (p a) -> m (p b) -> m (p b)
propagate' m (p a)
m1 m (p b)
m2 = m (p a) -> (a -> m (p b)) -> m (p b)
forall (p :: * -> *) (m :: * -> *) a b.
(Propagate p, Monad m) =>
m (p a) -> (a -> m (p b)) -> m (p b)
forall (m :: * -> *) a b.
Monad m =>
m (p a) -> (a -> m (p b)) -> m (p b)
propagate m (p a)
m1 (m (p b) -> a -> m (p b)
forall a b. a -> b -> a
const m (p b)
m2)