MutuallyDefinedList

MutuallyDefinedList.hs

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

module MutuallyDefinedList where

import Data.Type.Natural

-- Machinery

data Elem :: k -> [k] -> * where
  Here  :: Elem x (x ': xs)
  There :: Elem x xs -> Elem x (y ': xs)

newtype Tie (f :: * -> * -> (u -> *) -> [u] -> u -> *) l n ks k =
  Tie { unTie :: f l n (Tie f l n ks) ks k }

tie :: forall f l n r ks .
       (forall a b. (forall k. Elem k ks -> a k -> b k) ->
                     forall k . Elem k ks -> f l n a ks k -> f l n b ks k)
    -> (forall k . Elem k ks -> f l n r ks k -> r k)
    -> (forall k . Elem k ks -> Tie f l n ks k -> r k)
tie fmap phi elt = phi elt . rec elt . unTie
  where
    rec :: Elem k ks -> f l n (Tie f l n ks) ks k -> f l n r ks k
    rec = fmap $ tie fmap phi

-- Tree

data Tags = TTree | TForest

data ForestF l n (r :: Tags -> *) ks k where
  Leaf  :: l -> ForestF l n r (TTree ': TForest ': '[]) TTree
  Node  :: n -> r TForest -> ForestF l n r (TTree ': TForest ': '[]) TTree
  Nil   :: ForestF l n r (TTree ': TForest ': '[]) TForest
  Cons  :: r TTree -> r TForest -> ForestF l n r (TTree ': TForest ': '[]) TForest

type Tree   l n = Tie ForestF l n (TTree ': TForest ': '[]) TTree
type Forest l n = Tie ForestF l n (TTree ': TForest ': '[]) TForest

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 k . Elem k (TTree ': TForest ': '[]) -> a k -> b k) ->
              forall k. Elem k (TTree ': TForest ': '[]) ->
                        ForestF l n a (TTree ': TForest ': '[]) k ->
                        ForestF l n b (TTree ': TForest ': '[]) k
mapForestF f Here         (Leaf l)    = Leaf l
mapForestF f Here         (Node n ts) = Node n $ f (There Here) ts
mapForestF f (There Here) Nil         = Nil
mapForestF f (There Here) (Cons t ts) = Cons (f Here t) $ f (There Here) ts


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

numberLeaves :: forall l n. Tree l n -> Nat
numberLeaves = runCst . tie mapForestF alg Here
  where
    alg :: forall k . Elem k (TTree ': TForest ': '[]) ->
                      ForestF l n (Cst Nat) (TTree ': TForest ': '[]) k ->
                      Cst Nat k
    alg Here         (Leaf l)    = Cst $ S Z
    alg Here         (Node n ts) = Cst . runCst $ ts
    alg (There Here) Nil         = Cst Z
    alg (There Here) (Cons t ts) = Cst $ runCst t + runCst ts

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