Linear.hs

{-# LANGUAGE LinearTypes #-}

{- HLINT ignore "Avoid lambda using `infix`" -}
{- HLINT ignore "Use curry" -}
{- HLINT ignore "Use bimap" -}
{- HLINT ignore "Use tuple-section" -}

-- | The category of Haskell types and __linear functions__: the kind 'LINEAR' wraps 'Data.Kind.Type'
-- in 'L', and a morphism is a @a %1 -> b@ function. Symmetric monoidal closed with
-- @'L' a '**' 'L' b = 'L' (a, b)@. The categorical product is 'With', and only comonoid objects
-- (such as @'L' ('Ur' a)@) can be copied or discarded, so it is deliberately not
-- 'Proarrow.Category.Monoidal.CopyDiscard.CopyDiscard'.
module Proarrow.Category.Instance.Linear where

import Control.Exception (evaluate)
import Data.IORef (newIORef, readIORef, writeIORef)
import Data.Kind (Type)
import Data.Void (Void)
import System.IO.Unsafe (unsafeDupablePerformIO, unsafePerformIO)
import Unsafe.Coerce (unsafeCoerce)
import Prelude (Bool (..), Either (..), Eq (..), Show (..), error, showParen, showString, (&&), (>))
import Prelude qualified as P

import Proarrow.Category.Monoidal (Monoidal (..), MonoidalProfunctor (..), SymMonoidal (..))
import Proarrow.Category.Monoidal.Action (CoprodAction)
import Proarrow.Category.Monoidal.Closed (Closed (..))
import Proarrow.Category.Monoidal.Dialogue (Dialogue (..))
import Proarrow.Category.Monoidal.Distributive (Distributive (..))
import Proarrow.Category.Monoidal.IsoMix (IsoMix (..))
import Proarrow.Category.Monoidal.StarAutonomous (StarAutonomous (..))
import Proarrow.Category.Monoidal.Strength (Costrong (..))
import Proarrow.Colimit.BinaryCoproduct (Coprod (..), HasBinaryCoproducts (..))
import Proarrow.Colimit.Copower (Copowered (..))
import Proarrow.Colimit.Initial (HasInitialObject (..))
import Proarrow.Core (CAT, CategoryOf (..), Is, Profunctor (..), Promonad (..), UN, dimapDefault, type (+->))
import Proarrow.Functor (Functor (..), FunctorForRep (..))
import Proarrow.Limit.BinaryProduct (HasBinaryProducts (..))
import Proarrow.Limit.Power (Powered (..))
import Proarrow.Limit.Terminal (HasTerminalObject (..))
import Proarrow.Monoid (Comonoid (..))
import Proarrow.Profunctor.Corepresentable (Corep (..), Corepresentable (..))
import Proarrow.Profunctor.Instance.Composition ((:.:) (..))
import Proarrow.Profunctor.Representable (Rep (..))

type data LINEAR = L Type

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

unLinear :: (L a ~> L b) %1 -> (a %1 -> b)
unLinear (Linear f) = f

instance Profunctor Linear where
  dimap = dimapDefault
  r \\ Linear{} = r
instance Promonad Linear where
  {-# INLINE id #-}
  {-# INLINE (.) #-}
  id = Linear \x -> x
  Linear f . Linear g = Linear \x -> f (g x)

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

instance MonoidalProfunctor Linear where
  {-# INLINE one #-}
  {-# INLINE (**) #-}
  one = id
  Linear f ** Linear g = Linear \(x, y) -> (f x, g y)

-- | Tuples as monoidal tensor. Tuples are not the binary product in LINEAR.
instance Monoidal LINEAR where
  {-# INLINE withOb2 #-}
  {-# INLINE leftUnitor #-}
  {-# INLINE leftUnitorInv #-}
  {-# INLINE rightUnitor #-}
  {-# INLINE rightUnitorInv #-}
  {-# INLINE associator #-}
  {-# INLINE associatorInv #-}
  type Unit = L ()
  type L a ** L b = L (a, b)
  withOb2 r = r
  leftUnitor = Linear \((), x) -> x
  leftUnitorInv = Linear \x -> ((), x)
  rightUnitor = Linear \(x, ()) -> x
  rightUnitorInv = Linear \x -> (x, ())
  associator = Linear \((x, y), z) -> (x, (y, z))
  associatorInv = Linear \(x, (y, z)) -> ((x, y), z)

instance SymMonoidal LINEAR where
  {-# INLINE swap #-}
  swap = Linear \(x, y) -> (y, x)

instance Closed LINEAR where
  type a ~~> b = L (UN L a %1 -> UN L b)
  withObExp r = r
  curry (Linear f) = Linear \a b -> f (a, b)
  apply = Linear \(f, a) -> f a
  Linear f ^^^ Linear g = Linear \h x -> f (h (g x))

data family Forget :: LINEAR +-> Type
instance FunctorForRep Forget where
  type Forget @ a = UN L a
  fmap (Linear f) x = f x

-- | By creating the left adjoint to the forgetful functor,
-- we obtain the free-forgetful adjunction between Hask and LINEAR
instance Corepresentable (Rep Forget :: LINEAR +-> Type) where
  type Rep Forget %% a = L (Ur a)
  coindex (Rep f) = Linear \(Ur a) -> f a
  cotabulate (Linear f) = Rep \a -> f (Ur a)
  corepMap f = Linear \(Ur a) -> Ur (f a)

-- | Forget is a lax monoidal functor
instance MonoidalProfunctor (Rep Forget) where
  {-# INLINE one #-}
  {-# INLINE (**) #-}
  one = Rep \() -> ()
  Rep f ** Rep g = Rep \(x, y) -> (f x, g y)

-- | Forget is also a colax monoidal functor
instance MonoidalProfunctor (Corep Forget) where
  {-# INLINE one #-}
  {-# INLINE (**) #-}
  one = Corep id
  Corep f ** Corep g = Corep \(x, y) -> (f x, g y)

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

counitUr :: Ur a %1 -> a
counitUr (Ur a) = a

dupUr :: Ur a %1 -> Ur (Ur a)
dupUr (Ur a) = Ur (Ur a)

instance Functor Ur where
  map f (Ur a) = Ur (f a)

instance Comonoid (L (Ur a)) where
  counit = Linear \(Ur _) -> ()
  comult = Linear \(Ur a) -> (Ur a, Ur a)

-- | @L Bool@ is a comonoid: a @Bool@ is duplicated and discarded by case analysis, which is
-- linear (it consumes the input exactly once). The same holds for any finite, pattern-matchable
-- classical type. Only the @Bool@ instance is spelled out here.
instance Comonoid (L Bool) where
  counit = Linear \case True -> (); False -> ()
  comult = Linear \case True -> (True, True); False -> (False, False)

instance HasBinaryCoproducts LINEAR where
  type L a || L b = L (Either a b)
  withObCoprod r = r
  lft = Linear Left
  rgt = Linear Right
  Linear f ||| Linear g = Linear \case
    Left x -> f x
    Right y -> g y

instance HasInitialObject LINEAR where
  type InitialObject = L Void
  initiate = Linear \case {}

instance Costrong CoprodAction Linear where
  coact (Linear uxuy) = loop . Linear Right
    where
      loop = Linear \ux -> case uxuy ux of Left x -> unLinear loop (Left x); Right b -> b

data Top where
  Top :: a %1 -> Top
instance Show Top where
  showsPrec _ _ = showString "⊤"
instance Eq Top where
  _ == _ = True

data With a b where
  With :: x %1 -> (x %1 -> a) -> (x %1 -> b) -> With a b
instance (Show a, Show b) => Show (With a b) where
  showsPrec d (With x f g) = showParen (d > 9) (showString "mkWith " . showsPrec 10 (f x) . showString " " . showsPrec 10 (g x))
instance (Eq a, Eq b) => Eq (With a b) where
  With x fa fb == With y ga gb = (fa x == ga y) && (fb x == gb y)

urWith :: Ur (With a b) %1 -> (Ur a, Ur b)
urWith (Ur (With x f g)) = (Ur (f x), Ur (g x))

mkWith :: a -> b -> With a b
mkWith a b = With () (\() -> a) (\() -> b)

instance HasTerminalObject LINEAR where
  type TerminalObject = L Top
  terminate = Linear Top

instance HasBinaryProducts LINEAR where
  type L a && L b = L (With a b)
  withObProd r = r
  fst = Linear \(With x xa _) -> xa x
  snd = Linear \(With x _ xb) -> xb x
  Linear f &&& Linear g = Linear \x -> With x f g

instance Powered Type LINEAR where
  type L a ^ n = L (n -> a)
  withObPower r = r
  power f = Linear \x n -> unLinear (f n) x
  unpower (Linear f) n = Linear \x -> f x n

instance Copowered Type LINEAR where
  type n *. L a = L (Ur n, a)
  withObCopower r = r
  copower f = Linear \(Ur n, a) -> unLinear (f n) a
  uncopower (Linear f) n = Linear \x -> f (Ur n, x)

instance MonoidalProfunctor (Coprod Linear) where
  {-# INLINE one #-}
  {-# INLINE (**) #-}
  one = Coprod (Linear \x -> x)
  Coprod f ** Coprod g = Coprod (f +++ g)

instance Distributive LINEAR where
  distL = Linear \(a, ebc) -> case ebc of Left b -> Left (a, b); Right c -> Right (a, c)
  distR = Linear \(eab, c) -> case eab of Left a -> Left (a, c); Right b -> Right (b, c)
  absorbL = Linear \(_a, v) -> case v of {}
  absorbR = Linear \(v, _a) -> case v of {}

type Not a = a %1 -> ()

not :: (Not b %1 -> Not a) %1 -> a %1 -> b
not nbna a = dn \nb -> nbna nb a

not' :: (a %1 -> b) %1 -> Not b %1 -> Not a
not' ab nb a = nb (ab a)

newtype Par a b = Par (Not (Not a, Not b))

mkPar :: a %1 -> b %1 -> Par a b
mkPar a b = Par \(na, nb) -> case (na a, nb b) of ((), ()) -> ()

pairFst :: (a, b `Par` c) %1 -> (a, b) `Par` c
pairFst (a, Par f) = Par \(nab, nc) -> f (\b -> nab (a, b), nc)

pairSnd :: (a `Par` b, c) %1 -> a `Par` (b, c)
pairSnd (Par f, c) = Par \(na, nbc) -> f (na, \b -> nbc (b, c))

parAppL :: (a `Par` b) %1 -> Not a %1 -> b
parAppL (Par f) na = dn \nb -> f (na, nb)

parAppR :: (a `Par` b) %1 -> Not b %1 -> a
parAppR (Par f) nb = dn \na -> f (na, nb)

newtype Quest a = Quest (Not (Ur (Not a)))

notQuest :: Not (Quest a) %1 -> Ur (Not a)
notQuest nqa = dn \nuna -> nqa (Quest nuna)

unitQuest :: a %1 -> Quest a
unitQuest a = Quest \(Ur na) -> na a

multQuest :: Quest (Quest a) %1 -> Quest a
multQuest (Quest f) = Quest \(Ur na) -> f (Ur (\(Quest nuna) -> nuna (Ur na)))

questPar :: Par (Quest a) (Quest b) %1 -> Quest (Either a b)
questPar (Par f) = Quest (\(Ur g) -> f (\(Quest nuna) -> nuna (Ur (\a -> g (Left a))), \(Quest nunb) -> nunb (Ur (\b -> g (Right b)))))

-- LINEAR is not CompactClosed. And hence it is also not traced,
-- since any star autonomous category with a trace is compact closed.
instance Dialogue LINEAR where
  {-# INLINE withObDual #-}
  {-# INLINE dual #-}
  {-# INLINE linDist #-}
  {-# INLINE linDistInv #-}
  {-# INLINE doubleNegInv #-}
  type Dual (L a) = L (Not a)
  withObDual r = r
  dual (Linear f) = Linear (\nb a -> nb (f a))
  linDist (Linear f) = Linear (\a (b, c) -> f (a, b) c)
  linDistInv (Linear f) = Linear (\(a, b) c -> f a (b, c))
  doubleNegInv = Linear (\a na -> na a)

instance StarAutonomous LINEAR where
  {-# INLINE dualInv #-}
  {-# INLINE doubleNeg #-}
  dualInv (Linear f) = Linear (\b -> dn (\na -> f na b))
  doubleNeg = Linear dn

-- | The unit of par is @() %1 -> ()@, which has one value, just like @()@: apply it to @()@, or
-- give back the identity. So LINEAR is isomix, while tensor and par still differ.
instance IsoMix LINEAR where
  dualUnit = Linear (\f -> f ())
  dualUnitInv = Linear (\() u -> u)
  dualityCounit = Linear (\(na, a) -> na a)

-- | Double negation is possible with linear functions, though using `unsafePerformIO`.
-- Derived from https://gist.github.com/ant-arctica/7563282c57d9d1ce0c4520c543187932
dn :: Not (Not a) %1 -> a
-- One IO action that depends on the argument, so that optimisation can't share the reference
-- between calls by floating it out.
dn = unsafeLinear \nna ->
  unsafePerformIO
    ( newIORef (error "Linear.dn: the continuation was not called") P.>>= \ref ->
        evaluate (nna (unsafeLinear \x -> unsafeDupablePerformIO (writeIORef ref x))) P.>> readIORef ref
    )

unsafeLinear :: (a -> b) -> (a %1 -> b)
unsafeLinear = unsafeCoerce

unit :: L () ~> L (Par a (Not a))
unit = Linear \() -> Par \(na, nna) -> nna na

counit :: L (Not a, a) ~> L ()
counit = Linear \(na, a) -> na a

type p !~> q = forall a b. p a b %1 -> q a b

type NegComp :: (j +-> k) -> (i +-> j) -> (i +-> k)
data NegComp p q a c where
  NegComp :: (forall b. Par (p a b) (q b c)) %1 -> NegComp p q a c

newtype Neg p a b = Neg (Not (p b a))

getNeg :: Neg p a b %1 -> Not (p b a)
getNeg (Neg f) = f

conv1 :: NegComp p q !~> Neg (Neg q :.: Neg p)
conv1 (NegComp e) = Neg \(Neg nq :.: Neg np) -> case e of Par e' -> e' (np, nq)

conv2 :: Neg (Neg q :.: Neg p) !~> NegComp p q
conv2 (Neg f) = NegComp (Par (\(np, nq) -> f (Neg nq :.: Neg np)))

asCocat :: (Neg p :.: Neg p !~> Neg p) -> p !~> NegComp p p
asCocat comp p = NegComp (Par \(np1, np2) -> getNeg (comp (Neg np2 :.: Neg np1)) p)