{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE ScopedTypeVariables #-}
module Control.Monad.Flow(
    -- * Flow: Monad from a Category
    Flow,
    mkFlow,
    truncateFlow,

    toCoTravel,
    fromCoTravel,

    -- * Transformer version
    FlowT (..),
    truncateFlowT,
) where

import Control.Category ( Category(..), (>>>) )
import Control.Comonad
import Control.Monad (ap)
import Control.Monad.Co
import Data.Kind (Type)
import Data.Singletons
import Data.Singletons.Sigma
import Prelude hiding (id, (.))
import Control.Monad.Trans.State (State, StateT (..))

import Data.Functor.Identity
import Data.Functor ((<&>))
import Control.Monad.Trans.Class (MonadTrans (..))

import Control.Comonad.Travel

-- | @Flow@ can be thought of as a generalization of the @'State'@ Monad.
--
--   A value of @State s x@ describes:
--
--   - For each state @s1 :: s@, a next state @s2 :: s@
--   - For each state @s1 :: s@, a value @x1 :: x@
--
--   A value of @Flow cat@ describes:
--
--   - For each object @s1 :: k@, an arrow @f :: cat s1 s2@ starting from @s1@
--     and going to a next object @s2@
--   - For each object @s1 :: k@, a value @x1 :: x@
type Flow cat = FlowT cat Identity

mkFlow :: forall k (cat :: k -> k -> Type) x.
  (forall (a :: k). Sing a -> (x, Sigma k (TyCon (cat a))))
  -> Flow cat x
mkFlow :: forall k (cat :: k -> k -> *) x.
(forall (a :: k). Sing a -> (x, Sigma k (TyCon (cat a))))
-> Flow cat x
mkFlow forall (a :: k). Sing a -> (x, Sigma k (TyCon (cat a)))
f = (forall (a :: k). Sing a -> Identity (x, Sigma k (TyCon (cat a))))
-> FlowT cat Identity x
forall k (cat :: k -> k -> *) (m :: * -> *) x.
(forall (a :: k). Sing a -> m (x, Sigma k (TyCon (cat a))))
-> FlowT cat m x
MkFlowT ((x, Sigma k (TyCon (cat a)))
-> Identity (x, Sigma k (TyCon (cat a)))
forall a. a -> Identity a
Identity ((x, Sigma k (TyCon (cat a)))
 -> Identity (x, Sigma k (TyCon (cat a))))
-> (Sing a -> (x, Sigma k (TyCon (cat a))))
-> Sing a
-> Identity (x, Sigma k (TyCon (cat a)))
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Sing a -> (x, Sigma k (TyCon (cat a)))
forall (a :: k). Sing a -> (x, Sigma k (TyCon (cat a)))
f)

runFlow :: forall k (cat :: k -> k -> Type) x (a :: k).
  Flow cat x -> Sing a -> (x, Sigma k (TyCon (cat a)))
runFlow :: forall k (cat :: k -> k -> *) x (a :: k).
Flow cat x -> Sing a -> (x, Sigma k (TyCon (cat a)))
runFlow Flow cat x
mx Sing a
sa = Identity (x, Sigma k (TyCon (cat a)))
-> (x, Sigma k (TyCon (cat a)))
forall a. Identity a -> a
runIdentity (Identity (x, Sigma k (TyCon (cat a)))
 -> (x, Sigma k (TyCon (cat a))))
-> Identity (x, Sigma k (TyCon (cat a)))
-> (x, Sigma k (TyCon (cat a)))
forall a b. (a -> b) -> a -> b
$ Flow cat x
-> forall (a :: k). Sing a -> Identity (x, Sigma k (TyCon (cat a)))
forall k (cat :: k -> k -> *) (m :: * -> *) x.
FlowT cat m x
-> forall (a :: k). Sing a -> m (x, Sigma k (TyCon (cat a)))
runFlowT Flow cat x
mx Sing a
sa

-- | Forgets arrows of the underlying category from @Flow cat@,
--   converting it to a simple @State@ monad on objects of @cat@.
truncateFlow ::
  forall k (cat :: k -> k -> Type) x.
  (SingKind k) =>
  Flow cat x ->
  State (Demote k) x
truncateFlow :: forall k (cat :: k -> k -> *) x.
SingKind k =>
Flow cat x -> State (Demote k) x
truncateFlow Flow cat x
mx = (Demote k -> Identity (x, Demote k))
-> StateT (Demote k) Identity x
forall s (m :: * -> *) a. (s -> m (a, s)) -> StateT s m a
StateT \Demote k
a -> case Demote k -> SomeSing k
forall k. SingKind k => Demote k -> SomeSing k
toSing Demote k
a of
  SomeSing Sing a
sa -> Flow cat x
-> forall (a :: k). Sing a -> Identity (x, Sigma k (TyCon (cat a)))
forall k (cat :: k -> k -> *) (m :: * -> *) x.
FlowT cat m x
-> forall (a :: k). Sing a -> m (x, Sigma k (TyCon (cat a)))
runFlowT Flow cat x
mx Sing a
sa Identity (x, Sigma k (TyCon (cat a)))
-> ((x, Sigma k (TyCon (cat a))) -> (x, Demote k))
-> Identity (x, Demote k)
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \(x
x, Sing fst
sb :&: TyCon (cat a) @@ fst
_) -> (x
x, Sing fst -> Demote k
forall (a :: k). Sing a -> Demote k
forall k (a :: k). SingKind k => Sing a -> Demote k
fromSing Sing fst
sb)

-- | @'Co' ('Travel' cat)@ is isomorphic to @'Flow' cat@
fromCoTravel ::
  forall k (cat :: k -> k -> Type) x.
  (Category cat) =>
  Co (Travel cat) x ->
  Flow cat x
fromCoTravel :: forall k (cat :: k -> k -> *) x.
Category cat =>
Co (Travel cat) x -> Flow cat x
fromCoTravel Co (Travel cat) x
coTravel = (forall (a :: k). Sing a -> Identity (x, Sigma k (TyCon (cat a))))
-> FlowT cat Identity x
forall k (cat :: k -> k -> *) (m :: * -> *) x.
(forall (a :: k). Sing a -> m (x, Sigma k (TyCon (cat a))))
-> FlowT cat m x
MkFlowT ((forall (a :: k). Sing a -> Identity (x, Sigma k (TyCon (cat a))))
 -> FlowT cat Identity x)
-> (forall (a :: k).
    Sing a -> Identity (x, Sigma k (TyCon (cat a))))
-> FlowT cat Identity x
forall a b. (a -> b) -> a -> b
$ \Sing a
sa ->
    Co (Travel cat) x
-> Travel cat (x -> Identity (x, Sigma k (TyCon (cat a))))
-> Identity (x, Sigma k (TyCon (cat a)))
forall (w :: * -> *) a r. Functor w => Co w a -> w (a -> r) -> r
runCo Co (Travel cat) x
coTravel (Travel cat (x -> Identity (x, Sigma k (TyCon (cat a))))
 -> Identity (x, Sigma k (TyCon (cat a))))
-> Travel cat (x -> Identity (x, Sigma k (TyCon (cat a))))
-> Identity (x, Sigma k (TyCon (cat a)))
forall a b. (a -> b) -> a -> b
$ Sing a
-> (Sigma k (TyCon (cat a))
    -> x -> Identity (x, Sigma k (TyCon (cat a))))
-> Travel cat (x -> Identity (x, Sigma k (TyCon (cat a))))
forall k (cat :: k -> k -> *) r (a :: k).
Sing a -> (Sigma k (TyCon (cat a)) -> r) -> Travel cat r
MkTravel Sing a
sa (\Sigma k (TyCon (cat a))
key x
x -> (x, Sigma k (TyCon (cat a)))
-> Identity (x, Sigma k (TyCon (cat a)))
forall a. a -> Identity a
Identity (x
x, Sigma k (TyCon (cat a))
key))

-- | @'Co' ('Travel' cat)@ is isomorphic to @'Flow' cat@
toCoTravel ::
  forall k (cat :: k -> k -> Type) x.
  (Category cat) =>
  Flow cat x ->
  Co (Travel cat) x
toCoTravel :: forall k (cat :: k -> k -> *) x.
Category cat =>
Flow cat x -> Co (Travel cat) x
toCoTravel Flow cat x
mx = (forall r. Travel cat (x -> r) -> r) -> Co (Travel cat) x
forall (w :: * -> *) a.
Functor w =>
(forall r. w (a -> r) -> r) -> Co w a
co \(MkTravel Sing a
sa Sigma k (TyCon (cat a)) -> x -> r
body) ->
  case Flow cat x -> Sing a -> (x, Sigma k (TyCon (cat a)))
forall k (cat :: k -> k -> *) x (a :: k).
Flow cat x -> Sing a -> (x, Sigma k (TyCon (cat a)))
runFlow Flow cat x
mx Sing a
sa of
    (x
x, Sigma k (TyCon (cat a))
key) -> Sigma k (TyCon (cat a)) -> x -> r
body Sigma k (TyCon (cat a))
key x
x

---- FlowT ----

-- | Transformer version of @Flow@.
newtype FlowT (cat :: k -> k -> Type) m x = MkFlowT {
    forall k (cat :: k -> k -> *) (m :: * -> *) x.
FlowT cat m x
-> forall (a :: k). Sing a -> m (x, Sigma k (TyCon (cat a)))
runFlowT :: forall (a :: k). Sing a -> m (x, Sigma k (TyCon (cat a)))
  }
  deriving ((forall a b. (a -> b) -> FlowT cat m a -> FlowT cat m b)
-> (forall a b. a -> FlowT cat m b -> FlowT cat m a)
-> Functor (FlowT cat m)
forall a b. a -> FlowT cat m b -> FlowT cat m a
forall a b. (a -> b) -> FlowT cat m a -> FlowT cat m b
forall k (cat :: k -> k -> *) (m :: * -> *) a b.
Functor m =>
a -> FlowT cat m b -> FlowT cat m a
forall k (cat :: k -> k -> *) (m :: * -> *) a b.
Functor m =>
(a -> b) -> FlowT cat m a -> FlowT cat m b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
$cfmap :: forall k (cat :: k -> k -> *) (m :: * -> *) a b.
Functor m =>
(a -> b) -> FlowT cat m a -> FlowT cat m b
fmap :: forall a b. (a -> b) -> FlowT cat m a -> FlowT cat m b
$c<$ :: forall k (cat :: k -> k -> *) (m :: * -> *) a b.
Functor m =>
a -> FlowT cat m b -> FlowT cat m a
<$ :: forall a b. a -> FlowT cat m b -> FlowT cat m a
Functor)

instance (Category cat, Monad m) => Applicative (FlowT cat m) where
  pure :: forall a. a -> FlowT cat m a
pure a
x = (forall (a :: k). Sing a -> m (a, Sigma k (TyCon (cat a))))
-> FlowT cat m a
forall k (cat :: k -> k -> *) (m :: * -> *) x.
(forall (a :: k). Sing a -> m (x, Sigma k (TyCon (cat a))))
-> FlowT cat m x
MkFlowT ((forall (a :: k). Sing a -> m (a, Sigma k (TyCon (cat a))))
 -> FlowT cat m a)
-> (forall (a :: k). Sing a -> m (a, Sigma k (TyCon (cat a))))
-> FlowT cat m a
forall a b. (a -> b) -> a -> b
$ \Sing a
sa -> (a, Sigma k (TyCon (cat a))) -> m (a, Sigma k (TyCon (cat a)))
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a
x, Sing a
sa Sing a -> (TyCon (cat a) @@ a) -> Sigma k (TyCon (cat a))
forall s (a :: s ~> *) (fst :: s).
Sing fst -> (a @@ fst) -> Sigma s a
:&: cat a a
TyCon (cat a) @@ a
forall (a :: k). cat a a
forall {k} (cat :: k -> k -> *) (a :: k). Category cat => cat a a
id)
  <*> :: forall a b. FlowT cat m (a -> b) -> FlowT cat m a -> FlowT cat m b
(<*>) = FlowT cat m (a -> b) -> FlowT cat m a -> FlowT cat m b
forall (m :: * -> *) a b. Monad m => m (a -> b) -> m a -> m b
ap

instance (Category cat, Monad m) => Monad (FlowT cat m) where
  FlowT cat m a
mx >>= :: forall a b. FlowT cat m a -> (a -> FlowT cat m b) -> FlowT cat m b
>>= a -> FlowT cat m b
k = (forall (a :: k). Sing a -> m (b, Sigma k (TyCon (cat a))))
-> FlowT cat m b
forall k (cat :: k -> k -> *) (m :: * -> *) x.
(forall (a :: k). Sing a -> m (x, Sigma k (TyCon (cat a))))
-> FlowT cat m x
MkFlowT ((forall (a :: k). Sing a -> m (b, Sigma k (TyCon (cat a))))
 -> FlowT cat m b)
-> (forall (a :: k). Sing a -> m (b, Sigma k (TyCon (cat a))))
-> FlowT cat m b
forall a b. (a -> b) -> a -> b
$ \Sing a
sa ->
    do (a
y, Sing fst
sb :&: TyCon (cat a) @@ fst
f) <- FlowT cat m a
-> forall (a :: k). Sing a -> m (a, Sigma k (TyCon (cat a)))
forall k (cat :: k -> k -> *) (m :: * -> *) x.
FlowT cat m x
-> forall (a :: k). Sing a -> m (x, Sigma k (TyCon (cat a)))
runFlowT FlowT cat m a
mx Sing a
sa
       (b
z, Sing fst
sc :&: TyCon (cat fst) @@ fst
g) <- FlowT cat m b
-> forall (a :: k). Sing a -> m (b, Sigma k (TyCon (cat a)))
forall k (cat :: k -> k -> *) (m :: * -> *) x.
FlowT cat m x
-> forall (a :: k). Sing a -> m (x, Sigma k (TyCon (cat a)))
runFlowT (a -> FlowT cat m b
k a
y) Sing fst
sb
       (b, Sigma k (TyCon (cat a))) -> m (b, Sigma k (TyCon (cat a)))
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (b
z, Sing fst
sc Sing fst -> (TyCon (cat a) @@ fst) -> Sigma k (TyCon (cat a))
forall s (a :: s ~> *) (fst :: s).
Sing fst -> (a @@ fst) -> Sigma s a
:&: (cat a fst
TyCon (cat a) @@ fst
f cat a fst -> cat fst fst -> cat a fst
forall {k} (cat :: k -> k -> *) (a :: k) (b :: k) (c :: k).
Category cat =>
cat a b -> cat b c -> cat a c
>>> cat fst fst
TyCon (cat fst) @@ fst
g))

instance (Category cat) => MonadTrans (FlowT cat) where
  lift :: forall (m :: * -> *) a. Monad m => m a -> FlowT cat m a
lift m a
mx = (forall (a :: k). Sing a -> m (a, Sigma k (TyCon (cat a))))
-> FlowT cat m a
forall k (cat :: k -> k -> *) (m :: * -> *) x.
(forall (a :: k). Sing a -> m (x, Sigma k (TyCon (cat a))))
-> FlowT cat m x
MkFlowT ((forall (a :: k). Sing a -> m (a, Sigma k (TyCon (cat a))))
 -> FlowT cat m a)
-> (forall (a :: k). Sing a -> m (a, Sigma k (TyCon (cat a))))
-> FlowT cat m a
forall a b. (a -> b) -> a -> b
$ \Sing a
sa -> m a
mx m a
-> (a -> (a, Sigma k (TyCon (cat a))))
-> m (a, Sigma k (TyCon (cat a)))
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \a
x -> (a
x, Sing a
sa Sing a -> (TyCon (cat a) @@ a) -> Sigma k (TyCon (cat a))
forall s (a :: s ~> *) (fst :: s).
Sing fst -> (a @@ fst) -> Sigma s a
:&: cat a a
TyCon (cat a) @@ a
forall (a :: k). cat a a
forall {k} (cat :: k -> k -> *) (a :: k). Category cat => cat a a
id)

-- | Forgets arrows of the underlying category from @FlowT cat m@,
--   converting it to a simple @StateT@ on objects of @cat@.
truncateFlowT ::
  forall k (cat :: k -> k -> Type) m x.
  (SingKind k, Monad m) =>
  FlowT cat m x ->
  StateT (Demote k) m x
truncateFlowT :: forall k (cat :: k -> k -> *) (m :: * -> *) x.
(SingKind k, Monad m) =>
FlowT cat m x -> StateT (Demote k) m x
truncateFlowT FlowT cat m x
mx = (Demote k -> m (x, Demote k)) -> StateT (Demote k) m x
forall s (m :: * -> *) a. (s -> m (a, s)) -> StateT s m a
StateT \Demote k
a -> case Demote k -> SomeSing k
forall k. SingKind k => Demote k -> SomeSing k
toSing Demote k
a of
  SomeSing Sing a
sa -> FlowT cat m x
-> forall (a :: k). Sing a -> m (x, Sigma k (TyCon (cat a)))
forall k (cat :: k -> k -> *) (m :: * -> *) x.
FlowT cat m x
-> forall (a :: k). Sing a -> m (x, Sigma k (TyCon (cat a)))
runFlowT FlowT cat m x
mx Sing a
sa m (x, Sigma k (TyCon (cat a)))
-> ((x, Sigma k (TyCon (cat a))) -> (x, Demote k))
-> m (x, Demote k)
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \(x
x, Sing fst
sb :&: TyCon (cat a) @@ fst
_) -> (x
x, Sing fst -> Demote k
forall (a :: k). Sing a -> Demote k
forall k (a :: k). SingKind k => Sing a -> Demote k
fromSing Sing fst
sb)