SimpleCore.hs

{-# LANGUAGE DeriveFoldable, DeriveFunctor, DeriveTraversable #-}
module SimpleCore where

import Bound
import Bound.Name
import Bound.Var
import Control.Applicative
import Control.Comonad
import Control.Monad
import Data.Bifunctor
import Data.Foldable
import Data.Traversable
import Prelude.Extras

data Binder n a
  = Lam {binder :: Name n a}
  | Pi  {binder :: Name n a}
  deriving (Functor,Foldable,Traversable,Show,Eq)

instance Comonad (Binder n) where
  extract    = extract . binder
  extend f w = fmap (const (f w)) w

data Term n a
  = Var !a
  | Universe Integer
  | App !(Term n a) !(Term n a)
  | Bind !(Binder n (Term n a)) !(Scope (Name n ()) (Term n) a)
  deriving (Functor,Foldable,Traversable,Show,Eq)

instance Show n => Show1 (Term n)
instance Eq1 (Term n)

instance Applicative (Term n) where
  pure  = Var
  (<*>) = ap

instance Monad (Term n) where
  return = Var
  (>>=)  = bindTerm

bindTerm :: Term n a -> (a -> Term n b) -> Term n b
bindTerm (Var a)      f = f a
bindTerm (Universe i) _ = Universe i
bindTerm (App e1 e2)  f = App (bindTerm e1 f) (bindTerm e2 f)
bindTerm (Bind b s)   f = Bind (fmap (`bindTerm` f) b) (s >>>= f)

inferType :: Eq a => (a -> (Term n a,Maybe (Term n a))) -> Term n a -> Term n a
inferType ctx (Var a)      = fst $! ctx a
inferType _   (Universe i) = Universe (i + 1)
inferType ctx (App e1 e2)  = if equalTy ctx s te then instantiate1Name e2 t
                                                 else error "mismatch"
  where
    te      = inferType ctx e2
    (_,s,t) = inferPi   ctx e1
inferType ctx (Bind (Pi b) s) = Universe (max k1 k2)
  where
    t  = extract b
    k1 = inferUniverse ctx t
    k2 = inferUniverse ctx (instantiate1Name t s)
inferType ctx (Bind (Lam b) s) = Bind (Pi b) s'
  where
    ctx' = unvar (const (fmap F $ extract b,Nothing)) (bimap (fmap F) (fmap (fmap F)) . ctx)
    s'   = toScope $ inferType ctx' $ fromScope s

inferPi :: Eq a => (a -> (Term n a,Maybe (Term n a))) -> Term n a
        -> (n, Term n a,Scope  (Name n ()) (Term n) a)
inferPi ctx e = case normalize ctx (inferType ctx e) of
                  Bind (Pi b) s -> (name b, extract b, s)
                  _ -> error "function expected"

inferUniverse :: Eq a => (a -> (Term n a,Maybe (Term n a))) -> Term n a -> Integer
inferUniverse ctx t = case normalize ctx (inferType ctx t) of
                        Universe k -> k
                        _ -> error "type expected"

normalize :: (a -> (Term n a,Maybe (Term n a))) -> Term n a -> Term n a
normalize ctx (Var a)      = case snd (ctx a) of
                               Nothing -> Var a
                               Just e  -> normalize ctx e
normalize _   (Universe k) = Universe k
normalize ctx (App e1 e2)  = case normalize ctx e1 of
                               (Bind (Lam _) s) -> normalize ctx (instantiate1Name e2' s)
                               e1' -> App e1' e2'
  where
    e2' = normalize ctx e2
normalize ctx (Bind b s) = Bind (fmap (normalize ctx) b) s'
  where
    b'   = fmap (normalize ctx) b
    ctx' = unvar (const (fmap F $ extract b',Nothing)) (bimap (fmap F) (fmap (fmap F)) . ctx)
    s'   = toScope $ normalize ctx' $ fromScope s

equalTy :: Eq a => (a -> (Term n a,Maybe (Term n a))) -> Term n a -> Term n a -> Bool
equalTy ctx e1 e2 = e1' == e2'
  where
    e1' = normalize ctx e1
    e2' = normalize ctx e2