Simplex.hs

{-# LANGUAGE AllowAmbiguousTypes #-}
{-# OPTIONS_GHC -Wno-orphans #-}

-- | The __augmented simplex category__: objects are the finite ordinals (as type-level 'Nat's,
-- including the empty ordinal 'Z') and morphisms are order-preserving maps, built from the
-- constructors 'ZZ', 'Y' (skip a target) and 'X' (repeat a source). Ordinal sum makes it monoidal,
-- and it is the walking monoid: monoids in a monoidal category correspond to monoidal functors out
-- of it.
module Proarrow.Category.Instance.Simplex (module Proarrow.Category.Instance.Simplex, Nat (..)) where

import Data.Fin (Fin (..))
import Data.Kind (Type)
import Data.Type.Nat (Nat (..), SNatI, type Plus)
import Data.Vec.Lazy (Vec (..))
import Prelude (Eq, Show (..), (++), type (~))

import Data.Typeable (Typeable)
import Proarrow.Category.Instance.Opposite (OPPOSITE (..), Op (..))
import Proarrow.Category.Monoidal (Monoidal (..), MonoidalProfunctor (..), Strictly, associatorDefault)
import Proarrow.Colimit.Initial (HasInitialObject (..))
import Proarrow.Core (CAT, CategoryOf (..), Profunctor (..), Promonad (..), dimapDefault, obj, src, type (+->))
import Proarrow.Functor (FunctorForRep (..))
import Proarrow.Limit.Terminal (HasTerminalObject (..))
import Proarrow.Monoid (Monoid (..))

type n + m = Plus n m

data SNat :: Nat -> Type where
  SZ :: SNat Z
  SS :: (IsNat n) => SNat (S n)
instance Show (SNat n) where
  show SZ = "Z"
  show (SS @n') = "S" ++ show (singNat @n')

class (a + S b ~ S (a + b), Strictly a) => Rules a b
instance (a + S b ~ S (a + b), Strictly a) => Rules a b

class (forall b. Rules a b, SNatI a, Typeable a) => IsNat (a :: Nat) where singNat :: SNat a
instance IsNat Z where singNat = SZ
instance (IsNat a) => IsNat (S a) where singNat = SS

type Simplex :: CAT Nat
data Simplex a b where
  ZZ :: Simplex Z Z
  Y :: Simplex x y -> Simplex x (S y)
  X :: Simplex x (S y) -> Simplex (S x) (S y)
deriving instance Eq (Simplex a b)
deriving instance Show (Simplex a b)

suc :: Simplex a b -> Simplex (S a) (S b)
suc = X . Y

-- | The (augmented) simplex category is the category of finite ordinals and order preserving maps.
instance CategoryOf Nat where
  type (~>) = Simplex
  type Ob a = IsNat a

instance Promonad Simplex where
  id @a = case singNat @a of
    SZ -> ZZ
    SS -> suc id
  ZZ . f = f
  Y f . g = Y (f . g)
  X f . Y g = f . g
  X f . X g = X (X f . g)

instance Profunctor Simplex where
  dimap = dimapDefault
  r \\ ZZ = r
  r \\ Y f = r \\ f
  r \\ X f = r \\ f

instance HasInitialObject Nat where
  type InitialObject = Z
  initiate @a = case singNat @a of
    SZ -> ZZ
    SS @a' -> Y (initiate @_ @a')

instance HasTerminalObject Nat where
  type TerminalObject = S Z
  terminate @a = case singNat @a of
    SZ -> Y ZZ
    SS @n -> X (terminate @_ @n)

data family Forget :: Nat +-> Type
instance FunctorForRep Forget where
  type Forget @ n = Fin n
  fmap ZZ = id
  fmap (Y f) = FS . fmap @Forget f
  fmap (X f) = \case
    FZ -> FZ
    FS n -> fmap @Forget f n

data family Pick :: Type -> OPPOSITE Nat +-> Type
instance FunctorForRep (Pick a) where
  type (Pick a) @ OP n = Vec n a
  fmap (Op ZZ) VNil = VNil
  fmap (Op (Y f)) (_ ::: xs) = fmap @(Pick a) (Op f) xs
  fmap (Op (X f)) (x ::: xs) = x ::: fmap @(Pick a) (Op f) (x ::: xs)

instance MonoidalProfunctor Simplex where
  one = ZZ
  ZZ ** g = g
  Y f ** g = Y (f ** g)
  X f ** g = X (f ** g)

-- | Addition as monoidal tensor.
instance Monoidal Nat where
  type Unit = Z
  type a ** b = a + b
  withOb2 @a @b r = case singNat @a of
    SZ -> r
    SS @a' -> withOb2 @_ @a' @b r
  associator @a @b @c = associatorDefault @a @b @c
  associatorInv @a @b @c = associatorDefault @a @b @c

-- Not symmetric monoidal

instance Monoid Z where
  mempty = ZZ
  mappend = ZZ

instance Monoid (S Z) where
  mempty = Y ZZ
  mappend = X (X (Y ZZ))

data family Replicate :: k -> Nat +-> k
instance (Monoid m) => FunctorForRep (Replicate m) where
  type Replicate m @ Z = Unit
  type Replicate m @ S b = m ** (Replicate m @ b)
  fmap ZZ = one
  fmap (Y f) = let g = fmap @(Replicate m) f in (mempty @m ** g) . leftUnitorInv \\ g
  fmap (X (Y f)) = obj @m ** fmap @(Replicate m) f
  fmap (X (X @x f)) =
    let g = fmap @(Replicate m) (X f)
        b = fmap @(Replicate m) (src f)
    in g . (mappend @m ** b) . associatorInv @_ @m @m @(Replicate m @ x) \\ b