Subst.hs

{-# LANGUAGE DeriveFoldable #-}
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE DeriveTraversable #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE ScopedTypeVariables #-}
module Data.Subst where

import Data.Traversable
import Data.Foldable


data Var v = B | F v
    deriving (Functor, Foldable, Traversable)

infixl 1 //=

class (Traversable f, Monad (SubstTm f)) => Subst f where
    type SubstTm f :: * -> *
    (//=) :: Subst (SubstTm f) => f a -> (a -> SubstTm f b) -> f b

class (Subst f, Subst (SubstTm f)) => Subst' f

subst :: (Eq v, Subst' f) => v -> SubstTm f v -> f v -> f v
subst v m n = n //= \v' -> if v == v' then m else return v'

inst :: Subst' f => SubstTm f v -> f (Var v) -> f v
inst t s = s //= \v -> case v of
                           B    -> t
                           F v' -> return v'

nest :: Monad m => (t -> m v) -> Var t -> m (Var v)
nest _ B     = return B
nest f (F v) = f v >>= return . F

infixr 5 :<
infixl 5 :>

data Fwd m g v where
    F0   :: g v -> Fwd m g v
    (:>) :: m v -> Fwd m g (Var v) -> Fwd m g v
    deriving (Functor, Foldable, Traversable)

instance (Traversable g, Subst m, Subst g, SubstTm g ~ SubstTm m) => Subst (Fwd m g) where
    type SubstTm (Fwd m g) = SubstTm m

    F0 t     //= f = F0 (t //= f)
    ty :> te //= f = (ty //= f) :> (te //= nest f)

data Bwd m g v where
    B0   :: g v -> Bwd m g v
    (:<) :: Bwd m g v -> m v -> Bwd m g (Var v)

lookBwd :: Functor m => Bwd m g v -> v -> Maybe (m v)
lookBwd (B0 _)    _    = Nothing
lookBwd (bw :< _) (F v) = fmap (fmap F) (lookBwd bw v)
lookBwd (_  :< x) B     = Just (fmap F x)

data Proxy1 (f :: * -> *) v = Proxy1
    deriving (Functor, Foldable, Traversable)

instance Monad (Proxy1 m) where
    return _ = Proxy1
    _ >>= _  = Proxy1

instance Monad m => Subst (Proxy1 m) where
    type SubstTm (Proxy1 m) = m
    _ //= _ = Proxy1

type Zip m gl gr v = (Bwd m gl v, Fwd m gr v)

zipFwd :: Zip m gl gr v -> Maybe (Zip m gl gr (Var v))
zipFwd (_,  F0 _   ) = Nothing
zipFwd (bw, t :> fw) = Just (bw :< t, fw)

zipBwd :: Zip m gl gr (Var v) -> Zip m gl gr v
zipBwd (bw :< t, fw) = (bw, t :> fw)

data Tm v
    = V v
    | (Tm v) :@ (Tm v)
    | Lam (Tm (Var v))
    deriving (Functor, Foldable, Traversable)

instance Monad Tm where
  return = V

  V v    >>= f = f v
  m :@ n >>= f = (m >>= f) :@ (n >>= f)
  Lam s  >>= f = Lam (s >>= nest f)

instance Subst Tm where
    type SubstTm Tm = Tm
    (//=) = (>>=)