{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE PartialTypeSignatures #-}
module Data.InternalCategory.PolyComonad(
  Pos(..), Dir(..), identityDir, composeDir,
) where

import Data.Maybe (isJust)

import Control.Comonad

import Data.Type.Equality ((:~:)(..))
import Data.GADT.Compare (GEq(..))
import Data.Functor.Polynomial
import Data.Functor.Polynomial.Class

import Data.InternalCategory
import qualified Data.List

data Pos f where
  Position :: Tag f a -> Pos f

instance GEq (Tag f) => Eq (Pos f) where
  Position Tag f a
tag == :: Pos f -> Pos f -> Bool
== Position Tag f a
tag' = Maybe (a :~: a) -> Bool
forall a. Maybe a -> Bool
isJust (Maybe (a :~: a) -> Bool) -> Maybe (a :~: a) -> Bool
forall a b. (a -> b) -> a -> b
$ Tag f a -> Tag f a -> Maybe (a :~: a)
forall a b. Tag f a -> Tag f b -> Maybe (a :~: b)
forall k (f :: k -> *) (a :: k) (b :: k).
GEq f =>
f a -> f b -> Maybe (a :~: b)
geq Tag f a
tag Tag f a
tag'

data Dir f where
  Direction :: Tag f a -> a -> Dir f

domTag :: Dir f -> Pos f
domTag :: forall (f :: * -> *). Dir f -> Pos f
domTag (Direction Tag f a
tag a
_) = Tag f a -> Pos f
forall (f :: * -> *) a. Tag f a -> Pos f
Position Tag f a
tag

codTag :: (Comonad f, Polynomial f, GEq (Tag f)) => Dir f -> Pos f
codTag :: forall (f :: * -> *).
(Comonad f, Polynomial f, GEq (Tag f)) =>
Dir f -> Pos f
codTag Dir f
dir = case f (Dir f) -> Poly (Tag f) (Dir f)
forall x. f x -> Poly (Tag f) x
forall (f :: * -> *) x. Polynomial f => f x -> Poly (Tag f) x
toPoly (Dir f -> f (Dir f)
forall (f :: * -> *).
(Comonad f, Polynomial f, GEq (Tag f)) =>
Dir f -> f (Dir f)
extendDirection Dir f
dir) of
  P Tag f n
tag'' n -> Dir f
_ -> Tag f n -> Pos f
forall (f :: * -> *) a. Tag f a -> Pos f
Position Tag f n
tag''

identityDir :: forall f. (Comonad f, Polynomial f) => Pos f -> Dir f
identityDir :: forall (f :: * -> *). (Comonad f, Polynomial f) => Pos f -> Dir f
identityDir Pos f
pos = f (Dir f) -> Dir f
forall a. f a -> a
forall (w :: * -> *) a. Comonad w => w a -> a
extract (Pos f -> f (Dir f)
forall (f :: * -> *). Polynomial f => Pos f -> f (Dir f)
unfoldPosition Pos f
pos)

composeDir :: (Comonad f, Polynomial f, GEq (Tag f)) => Dir f -> Dir f -> Maybe (Dir f)
composeDir :: forall (f :: * -> *).
(Comonad f, Polynomial f, GEq (Tag f)) =>
Dir f -> Dir f -> Maybe (Dir f)
composeDir Dir f
dir1 (Direction Tag f a
tag2 a
a2) =
  case f (Dir f) -> Poly (Tag f) (Dir f)
forall x. f x -> Poly (Tag f) x
forall (f :: * -> *) x. Polynomial f => f x -> Poly (Tag f) x
toPoly (Dir f -> f (Dir f)
forall (f :: * -> *).
(Comonad f, Polynomial f, GEq (Tag f)) =>
Dir f -> f (Dir f)
extendDirection Dir f
dir1) of
    P Tag f n
tag2' n -> Dir f
rep -> case Tag f a -> Tag f n -> Maybe (a :~: n)
forall a b. Tag f a -> Tag f b -> Maybe (a :~: b)
forall k (f :: k -> *) (a :: k) (b :: k).
GEq f =>
f a -> f b -> Maybe (a :~: b)
geq Tag f a
tag2 Tag f n
tag2' of
      -- Case: cod dir1 == dom dir2
      Just a :~: n
Refl -> Dir f -> Maybe (Dir f)
forall a. a -> Maybe a
Just (n -> Dir f
rep a
n
a2)
      -- Case: cod dir1 /= dom dir2
      Maybe (a :~: n)
Nothing -> Maybe (Dir f)
forall a. Maybe a
Nothing

unfoldPosition :: (Polynomial f) => Pos f -> f (Dir f)
unfoldPosition :: forall (f :: * -> *). Polynomial f => Pos f -> f (Dir f)
unfoldPosition (Position Tag f a
tag) = Poly (Tag f) (Dir f) -> f (Dir f)
forall x. Poly (Tag f) x -> f x
forall (f :: * -> *) x. Polynomial f => Poly (Tag f) x -> f x
fromPoly (Poly (Tag f) (Dir f) -> f (Dir f))
-> Poly (Tag f) (Dir f) -> f (Dir f)
forall a b. (a -> b) -> a -> b
$ Tag f a -> (a -> Dir f) -> Poly (Tag f) (Dir f)
forall (tag :: * -> *) n x. tag n -> (n -> x) -> Poly tag x
P Tag f a
tag (Tag f a -> a -> Dir f
forall (f :: * -> *) a. Tag f a -> a -> Dir f
Direction Tag f a
tag)

extendDirection :: forall f. (Comonad f, Polynomial f, GEq (Tag f)) => Dir f -> f (Dir f)
extendDirection :: forall (f :: * -> *).
(Comonad f, Polynomial f, GEq (Tag f)) =>
Dir f -> f (Dir f)
extendDirection (Direction Tag f a
tag a
a) = case f (f (Dir f)) -> Poly (Tag f) (f (Dir f))
forall x. f x -> Poly (Tag f) x
forall (f :: * -> *) x. Polynomial f => f x -> Poly (Tag f) x
toPoly (f (Dir f) -> f (f (Dir f))
forall a. f a -> f (f a)
forall (w :: * -> *) a. Comonad w => w a -> w (w a)
duplicate (forall (f :: * -> *). Polynomial f => Pos f -> f (Dir f)
unfoldPosition @f (Tag f a -> Pos f
forall (f :: * -> *) a. Tag f a -> Pos f
Position Tag f a
tag))) of
  P Tag f n
tag' n -> f (Dir f)
rep -> case Tag f a -> Tag f n -> Maybe (a :~: n)
forall a b. Tag f a -> Tag f b -> Maybe (a :~: b)
forall k (f :: k -> *) (a :: k) (b :: k).
GEq f =>
f a -> f b -> Maybe (a :~: b)
geq Tag f a
tag Tag f n
tag' of
    Maybe (a :~: n)
Nothing -> [Char] -> f (Dir f)
forall a. HasCallStack => [Char] -> a
error [Char]
"unreachable; duplicate shouldn't change the outer shape"
    Just a :~: n
Refl -> n -> f (Dir f)
rep a
n
a

instance (Comonad f, Polynomial f, GEq (Tag f)) => IQuiver (Pos f) (Dir f) where
  src :: Dir f -> Pos f
src = Dir f -> Pos f
forall (f :: * -> *). Dir f -> Pos f
domTag
  tgt :: Dir f -> Pos f
tgt = Dir f -> Pos f
forall (f :: * -> *).
(Comonad f, Polynomial f, GEq (Tag f)) =>
Dir f -> Pos f
codTag

instance (Comonad f, Polynomial f, GEq (Tag f)) => ICategory (Pos f) (Dir f) where
  foldPath :: Path (Pos f) (Dir f) -> Dir f
foldPath (Path Pos f
fstTag [Dir f]
directions Pos f
_) = (Dir f -> Dir f -> Dir f) -> Dir f -> [Dir f] -> Dir f
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
Data.List.foldl' Dir f -> Dir f -> Dir f
forall {f :: * -> *}.
(Comonad f, Polynomial f, GEq (Tag f)) =>
Dir f -> Dir f -> Dir f
step (Pos f -> Dir f
forall (f :: * -> *). (Comonad f, Polynomial f) => Pos f -> Dir f
identityDir Pos f
fstTag) [Dir f]
directions
    where
      step :: Dir f -> Dir f -> Dir f
step Dir f
dir1 Dir f
dir2 = case Dir f -> Dir f -> Maybe (Dir f)
forall (f :: * -> *).
(Comonad f, Polynomial f, GEq (Tag f)) =>
Dir f -> Dir f -> Maybe (Dir f)
composeDir Dir f
dir1 Dir f
dir2 of
        Maybe (Dir f)
Nothing -> [Char] -> Dir f
forall a. HasCallStack => [Char] -> a
error [Char]
"Invalid path"
        Just Dir f
dir -> Dir f
dir