Folds and unfolds using squares

FixSquares.hs

-- Using https://hackage.haskell.org/package/squares
{-# LANGUAGE GADTs #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE StandaloneKindSignatures #-}
module FixSquares where

import Data.Kind (Type)
import Data.Profunctor
import Data.Functor.Adjunction (Adjunction)
import Data.Functor.Adjunction.Square
import Data.Square
import Data.Functor.Square
import Data.Functor.Compose
import Data.Functor.Const


fromExtNat :: (Profunctor p, Functor x) => (forall a. p (x a) (x a)) -> Square '[] '[p] '[x] '[x]
fromExtNat p = mkSquare $ \f -> rmap (fmap f) p

-- +--x--+
-- |  v  |
-- |  @--p
-- |  v  |
-- +--x--+
type Alg' p x = Square '[] p x x
type Alg p x = Alg' p '[x]
type AlgArr' x y = Square '[] '[] x y
type AlgArr x y = AlgArr' '[x] '[y]


newtype Fix f a = Fix { runFix :: f (Fix f a) } deriving (Functor)
deriving instance Show (f (Fix f a)) => Show (Fix f a)

runFixSq :: Functor f => Square '[] '[] '[Fix f] '[Fix f, f]
runFixSq =
  funNat (Compose . runFix)
  ===
  fromCompose

fixSq :: Functor f => Square '[] '[] '[Fix f, f] '[Fix f]
fixSq =
  toCompose
  ===
  funNat (Fix . getCompose)


class HasInitialAlg p where
  type InitCarrier p :: [Type -> Type]
  initAlg :: Alg' p (InitCarrier p)
  isInit :: Functor y => Alg p y -> AlgArr' (InitCarrier p) '[y]

-- +----µ f----+
-- |     v     |
-- |     @     |
-- |   /   \   |
-- |  v     v  |
-- +-µ f-+--f--+
-- |  v  |  v  |
-- |  @  |  |  |
-- |  v  |  v  |
-- +--x--+--f--+
-- |  v  |  v  |
-- |  @-<f<-/  |
-- |  v  |     |
-- +--x--+-----+
cata :: Functor f => Alg '[Costar f] x -> AlgArr (Fix f) x
cata alg =
  runFixSq
  ===
  cata alg ||| funId
  ===
  alg ||| toLeft

fromPlainAlg :: Functor f => (f a -> a) -> Alg '[Costar f] (Const a)
fromPlainAlg alg = fromExtNat $ Costar $ Const . alg . fmap getConst

cataPlain :: Functor f => (f a -> a) -> Fix f () -> a
cataPlain = (getConst .) . (`runSquare` id) . cata . fromPlainAlg

instance Functor f => HasInitialAlg '[Costar f] where
  type InitCarrier '[Costar f] = '[Fix f]
  -- +--µ f---+
  -- +   v    |
  -- +   | /-<f
  -- +   \ /  |
  -- +    @   |
  -- +    v   |
  -- +---µ f--+
  initAlg =
    funId ||| fromRight
    ===
    fixSq
  isInit = cata


class HasTerminalAlg p where
  type TermCarrier p :: [Type -> Type]
  termAlg :: Alg' p (TermCarrier p)
  isTerm :: Alg p y -> AlgArr' '[y] (TermCarrier p)

-- +--x--+-----+
-- |  v  |     |
-- |  @->f>-\  |
-- |  v  |  v  |
-- +--x--+--f--+
-- |  v  |  v  |
-- |  @  |  |  |
-- |  v  |  v  |
-- +-µ f-+--f--+
-- |  v     v  |
-- |   \   /   |
-- |     @     |
-- |     v     |
-- +----µ f----+
ana :: Functor f => Alg '[Star f] x -> AlgArr x (Fix f)
ana coalg =
  coalg ||| fromLeft
  ===
  ana coalg ||| funId
  ===
  fixSq

fromPlainCoalg :: Functor f => (a -> f a) -> Alg '[Star f] (Const a)
fromPlainCoalg coalg = fromExtNat $ Star $ fmap Const . coalg . getConst

anaPlain :: Functor f => (a -> f a) -> a -> Fix f ()
anaPlain = (. Const) . (`runSquare` id) . ana . fromPlainCoalg

instance Functor f => HasTerminalAlg '[Star f] where
  type TermCarrier '[Star f] = '[Fix f]
  -- +---µ f--+
  -- +    v   |
  -- +    @   |
  -- +   / \  |
  -- +   | \->f
  -- +   v    |
  -- +--µ f---+
  termAlg =
    runFixSq
    ===
    funId ||| toRight
  isTerm = ana

-- +--x--+-----+
-- |  v  |     |
-- |  @->f>-\  |
-- |  v  |  v  |
-- +--x--+--f--+
-- |  v  |  v  |
-- |  @  |  |  |
-- |  v  |  v  |
-- +--y--+--f--+
-- |  v  |  v  |
-- |  @-<f<-/  |
-- |  v  |     |
-- +--y--+-----+
hylo :: Functor f => Alg '[Star f] x -> Alg '[Costar f] y -> AlgArr x y
hylo coalg alg =
  coalg ||| fromLeft
  ===
  hylo coalg alg ||| funId
  ===
  alg ||| toLeft

hyloPlain :: Functor f => (a -> f a) -> (f b -> b) -> a -> b
hyloPlain coalg alg = getConst . runSquare (hylo (fromPlainCoalg coalg) (fromPlainAlg alg)) id . Const




-- +--µ f--+-l-+
-- |   v   | v |
-- |   @   | | |
-- |  / \  | | |
-- | v   v | v |
-- +-x-+-r-+-l-+
-- | v | v   v |
-- | | | \-@-/ |
-- | v |       |
-- +-x-+-------+
adjointCata
  :: (Adjunction l r, Functor f, Functor x)
  => Alg '[Costar r, Costar f, Costar l] x
  -> AlgArr' '[Fix f, l] '[x]
adjointCata alg =
  (cata (adjointAlg alg)
   ===
   fromCompose) ||| funId
  ===
  funId ||| counit

-- +--x---+--------r--+
-- |  v   |        v  |
-- |  | /<r<-------/  |
-- |  |/  +-----------+
-- |  @--<f<---------<f
-- |  |\  +-----+-----+
-- |  | \<l<-@->r>-\  |
-- |  v   |     |  v  |
-- +--x---+-----+--r--+
adjointAlg
  :: (Adjunction l r, Functor f, Functor x)
  => Alg '[Costar r, Costar f, Costar l] x
  -> Alg '[Costar f] (Compose r x)
adjointAlg alg =
  fromCompose
  ===
  alg ||| (toLeft
            ===
            proId
            ===
            leftAdjunct ||| fromLeft)
  ===
  toCompose


-- +-x-+-------+
-- | v |       |
-- | | | /-@-\ |
-- | v | v   v |
-- +-x-+-l-+-r-+
-- | v   v | v |
-- |  \ /  | | |
-- |   @   | | |
-- |   v   | v |
-- +--µ f--+-r-+
adjointAna
  :: (Adjunction l r, Functor f, Functor x)
  => Alg '[Star r, Star f, Star l] x
  -> AlgArr' '[x] '[Fix f, r]
adjointAna coalg =
  funId ||| unit
  ===
  (toCompose
   ===
   ana (adjointCoalg coalg)) ||| funId

-- +--x---+-----+--l--+
-- |  v   |     |  v  |
-- |  | />r>-@-<l<-/  |
-- |  |/  +-----+-----+
-- |  @-->f>--------->f
-- |  |\  +-----------+
-- |  | \>l>-------\  |
-- |  v   |        v  |
-- +--x---+--------l--+
adjointCoalg
  :: (Adjunction l r, Functor f, Functor x)
  => Alg '[Star r, Star f, Star l] x
  -> Alg '[Star f] (Compose l x)
adjointCoalg coalg =
  fromCompose
  ===
  coalg ||| (rightAdjunct ||| toLeft
            ===
            proId
            ===
            fromLeft)
  ===
  toCompose


-- +-x-+-------+
-- | v |       |
-- | | | /-@-\ |
-- | v | v   v |
-- +-x-+-l-+-m-+
-- | v   v | v |
-- |  \ /  | | |
-- |   @   | | |
-- |  / \  | | |
-- | v   v | v |
-- +-x-+-r-+-m-+
-- | v | v   v |
-- | | | \-@-/ |
-- | v |       |
-- +-x-+-------+
adjointHylo
  :: (Adjunction l m, Adjunction m r, Functor f, Functor x, Functor y)
  => Alg '[Star m, Star f, Star l] x
  -> Alg '[Costar r, Costar f, Costar m] y
  -> AlgArr x y
adjointHylo coalg alg =
  funId ||| unit
  ===
  (toCompose
   ===
   hylo (adjointCoalg coalg) (adjointAlg alg)
   ===
   fromCompose) ||| funId
  ===
  funId ||| counit




toSquare :: Functor x => x () -> Square '[] '[] '[] '[x]
toSquare x = mkSquare $ \f b -> f b <$ x

fromSquare :: Square '[] '[] '[] '[x] -> x ()
fromSquare sq = runSquare sq id ()

toConstSquare :: a -> Square '[] '[] '[] '[Const a]
toConstSquare a = toSquare $ Const a

fromConstSquare :: Square '[] '[] '[] '[Const a] -> a
fromConstSquare = getConst . fromSquare


data NatF n = Z | S n deriving (Functor, Show)
type Nat = Fix NatF ()

toIntAlg :: Alg '[Costar NatF] (Const Int)
toIntAlg = fromPlainAlg $ \cases
  Z -> 0
  (S n) -> 1 + n

toInt :: Nat -> Int
toInt n = fromConstSquare $
  toSquare n
  ===
  cata toIntAlg

fromIntCoalg :: Alg '[Star NatF] (Const Int)
fromIntCoalg = fromPlainCoalg $ \cases
  0 -> Z
  n -> S (n - 1)

fromInt :: Int -> Nat
fromInt i = fromSquare $
  toConstSquare i
  ===
  ana fromIntCoalg