{-# LANGUAGE LinearTypes #-}

module Proarrow.Category.Instance.Linear where

import Data.Kind (Type)
import Data.Void (Void)
import Prelude (Either (..))

import Proarrow.Adjunction (Adjunction (..))
import Proarrow.Category.Monoidal (Monoidal (..), MonoidalProfunctor (..), SymMonoidal (..))
import Proarrow.Core (CAT, CategoryOf (..), Is, Profunctor (..), Promonad (..), UN, dimapDefault, type (+->))
import Proarrow.Functor (Functor (..))
import Proarrow.Object.BinaryCoproduct (HasBinaryCoproducts (..))
import Proarrow.Object.Exponential (Closed (..))
import Proarrow.Object.Initial (HasInitialObject (..))
import Proarrow.Profunctor.Composition ((:.:) (..))
import Proarrow.Profunctor.Representable (Representable (..), dimapRep, RepCostar (..))
import Proarrow.Object.BinaryProduct (HasBinaryProducts (..))
import Proarrow.Object.Terminal (HasTerminalObject (..))
import Proarrow.Monoid (Comonoid (..))

type data LINEAR = L Type
type instance UN L (L a) = a

type Linear :: CAT LINEAR
data Linear a b where
  Linear :: (a %1 -> b) -> Linear (L a) (L b)

instance Profunctor Linear where
  dimap :: forall (c :: LINEAR) (a :: LINEAR) (b :: LINEAR) (d :: LINEAR).
(c ~> a) -> (b ~> d) -> Linear a b -> Linear c d
dimap = (c ~> a) -> (b ~> d) -> Linear a b -> Linear c d
Linear c a -> Linear b d -> Linear a b -> Linear c d
forall {k} (p :: PRO k k) (c :: k) (a :: k) (b :: k) (d :: k).
Promonad p =>
p c a -> p b d -> p a b -> p c d
dimapDefault
  (Ob a, Ob b) => r
r \\ :: forall (a :: LINEAR) (b :: LINEAR) r.
((Ob a, Ob b) => r) -> Linear a b -> r
\\ Linear{} = r
(Ob a, Ob b) => r
r
instance Promonad Linear where
  id :: forall (a :: LINEAR). Ob a => Linear a a
id = (UN L a %1 -> UN L a) -> Linear (L (UN L a)) (L (UN L a))
forall x b. (x %1 -> b) -> Linear (L x) (L b)
Linear \UN L a
x -> UN L a
x
  Linear a %1 -> b
f . :: forall (b :: LINEAR) (c :: LINEAR) (a :: LINEAR).
Linear b c -> Linear a b -> Linear a c
. Linear a %1 -> b
g = (a %1 -> b) -> Linear (L a) (L b)
forall x b. (x %1 -> b) -> Linear (L x) (L b)
Linear \a
x -> a %1 -> b
f (a %1 -> b
g a
x)

-- | Category of linear functions.
instance CategoryOf LINEAR where
  type (~>) = Linear
  type Ob (a :: LINEAR) = Is L a

instance MonoidalProfunctor Linear where
  par0 :: Linear Unit Unit
par0 = Linear Unit Unit
Linear (L ()) (L ())
forall {k} (p :: PRO k k) (a :: k). (Promonad p, Ob a) => p a a
forall (a :: LINEAR). Ob a => Linear a a
id
  Linear a %1 -> b
f par :: forall (x1 :: LINEAR) (x2 :: LINEAR) (y1 :: LINEAR) (y2 :: LINEAR).
Linear x1 x2 -> Linear y1 y2 -> Linear (x1 ** y1) (x2 ** y2)
`par` Linear a %1 -> b
g = ((a, a) %1 -> (b, b)) -> Linear (L (a, a)) (L (b, b))
forall x b. (x %1 -> b) -> Linear (L x) (L b)
Linear \(a
x, a
y) -> (a %1 -> b
f a
x, a %1 -> b
g a
y)

instance Monoidal LINEAR where
  type Unit = L ()
  type L a ** L b = L (a, b)
  leftUnitor :: forall (a :: LINEAR). Ob a => (Unit ** a) ~> a
leftUnitor = (((), UN L a) %1 -> UN L a) -> Linear (L ((), UN L a)) (L (UN L a))
forall x b. (x %1 -> b) -> Linear (L x) (L b)
Linear \((), UN L a
x) -> UN L a
x
  leftUnitorInv :: forall (a :: LINEAR). Ob a => a ~> (Unit ** a)
leftUnitorInv = (UN L a %1 -> ((), UN L a)) -> Linear (L (UN L a)) (L ((), UN L a))
forall x b. (x %1 -> b) -> Linear (L x) (L b)
Linear \UN L a
x -> ((), UN L a
x)
  rightUnitor :: forall (a :: LINEAR). Ob a => (a ** Unit) ~> a
rightUnitor = ((UN L a, ()) %1 -> UN L a) -> Linear (L (UN L a, ())) (L (UN L a))
forall x b. (x %1 -> b) -> Linear (L x) (L b)
Linear \(UN L a
x, ()) -> UN L a
x
  rightUnitorInv :: forall (a :: LINEAR). Ob a => a ~> (a ** Unit)
rightUnitorInv = (UN L a %1 -> (UN L a, ())) -> Linear (L (UN L a)) (L (UN L a, ()))
forall x b. (x %1 -> b) -> Linear (L x) (L b)
Linear \UN L a
x -> (UN L a
x, ())
  associator :: forall (a :: LINEAR) (b :: LINEAR) (c :: LINEAR).
(Ob a, Ob b, Ob c) =>
((a ** b) ** c) ~> (a ** (b ** c))
associator = (((UN L a, UN L b), UN L c) %1 -> (UN L a, (UN L b, UN L c)))
-> Linear
     (L ((UN L a, UN L b), UN L c)) (L (UN L a, (UN L b, UN L c)))
forall x b. (x %1 -> b) -> Linear (L x) (L b)
Linear \((UN L a
x, UN L b
y), UN L c
z) -> (UN L a
x, (UN L b
y, UN L c
z))
  associatorInv :: forall (a :: LINEAR) (b :: LINEAR) (c :: LINEAR).
(Ob a, Ob b, Ob c) =>
(a ** (b ** c)) ~> ((a ** b) ** c)
associatorInv = ((UN L a, (UN L b, UN L c)) %1 -> ((UN L a, UN L b), UN L c))
-> Linear
     (L (UN L a, (UN L b, UN L c))) (L ((UN L a, UN L b), UN L c))
forall x b. (x %1 -> b) -> Linear (L x) (L b)
Linear \(UN L a
x, (UN L b
y, UN L c
z)) -> ((UN L a
x, UN L b
y), UN L c
z)

instance SymMonoidal LINEAR where
  swap' :: forall (a :: LINEAR) (a' :: LINEAR) (b :: LINEAR) (b' :: LINEAR).
(a ~> a') -> (b ~> b') -> (a ** b) ~> (b' ** a')
swap' (Linear a %1 -> b
f) (Linear a %1 -> b
g) = ((a, a) %1 -> (b, b)) -> Linear (L (a, a)) (L (b, b))
forall x b. (x %1 -> b) -> Linear (L x) (L b)
Linear \(a
x, a
y) -> (a %1 -> b
g a
y, a %1 -> b
f a
x)

-- instance TracedMonoidalProfunctor Linear where
--   trace (Linear f) = Linear \x -> let !(y, u) = f (x, u) in y

instance Closed LINEAR where
  type a ~~> b = L (UN L a %1 -> UN L b)
  curry' :: forall (a :: LINEAR) (b :: LINEAR) (c :: LINEAR).
Obj a -> Obj b -> ((a ** b) ~> c) -> a ~> (b ~~> c)
curry' Linear{} Linear{} (Linear a %1 -> b
f) = (b %1 -> b %1 -> b) -> Linear (L b) (L (b %1 -> b))
forall x b. (x %1 -> b) -> Linear (L x) (L b)
Linear \b
a b
b -> a %1 -> b
f (b
a, b
b)
  uncurry' :: forall (b :: LINEAR) (c :: LINEAR) (a :: LINEAR).
Obj b -> Obj c -> (a ~> (b ~~> c)) -> (a ** b) ~> c
uncurry' Linear{} Linear{} (Linear a %1 -> b
f) = ((a, b) %1 -> b) -> Linear (L (a, b)) (L b)
forall x b. (x %1 -> b) -> Linear (L x) (L b)
Linear \(a
a, b
b) -> a %1 -> b
f a
a b
b
  Linear a %1 -> b
f ^^^ :: forall (b :: LINEAR) (y :: LINEAR) (x :: LINEAR) (a :: LINEAR).
(b ~> y) -> (x ~> a) -> (a ~~> b) ~> (x ~~> y)
^^^ Linear a %1 -> b
g = ((b %1 -> a) %1 -> a %1 -> b)
-> Linear (L (b %1 -> a)) (L (a %1 -> b))
forall x b. (x %1 -> b) -> Linear (L x) (L b)
Linear \b %1 -> a
h a
x -> a %1 -> b
f (b %1 -> a
h (a %1 -> b
g a
x))

type Forget :: LINEAR +-> Type
data Forget a b where
  Forget :: (a -> b) -> Forget a (L b)
instance Profunctor Forget where
  dimap :: forall c a (b :: LINEAR) (d :: LINEAR).
(c ~> a) -> (b ~> d) -> Forget a b -> Forget c d
dimap = (c ~> a) -> (b ~> d) -> Forget a b -> Forget c d
forall {j} {k} (p :: j +-> k) (a :: k) (b :: j) (c :: k) (d :: j).
Representable p =>
(c ~> a) -> (b ~> d) -> p a b -> p c d
dimapRep
  (Ob a, Ob b) => r
r \\ :: forall a (b :: LINEAR) r. ((Ob a, Ob b) => r) -> Forget a b -> r
\\ Forget{} = r
(Ob a, Ob b) => r
r
instance Representable Forget where
  type Forget % a = UN L a
  index :: forall a (b :: LINEAR). Forget a b -> a ~> (Forget % b)
index (Forget a -> b
f) = a ~> (Forget % b)
a -> b
f
  tabulate :: forall (b :: LINEAR) a. Ob b => (a ~> (Forget % b)) -> Forget a b
tabulate = (a ~> (Forget % b)) -> Forget a b
(a -> UN L b) -> Forget a (L (UN L b))
forall a x. (a -> x) -> Forget a (L x)
Forget
  repMap :: forall (a :: LINEAR) (b :: LINEAR).
(a ~> b) -> (Forget % a) ~> (Forget % b)
repMap (Linear a %1 -> b
f) UN L a
x = a %1 -> b
f a
UN L a
x
-- | Forget is a lax monoidal functor
instance MonoidalProfunctor Forget where
  par0 :: Forget Unit Unit
par0 = (() -> ()) -> Forget () (L ())
forall a x. (a -> x) -> Forget a (L x)
Forget \() -> ()
  Forget x1 -> b
f par :: forall x1 (x2 :: LINEAR) y1 (y2 :: LINEAR).
Forget x1 x2 -> Forget y1 y2 -> Forget (x1 ** y1) (x2 ** y2)
`par` Forget y1 -> b
g = ((x1, y1) -> (b, b)) -> Forget (x1, y1) (L (b, b))
forall a x. (a -> x) -> Forget a (L x)
Forget \(x1
x, y1
y) -> (x1 -> b
f x1
x, y1 -> b
g y1
y)
-- | Forget is also a colax monoidal functor
instance MonoidalProfunctor (RepCostar Forget) where
  par0 :: RepCostar Forget Unit Unit
par0 = ((Forget % L ()) ~> ()) -> RepCostar Forget (L ()) ()
forall {k} {j} (a :: k) (p :: k +-> j) (b :: j).
Ob a =>
((p % a) ~> b) -> RepCostar p a b
RepCostar (Forget % L ()) ~> ()
() -> ()
forall a. Ob a => a -> a
forall {k} (p :: PRO k k) (a :: k). (Promonad p, Ob a) => p a a
id
  RepCostar (Forget % x1) ~> x2
f par :: forall (x1 :: LINEAR) x2 (y1 :: LINEAR) y2.
RepCostar Forget x1 x2
-> RepCostar Forget y1 y2 -> RepCostar Forget (x1 ** y1) (x2 ** y2)
`par` RepCostar (Forget % y1) ~> y2
g = ((Forget % L (UN L x1, UN L y1)) ~> (x2, y2))
-> RepCostar Forget (L (UN L x1, UN L y1)) (x2, y2)
forall {k} {j} (a :: k) (p :: k +-> j) (b :: j).
Ob a =>
((p % a) ~> b) -> RepCostar p a b
RepCostar \(UN L x1
x, UN L y1
y) -> ((Forget % x1) ~> x2
UN L x1 -> x2
f UN L x1
x, (Forget % y1) ~> y2
UN L y1 -> y2
g UN L y1
y)

data Ur a where
  Ur :: a -> Ur a

instance Functor Ur where
  map :: forall a b. (a ~> b) -> Ur a ~> Ur b
map a ~> b
f (Ur a
a) = b -> Ur b
forall a. a -> Ur a
Ur (a ~> b
a -> b
f a
a)

instance Comonoid (L (Ur a)) where
  counit :: L (Ur a) ~> Unit
counit = (Ur a %1 -> ()) -> Linear (L (Ur a)) (L ())
forall x b. (x %1 -> b) -> Linear (L x) (L b)
Linear \(Ur a
_) -> ()
  comult :: L (Ur a) ~> (L (Ur a) ** L (Ur a))
comult = (Ur a %1 -> (Ur a, Ur a)) -> Linear (L (Ur a)) (L (Ur a, Ur a))
forall x b. (x %1 -> b) -> Linear (L x) (L b)
Linear \(Ur a
a) -> (a -> Ur a
forall a. a -> Ur a
Ur a
a, a -> Ur a
forall a. a -> Ur a
Ur a
a)

type Free :: Type +-> LINEAR
data Free a b where
  Free :: (a %1 -> Ur b) -> Free (L a) b
instance Profunctor Free where
  dimap :: forall (c :: LINEAR) (a :: LINEAR) b d.
(c ~> a) -> (b ~> d) -> Free a b -> Free c d
dimap = (c ~> a) -> (b ~> d) -> Free a b -> Free c d
forall {j} {k} (p :: j +-> k) (a :: k) (b :: j) (c :: k) (d :: j).
Representable p =>
(c ~> a) -> (b ~> d) -> p a b -> p c d
dimapRep
  (Ob a, Ob b) => r
r \\ :: forall (a :: LINEAR) b r. ((Ob a, Ob b) => r) -> Free a b -> r
\\ Free{} = r
(Ob a, Ob b) => r
r
instance Representable Free where
  type Free % a = L (Ur a)
  index :: forall (a :: LINEAR) b. Free a b -> a ~> (Free % b)
index (Free a %1 -> Ur b
f) = (a %1 -> Ur b) -> Linear (L a) (L (Ur b))
forall x b. (x %1 -> b) -> Linear (L x) (L b)
Linear a %1 -> Ur b
f
  tabulate :: forall b (a :: LINEAR). Ob b => (a ~> (Free % b)) -> Free a b
tabulate (Linear a %1 -> b
f) = (a %1 -> Ur b) -> Free (L a) b
forall x b. (x %1 -> Ur b) -> Free (L x) b
Free a %1 -> b
a %1 -> Ur b
f
  repMap :: forall a b. (a ~> b) -> (Free % a) ~> (Free % b)
repMap a ~> b
f = (Ur a %1 -> Ur b) -> Linear (L (Ur a)) (L (Ur b))
forall x b. (x %1 -> b) -> Linear (L x) (L b)
Linear \(Ur a
a) -> b -> Ur b
forall a. a -> Ur a
Ur (a ~> b
a -> b
f a
a)
-- | Free is a lax monoidal functor
instance MonoidalProfunctor Free where
  par0 :: Free Unit Unit
par0 = (() %1 -> Ur ()) -> Free (L ()) ()
forall x b. (x %1 -> Ur b) -> Free (L x) b
Free \() -> () -> Ur ()
forall a. a -> Ur a
Ur ()
  Free a %1 -> Ur x2
f par :: forall (x1 :: LINEAR) x2 (y1 :: LINEAR) y2.
Free x1 x2 -> Free y1 y2 -> Free (x1 ** y1) (x2 ** y2)
`par` Free a %1 -> Ur y2
g = ((a, a) %1 -> Ur (x2, y2)) -> Free (L (a, a)) (x2, y2)
forall x b. (x %1 -> Ur b) -> Free (L x) b
Free \(a
x, a
y) -> case (a %1 -> Ur x2
f a
x, a %1 -> Ur y2
g a
y) of (Ur x2
a, Ur y2
b) -> (x2, y2) -> Ur (x2, y2)
forall a. a -> Ur a
Ur (x2
a, y2
b)
-- | Free is also a colax monoidal functor
instance MonoidalProfunctor (RepCostar Free) where
  par0 :: RepCostar Free Unit Unit
par0 = ((Free % ()) ~> L ()) -> RepCostar Free () (L ())
forall {k} {j} (a :: k) (p :: k +-> j) (b :: j).
Ob a =>
((p % a) ~> b) -> RepCostar p a b
RepCostar ((Ur () %1 -> ()) -> Linear (L (Ur ())) (L ())
forall x b. (x %1 -> b) -> Linear (L x) (L b)
Linear (\(Ur ()
_) -> ()))
  RepCostar (Linear a %1 -> b
f) par :: forall x1 (x2 :: LINEAR) y1 (y2 :: LINEAR).
RepCostar Free x1 x2
-> RepCostar Free y1 y2 -> RepCostar Free (x1 ** y1) (x2 ** y2)
`par` RepCostar (Linear a %1 -> b
g) = ((Free % (x1, y1)) ~> L (b, b))
-> RepCostar Free (x1, y1) (L (b, b))
forall {k} {j} (a :: k) (p :: k +-> j) (b :: j).
Ob a =>
((p % a) ~> b) -> RepCostar p a b
RepCostar ((Ur (x1, y1) %1 -> (b, b)) -> Linear (L (Ur (x1, y1))) (L (b, b))
forall x b. (x %1 -> b) -> Linear (L x) (L b)
Linear \(Ur (x1
x, y1
y)) -> (a %1 -> b
f (x1 -> Ur x1
forall a. a -> Ur a
Ur x1
x), a %1 -> b
g (y1 -> Ur y1
forall a. a -> Ur a
Ur y1
y)))

instance Adjunction Free Forget where
  unit :: forall a. Ob a => (:.:) Forget Free a a
unit = (a -> Ur a) -> Forget a (L (Ur a))
forall a x. (a -> x) -> Forget a (L x)
Forget a -> Ur a
forall a. a -> Ur a
Ur Forget a (L (Ur a)) -> Free (L (Ur a)) a -> (:.:) Forget Free a a
forall {j} {k} {i} (p :: j +-> k) (a :: k) (b :: j) (q :: i +-> j)
       (c :: i).
p a b -> q b c -> (:.:) p q a c
:.: (Ur a %1 -> Ur a) -> Free (L (Ur a)) a
forall x b. (x %1 -> Ur b) -> Free (L x) b
Free \Ur a
x -> Ur a
x
  counit :: (Free :.: Forget) :~> (~>)
counit (Free a %1 -> Ur b
f :.: Forget b -> b
g) = (a %1 -> b) -> Linear (L a) (L b)
forall x b. (x %1 -> b) -> Linear (L x) (L b)
Linear \a
a -> case a %1 -> Ur b
f a
a of Ur b
b -> b -> b
g b
b

instance HasBinaryCoproducts LINEAR where
  type L a || L b = L (Either a b)
  lft :: forall (a :: LINEAR) (b :: LINEAR). (Ob a, Ob b) => a ~> (a || b)
lft = (UN L a %1 -> Either (UN L a) (UN L b))
-> Linear (L (UN L a)) (L (Either (UN L a) (UN L b)))
forall x b. (x %1 -> b) -> Linear (L x) (L b)
Linear UN L a -> Either (UN L a) (UN L b)
forall a b. a -> Either a b
Left
  rgt :: forall (a :: LINEAR) (b :: LINEAR). (Ob a, Ob b) => b ~> (a || b)
rgt = (UN L b %1 -> Either (UN L a) (UN L b))
-> Linear (L (UN L b)) (L (Either (UN L a) (UN L b)))
forall x b. (x %1 -> b) -> Linear (L x) (L b)
Linear UN L b -> Either (UN L a) (UN L b)
forall a b. b -> Either a b
Right
  Linear a %1 -> b
f ||| :: forall (x :: LINEAR) (a :: LINEAR) (y :: LINEAR).
(x ~> a) -> (y ~> a) -> (x || y) ~> a
||| Linear a %1 -> b
g = (Either a a %1 -> b) -> Linear (L (Either a a)) (L b)
forall x b. (x %1 -> b) -> Linear (L x) (L b)
Linear \case
    Left a
x -> a %1 -> b
f a
x
    Right a
y -> a %1 -> b
g a
y

instance HasInitialObject LINEAR where
  type InitialObject = L Void
  initiate :: forall (a :: LINEAR). Ob a => InitialObject ~> a
initiate = (Void %1 -> UN L a) -> Linear (L Void) (L (UN L a))
forall x b. (x %1 -> b) -> Linear (L x) (L b)
Linear \case {}

data Top where
  Top :: a %1 -> Top

data With a b where
  With :: x %1 -> (x %1 -> a) -> (x %1 -> b) -> With a b

instance HasTerminalObject LINEAR where
  type TerminalObject = L Top
  terminate :: forall (a :: LINEAR). Ob a => a ~> TerminalObject
terminate = (UN L a %1 -> Top) -> Linear (L (UN L a)) (L Top)
forall x b. (x %1 -> b) -> Linear (L x) (L b)
Linear UN L a -> Top
forall x. x -> Top
Top

instance HasBinaryProducts LINEAR where
  type L a && L b = L (With a b)
  fst :: forall (a :: LINEAR) (b :: LINEAR). (Ob a, Ob b) => (a && b) ~> a
fst = (With (UN L a) (UN L b) %1 -> UN L a)
-> Linear (L (With (UN L a) (UN L b))) (L (UN L a))
forall x b. (x %1 -> b) -> Linear (L x) (L b)
Linear \(With x
x x %1 -> UN L a
xa x %1 -> UN L b
_) -> x %1 -> UN L a
xa x
x
  snd :: forall (a :: LINEAR) (b :: LINEAR). (Ob a, Ob b) => (a && b) ~> b
snd = (With (UN L a) (UN L b) %1 -> UN L b)
-> Linear (L (With (UN L a) (UN L b))) (L (UN L b))
forall x b. (x %1 -> b) -> Linear (L x) (L b)
Linear \(With x
x x %1 -> UN L a
_ x %1 -> UN L b
xb) -> x %1 -> UN L b
xb x
x
  Linear a %1 -> b
f &&& :: forall (a :: LINEAR) (x :: LINEAR) (y :: LINEAR).
(a ~> x) -> (a ~> y) -> a ~> (x && y)
&&& Linear a %1 -> b
g = (a %1 -> With b b) -> Linear (L a) (L (With b b))
forall x b. (x %1 -> b) -> Linear (L x) (L b)
Linear \a
x -> a -> (a %1 -> b) -> (a %1 -> b) -> With b b
forall x a b. x -> (x %1 -> a) -> (x %1 -> b) -> With a b
With a
x a %1 -> b
f a %1 -> b
a %1 -> b
g