CofreeWZap.hs

{-# LANGUAGE
    FunctionalDependencies,
    RankNTypes,
    NoMonomorphismRestriction,
    TypeOperators,
    TypeFamilies,
    TypeFamilyDependencies,
    TypeApplications,
    InstanceSigs,
    PolyKinds,
    KindSignatures,
    GADTs,
    DataKinds,
    MultiParamTypeClasses,
    FlexibleContexts,
    FlexibleInstances,
    UndecidableInstances,
    GeneralizedNewtypeDeriving,
    StandaloneDeriving,
    TupleSections,
    DerivingStrategies,
    ScopedTypeVariables,
    LambdaCase,
    AllowAmbiguousTypes
#-}

import Prelude hiding (read)

import Control.Applicative
import Control.Comonad
import Control.Comonad.Cofree
import Control.Concurrent.MVar
import Control.Exception (Exception(..), SomeException, throwIO)
import Control.Monad ((>=>))
import Control.Monad.Except (ExceptT(..))
import Control.Monad.Free (Free(..))
import Control.Monad.IO.Class (MonadIO(..))
import Control.Monad.Reader (ReaderT(..))
import Control.Monad.State (StateT(..))
import Control.Monad.Trans.Class (MonadTrans(..))
import Control.Monad.Writer (WriterT(..))
import Data.Bifunctor
import Data.Coerce
import Data.Functor.Adjunction
import Data.Functor.Compose
import Data.Functor.Const
import Data.Functor.Coyoneda
import Data.Functor.Identity
import Data.Functor.Product
import Data.Functor.Sum
import Data.Functor.Sum
import Data.Kind
import Data.List (isPrefixOf)
import Data.Map (Map)
import Data.Proxy
import Data.Text(Text)
import Data.Unique
import Data.Void
import GHC.TypeLits
import qualified Control.Exception as IO
import qualified Control.Monad.Except as MTL
import qualified Control.Monad.Free as Free
import qualified Control.Monad.Reader as MTL
import qualified Control.Monad.State as MTL
import qualified Control.Monad.Writer as MTL
import qualified Data.ByteString.Lazy as LBS
import qualified Data.Map as Map
import qualified Data.OpenUnion as OU
import qualified Prelude
import qualified System.Directory as D
import qualified System.Environment as D
import qualified System.Process as D




type f ~> g = forall a . f a -> g a
type (:+:) f g = Sum f g
type (:*:) f g = Product f g
type (:.:) f g = Compose f g

-- Like Data.Functor.Compose, but more polymorphic
type (:..) f g x = f (g x)


infixr 0 ~>
infixr 3 :+:
infixr 4 :*:

class f :< g where
  liftTo :: f ~> g

liftTo2 :: f :< g => (a -> b -> f c) -> a -> b -> g c
liftTo2 f x y = liftTo $ f x y


data Reads a :: Type -> Type where
  Read :: Reads a a

read :: Reads a :< f => f a
read = liftTo Read

data Writes w :: Type -> Type where
  Write :: w -> Writes w ()

write :: Writes w :< f => w -> f ()
write x = liftTo (Write x)

data State s :: Type -> Type where
  Get :: State s s
  Put :: s -> State s ()

get :: State s :< f => f s
get = liftTo Get

put :: State s :< f => s -> f ()
put x = liftTo (Put x)

modify :: (Monad m, State s :< m) => (s -> s) -> m ()
modify f = do
  s <- get
  put $ f s

data Throws e :: Type -> Type where
  Throw :: e -> Throws e a

throw :: Throws e :< f => e -> f a
throw x = liftTo (Throw x)


data Foo = Foo Int
data Bar = Bar String

exManyErrors :: (Throws Foo :< f, Throws Bar :< f) => Int -> f a
exManyErrors x = do
  if (x > 2)
    then
      throw $ Foo 1
    else
      throw $ Bar "bad!"


data ReadEnv :: Symbol -> Type -> Type where
  ReadEnv :: ReadEnv s (Maybe String)

instance KnownSymbol s => ReadEnv s :< IO where
  liftTo ReadEnv = D.lookupEnv $ symbolVal (Proxy @s)

readEnv :: forall s f . ReadEnv s :< f => f (Maybe String)
readEnv = liftTo $ ReadEnv @s

exEnv :: (Monad f, ReadEnv "HOME" :< f, Writes String :< f) => f ()
exEnv = do
  homeDir <- readEnv @"HOME"
  maybe (pure ()) write homeDir

-- throw x `catch` q = q x
-- p `catch` throw = p
class Throws e :< f => Catch e f where
  catch :: f a -> (e -> f a) -> f a

instance Catch e (Either e) where
  catch (Left e) h = h e
  catch x        _ = x

-- local f read = f <$> read
class Reads e :< f => Local e f where
  local :: (e -> e) -> f a -> f a

-- censor f write = write . f
class Writes e :< f => Censor e f where
  censor :: (e -> e) -> f a -> f a

exHalfCapacity :: (Monad f, Local Double f, Censor String f) => f ()
exHalfCapacity = do
  showCapacity
  local @Double (/ 2) $ censor ("Running at half capacity: " <> ) $ do
    showCapacity
  showCapacity
  where
    showCapacity = do
      numberOfCapabilities <- read @Double
      write $ "Available capabilities: " <> show numberOfCapabilities

data Cut = Cut

cut :: (Functor f, Throws Cut :< f) => f a
cut = vacuous $ liftTo $ Throw Cut

choice_ :: Catch Cut f => f a -> f a -> f a
p `choice_` q = p `catch` (\Cut -> q)

instance Throws e :< Either e where
  liftTo (Throw e) = Left e
-- | Fold for Sums.
foldSum :: (f a -> b) -> (g a -> b) -> Sum f g a -> b
foldSum f _ (InL a) = f a
foldSum _ g (InR b) = g b

-- |
-- Special case of foldSum, more in line with the categorical definition.
-- We recognize this as the 'either' function lifted to endofunctors
foldCoproduct :: (f ~> h) -> (g ~> h) -> (f :+: g) ~> h
foldCoproduct = foldSum

type NT3 f g h = forall x . f x -> g x -> h x
-- (dually:)
foldProduct :: (NT3 f g h) -> (f :*: g) ~> h
foldProduct fgh (x `Pair` y) = fgh x y



class Member f g where
  inj :: forall a. f a -> g a
  prj :: forall a. g a -> Maybe (f a)

instance Member f f where
  inj = id
  prj = Just

instance  {-# OVERLAPPING #-} Member f (Sum f g) where
  inj = InL
  prj = foldSum Just (const Nothing)

instance  {-# OVERLAPPING #-} (Member f g) => Member f (Sum h g) where
  inj = InR . inj
  prj = foldSum (const Nothing) prj



-- newtype Freer f = Free . Coyoneda f
newtype Freer f a = Freer { runFreer :: ((Free :.. Coyoneda) f) a }
  deriving (Functor, Applicative, Monad)

instance MonadTrans Freer where
  lift = Freer . lift . liftCoyoneda

instance Member a f => a :< Freer f where
  liftTo = liftFreer . inj

liftFreer :: g ~> Freer g
liftFreer = Freer . Free.liftF . liftCoyoneda

foldFreer :: Monad g => (f ~> g) -> Freer f ~> g
foldFreer f = Free.foldFree (foldCoyoneda f) . runFreer

foldCoyoneda :: Functor g => (f ~> g) -> Coyoneda f ~> g
foldCoyoneda f = lowerCoyoneda . hoistCoyoneda f



-- newtype Cofreer' f = Cofree . Coyoneda f
newtype Cofreer' f a = Cofreer' { runCofreer' :: Cofree (Coyoneda f) a }
  deriving (Functor)

instance Comonad (Cofreer' f) where
  extract = extract . runCofreer'
  duplicate = Cofreer' . fmap Cofreer' . duplicate . runCofreer'



newtype Zap f g = Zap { zapWith :: forall a b c. (a -> b -> c) -> f a -> g b -> c }

zap :: Zap f g -> f (a -> b) -> g a -> b
zap z = zapWith z id

zapPair :: Zap f g -> f a -> g b -> (a, b)
zapPair z = zapWith z (,)

toZap :: Functor f => (forall a b . f (a -> b) -> g a -> b) -> Zap f g
toZap v = Zap $ \u f g -> v (u <$> f) g

flipZap :: Zap f g -> Zap g f
flipZap (Zap z) = Zap $ \f a b -> z (flip f) b a

contramapZap :: (f ~> g) -> (h ~> i) -> Zap g i -> Zap f h
contramapZap fg hi (Zap z) = Zap $ \u f h -> z u (fg f) (hi h)

strength :: Functor f => a -> f b -> f (a, b)
strength = fmap . (,)

zapAdjunctionL :: Adjunction f g => Zap f g
zapAdjunctionL = flipZap zapAdjunctionR

zapAdjunctionR :: Adjunction f g => Zap g f
zapAdjunctionR = Zap $ \f a b -> uncurry (flip f) $ rightAdjunct (uncurry (flip strength)) $ strength a b

prodZ :: Zap f h -> Zap g i -> Zap (f :+: g) (h :*: i)
prodZ (Zap fh) (Zap gi) = Zap $ \z fg (h `Pair` i) -> case fg of
  InL f -> fh z f h
  InR g -> gi z g i

coprodZ :: Zap f h -> Zap g i -> Zap (f :*: g) (h :+: i)
coprodZ f g = flipZap $ prodZ (flipZap f) (flipZap g)

idZ :: Zap Identity Identity
idZ = zapAdjunctionR

constZ :: Zap (Const Void) (Const a)
constZ = Zap $ \r (Const v) _ -> r (absurd v) (absurd v)

newtype Handles e a = Handles (e -> Const Void a)

handlesZ :: Zap (Handles e) (Throws e)
handlesZ = Zap $ \r (Handles h) (Throw e) -> r (absurd $ getConst $ h e) (absurd $ getConst $ h e)

curryZ :: Zap (Reads e) (Writes e)
curryZ = Zap $ \z Read (Write x) -> z x ()

composeZap :: Zap f g -> Zap h i -> Zap (f :.: h) (g :.: i)
composeZap (Zap u) (Zap v) = Zap $ \f (Compose a) (Compose b) -> u (v f) a b

freeZ :: forall f g . Zap f g -> Zap (Free f) (Cofree g)
freeZ (Zap zap) = Zap pair
  where
    pair :: forall a b c. (a -> b -> c) -> (Free f) a -> (Cofree g) b -> c
    pair p (Pure x)  (a :< _)  = p x a
    pair p (Free gs) (_ :< fs) = zap (pair p) gs fs

coyonedaZ :: Zap f g -> Zap (Coyoneda f) (Coyoneda g)
coyonedaZ (Zap fg) = Zap $ \u (Coyoneda f fs) (Coyoneda g gs) -> fg (\x y -> u (f x) (g y)) fs gs

freerZ :: forall f g . Zap f g -> Zap (Freer f) (Cofreer' g)
freerZ = contramapZap runFreer runCofreer' . zapX
  where
    zapX :: Zap f g -> Zap ((Free :.. Coyoneda) f) ((Cofree :.. Coyoneda) g)
    zapX  = freeZ . coyonedaZ


class HasCo (f :: Type -> Type) where
  type Co f = (cf :: Type -> Type) | cf -> f

-- TODO: Co (State x) = Store x
-- type State s a = Reader s . Writer s
-- type Store s a = Writer s . Reader s

instance HasCo (Reads x) where
  type Co (Reads x) = Writes x

instance HasCo (Writes x) where
  type Co (Writes x) = Reads x

instance HasCo (Throws x) where
  type Co (Throws x) = Handles x

instance (HasCo f, HasCo g) => HasCo (f :+: g) where
  type Co (f :+: g) = Co f :*: Co g

instance (HasCo f, HasCo g) => HasCo (f :*: g) where
  type Co (f :*: g) = Co f :+: Co g


class Functor (Lift f) => Liftable (f :: Type -> Type) where
  type Lift f :: Type -> Type
  prom :: f a -> Lift f a


instance Liftable (Reads e) where
  type Lift (Reads e) = Coyoneda (Reads e)
  prom = liftCoyoneda

instance Liftable (Writes e) where
  type Lift (Writes e) = Coyoneda (Writes e)
  prom = liftCoyoneda

instance Liftable (Throws e) where
  type Lift (Throws e) = Coyoneda (Throws e)
  prom = liftCoyoneda

instance Liftable (Handles e) where
  type Lift (Handles e) = Coyoneda (Handles e)
  prom = liftCoyoneda

instance (Liftable f, Liftable g) => Liftable (f :.: g) where
  type Lift (f :.: g) = Lift f :.: Lift g
  prom (Compose x) = Compose (fmap prom $ prom x)

instance (Liftable f, Liftable g) => Liftable (f :*: g) where
  type Lift (f :*: g) = Lift f :*: Lift g
  prom (f `Pair` g) = Pair (prom f) (prom g)

instance (Liftable f, Liftable g) => Liftable (f :+: g) where
  type Lift (f :+: g) = Lift f :+: Lift g
  prom (InL a) = InL (prom a)
  prom (InR a) = InR (prom a)



instance Liftable f => Liftable (Freer f) where
  type Lift (Freer f) = Free (Lift f)
  prom (Freer f) = Free.hoistFree g f where g = lowerCoyoneda . hoistCoyoneda prom



class HasZap f g where
  dazap :: Zap f g

instance HasZap f g => HasZap (Coyoneda f) (Coyoneda g) where
  dazap = coyonedaZ dazap

instance (HasZap f h, HasZap g i) => HasZap (f :+: g) (h :*: i) where
  dazap = prodZ dazap dazap

instance (HasZap f g, HasZap h i) => HasZap (f :.: h) (g :.: i) where
  dazap = composeZap dazap dazap

instance HasZap (Writes e) (Reads e) where
  dazap = flipZap curryZ

instance HasZap (Reads e) (Writes e) where
  dazap = curryZ

instance HasZap (Throws e) (Handles e) where
  dazap = flipZap handlesZ

instance HasZap (Throws e) (Const Void) where
  dazap = Zap $ \u (Throw x) (Const v) -> absurd v

type Colift f      = Lift (Co f)
type Program f     = Freer f
type Interpreter f = Cofree (Colift f)

interpret ::
  ( Liftable f
  , HasZap (Lift f) (Colift f)
  )
  => Program f a
  -> Interpreter f b
  -> (a, b)
interpret f = zapWith (freeZ dazap) (,) (prom f)


ex4 :: (Monad f, Reads [String] :< f, Writes String :< f) => f Int
ex4 = do
  write $ "foo"
  write $ "bar"
  msgs <- read @[String]
  pure $ length $ concat msgs

type F = Writes String :+: Reads [String]

i1 :: Interpreter F ()
i1 = () :< (fmap (const i1) (liftCoyoneda Read) `Pair` fmap (const i1) (liftCoyoneda (Write [])))


i3 :: Interpreter F [String]
i3 = go []
  where
    go xs = xs :< fmap go (fmap (\l -> xs ++ [l]) (liftCoyoneda Read) `Pair` fmap (const xs) (liftCoyoneda (Write xs)))

i4 :: Interpreter F [String]
i4 = go []
  where
    go xs = xs :< fmap go (fmap (: xs) (liftCoyoneda Read) `Pair` fmap (const xs) (liftCoyoneda (Write xs)))

exB :: (Monad f, Writes Integer :< f, Throws String :< f) => f ()
exB = do
  write 22
  write 23
  throw "Bad!"

type G = Writes Integer :+: Throws String

i5 :: Interpreter G Integer
i5 = go 0
  where
    go s = s :< fmap go ((fmap (+ s) . liftCoyoneda) Read `Pair` liftCoyoneda (Handles (\e -> Const $ error $ "Failed: " <> e)))

main = do
  print $ interpret @F ex4 i1
  -- >> (0, ())

  print $ interpret @F ex4 i3
  -- >> (6,["foo","bar"])

  print $ interpret @F ex4 i4
  -- >> (6,["bar","foo"])

  print $ interpret @G exB i5
  -- >> (*** Exception: Failed: Bad!