{-# LANGUAGE RankNTypes #-}

module Data.Stream.Recursive where

-- base
import Control.Applicative (Alternative (..))
import Data.Function ((&))
import Data.Functor ((<&>))

-- mmorph
import Control.Monad.Morph (MFunctor (..))

-- automaton
import Data.Stream.Result

{- | A stream transformer in recursive encoding.

One step of the stream transformer performs a monadic action and results in an output and a new stream.
-}
newtype Recursive m a = Recursive {forall (m :: Type -> Type) a.
Recursive m a -> m (Result (Recursive m a) a)
getRecursive :: m (Result (Recursive m a) a)}

instance MFunctor Recursive where
  hoist :: forall (m :: Type -> Type) (n :: Type -> Type) b.
Monad m =>
(forall a. m a -> n a) -> Recursive m b -> Recursive n b
hoist = (forall x. m x -> n x) -> Recursive m b -> Recursive n b
forall (f :: Type -> Type) (g :: Type -> Type) a.
Functor f =>
(forall x. f x -> g x) -> Recursive f a -> Recursive g a
hoist'

{- | Hoist a stream along a monad morphism, by applying said morphism to the step function.

This is like @mmorph@'s 'hoist', but it doesn't require a 'Monad' constraint on @m2@.
-}
hoist' :: (Functor f) => (forall x. f x -> g x) -> Recursive f a -> Recursive g a
hoist' :: forall (f :: Type -> Type) (g :: Type -> Type) a.
Functor f =>
(forall x. f x -> g x) -> Recursive f a -> Recursive g a
hoist' forall x. f x -> g x
morph = Recursive f a -> Recursive g a
go
  where
    go :: Recursive f a -> Recursive g a
go Recursive {f (Result (Recursive f a) a)
getRecursive :: forall (m :: Type -> Type) a.
Recursive m a -> m (Result (Recursive m a) a)
getRecursive :: f (Result (Recursive f a) a)
getRecursive} = g (Result (Recursive g a) a) -> Recursive g a
forall (m :: Type -> Type) a.
m (Result (Recursive m a) a) -> Recursive m a
Recursive (g (Result (Recursive g a) a) -> Recursive g a)
-> g (Result (Recursive g a) a) -> Recursive g a
forall a b. (a -> b) -> a -> b
$ f (Result (Recursive g a) a) -> g (Result (Recursive g a) a)
forall x. f x -> g x
morph (f (Result (Recursive g a) a) -> g (Result (Recursive g a) a))
-> f (Result (Recursive g a) a) -> g (Result (Recursive g a) a)
forall a b. (a -> b) -> a -> b
$ (Recursive f a -> Recursive g a)
-> Result (Recursive f a) a -> Result (Recursive g a) a
forall s1 s2 a. (s1 -> s2) -> Result s1 a -> Result s2 a
mapResultState Recursive f a -> Recursive g a
go (Result (Recursive f a) a -> Result (Recursive g a) a)
-> f (Result (Recursive f a) a) -> f (Result (Recursive g a) a)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> f (Result (Recursive f a) a)
getRecursive

instance (Functor m) => Functor (Recursive m) where
  fmap :: forall a b. (a -> b) -> Recursive m a -> Recursive m b
fmap a -> b
f Recursive {m (Result (Recursive m a) a)
getRecursive :: forall (m :: Type -> Type) a.
Recursive m a -> m (Result (Recursive m a) a)
getRecursive :: m (Result (Recursive m a) a)
getRecursive} = m (Result (Recursive m b) b) -> Recursive m b
forall (m :: Type -> Type) a.
m (Result (Recursive m a) a) -> Recursive m a
Recursive (m (Result (Recursive m b) b) -> Recursive m b)
-> m (Result (Recursive m b) b) -> Recursive m b
forall a b. (a -> b) -> a -> b
$ (a -> b) -> Result (Recursive m b) a -> Result (Recursive m b) b
forall a b.
(a -> b) -> Result (Recursive m b) a -> Result (Recursive m b) b
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f (Result (Recursive m b) a -> Result (Recursive m b) b)
-> (Result (Recursive m a) a -> Result (Recursive m b) a)
-> Result (Recursive m a) a
-> Result (Recursive m b) b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Recursive m a -> Recursive m b)
-> Result (Recursive m a) a -> Result (Recursive m b) a
forall s1 s2 a. (s1 -> s2) -> Result s1 a -> Result s2 a
mapResultState ((a -> b) -> Recursive m a -> Recursive m b
forall a b. (a -> b) -> Recursive m a -> Recursive m b
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f) (Result (Recursive m a) a -> Result (Recursive m b) b)
-> m (Result (Recursive m a) a) -> m (Result (Recursive m b) b)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> m (Result (Recursive m a) a)
getRecursive

instance (Applicative m) => Applicative (Recursive m) where
  pure :: forall a. a -> Recursive m a
pure a
a = Recursive m a
go
    where
      go :: Recursive m a
go = m (Result (Recursive m a) a) -> Recursive m a
forall (m :: Type -> Type) a.
m (Result (Recursive m a) a) -> Recursive m a
Recursive (m (Result (Recursive m a) a) -> Recursive m a)
-> m (Result (Recursive m a) a) -> Recursive m a
forall a b. (a -> b) -> a -> b
$! Result (Recursive m a) a -> m (Result (Recursive m a) a)
forall a. a -> m a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (Result (Recursive m a) a -> m (Result (Recursive m a) a))
-> Result (Recursive m a) a -> m (Result (Recursive m a) a)
forall a b. (a -> b) -> a -> b
$! Recursive m a -> a -> Result (Recursive m a) a
forall s a. s -> a -> Result s a
Result Recursive m a
go a
a

  Recursive m (Result (Recursive m (a -> b)) (a -> b))
mf <*> :: forall a b. Recursive m (a -> b) -> Recursive m a -> Recursive m b
<*> Recursive m (Result (Recursive m a) a)
ma = m (Result (Recursive m b) b) -> Recursive m b
forall (m :: Type -> Type) a.
m (Result (Recursive m a) a) -> Recursive m a
Recursive (m (Result (Recursive m b) b) -> Recursive m b)
-> m (Result (Recursive m b) b) -> Recursive m b
forall a b. (a -> b) -> a -> b
$! (\(Result Recursive m (a -> b)
cf a -> b
f) (Result Recursive m a
ca a
a) -> Recursive m b -> b -> Result (Recursive m b) b
forall s a. s -> a -> Result s a
Result (Recursive m (a -> b)
cf Recursive m (a -> b) -> Recursive m a -> Recursive m b
forall a b. Recursive m (a -> b) -> Recursive m a -> Recursive m b
forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> Recursive m a
ca) (b -> Result (Recursive m b) b) -> b -> Result (Recursive m b) b
forall a b. (a -> b) -> a -> b
$! a -> b
f a
a) (Result (Recursive m (a -> b)) (a -> b)
 -> Result (Recursive m a) a -> Result (Recursive m b) b)
-> m (Result (Recursive m (a -> b)) (a -> b))
-> m (Result (Recursive m a) a -> Result (Recursive m b) b)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> m (Result (Recursive m (a -> b)) (a -> b))
mf m (Result (Recursive m a) a -> Result (Recursive m b) b)
-> m (Result (Recursive m a) a) -> m (Result (Recursive m b) b)
forall a b. m (a -> b) -> m a -> m b
forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> m (Result (Recursive m a) a)
ma

-- | Constantly perform the same effect, without remembering a state.
constM :: (Functor m) => m a -> Recursive m a
constM :: forall (m :: Type -> Type) a. Functor m => m a -> Recursive m a
constM m a
ma = Recursive m a
go
  where
    go :: Recursive m a
go = m (Result (Recursive m a) a) -> Recursive m a
forall (m :: Type -> Type) a.
m (Result (Recursive m a) a) -> Recursive m a
Recursive (m (Result (Recursive m a) a) -> Recursive m a)
-> m (Result (Recursive m a) a) -> Recursive m a
forall a b. (a -> b) -> a -> b
$ Recursive m a -> a -> Result (Recursive m a) a
forall s a. s -> a -> Result s a
Result Recursive m a
go (a -> Result (Recursive m a) a)
-> m a -> m (Result (Recursive m a) a)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> m a
ma

instance (Alternative m) => Alternative (Recursive m) where
  empty :: forall a. Recursive m a
empty = m a -> Recursive m a
forall (m :: Type -> Type) a. Functor m => m a -> Recursive m a
constM m a
forall a. m a
forall (f :: Type -> Type) a. Alternative f => f a
empty

  Recursive m (Result (Recursive m a) a)
ma1 <|> :: forall a. Recursive m a -> Recursive m a -> Recursive m a
<|> Recursive m (Result (Recursive m a) a)
ma2 = m (Result (Recursive m a) a) -> Recursive m a
forall (m :: Type -> Type) a.
m (Result (Recursive m a) a) -> Recursive m a
Recursive (m (Result (Recursive m a) a) -> Recursive m a)
-> m (Result (Recursive m a) a) -> Recursive m a
forall a b. (a -> b) -> a -> b
$ m (Result (Recursive m a) a)
ma1 m (Result (Recursive m a) a)
-> m (Result (Recursive m a) a) -> m (Result (Recursive m a) a)
forall a. m a -> m a -> m a
forall (f :: Type -> Type) a. Alternative f => f a -> f a -> f a
<|> m (Result (Recursive m a) a)
ma2

instance (Foldable m) => Foldable (Recursive m) where
  foldMap :: forall m a. Monoid m => (a -> m) -> Recursive m a -> m
foldMap a -> m
f Recursive {m (Result (Recursive m a) a)
getRecursive :: forall (m :: Type -> Type) a.
Recursive m a -> m (Result (Recursive m a) a)
getRecursive :: m (Result (Recursive m a) a)
getRecursive} = (Result (Recursive m a) a -> m)
-> m (Result (Recursive m a) a) -> m
forall m a. Monoid m => (a -> m) -> m a -> m
forall (t :: Type -> Type) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (\(Result Recursive m a
recursive a
a) -> a -> m
f a
a m -> m -> m
forall a. Semigroup a => a -> a -> a
<> (a -> m) -> Recursive m a -> m
forall m a. Monoid m => (a -> m) -> Recursive m a -> m
forall (t :: Type -> Type) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap a -> m
f Recursive m a
recursive) m (Result (Recursive m a) a)
getRecursive

instance (Traversable m) => Traversable (Recursive m) where
  traverse :: forall (f :: Type -> Type) a b.
Applicative f =>
(a -> f b) -> Recursive m a -> f (Recursive m b)
traverse a -> f b
f = Recursive m a -> f (Recursive m b)
go
    where
      go :: Recursive m a -> f (Recursive m b)
go Recursive {m (Result (Recursive m a) a)
getRecursive :: forall (m :: Type -> Type) a.
Recursive m a -> m (Result (Recursive m a) a)
getRecursive :: m (Result (Recursive m a) a)
getRecursive} = (m (Result (Recursive m a) a)
getRecursive m (Result (Recursive m a) a)
-> (m (Result (Recursive m a) a)
    -> f (m (Result (Recursive m b) b)))
-> f (m (Result (Recursive m b) b))
forall a b. a -> (a -> b) -> b
& (Result (Recursive m a) a -> f (Result (Recursive m b) b))
-> m (Result (Recursive m a) a) -> f (m (Result (Recursive m b) b))
forall (t :: Type -> Type) (f :: Type -> Type) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: Type -> Type) a b.
Applicative f =>
(a -> f b) -> m a -> f (m b)
traverse (\(Result Recursive m a
cont a
a) -> (Recursive m b -> b -> Result (Recursive m b) b)
-> b -> Recursive m b -> Result (Recursive m b) b
forall a b c. (a -> b -> c) -> b -> a -> c
flip Recursive m b -> b -> Result (Recursive m b) b
forall s a. s -> a -> Result s a
Result (b -> Recursive m b -> Result (Recursive m b) b)
-> f b -> f (Recursive m b -> Result (Recursive m b) b)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> f b
f a
a f (Recursive m b -> Result (Recursive m b) b)
-> f (Recursive m b) -> f (Result (Recursive m b) b)
forall a b. f (a -> b) -> f a -> f b
forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> Recursive m a -> f (Recursive m b)
go Recursive m a
cont)) f (m (Result (Recursive m b) b))
-> (m (Result (Recursive m b) b) -> Recursive m b)
-> f (Recursive m b)
forall (f :: Type -> Type) a b. Functor f => f a -> (a -> b) -> f b
<&> m (Result (Recursive m b) b) -> Recursive m b
forall (m :: Type -> Type) a.
m (Result (Recursive m a) a) -> Recursive m a
Recursive

-- | Like 'fmap' or 'rmap', but the postcomposed function may have an effect in @m@.
mmap :: (Monad m) => (a -> m b) -> Recursive m a -> Recursive m b
mmap :: forall (m :: Type -> Type) a b.
Monad m =>
(a -> m b) -> Recursive m a -> Recursive m b
mmap a -> m b
f Recursive {m (Result (Recursive m a) a)
getRecursive :: forall (m :: Type -> Type) a.
Recursive m a -> m (Result (Recursive m a) a)
getRecursive :: m (Result (Recursive m a) a)
getRecursive} = m (Result (Recursive m b) b) -> Recursive m b
forall (m :: Type -> Type) a.
m (Result (Recursive m a) a) -> Recursive m a
Recursive (m (Result (Recursive m b) b) -> Recursive m b)
-> m (Result (Recursive m b) b) -> Recursive m b
forall a b. (a -> b) -> a -> b
$ do
  Result Recursive m a
recursive a
a <- m (Result (Recursive m a) a)
getRecursive
  b
b <- a -> m b
f a
a
  Result (Recursive m b) b -> m (Result (Recursive m b) b)
forall a. a -> m a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (Result (Recursive m b) b -> m (Result (Recursive m b) b))
-> Result (Recursive m b) b -> m (Result (Recursive m b) b)
forall a b. (a -> b) -> a -> b
$ Recursive m b -> b -> Result (Recursive m b) b
forall s a. s -> a -> Result s a
Result ((a -> m b) -> Recursive m a -> Recursive m b
forall (m :: Type -> Type) a b.
Monad m =>
(a -> m b) -> Recursive m a -> Recursive m b
mmap a -> m b
f Recursive m a
recursive) b
b