{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE PartialTypeSignatures #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE ScopedTypeVariables #-}
{- |

Implementation of
"Directed Containers as Categories," D. Ahman and T. Uustalu, EPTCS 207, 2016, pp. 89-98

https://arxiv.org/abs/1604.01187

-}
module Control.Comonad.Travel
  (
    -- * Travel: Comonad from a Category
    Travel (..),
    TravelTag (..),

    truncateTravel,
  )
where

import Control.Category ( Category(..) )
import Control.Comonad
import Data.Kind (Type)
import Data.Singletons
import Data.Singletons.Sigma
import Prelude hiding (id, (.))

import Data.Functor.Polynomial
import Data.Functor.Polynomial.Class
import Data.Singletons.Decide
import Data.Type.Equality
import Data.GADT.Compare (GEq (..))
import Control.Comonad.Store (Store)
import Control.Comonad.Store.Class (ComonadStore(..))

-- | Make 'Polynomial' 'Comonad' from 'Category'.
--
-- @Travel cat@ is a polynomial functor for any type constructor
-- @cat@.
--
-- @
-- Travel cat r ~ ∑a. ((∑b. cat a b) -> r)
-- @
--
-- @Travel cat@ is a @Comonad@ if @cat@ is a @Category@.
type Travel :: (k -> k -> Type) -> Type -> Type
data Travel (cat :: k -> k -> Type) r where
  MkTravel ::
    forall k (cat :: k -> k -> Type) (r :: Type) (a :: k).
    Sing a ->
    (Sigma k (TyCon (cat a)) -> r) ->
    Travel cat r

instance Functor (Travel cat) where
  fmap :: forall a b. (a -> b) -> Travel cat a -> Travel cat b
fmap a -> b
f (MkTravel Sing a
sa Sigma k (TyCon (cat a)) -> a
k) = Sing a -> (Sigma k (TyCon (cat a)) -> b) -> Travel cat b
forall k (cat :: k -> k -> *) r (a :: k).
Sing a -> (Sigma k (TyCon (cat a)) -> r) -> Travel cat r
MkTravel Sing a
sa (a -> b
f (a -> b)
-> (Sigma k (TyCon (cat a)) -> a) -> Sigma k (TyCon (cat a)) -> b
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
. Sigma k (TyCon (cat a)) -> a
k)

instance (Category cat) => Comonad (Travel (cat :: k -> k -> Type)) where
  extract :: forall r. Travel cat r -> r
  extract :: forall a. Travel cat a -> a
extract (MkTravel Sing a
sa Sigma k (TyCon (cat a)) -> r
g) = Sigma k (TyCon (cat a)) -> r
g (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)

  duplicate :: forall r. Travel cat r -> Travel cat (Travel cat r)
  duplicate :: forall a. Travel cat a -> Travel cat (Travel cat a)
duplicate (MkTravel Sing a
sa Sigma k (TyCon (cat a)) -> r
k) = Sing a
-> (Sigma k (TyCon (cat a)) -> Travel cat r)
-> Travel cat (Travel cat r)
forall k (cat :: k -> k -> *) r (a :: k).
Sing a -> (Sigma k (TyCon (cat a)) -> r) -> Travel cat r
MkTravel Sing a
sa \(Sing fst
sb :&: TyCon (cat a) @@ fst
ab) -> Sing fst -> (Sigma k (TyCon (cat fst)) -> r) -> Travel cat r
forall k (cat :: k -> k -> *) r (a :: k).
Sing a -> (Sigma k (TyCon (cat a)) -> r) -> Travel cat r
MkTravel Sing fst
sb \(Sing fst
sc :&: TyCon (cat fst) @@ fst
bc) -> Sigma k (TyCon (cat a)) -> r
k (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 fst fst
TyCon (cat fst) @@ fst
bc cat fst fst -> cat a fst -> cat a fst
forall (b :: k) (c :: k) (a :: k). cat b c -> cat a b -> cat a c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. cat a fst
TyCon (cat a) @@ fst
ab))

type TravelTag :: (k -> k -> Type) -> Type -> Type
data TravelTag cat x where
  TravelFrom
    :: forall k (cat :: k -> k -> Type) (a :: k).
       Sing a -> TravelTag cat (Sigma k (TyCon (cat a)))

instance SDecide k => TestEquality (TravelTag (cat :: k -> k -> Type)) where
  testEquality :: forall a b. TravelTag cat a -> TravelTag cat b -> Maybe (a :~: b)
testEquality (TravelFrom Sing a
sa) (TravelFrom Sing a
sb) = case Sing a
sa Sing a -> Sing a -> Decision (a :~: a)
forall (a :: k) (b :: k). Sing a -> Sing b -> Decision (a :~: b)
forall k (a :: k) (b :: k).
SDecide k =>
Sing a -> Sing b -> Decision (a :~: b)
%~ Sing a
sb of
    Proved a :~: a
Refl -> (a :~: b) -> Maybe (a :~: b)
forall a. a -> Maybe a
Just a :~: a
a :~: b
forall {k} (a :: k). a :~: a
Refl
    Disproved Refuted (a :~: a)
_ -> Maybe (a :~: b)
forall a. Maybe a
Nothing

instance SDecide k => GEq (TravelTag (cat :: k -> k -> Type)) where
  geq :: forall a b. TravelTag cat a -> TravelTag cat b -> Maybe (a :~: b)
geq = TravelTag cat a -> TravelTag cat b -> Maybe (a :~: b)
forall a b. TravelTag cat a -> TravelTag cat b -> Maybe (a :~: b)
forall {k} (f :: k -> *) (a :: k) (b :: k).
TestEquality f =>
f a -> f b -> Maybe (a :~: b)
testEquality

instance Polynomial (Travel cat) where
  type Tag (Travel cat) = TravelTag cat

  toPoly :: forall x. Travel cat x -> Poly (Tag (Travel cat)) x
toPoly (MkTravel Sing a
sa Sigma k (TyCon (cat a)) -> x
f) = TravelTag cat (Sigma k (TyCon (cat a)))
-> (Sigma k (TyCon (cat a)) -> x) -> Poly (TravelTag cat) x
forall (tag :: * -> *) n x. tag n -> (n -> x) -> Poly tag x
P (Sing a -> TravelTag cat (Sigma k (TyCon (cat a)))
forall k (cat :: k -> k -> *) (a :: k).
Sing a -> TravelTag cat (Sigma k (TyCon (cat a)))
TravelFrom Sing a
sa) Sigma k (TyCon (cat a)) -> x
f
  fromPoly :: forall x. Poly (Tag (Travel cat)) x -> Travel cat x
fromPoly (P (TravelFrom Sing a
sa) n -> x
f) = Sing a -> (Sigma k (TyCon (cat a)) -> x) -> Travel cat x
forall k (cat :: k -> k -> *) r (a :: k).
Sing a -> (Sigma k (TyCon (cat a)) -> r) -> Travel cat r
MkTravel Sing a
sa n -> x
Sigma k (TyCon (cat a)) -> x
f

-- | Constructs @Travel cat@ from @Store (Demote k)@,
--   which can be thought of as a @Travel cat@ but its arrows are
--   forgotten.
truncateTravel :: forall k (cat :: k -> k -> Type) (r :: Type).
  (SingKind k) => Store (Demote k) r -> Travel cat r
truncateTravel :: forall k (cat :: k -> k -> *) r.
SingKind k =>
Store (Demote k) r -> Travel cat r
truncateTravel Store (Demote k) r
wr = case Demote k -> SomeSing k
forall k. SingKind k => Demote k -> SomeSing k
toSing (Store (Demote k) r -> Demote k
forall a. StoreT (Demote k) Identity a -> Demote k
forall s (w :: * -> *) a. ComonadStore s w => w a -> s
pos Store (Demote k) r
wr) of
  SomeSing Sing a
sa -> Sing a -> (Sigma k (TyCon (cat a)) -> r) -> Travel cat r
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)) -> r) -> Travel cat r)
-> (Sigma k (TyCon (cat a)) -> r) -> Travel cat r
forall a b. (a -> b) -> a -> b
$ \(Sing fst
sb :&: TyCon (cat a) @@ fst
_) -> Demote k -> Store (Demote k) r -> r
forall a. Demote k -> StoreT (Demote k) Identity a -> a
forall s (w :: * -> *) a. ComonadStore s w => s -> w a -> a
peek (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) Store (Demote k) r
wr