Example for recursion schemes for mutually recursive data

IxFix.hs

{-# LANGUAGE
  UndecidableInstances, RankNTypes, TypeOperators, TypeFamilies,
  StandaloneDeriving, DataKinds, PolyKinds, DeriveFunctor, DeriveFoldable,
  DeriveTraversable, LambdaCase, PatternSynonyms, TemplateHaskell #-}

import Control.Monad
import Control.Applicative
import Data.Singletons.TH


-- type synonyms
--------------------------------------------------------------------------------
infixr 5 ~>
infixr 5 .~>
infixr 5 ~>.

type f ~>  g = forall i. f i -> g i  -- Natural transformation
type f .~> g = forall i. f -> g i   -- Constant on the left
type f ~>. g = forall i. f i -> g   -- Constant on the right

-- K and I
--------------------------------------------------------------------------------
newtype K a b = K a deriving
  (Eq, Show, Functor, Foldable, Traversable)

newtype I a = I a deriving
  (Eq, Show, Functor, Foldable, Traversable)

instance Applicative I where
  pure = I
  I f <*> I a = I (f a)

instance Monoid a => Applicative (K a) where
  pure _ = K mempty
  K a <*> K a' = K (mappend a a')

getK (K a) = a
getI (I a) = a

-- Indexed classses
--------------------------------------------------------------------------------

class IxFunctor (f :: (k -> *) -> (k -> *)) where
  imap :: (a ~> b) -> (f a ~> f b)

class IxFunctor t => IxTraversable t where
  -- the type would be "(a ~> (f . b)) -> (t a ~> (f . t b))" if we had the composition
  -- function on type level.
  itraverse :: Applicative f => (forall i. a i -> f (b i)) -> (forall i. t a i -> f (t b i))

class IxFoldable t where
  iFoldMap :: Monoid m => (a ~>. m) -> (t a ~>. m)

imapDefault :: IxTraversable t => (a ~> b) -> (t a ~> t b)
imapDefault f = getI . itraverse (I . f)

iFoldMapDefault :: (IxTraversable t, Monoid m) => (a ~>. m) -> (t a ~>. m)
iFoldMapDefault f = getK . itraverse (K . f)

-- Indexed fixpoint
--------------------------------------------------------------------------------
newtype IxFix f i = In (f (IxFix f) i)
out (In f) = f

deriving instance Show (f (IxFix f) t) => Show (IxFix f t)
deriving instance Eq (f (IxFix f) t) => Eq (IxFix f t)
deriving instance Ord (f (IxFix f) t) => Ord (IxFix f t)

-- Folds returning an indexed type
--------------------------------------------------------------------------------
cata :: IxFunctor f => (f a ~> a) -> (IxFix f ~> a)
cata phi = phi . imap (cata phi) . out

cataM :: (Monad m, IxTraversable t) => (forall i. t a i -> m (a i)) -> (forall i. IxFix t i -> m (a i))
cataM phi = phi <=< itraverse (cataM phi) . out

-- Folds returning constant types
--------------------------------------------------------------------------------
cata' :: IxFunctor f => (f (K a) ~>. a) -> (IxFix f ~>. a)
cata' phi = getK . cata (K . phi)

cataM' :: (Monad m, IxTraversable t) => (t (K a) ~>. m a) -> (IxFix t ~>. m a)
cataM' phi = phi <=< itraverse (fmap K . cataM' phi) . out


-- EXAMPLE TIME
--------------------------------------------------------------------------------

data ExprType = ExprTy | PatTy

data ExprF (k :: ExprType -> *) (i :: ExprType) where
  
  -- "Expr" tagged family of constructors
  EIntF :: Int -> ExprF k ExprTy
  ELamF :: k PatTy -> k ExprTy -> ExprF k ExprTy

  -- "PatTy" tagged familfy of constructors
  PAsF  :: String -> k ExprTy -> ExprF k PatTy

type Expr = IxFix ExprF

pattern EInt i        = In (EIntF i)
pattern ELam pat expr = In (ELamF pat expr)
pattern PAs str expr  = In (PAsF str expr)

deriving instance (Show (k PatTy), Show (k ExprTy)) => Show (ExprF k i)
deriving instance (Eq (k PatTy), Eq (k ExprTy))     => Eq (ExprF k i)

instance IxFunctor ExprF where
  imap = imapDefault
   
instance IxFoldable ExprF where
  iFoldMap = iFoldMapDefault

instance IxTraversable ExprF where
  itraverse f = \case
    EIntF i -> pure (EIntF i)
    ELamF pat expr -> ELamF <$> f pat <*> f expr
    PAsF str expr -> PAsF str <$> f expr
    
-- Some use cases
--------------------------------------------------------------------------------
expr = ELam (PAs "foo" (ELam (PAs "bar" (EInt 0)) (EInt 0))) (EInt 0)

-- Non-indexed return
----------------------------------------
patStrings :: Expr i -> [String] 
patStrings = cata' $ \case
  EIntF _ -> []
  PAsF str (K strs) -> str : strs
  ELamF (K strs) (K strs') -> strs ++ strs'

-- patStrings expr == ["foo","bar"]

-- indexed return
-- ----------------------------------------
incrementInts :: Expr i -> Expr i 
incrementInts = cata $ \case
  EIntF i -> EInt (i + 1)
  other   -> In other

-- incrementInts expr ==
--   In (ELamF (In (PAsF "foo" (In (ELamF (In (PAsF "bar" (In (EIntF 1)))) (In (EIntF 1)))))) (In (EIntF 1)))  


-- return with elimination on type tags (warning: rampant singletons usage!)
----------------------------------------

-- Reify TyFun application
data At (f :: TyFun k * -> *) (x :: k) where
  At :: f @@ x -> At f x

-- type level case
type family Case (cases :: [(key, val)]) (k :: key) where
  Case ('(k  , v) ': kvs) k = v
  Case ('(k' , v) ': kvs) k = Case kvs k
  Case '[]                k = Error "Case: key not found"   
-- Note: Lookup in "singletons" is bugged out as of 2.0.1
-- If we had Lookup, then we could have written Case the following way:
-- type family Case cases key where
--    Case cases key = FromJust (Lookup key cases)

$(genDefunSymbols [''Case])  

foldFoo :: Expr i -> At (CaseSym1 ['(ExprTy, Int), '(PatTy, Bool)]) i
foldFoo = cata $ \case
  EIntF i        -> At i
  PAsF str _     -> At False
  ELamF pat expr -> expr