Core.hs

{-# LANGUAGE DeriveFunctor, DeriveFoldable, DeriveTraversable,
             ScopedTypeVariables, OverloadedStrings, TypeSynonymInstances,
             FlexibleInstances, GADTs #-}
module Core where

import Bound
import Bound.Name
import Bound.Var
import Control.Applicative
import Control.Comonad
import Control.Monad
import Control.Monad.Trans.Class
import Data.Bifunctor
import Data.Foldable
import qualified Data.Text.Lazy as L
import qualified Data.Text.Encoding as E
import Data.Traversable
import GHC.Exts
import Prelude.Extras
import Text.PrettyPrint.Leijen.Text

-- GHC API
import SrcLoc
import FastString

import Debug.Trace

-- Type definitions and class instances
instance (Pretty b, Pretty a) => Pretty (Var a b) where
  pretty (B a) = pretty a
  pretty (F a) = pretty a

instance Pretty n => Pretty (Name n a) where
  pretty (Name n _) = pretty n

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)

instance (Eq a,Pretty n, Show n, Pretty a) => Pretty (Term n a) where
  pretty = pretty' 0
    where
      pretty' :: (Eq b, Show n, Pretty n, Pretty b) => Int -> Term n b -> Doc
      pretty' _ (Var a)          = pretty a
      pretty' _ (Universe _)     = char '*'
      pretty' n (App f a)        = parensIf (n > 0)
                                    (pretty' n f <+> pretty' (n + 1) a)
      pretty' n (Bind (Pi b) s)  = text "∀" <> parens ((pretty $ name b) <+>
                                    ":" <+> (pretty' n $ extract b)) <>
                                    char '.' <> pretty' n (unscope s)
      pretty' n (Bind (Lam b) s) = parens ((pretty $ name b) <+> ":" <+>
                                    (pretty' n $ extract b)) <+>
                                    text "->" <+> pretty' n (unscope s)

parensIf :: Bool -> Doc -> Doc
parensIf b a = if b then parens a else a

ppr :: Pretty a => a -> IO ()
ppr = putDoc . (<> "\n") . pretty

type LVar = Located FastString

instance Show LVar where
  show = unpackFS . unLoc

instance IsString LVar where
  fromString = noLoc . fsLit

instance Pretty LVar where
  pretty = text . L.fromStrict . E.decodeUtf8 . fastStringToByteString . unLoc

type CoreTerm = Term LVar LVar

instance IsString CoreTerm where
  fromString = Var . fromString

type Type n a = Term n a
type CoreType = CoreTerm

type ScopedTerm n a = Scope (Name n ()) (Term n) a
type ScopedType n a = ScopedTerm n a

-- Existential Pair
data EPair a where
  EP :: b -> (b -> a) -> EPair a

eapply :: (a -> EPair a) -> a -> a
eapply f e = case f e of EP b g -> g b

startEP :: Term n a -> EPair (Term n a)
startEP e = EP e id

-- type inference
inferType :: forall n a . Eq a => Show n => Show a
          => (a -> Type n a)
          -> (Term n a -> EPair (Term n a))
          -> Term n a
          -> Type n a
inferType ctx dist (Var a)      = ctx a
inferType _   dist (Universe i) = Universe (i + 1)
inferType ctx dist (App e1 e2)  = if s == te then instantiate1Name e2 t
                                             else error "mismatch"
  where
    te      = inferType ctx dist e2
    (_,s,t) = inferPi   ctx dist e1
inferType ctx dist (Bind (Pi b) s) = Universe (max k1 k2)
  where
    t  = extract b
    k1 = inferUniverse ctx dist t
    k2 = inferUniverse ctx dist (instantiate1Name t s)
inferType ctx distEP (Bind (Lam b) s) =
    Bind (Pi b) (Scope . inferType ctx' distEP' $ unscope s)
  where
    ctx' = unvar (const (unscope $ lift $ extract b))
                 (unscope . lift . inferType ctx distEP)

    -- Distribute 'Var F' to leafs.
    -- Unlike 'fromScope' which does has the type:
    -- 'Term n (Var b (Term n a)) -> Term n (Var b a)'
    --
    -- 'eapply distEP'' will the type:
    -- 'Term n (Var b (Term n a)) -> Term n (Var b (Term n a))'
    --
    -- It will, for example, convert:
    -- 'Var (F (App (Var (B ())) (Var (F (Var (B ())))))'
    -- To:
    -- 'App (Var (F (Var (B ()))))) (Var (F (Var (F (Var (B ()))))))'
    distEP' (Var (F a)) = case distEP a of EP b g -> EP b (fmap (F . Var) . g)
    distEP' e           = EP e id

inferPi :: Eq a => Show n => Show a
        => (a -> Type n a)
        -> (Term n a -> EPair (Term n a))
        -> Term n a
        -> (n, Type n a, ScopedType n a)
inferPi ctx distEP tm = case eapply distEP $ inferType ctx distEP tm of
  Bind (Pi b) s -> (name b, extract b, s)
  t             -> error ("Function expected: " ++ show t)

inferUniverse :: Eq a => Show n => Show a
              => (a -> Type n a)
              -> (Term n a -> EPair (Term n a))
              -> Type n a
              -> Integer
inferUniverse ctx distEP ty = case eapply distEP $ inferType ctx distEP ty of
  Universe i -> i
  t          -> error ("Type expected" ++ show t)

-- Smart builders
uni :: CoreTerm
uni = Universe 0

lam :: (LVar,CoreTerm) -> CoreTerm -> CoreTerm
lam (v,b) e = Bind (Lam (Name v b)) (abstract1Name v e)

pi_ :: (LVar,CoreTerm) -> CoreTerm -> CoreTerm
pi_ (v,b) e = Bind (Pi (Name v b)) (abstract1Name v e)


-- Examples
id_ :: CoreTerm
id_ = lam ("a",uni) $ lam ("x","a") "x"

dollar :: CoreTerm
dollar = lam ("a",uni) $ lam ("b",uni)
       $ lam ("f",pi_ ("_","a") "b")
       $ lam ("x","a")
       $ App "f" "x"

comp :: CoreTerm
comp = lam ("a",uni) $ lam ("b",uni) $ lam ("c",uni)
     $ lam ("f",pi_ ("_","a") "b")
     $ lam ("g",pi_ ("_","b") "c")
     $ lam ("x","a")
     $ App "g" (App "f" "x")