Categorical soup. CCC compilation a la the work of Conal Elliot, Chris Penner, and Phil Freeman

ccc.hs

{-# LANGUAGE NoImplicitPrelude #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE OverloadedStrings #-}


import Control.Category
import Data.Text (Text)
import Data.Void
import Prelude hiding ((.), id, curry, uncurry)
import Prelude qualified
import Prettyprinter qualified as P

--------------------------------------------------------------------------------

main = print $ renderJS collatzStep 
-- renderJS isPalindrome
-- (times10 >>> fork @_ @(,) times10 times10)

--------------------------------------------------------------------------------

class (Category cat1, Category cat2, Category cat3) => 
  GBifunctor cat1 cat2 cat3 t | t cat3 -> cat1 cat2 where 
  gbimap :: cat1 a b -> cat2 c d -> cat3 (a `t` c) (b `t` d) 

second :: (Category cat, Category cat') => 
  cat a b -> cat' (q c a) (q c b) 
second = undefined

grmap :: 
  GBifunctor cat1 cat2 cat3 t => 
  cat2 c d -> cat3 (a `t` c) (a `t` d) 
grmap = gbimap id

glmap :: 
  GBifunctor cat1 cat2 cat3 t => 
  cat1 a b -> cat3 (a `t` c) (b `t` c)
glmap = flip gbimap id

instance GBifunctor (->) (->) (->) (,) where
  gbimap f g (a, c) = (f a, g c)

instance GBifunctor (->) (->) (->) Either where
  gbimap f g = either (Left . f) (Right . g)
  
--------------------------------------------------------------------------------

data Iso cat a b = Iso {fwd :: cat a b, bwd :: cat b a}

instance Category cat => Category (Iso cat) where
  id :: Iso cat a a
  id = Iso id id

  (.) :: Iso cat b c -> Iso cat a b -> Iso cat a c
  bc . ab = Iso (fwd bc . fwd ab) (bwd ab . bwd bc)

class (Category cat, GBifunctor cat cat cat t) => Associative cat t where 
  assoc :: Iso cat (a `t` (b `t` c)) ((a `t` b) `t` c) 

class Associative cat t => Tensor cat t i | t -> i where 
  unitl :: Iso cat (i `t` a) a 
  unitr :: Iso cat (a `t` i) a 
  
--------------------------------------------------------------------------------

class Associative cat t => Symmetric cat t where
  swap :: cat (a `t` b) (b `t` a) 
  
class Symmetric cat t => Semicartesian cat t where 
  split :: cat a (a `t` a) 
  fork :: cat a x -> cat a y -> cat a (x `t` y) 
  
class Symmetric cat t => Semicocartesian cat t where 
  merge :: cat (a `t` a) a 
  fuse :: cat x a -> cat y a -> cat (x `t` y) a 
  
class Semicartesian cat t => Cartesian cat t i | i -> t, t -> i where
  kill :: cat a i 
  
  projl :: cat (x `t` y) x 
  projr :: cat (x `t` y) y 

  unfork :: cat a (x `t` y) -> (cat a x, cat a y)
  unfork h = (h >>> projl, h >>> projr)

first' :: 
  (Cartesian cat t i) => 
  cat a b -> cat (a `t` x) (b `t` x)
first' f = fork (projl >>> f) projr

second' :: 
  (Cartesian cat t i) => 
  cat a b -> cat (x `t` a) (x `t` b)
second' f = fork projl (projr >>> f)

(***) ::
  (Cartesian cat t i) =>
  cat a b -> cat a' b' -> cat (a `t` a') (b `t` b')
(***) f g = first' f >>> second' g

strong ::
  (Cartesian cat (,) ()) =>
  cat (a, b) r -> cat a b -> cat a r
strong f g = split >>> second' g >>> f

class (Semicocartesian cat t) => Cocartesian cat t i | i -> t, t -> i where
  spawn :: cat i a 
  incll :: cat x (x `t` y)
  inclr :: cat y (x `t` y) 
  unfuse :: cat (x `t` y) a -> (cat x a, cat y a)

class Cartesian cat t i => Closed cat t i f | cat -> f where
  curry :: cat (a `t` b) c -> cat a (b `f` c)
  uncurry :: (cat a (b `f` c)) -> cat (a `t` b) c

apply :: Closed cat t i f => cat ((a `f` b) `t` a) b
apply = uncurry id

class (Cartesian cat t i, Cocartesian cat t' i') => 
  Distrib cat t i t' i' | i -> t, t -> i, t' -> i', i' -> t' where
  distl ::  cat (a `t` (u `t'` v)) ((a `t` u) `t'` (a `t` v))

distr :: forall cat t i t' i' a b c.
  (Cartesian cat t i, Cocartesian cat t' i') => 
  cat ((a `t` b) `t'` (a `t` c)) (a `t` (b `t'` c))
distr = fuse (grmap $ incll @cat @t') (grmap $ inclr @cat @t')

left' ::
  (Cocartesian cat t i) =>
  cat a b -> cat (a `t` x) (b `t` x)
left' f = fuse (f >>> incll) inclr

right' ::
  (Cocartesian cat t i) =>
  cat a b -> cat (x `t` a) (x `t` b)
right' f = fuse incll (f >>> inclr)

(+++) ::
  (Cocartesian cat t i) =>
  cat l l' -> cat r r' -> cat (l `t` r) (l' `t` r')
(+++) f g = left' f >>> right' g

type GBool = Either () ()

tag ::
  (Distrib cat (,) () Either Void) => 
  cat (GBool, a) (Either a a)
tag = swap >>> distl >>> fuse (projl >>> incll) (projl >>> inclr)

matchOn :: 
  (Distrib cat (,) () Either Void) =>
  cat a GBool -> cat a (Either a a)
matchOn p = split >>> first' p >>> tag

class Associative cat t => Traced cat t where
  rightTraced :: cat (a `t` c) (b `t` c) -> cat a b
  leftTraced :: cat (c `t` a) (c `t` b) -> cat a b

--------------------------------------------------------------------------------

class Primitives cat where
  reverseString :: cat String String
  eq :: Eq a => cat (a, a) Bool
  fromBool :: cat Bool GBool
  toBool :: cat GBool Bool

isPalindrome :: 
  (Cartesian cat (,) (), Primitives cat) => cat String Bool 
isPalindrome = split >>> first' reverseString >>> eq

--------------------------------------------------------------------------------

class Numerics cat where
  num :: Int -> cat a Int
  negate' :: cat Int Int
  add :: cat (Int, Int) Int
  mult :: cat (Int, Int) Int
  div' :: cat (Int, Int) Int
  mod' :: cat (Int, Int) Int

mod2 :: forall cat.
  (Cartesian cat (,) (), Numerics cat) =>
  cat Int Int
mod2 = strong mod' (num 2)

isEven :: forall cat.
  (Cartesian cat (,) (), Numerics cat, Primitives cat) =>
  cat Int GBool
isEven = mod2 >>> strong eq (num 0) >>> fromBool

--------------------------------------------------------------------------------

instance Associative (->) (,) where
  assoc = Iso (\(a, (b, c)) -> ((a, b), c)) (\((a, b),c ) -> (a, (b, c))) 

instance Associative (->) Either where
  assoc = 
    Iso 
      (either (Left . Left) (either (Left . Right) Right)) 
      (either (either Left (Right . Left)) (Right . Right))
  
instance Symmetric (->) (,) where
  swap = \(a, b) -> (b, a)

instance Symmetric (->) Either where
  swap = either Right Left
  
instance Semicartesian (->) (,) where
  split = \a -> (a, a)
  fork = \f g a -> (f a, g a)
  
instance Semicocartesian (->) Either where
  merge = either id id
  fuse = \f g -> either f g
  
instance Cartesian (->) (,) () where
  kill = const ()
  projl = fst
  projr = snd

instance Closed (->) (,) () (->) where
  curry = Prelude.curry
  uncurry = Prelude.uncurry
  
instance Cocartesian (->) Either Void where
  spawn = absurd
  inclr = Right
  incll = Left
  unfuse = \f -> (f . Left, f . Right)

instance Traced (->) Either where
  rightTraced f a =
    let go = either id (go . Right) . f
     in go (Left a)
  
  leftTraced f a =
    let go = either (go . Left) id . f
     in go (Right a)
  
instance Distrib (->) (,) () Either Void where
  distl = uncurry $ \a -> either (Left . (a,)) (Right . (a,))

instance Primitives (->) where
  reverseString = reverse
  eq = uncurry (==)
  fromBool p = if p then Right () else Left ()
  toBool = either (const False) (const True)
  
instance Numerics (->) where
  num = const
  negate' = negate
  add = uncurry (+)
  mult = uncurry (*)
  div' = uncurry div
  mod' = uncurry mod

collatzStep :: forall cat.
  ( Distrib cat (,) () Either Void
  , Numerics cat
  , Primitives cat) =>
  cat Int Int
collatzStep = matchOn isEven >>> (onOdds +++ onEvens) >>> merge @_ @Either
  where
    onOdds :: cat Int Int
    onOdds = strong mult (num 3) >>> strong add (num 1)
    
    onEvens :: cat Int Int
    onEvens = strong div' (num 2)

--------------------------------------------------------------------------------

newtype JSFunc a b = JSFunc { renderJS :: forall ann. P.Doc ann }

jsLeft :: JSFunc a (Either a x)
jsLeft = JSFunc $ "((a) => ({tag: 'Left', value: a})"

jsRight :: JSFunc a (Either x x)
jsRight = JSFunc $ "((a) => ({tag: 'Right', value: a})"

instance GBifunctor JSFunc JSFunc JSFunc (,) where
  gbimap f g = 
    JSFunc $
      P.vsep $
        [ "(([a, c]) => {"
        , P.indent 2 $ P.vsep
          [ "const f =" P.<+> renderJS f
          , "const g =" P.<+> renderJS g
          , "return [f(a), g(c)]"
          ]
        , "})"
        ]
        
instance GBifunctor JSFunc JSFunc JSFunc Either where
  gbimap f g =
    JSFunc $
      P.vsep $
        [ "((x) => {"
        , P.indent 2 $ P.vsep
          [ "const f =" P.<+> renderJS f
          , "const g =" P.<+> renderJS g
          , "return (x.tag === 'Left' ? f(x) : g(x))"
          ]
        , "})"
        ]    
  
instance Associative JSFunc (,) where
  assoc = Iso from to
    where
      from = JSFunc $
        P.vsep $
          [ "(([a, [b, c]]) => {"
          , P.indent 2 $ P.vsep
            [ "return [[a, b], c]" ]
          , "})"
          ]
      to = JSFunc $
        P.vsep $
          [ "(([[a, b, c]) => {"
          , P.indent 2 $ P.vsep
            [ "return [a, [b, c]]" ]
          , "})"
          ]

-- TODO: Should flip around and define Cocartesian in terms of Associtive.
instance Associative JSFunc Either where
  assoc = Iso from to
    where
      from :: JSFunc (Either a (Either b c)) (Either (Either a b) c)
      from = fuse (incll >>> incll) (fuse (inclr >>> incll) inclr)

      to :: JSFunc (Either (Either a b) c) (Either a (Either b c))
      to = fuse (fuse incll (inclr >>> grmap incll)) (inclr >>> inclr)
          
instance Category JSFunc where
  id = JSFunc $ P.parens "((x) => x)"
  
  f . g = 
    JSFunc $
      P.vsep $
        [ "((x) => {"
        , P.indent 2 $ P.vsep 
          [ "const f =" P.<+> renderJS f
          , "const g =" P.<+> renderJS g
          , "return f(g(x))"
          ]
        , "})"
        ]
        
instance Symmetric JSFunc (,) where
  swap = JSFunc "(([x, y]) => [y, x])"

instance Semicartesian JSFunc (,) where
  split = JSFunc "((x) => [x, x])"
  fork f g = 
    JSFunc $
      P.vsep $
        [ "((a) => {"
        , P.indent 2 $ P.vsep
          [ "const f =" P.<+> renderJS f
          , "const g =" P.<+> renderJS g
          , "return [f(a), g(a)]"
          ]
        , "})"
        ]
        
instance Cartesian JSFunc (,) () where
  kill = JSFunc "((a) => null)"
  projl = JSFunc "(([x, y]) => x)"
  projr = JSFunc "(([x, y]) => y)"
  
instance Symmetric JSFunc Either where
  swap = 
    JSFunc $
      P.vsep $
        ["((t) => {"
        , P.indent 2 $ P.vsep
          [ "t.tag = t.tag === 'Left' ? 'Right' : 'Left'"
          , "return t"
          ]
        , "})"
        ]

instance Semicocartesian JSFunc Either where
  merge = JSFunc "((t) => t.value)"
  fuse f g =
    JSFunc $
      P.hsep
        [ "((t) => t.tag === 'Left' ?"
        , renderJS f <> "(t.value)"
        , ":"
        , renderJS g <> "(t.value)"
        , ")"
        ]

instance Cocartesian JSFunc Either Void where
  spawn = JSFunc "((x) => undefined)"
  incll = JSFunc "((x) => ({ tag: 'Left', value: x }))"
  inclr = JSFunc "((y) => ({ tag: 'Right', value: y }))"
  unfuse f = 
   (JSFunc $ "((x) =>" P.<+> renderJS f <> "({tag: 'Left', value: x})",
    JSFunc $ "((y) =>" P.<+> renderJS f <> "({tag: 'Right', value: y})") 

instance Distrib JSFunc (,) () Either Void where
  distl =
    JSFunc $
      P.vsep
        [ "(([a, uv]) => {"
        , P.indent 2 $ P.vsep
            [ "uv.value = [a, uv.value]"
            , "return uv"
            ]
        , "})"
        ]

times10 :: JSFunc Int Int       
times10 = JSFunc "((x) => 10 * x)"

instance Primitives JSFunc where
  reverseString = JSFunc "(s => s.split('').reverse().join(''))"
  eq = JSFunc "(([x, y]) => x === y)"
  fromBool = JSFunc "((b) => ({ tag: b ? 'Right' : 'Left', value: null }))"
  toBool = JSFunc "((b) => b.tag === 'Right' ? true : false) "

instance Numerics JSFunc where
  num n = JSFunc $ "((x) => (" <> P.pretty n <> "))"
  negate' = JSFunc "((x) => (-x))"
  add = JSFunc "(([x, y]) => (x + y))"
  mult = JSFunc "(([x, y]) => (x * y))"
  div' = JSFunc "(([x, y]) => (x / y))"
  mod' = JSFunc "(([x, y]) => (x % y))"