Mutually Defined Datatypes

MutuallDefined.hs

{-# LANGUAGE KindSignatures      #-}
{-# LANGUAGE GADTs               #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE PolyKinds           #-}
{-# LANGUAGE DataKinds           #-}
{-# LANGUAGE RankNTypes          #-}

module MutuallyDefined where

import Data.Type.Natural

data ForestF l n x idx  where
  Leaf  :: l -> ForestF l n x Z
  Node  :: n -> x (S Z) -> ForestF l n x Z
  Nil   :: ForestF l n x (S Z)
  Cons  :: x Z -> x (S Z) -> ForestF l n x (S Z)

newtype Tie (f :: * -> * -> (Nat -> *) -> Nat -> *) l n idx =
  Tie { unTie :: f l n (Tie f l n) idx }

type Tree   l n = Tie ForestF l n Z
type Forest l n = Tie ForestF l n (S Z)

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 t = Tie . Cons t

mapForestF :: forall l n a b . (forall i . SNat i -> a i -> b i) ->
              forall i. SNat i -> ForestF l n a i -> ForestF l n b i
mapForestF f SZ      (Leaf l)    = Leaf l
mapForestF f SZ      (Node n ts) = Node n $ f (SS SZ) ts
mapForestF f (SS SZ) Nil         = Nil
mapForestF f (SS SZ) (Cons t ts) = Cons (f SZ t) $ f (SS SZ) ts

tie :: forall f l n r .
       (forall a b . (forall i . SNat i -> a i -> b i) -> forall i. SNat i -> f l n a i -> f l n b i)
    -> (forall i . SNat i -> f l n r i -> r i)
    -> (forall i . SNat i -> Tie f l n i -> r i)
tie fmap phi snat = phi snat . rec snat . unTie
  where
    rec :: forall i . SNat i -> f l n (Tie f l n) i -> f l n r i
    rec = fmap $ tie fmap phi

newtype Cst k (a :: Nat) = Cst { runCst :: k }

numberLeaves :: forall l n. Tree l n -> Nat
numberLeaves = runCst . tie mapForestF alg SZ
  where
    alg :: forall i . SNat i -> ForestF l n (Cst Nat) i -> Cst Nat i
    alg SZ      (Leaf l)    = Cst $ S Z
    alg SZ      (Node n ts) = Cst . runCst $ ts
    alg (SS SZ) Nil         = Cst Z
    alg (SS SZ) (Cons t ts) = Cst $ runCst t + runCst ts

example :: Tree () ()
example = node () $ cons (leaf ()) $ cons (leaf ()) $ cons (node () nil) $ cons (leaf ()) nil

Nota: Agda version which is slightly nicer (we statically know we have treated all the possible cases!).