Adjoint folds

adjointFold.hs

-- http://www.cs.ox.ac.uk/ralf.hinze/SSGIP10/AdjointFolds.pdf
-- http://www.cs.ox.ac.uk/ralf.hinze/publications/SCP-78-11.pdf
-- https://www.cs.ox.ac.uk/people/nicolas.wu/papers/URS.pdf

{-# LANGUAGE
    MultiParamTypeClasses
  , FunctionalDependencies
  , GADTs
  , PolyKinds
  , DataKinds
  , ConstraintKinds
  , QuantifiedConstraints
  , TypeFamilies
  , RankNTypes
  , FlexibleContexts
  , FlexibleInstances
  , UndecidableInstances
  , DeriveFunctor
  , TypeOperators
  , LambdaCase
  , ScopedTypeVariables
  #-}
import Prelude hiding (Functor, (.), fmap, id, mapM, map)
import qualified Prelude as P

import Control.Category

import Control.Applicative (Const(..))
import Control.Arrow ((&&&), (***), (|||), (+++), Kleisli(..), arr)
import Control.Monad (liftM)
import Control.Monad.Free (Free(..), wrap)
import Control.Comonad (Comonad(..), Cokleisli(..), (=>>))
import Control.Comonad.Cofree (Cofree(..), unwrap)
import Data.Distributive
import Data.Traversable
import Data.Functor.Identity
import Data.Functor.Kan.Lan
import Data.Functor.Kan.Ran
import Data.Functor.Rep
import Data.Kind (Constraint, Type)


class (Category dom, Category cod) => Functor f dom cod | f -> dom cod where
  fmap :: dom a b -> cod (f a) (f b)

type family Obj (c :: k -> k -> Type) (a :: k) :: Constraint
type instance Obj (->) a = ()
type instance Obj (Kleisli m) a = ()
type instance Obj (Cokleisli w) a = ()

class Category k => FunctorDefault (k :: x -> x -> Type) where
  type IsFunctor k (f :: x -> x) :: Constraint
  map :: IsFunctor k f => k a b -> k (f a) (f b)

class FunctorDefault k => Fixable (k :: x -> x -> Type) where
  type Fix k (f :: x -> x) :: x
  fix :: k (f (Fix k f)) (Fix k f)
  unfix :: k (Fix k f) (f (Fix k f))

class (Functor f d c, Functor u c d) => Adjunction (f :: y -> x) (u :: x -> y) (c :: x -> x -> Type) (d :: y -> y -> Type) | f -> u c d, u -> f c d where
  leftAdjunct  :: Obj d a => c (f a) b -> d a (u b)
  rightAdjunct :: Obj c b => d a (u b) -> c (f a) b
  unit :: Obj d a => d a (u (f a))
  unit = leftAdjunct id
  counit :: Obj c b => c (f (u b)) b
  counit = rightAdjunct id


hylo :: (Fixable k, IsFunctor k f) => k (f b) b -> k a (f a) -> k a b
hylo alg coa = c where c = alg . map c . coa

ana  :: (Fixable k, IsFunctor k f) => k a (f a) -> k a (Fix k f)
ana  coa = hylo fix coa

cata :: (Fixable k, IsFunctor k f) => k (f b) b -> k (Fix k f) b
cata alg = hylo alg unfix


adjointFold :: (Adjunction l r c d, Fixable d, IsFunctor d f, Obj c b, Obj d (f (r b)))
            => c (l (f (r b))) b -> c (l (Fix d f)) b
adjointFold = rightAdjunct . cata . leftAdjunct

adjointUnfold :: (Adjunction l r d c, Fixable d, IsFunctor d f, Obj c a, Obj d (f (l a)))
              => c a (r (f (l a))) -> c a (r (Fix d f))
adjointUnfold = leftAdjunct . ana . rightAdjunct

adjointRefold :: (Adjunction l m d c, Adjunction m r c d, Fixable d, IsFunctor d f, Obj c a, Obj d (f (l a)), Obj c b, Obj d (f (r b)))
              => c (m (f (r b))) b -> c a (m (f (l a))) -> c a b
adjointRefold alg coalg = rightAdjunct (hylo (leftAdjunct alg) (rightAdjunct coalg)) . unit



newtype FixF f = In { out :: f (FixF f) }

instance Fixable (->) where
  type Fix (->) f = FixF f
  fix = In
  unfix = out

instance FunctorDefault (->) where
  type IsFunctor (->) f = P.Functor f
  map = P.fmap


instance Functor Identity (->) (->) where
  fmap = P.fmap

instance Adjunction Identity Identity (->) (->) where
  leftAdjunct f = Identity . f . Identity
  rightAdjunct f = runIdentity . f . runIdentity

fold :: P.Functor f => (f b -> b) -> FixF f -> b
fold alg = adjointFold (alg . P.fmap runIdentity . runIdentity) . Identity

unfold :: P.Functor f => (a -> f a) -> a -> FixF f
unfold coalg = runIdentity . adjointUnfold (Identity . P.fmap Identity . coalg)

refold :: P.Functor f => (f b -> b) -> (a -> f a) -> a -> b
refold alg coalg = adjointRefold (alg . P.fmap runIdentity . runIdentity) (Identity . P.fmap Identity . coalg)


instance Functor ((->) e) (->) (->) where
  fmap = P.fmap
instance Functor ((,) e) (->) (->) where
  fmap = P.fmap

instance Adjunction ((,) e) ((->) e) (->) (->) where
  leftAdjunct f a e = f (e, a)
  rightAdjunct f (e, a) = index (f a) e

paramFold :: P.Functor f => (a -> f (a -> b) -> b) -> a -> FixF f -> b
paramFold f = curry (adjointFold (uncurry f))

paramUnfold :: P.Functor f => (b -> a -> f (a, b)) -> b -> a -> FixF f
paramUnfold = adjointUnfold


instance Monad m => Fixable (Kleisli m) where
  type Fix (Kleisli m) f = FixF f
  fix = arr In
  unfix = arr out

instance Monad m => FunctorDefault (Kleisli m) where
  type IsFunctor (Kleisli m) f = Traversable f
  map = Kleisli . mapM . runKleisli

newtype KleisliAdjF (m :: Type -> Type) a = KAF { unKAF :: a }
instance Monad m => Functor (KleisliAdjF m) (->) (Kleisli m) where
  fmap f = Kleisli (return . KAF . f . unKAF)

newtype KleisliAdjG m a = KAG { unKAG :: m a }
instance Monad m => Functor (KleisliAdjG m) (Kleisli m) (->) where
  fmap (Kleisli f) = KAG . (>>= f) . unKAG

instance Monad m => Adjunction (KleisliAdjF m) (KleisliAdjG m) (Kleisli m) (->) where
  leftAdjunct (Kleisli f) = KAG . f . KAF
  rightAdjunct f = Kleisli $ unKAG . f . unKAF

unfoldM :: (Traversable f, Monad m) => (b -> m (f b)) -> b -> m (FixF f)
unfoldM coalg = unKAG . adjointUnfold (KAG . liftM (P.fmap KAF) . coalg)


instance Comonad w => Fixable (Cokleisli w) where
  type Fix (Cokleisli w) f = FixF f
  fix = arr In
  unfix = arr out

instance Comonad w => FunctorDefault (Cokleisli w) where
  type IsFunctor (Cokleisli w) t = Distributive t
  map = Cokleisli . cotraverse . runCokleisli

newtype CokleisliAdjF (w :: Type -> Type) a = CoKAF { unCoKAF :: a }
instance Comonad w => Functor (CokleisliAdjF w) (->) (Cokleisli w) where
  fmap f = Cokleisli (CoKAF . f . unCoKAF . extract)

newtype CokleisliAdjG w a = CoKAG { unCoKAG :: w a }
instance Comonad w => Functor (CokleisliAdjG w) (Cokleisli w) (->) where
  fmap (Cokleisli f) = CoKAG . (=>> f) . unCoKAG

instance Comonad w => Adjunction (CokleisliAdjG w) (CokleisliAdjF w) (->) (Cokleisli w) where
  leftAdjunct f = Cokleisli $ CoKAF . f . CoKAG
  rightAdjunct (Cokleisli f) = unCoKAF . f . unCoKAG

foldW :: (Distributive f, Comonad w) => (w (f b) -> b) -> w (FixF f) -> b
foldW alg = adjointFold (alg . P.fmap (P.fmap unCoKAF) . unCoKAG) . CoKAG



newtype Nat f g = Nat { unNat :: forall a. f a -> g a }
type instance Obj Nat a = P.Functor a
instance Category Nat where
  id = Nat id
  Nat f . Nat g = Nat (f . g)

instance FunctorDefault Nat where
  type IsFunctor Nat f = Functor f Nat Nat
  map = fmap

instance Functor (Ran j) Nat Nat where
  fmap (Nat f) = Nat $ \(Ran h) -> Ran (f . h)
instance Functor (Lan j) Nat Nat where
  fmap (Nat f) = Nat $ \(Lan h g) -> Lan h (f g)

newtype Pre j f a = Pre { unPre :: f (j a) }
instance Functor (Pre j) Nat Nat where
  fmap (Nat f) = Nat $ \(Pre fja) -> Pre (f fja)

instance Adjunction (Pre j) (Ran j) Nat Nat where
  leftAdjunct (Nat n) = Nat $ \fb -> Ran $ \bjx -> n (Pre (bjx <$> fb))
  rightAdjunct (Nat n) = Nat $ \(Pre fja) -> runRan (n fja) id

instance Adjunction (Lan j) (Pre j) Nat Nat where
  leftAdjunct (Nat n) = Nat $ Pre . n . Lan id
  rightAdjunct (Nat n) = Nat $ \(Lan jyx ay) -> jyx <$> unPre (n ay)

newtype FixH f a = InH { outH :: f (FixH f) a }
instance P.Functor (f (FixH f)) => P.Functor (FixH f) where
  fmap f = InH . P.fmap f . outH

instance Fixable Nat where
  type Fix Nat f = FixH f
  fix = Nat InH
  unfix = Nat outH

gnestedFold :: (Functor f Nat Nat, forall i. P.Functor i => P.Functor (f i), P.Functor b)
            => (forall y. f (Ran j b) (j y) -> b y) -> FixH f (j x) -> b x
gnestedFold alg = unNat (adjointFold (Nat $ alg . unPre)) . Pre

gnestedUnfold :: (Functor f Nat Nat, forall i. P.Functor i => P.Functor (f i), P.Functor a)
              => (forall y. a y -> f (Lan j a) (j y)) -> a x -> FixH f (j x)
gnestedUnfold coalg = unPre . unNat (adjointUnfold (Nat $ Pre . coalg))

gnestedRefold :: (Functor f Nat Nat, forall i. P.Functor i => P.Functor (f i), P.Functor a, P.Functor b)
              => (forall y. f (Ran j b) (j y) -> b y) -> (forall y. a y -> f (Lan j a) (j y)) -> a x -> b x
gnestedRefold alg coalg = unNat (adjointRefold (Nat $ alg . unPre) (Nat $ Pre . coalg))

newtype Rsh j b x = Rsh { unRsh :: (x -> j) -> b }
ran2rsh :: Nat (Ran (Const a) (Const b)) (Rsh a b)
ran2rsh = Nat $ \(Ran f) -> Rsh (\g -> getConst (f (Const . g)))

nestedFold :: (Functor f Nat Nat, forall i. P.Functor i => P.Functor (f i), P.Functor (f (FixH f)))
           => (f (Rsh j b) j -> b) -> FixH f j -> b
nestedFold alg = getConst . gnestedFold (Const . alg . unNat (fmap ran2rsh) . P.fmap getConst) . P.fmap Const

data Lsh j a x = Lsh (j -> x) a
lsh2lan :: Nat (Lsh j a) (Lan (Const j) (Const a))
lsh2lan = Nat $ \(Lsh f a) -> Lan (f . getConst) (Const a)
nestedUnfold :: (Functor f Nat Nat, forall i. P.Functor i => P.Functor (f i), P.Functor (f (FixH f)))
             => (a -> f (Lsh j a) j) -> a -> FixH f j
nestedUnfold coalg = P.fmap getConst . gnestedUnfold (P.fmap Const . unNat (fmap lsh2lan) . coalg . getConst) . Const

nestedRefold :: (Functor f Nat Nat, forall i. P.Functor i => P.Functor (f i))
             => (f (Rsh j b) j -> b) -> (a -> f (Lsh j a) j) -> a -> b
nestedRefold alg coalg =
  getConst . gnestedRefold
      (Const . alg . unNat (fmap ran2rsh) . P.fmap getConst)
      (P.fmap Const . unNat (fmap lsh2lan) . coalg . getConst)
  . Const

data Pair a = Pair a a deriving P.Functor
data PerfectF r a = Zero a | Succ (r (Pair a)) deriving P.Functor
type Perfect = FixH PerfectF
instance Functor PerfectF Nat Nat where
  fmap n = Nat $ \case
    Zero a -> Zero a
    Succ t -> Succ $ unNat n t

instance Show a => Show (Perfect a) where
  show = nestedFold f . P.fmap show where
    f (Zero i) = i
    f (Succ (Rsh g)) = g (\(Pair a b) -> "(" ++ a ++ ", " ++ b ++ ")")

fromDepth :: Int -> Perfect Int
fromDepth = nestedUnfold f where
  f 0 = Zero 0
  f b = Succ (Lsh (\j -> Pair (2 * j) (2 * j + 1)) (b - 1))



class Comonad w => Coalgebra w a where
  coalgebra :: a -> w a
newtype CoEMMor (w :: Type -> Type) a b = CoEMMor (a -> b)
type instance Obj (CoEMMor w) a = Coalgebra w a
instance Category (CoEMMor w) where
  id = CoEMMor id
  CoEMMor f . CoEMMor g = CoEMMor (f . g)

instance Fixable (CoEMMor w) where
  type Fix (CoEMMor w) f = FixF f
  fix = CoEMMor In
  unfix = CoEMMor out

instance FunctorDefault (CoEMMor w) where
  type IsFunctor (CoEMMor w) f = IsFunctor (->) f
  map (CoEMMor f) = CoEMMor (map f)


newtype CofreeCoalg w a = CofreeCoalg { unCofreeCoalg :: w a } deriving P.Functor
instance Comonad w => Functor (CofreeCoalg w) (->) (CoEMMor w) where
  fmap f = CoEMMor (P.fmap f)

newtype ForgetCoalg w a = ForgetCoalg { unForgetCoalg :: a } deriving P.Functor
instance Comonad w => Functor (ForgetCoalg w) (CoEMMor w) (->) where
  fmap (CoEMMor f) = P.fmap f

instance Comonad w => Adjunction (ForgetCoalg w) (CofreeCoalg w) (->) (CoEMMor w) where
  leftAdjunct f = CoEMMor $ CofreeCoalg . P.fmap (f . ForgetCoalg) . coalgebra
  rightAdjunct (CoEMMor f) = extract . unCofreeCoalg . f . unForgetCoalg

gcata :: (Coalgebra w (f (CofreeCoalg w b)), P.Functor f) => (f (w b) -> b) -> FixF f -> b
gcata alg = adjointFold (alg . P.fmap unCofreeCoalg . unForgetCoalg) . ForgetCoalg

distHisto :: P.Functor f => f (Cofree f a) -> Cofree f (f a)
distHisto fcfa = (extract <$> fcfa) :< (distHisto . unwrap <$> fcfa)

instance P.Functor f => Coalgebra (Cofree f) (f (CofreeCoalg (Cofree f) b)) where
  coalgebra = distHisto . P.fmap (extend CofreeCoalg . unCofreeCoalg)

histo :: P.Functor f => (f (Cofree f b) -> b) -> FixF f -> b
histo = gcata


class (Monad m, P.Functor m) => Algebra m a where
  algebra :: m a -> a
newtype EMMor (m :: Type -> Type) a b = EMMor (a -> b)
type instance Obj (EMMor m) a = Algebra m a
instance Category (EMMor m) where
  id = EMMor id
  EMMor f . EMMor g = EMMor (f . g)

instance Fixable (EMMor m) where
  type Fix (EMMor m) f = FixF f
  fix = EMMor In
  unfix = EMMor out

instance FunctorDefault (EMMor m) where
  type IsFunctor (EMMor m) f = IsFunctor (->) f
  map (EMMor f) = EMMor (map f)

newtype FreeAlg m a = FreeAlg { unFreeAlg :: m a } deriving P.Functor
instance P.Functor m => Functor (FreeAlg m) (->) (EMMor m) where
  fmap f = EMMor (P.fmap f)

newtype ForgetAlg m a = ForgetAlg { unForgetAlg :: a } deriving P.Functor
instance Functor (ForgetAlg m) (EMMor m) (->) where
  fmap (EMMor f) = P.fmap f

instance (Monad m, P.Functor m) => Adjunction (FreeAlg m) (ForgetAlg m) (EMMor m) (->) where
  leftAdjunct (EMMor f) = ForgetAlg . f . FreeAlg . return
  rightAdjunct f = EMMor $ algebra . P.fmap (unForgetAlg . f) . unFreeAlg

gana :: (Algebra m (f (FreeAlg m b)), P.Functor f) => (b -> f (m b)) -> b -> FixF f
gana coalg = unForgetAlg . adjointUnfold (ForgetAlg . P.fmap FreeAlg . coalg)

distFutu :: P.Functor f => Free f (f a) -> f (Free f a)
distFutu (Pure fa) = return <$> fa
distFutu (Free fmfa) = wrap . distFutu <$> fmfa

instance P.Functor f => Algebra (Free f) (f (FreeAlg (Free f) b)) where
  algebra = P.fmap (FreeAlg . (>>= unFreeAlg)) . distFutu

futu :: P.Functor f => (b -> f (Free f b)) -> b -> FixF f
futu = gana


{-
newtype DiEMMor (m :: Type -> Type) (w :: Type -> Type) a b = DiEMMor (a -> b)
type instance Obj (DiEMMor m w) a = (Algebra m a, Coalgebra w a)
instance Category (DiEMMor m w) where
  id = DiEMMor id
  DiEMMor f . DiEMMor g = DiEMMor (f . g)

instance Fixable (DiEMMor m w) where
  type Fix (DiEMMor m w) f = FixF f
  fix = DiEMMor In
  unfix = DiEMMor out

instance FunctorDefault (DiEMMor m w) where
  type IsFunctor (DiEMMor m w) f = IsFunctor (->) f
  map (DiEMMor f) = DiEMMor (map f)

newtype FreeDialg m w a = FreeDialg { unFreeDialg :: m a } deriving P.Functor
instance P.Functor m => Functor (FreeDialg m w) (->) (DiEMMor m w) where
  fmap f = DiEMMor (P.fmap f)
instance P.Functor f => Coalgebra (Cofree f) (f (FreeDialg (Free f) (Cofree f) b)) where
  coalgebra fma = fma :< _

newtype CofreeDialg m w a = CofreeDialg { unCofreeDialg :: w a } deriving P.Functor
instance P.Functor w => Functor (CofreeDialg m w) (->) (DiEMMor m w) where
  fmap f = DiEMMor (P.fmap f)

newtype ForgetDialg m w a = ForgetDialg { unForgetDialg :: a } deriving P.Functor
instance Functor (ForgetDialg m w) (DiEMMor m w) (->) where
  fmap (DiEMMor f) = P.fmap f

instance (Monad m, P.Functor m) => Adjunction (FreeDialg m w) (ForgetDialg m w) (DiEMMor m w) (->) where
  leftAdjunct (DiEMMor f) = ForgetDialg . f . FreeDialg . return
  rightAdjunct f = DiEMMor $ algebra . P.fmap (unForgetDialg . f) . unFreeDialg

instance Comonad w => Adjunction (ForgetDialg m w) (CofreeDialg m w) (->) (DiEMMor m w) where
  leftAdjunct f = DiEMMor $ CofreeDialg . P.fmap (f . ForgetDialg) . coalgebra
  rightAdjunct (DiEMMor f) = extract . unCofreeDialg . f . unForgetDialg

ghylo :: forall m w f a b. (Algebra m (f (FreeDialg m w a)), Coalgebra w (f (CofreeDialg m w b)), P.Functor f)
      => (f (w b) -> b) -> (a -> f (m a)) -> a -> b
ghylo alg coalg = adjointRefold alg' coalg'
  where
    alg' :: ForgetDialg m w (f (CofreeDialg m w b)) -> b
    alg' = alg . P.fmap unCofreeDialg . unForgetDialg
    coalg' :: a -> ForgetDialg m w (f (FreeDialg m w a))
    coalg' = ForgetDialg . P.fmap FreeDialg . coalg
-}


data Two = P0 | P1
type (:**:) :: (x -> x -> Type) -> (y -> y -> Type) -> (Two -> x) -> (Two -> y) -> Type
data (:**:) k0 k1 a b where
  (:**:) :: k0 (a 'P0) (b 'P0) -> k1 (a 'P1) (b 'P1) -> (k0 :**: k1) a b

type instance Obj (k0 :**: k1) a = (Obj k0 (a 'P0), Obj k1 (a 'P1))
instance (Category k0, Category k1) => Category (k0 :**: k1) where
  id = id :**: id
  (f0 :**: f1) . (g0 :**: g1) = (f0 . g0) :**: (f1 . g1)

newtype Prod a = Prod { unProd :: (a 'P0, a 'P1) }
instance Functor Prod ((->) :**: (->)) (->) where
  fmap (f :**: g) = Prod . (f *** g) . unProd

newtype Coprod a = Coprod { unCoprod :: Either (a 'P0) (a 'P1) }
instance Functor Coprod ((->) :**: (->)) (->) where
  fmap (f :**: g) = Coprod . (f +++ g) . unCoprod

newtype Diag a (b :: Two) = Diag { unDiag :: a }
instance Functor Diag (->) ((->) :**: (->)) where
  fmap f = (\(Diag a) -> Diag (f a)) :**: (\(Diag a) -> Diag (f a))

instance Adjunction Diag Prod ((->) :**: (->)) (->) where
  leftAdjunct (fab1 :**: fab2) = Prod . (fab1 . Diag &&& fab2 . Diag)
  rightAdjunct f = (fst . unProd . f . unDiag) :**: (snd . unProd . f . unDiag)

instance Adjunction Coprod Diag (->) ((->) :**: (->)) where
  leftAdjunct f = (Diag . f . Coprod . Left) :**: (Diag . f . Coprod . Right)
  rightAdjunct (fab1 :**: fab2) = (unDiag . fab1 ||| unDiag . fab2) . unCoprod

data Proj :: Type -> Type -> Two -> Type where
  Proj0 :: { unProj0 :: a } -> Proj a b 'P0
  Proj1 :: { unProj1 :: b } -> Proj a b 'P1

mutu :: forall f a b. P.Functor f => (f (a, b) -> a, f (a, b) -> b) -> (FixF f -> a, FixF f -> b)
mutu (algA, algB) = case adjointFold ((Proj0 . algA . conv) :**: (Proj1 . algB . conv)) of
    f :**: g -> (unProj0 . f . Diag, unProj1 . g . Diag)
  where
    conv = P.fmap ((unProj0 *** unProj1) . unProd) . unDiag

comutu :: P.Functor f => (a -> f (Either a b), b -> f (Either a b)) -> (a -> FixF f, b -> FixF f)
comutu (coalgA, coalgB) = case adjointUnfold ((conv . coalgA . unProj0) :**: (conv . coalgB . unProj1)) of
    f :**: g -> (unDiag . f . Proj0, unDiag . g . Proj1)
  where
    conv = Diag . P.fmap (Coprod . (Proj0 +++ Proj1))

zygo :: P.Functor f => (f b -> b) -> (f (a, b) -> a) -> FixF f -> a
zygo f alg = fst $ mutu (alg, f . P.fmap snd)

para :: P.Functor f => (f (a, FixF f) -> a) -> FixF f -> a
para = zygo In

apo :: P.Functor f => (a -> f (Either a (FixF f))) -> a -> FixF f
apo alg = fst $ comutu (alg, P.fmap Right . out)