Mutually defined data types with recursion principle

MutuallyDefined.hs

{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeOperators #-}
 
module MutuallyDefined where

import Data.Maybe (fromJust)

-- Really we want to only quantify over whatever the kind of the
-- argument to f and g is, but I don't think there's a sane way to do
-- that in Haskell.
newtype f ~> g = Nat { app :: forall i . f i -> g i }

class Functor1 k where
  fmap1 :: (f ~> g) -> (k f ~> k g)

newtype Tie (f :: (k -> *) -> (k -> *)) (i :: k) =
  Tie { unTie :: f (Tie f) i }

tie :: Functor1 k => (k r ~> r) -> (Tie k ~> r)
tie phi = go where go = Nat $ app phi . app (fmap1 go) . unTie

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

newtype Zero f x (i :: ()) = Zero { unZero :: f (x i) }

newtype Tie0 (f :: * -> *) = Tie0 { unTie0 :: Tie (Zero f) '() }

tie0 :: Functor f => (f r -> r) -> (Tie0 f -> r)
tie0 phi = phi . fmap (tie0 phi . Tie0) . unZero . unTie . unTie0

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

data Count = One | Many
 
data ForestF l n (x :: Count -> *) (idx :: Count) where
  Leaf  :: l -> ForestF l n x One
  Node  :: n -> x Many -> ForestF l n x One
  Nil   :: ForestF l n x Many
  Cons  :: x One -> x Many -> ForestF l n x Many

instance Functor1 (ForestF l n) where
  fmap1 f = Nat $ \case
    Leaf l     -> Leaf l
    Node n a   -> Node n (app f a)
    Nil        -> Nil
    Cons a1 a2 -> Cons (app f a1) (app f a2)
                   
type Tree   l n = Tie (ForestF l n) One
type Forest l n = Tie (ForestF l n) Many
 
leaf :: l -> Tree l n
leaf = Tie . Leaf
 
node :: n -> Forest l n -> Tree l n
node n = Tie . Node n

nil :: Forest l n
nil = Tie Nil

cons :: Tree l n -> Forest l n -> Forest l n
cons a x = Tie (Cons a x)

newtype K a b = K { getK :: a } deriving (Show, Functor)

count :: ForestF l n (K Int) ~> K Int
count = Nat $ \case
  Leaf _           -> K 1
  Node _ (K n)     -> K (n + 1)
  Nil              -> K 0
  Cons (K n) (K m) -> K (n + m)

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

data Ty = Bool | Arr Ty Ty deriving (Eq, Show)

data Hypothesis = Infer | Check
                                    
data ExprF (x :: Hypothesis -> *) (i :: Hypothesis) where
  Var    ::                           Int -> ExprF x Infer
  App    ::            x Infer -> x Check -> ExprF x Infer
  Annot  ::            x Check ->      Ty -> ExprF x Infer
  If     :: x Check -> x Infer -> x Infer -> ExprF x Infer
  ETrue  ::                                  ExprF x Infer
  EFalse ::                                  ExprF x Infer
  Lam    ::                       x Check -> ExprF x Check
  CI     ::                       x Infer -> ExprF x Check

instance Functor1 ExprF where
  fmap1 phi = go where
    go = Nat $ \case
      Var i     -> Var i
      App f a   -> App (app phi f) (app phi a)
      Annot e t -> Annot (app phi e) t
      If b t e  -> If (app phi b) (app phi t) (app phi e)
      ETrue     -> ETrue
      EFalse    -> EFalse
      Lam e     -> Lam (app phi e)
      CI e      -> CI (app phi e)

type Expr = Tie ExprF

data Val = VTrue
         | VFalse
         | VLam (Val -> Val)

vapp :: Val -> Val -> Val
vapp (VLam f) = f

seek :: Int -> [a] -> Maybe a
seek n []       = Nothing
seek 0 (x : _)  = Just x
seek n (_ : xs) = seek (n-1) xs

eval1 :: ExprF (K ([Val] -> Val)) ~> K ([Val] -> Val)
eval1 = Nat $ \x -> K $ \env -> case x of
  Var i     -> fromJust (seek i env)
  App f a   -> getK f env `vapp` getK a env
  Annot e _ -> getK e env
  If b t e  -> case getK b env of
    VTrue  -> getK t env
    VFalse -> getK e env
  ETrue     -> VTrue
  EFalse    -> VFalse
  Lam e     -> VLam $ \v -> getK e (v:env)
  CI e      -> getK e env

evalOpen :: [Val] -> Expr i -> Val
evalOpen env e = getK (tie eval1 `app` e) env

eval :: Expr i -> Val
eval = evalOpen []

Now if I could only figure out a way to push the bidirectional type checker through. It would probably require putting a typeclass constraint on the index quantified over with (~>).

You could go Nat-indexed for a better-structured definition of Tie (see updated gist) but you are giving up on the coverage checking as a side effect. :/ Might be worth investigating type-level lists à la @jonsterling's Vinyl too.

Your comment about typechecking reminded me of Brian McKenna's blog post about type inference using the cofree-comonad.

I could not resist doing the example with type-level lists: https://gist.github.com/gallais/d125853cea257f051578

Haha, once you go higher-typed you don't go back.