TypeTheory.hs

-- TypeTheory.hs

module Main where

import Control.Applicative
import Data.Unique
import qualified Data.Set as Set

-- Types

data Kind = KBase
          | KArrow Kind Kind
          deriving (Eq)

data TypeExpr = TEVar TypeVarName
              | TEAbs TypeVarName Kind TypeExpr
              | TEApp TypeExpr TypeExpr
              | TEPi Kind
              | TEObj [SigExpr]

data Type = TConst TypeConst
          | TVar Unique
          | TAbs Unique Type
          | TApp Type Type
          | TObj (Set.Set Sig)

data TypeConst = TCInteger
               | TCString
               | TCArrow
               | TCPi Kind
               deriving (Eq)

data SigExpr = SigExpr VarName TypeExpr
data Sig     = Sig VarName Type deriving (Eq)

data Expr = EConst ExprConst
          | EVar VarName
          | EAbs VarName TypeExpr Expr
          | EApp Expr Expr
          | ETypeAbs TypeVarName Kind Expr
          | ETypeApp Expr TypeExpr

data ExprConst = ECInteger Integer
               | ECString String

data Value = VLiteral ExprConst
           | VClosure Env VarName Expr

type TypeVarName = String
type VarName     = String
type TypeContext = [(TypeVarName, Kind)]
type TypeEnv     = [(TypeVarName, Type)]
type Context     = [(VarName, Type)]
type Env         = [(VarName, Value)]

instance Eq Type where
  TConst tc1 == TConst tc2 = tc1 == tc2
  TVar u1    == TVar u2    = u1 == u2
  TAbs u1 t1 == TAbs u2 t2 = u1 == u2 || replace t1 u1 (TVar u2) == t2
  TApp t1 t2 == TApp t3 t4 = t1 == t3 && t2 == t4
  TObj sigs1 == TObj sigs2 = sigs1 == sigs2
  _          == _          = False

instance Ord Sig where
  compare (Sig vn1 _) (Sig vn2 _) = compare vn1 vn2

replace :: Type -> Unique -> Type -> Type
replace t@(TConst _) _ _   = t
replace t@(TVar u) u' t'   = if u == u' then t' else t
replace (TAbs u t) u' t'   = TAbs u (replace t u' t')
replace (TApp t1 t2) u' t' = TApp (replace t1 u' t') (replace t2 u' t')
replace (TObj sigs) u' t'  = TObj (Set.map (replaceSig u' t') sigs)

replaceSig :: Unique -> Type -> Sig -> Sig
replaceSig u t' (Sig vn t) = Sig vn (replace t u t')

-- Kind Checking

checkTypeExpr :: TypeContext -> TypeExpr -> IO Kind
checkTypeExpr tctxt (TEVar tvn)      = case lookup tvn tctxt of
                                         (Just k) -> return k
                                         Nothing  -> unknownTypeVariable tvn
checkTypeExpr tctxt (TEAbs tvn k te) = KArrow k
                                         <$> checkTypeExpr ((tvn, k):tctxt) te
checkTypeExpr tctxt (TEApp te1 te2)  = do k1 <- checkTypeExpr tctxt te1
                                          k2 <- checkTypeExpr tctxt te2
                                          checkTypeExprApp k1 k2
checkTypeExpr tctxt (TEPi k)         = return (KArrow (KArrow k KBase) KBase)
checkTypeExpr tctxt (TEObj _)        = return KBase

checkTypeExprApp :: Kind -> Kind -> IO Kind
checkTypeExprApp (KArrow k1 k2) k3 = if k3 == k1
                                        then return k2
                                        else unexpectedKind k3 k1
checkTypeExprApp k _               = isNotAFunctionKind k

-- Type Evaluation

evalTypeExpr :: TypeEnv -> TypeExpr -> IO Type
evalTypeExpr tenv (TEVar tvn)      = case lookup tvn tenv of
                                       (Just t) -> return t
                                       Nothing  -> unboundTypeVariable tvn
evalTypeExpr tenv (TEAbs tvn _ te) = do u <- newUnique
                                        let tenv' = (tvn, TVar u):tenv
                                        TAbs u <$> evalTypeExpr tenv' te
evalTypeExpr tenv (TEApp te1 te2)  = do t1 <- evalTypeExpr tenv te1
                                        t2 <- evalTypeExpr tenv te2
                                        evalTypeExprApp t1 t2
evalTypeExpr tenv (TEPi k)         = return (TConst (TCPi k))
evalTypeExpr tenv (TEObj siges)    = (TObj . Set.fromList)
                                       <$> mapM (evalSigExpr tenv) siges

evalTypeExprApp :: Type -> Type -> IO Type
evalTypeExprApp (TAbs u t1) t2 = return (replace t1 u t2)
evalTypeExprApp t1 t2          = return (TApp t1 t2)

evalSigExpr :: TypeEnv -> SigExpr -> IO Sig
evalSigExpr tenv (SigExpr vn te) = Sig vn <$> evalTypeExpr tenv te

-- Type Checking

checkExpr :: TypeContext -> TypeEnv -> Context -> Expr -> IO Type
checkExpr tctxt tenv ctxt = checkExpr'
  where
    checkExpr' (EConst ec)        = return (TConst (exprConstToTypeConst ec))
    checkExpr' (EVar vn)          = case lookup vn ctxt of
                                      (Just t) -> return t
                                      Nothing  -> unknownVariable vn
    checkExpr' (EAbs vn te e)     = do checkParameterKind
                                         =<< checkTypeExpr tctxt te
                                       t1 <- evalTypeExpr tenv te
                                       let ctxt' = (vn, t1):ctxt
                                       t2 <- checkExpr tctxt tenv ctxt' e
                                       return (makeFunctionType t1 t2)
    checkExpr' (EApp e1 e2)       = do t1 <- checkExpr' e1
                                       t2 <- checkExpr' e2
                                       checkExprApp t1 t2
    checkExpr' (ETypeAbs tvn k e) = do u <- newUnique
                                       let tctxt' = (tvn, k):tctxt
                                       let tenv' = (tvn, TVar u):tenv
                                       t <- checkExpr tctxt' tenv' ctxt e
                                       return (makePolymorphicType u k t)
    checkExpr' (ETypeApp e te)    = do t1 <- checkExpr' e
                                       k <- checkTypeExpr tctxt te
                                       t2 <- evalTypeExpr tenv te
                                       checkExprTypeApp t1 k t2

exprConstToTypeConst :: ExprConst -> TypeConst
exprConstToTypeConst (ECInteger _) = TCInteger
exprConstToTypeConst (ECString _)  = TCString

checkParameterKind :: Kind -> IO ()
checkParameterKind KBase = return ()
checkParameterKind k     = isNotABaseKind k

makeFunctionType :: Type -> Type -> Type
makeFunctionType t1 t2 = TApp (TApp (TConst TCArrow) t1) t2

checkExprApp :: Type -> Type -> IO Type
checkExprApp (TApp (TApp (TConst TCArrow) t1) t2) t3 = checkExprApp'
  where
    checkExprApp' = if t3 == t1 then return t2 else unexpectedType t3 t1
checkExprApp t _  = isNotAFunctionType t

makePolymorphicType :: Unique -> Kind -> Type -> Type
makePolymorphicType u k t = TApp (TConst (TCPi k)) (TAbs u t)

checkExprTypeApp :: Type -> Kind -> Type -> IO Type
checkExprTypeApp (TApp (TConst (TCPi k1)) (TAbs u t1)) k2 t2 = checkExprTypeApp'
  where
    checkExprTypeApp'  = if k2 == k1
                            then evalType t1 u t2
                            else unexpectedKind k2 k1
checkExprTypeApp t _ _ = isNotAPolymorphicType t

evalType :: Type -> Unique -> Type -> IO Type
evalType t@(TConst _) _ _   = return t
evalType t@(TVar u) u' t'   = return (if u == u' then t' else t)
evalType (TAbs u t) u' t'   = TAbs u <$> evalType t u' t'
evalType (TApp t1 t2) u' t' = do t1 <- evalType t1 u' t'
                                 t2 <- evalType t2 u' t'
                                 evalTypeExprApp t1 t2
evalType (TObj sigs) u' t'  = (TObj . Set.fromList)
                                <$> mapM (evalSig u' t') (Set.toList sigs)

evalSig :: Unique -> Type -> Sig -> IO Sig
evalSig u t' (Sig vn t) = Sig vn <$> evalType t u t'

-- Evaluation

evalExpr :: Env -> Expr -> IO Value
evalExpr _ (EConst ec)        = return (VLiteral ec)
evalExpr env (EVar vn)        = case lookup vn env of
                              (Just v) -> return v
                              Nothing  -> unboundVariable vn
evalExpr env (EAbs vn _ e)    = return (VClosure env vn e)
evalExpr env (EApp e1 e2)     = do v1 <- evalExpr env e1
                                   v2 <- evalExpr env e2
                                   evalExprApp v1 v2
evalExpr env (ETypeAbs _ _ e) = evalExpr env e
evalExpr env (ETypeApp e _)   = evalExpr env e

evalExprApp :: Value -> Value -> IO Value
evalExprApp (VClosure env vn e) v = evalExpr ((vn, v):env) e
evalExprApp v _                   = isNotAFunction v

-- Show instances

instance Show Kind where
  show KBase            = "*"
  show (KArrow k1 k2)   = "(" ++ show k1 ++ " -> " ++ show k2 ++ ")"

instance Show TypeExpr where
  show (TEVar tvn)      = tvn
  show (TEAbs tvn k te) = "^" ++ tvn ++ ":" ++ show k ++ "{" ++ show te ++ "}"
  show (TEApp (TEApp (TEVar "Arrow") te1) te2)
                        = "(" ++ show te1 ++ " => " ++ show te2 ++ ")"
  show (TEApp (TEPi k1) (TEAbs tvn k2 te))
    | k1 == k2          = "~" ++ tvn ++ ":" ++ show k1 ++ "{" ++ show te ++ "}"
  show (TEApp te1 te2)  = show te1 ++ "(" ++ show te2 ++ ")"
  show (TEPi k)         = "pi[" ++ show k ++ "]"
  show (TEObj siges)    = "obj" ++ show siges

instance Show Type where
  show (TConst tc)  = show tc
  show (TVar u)     = show u
  show (TAbs u t)   = "^" ++ show u ++ "{" ++ show t ++ "}"
  show (TApp (TApp (TConst TCArrow) t1) t2)
                    = "(" ++ show t1 ++ " => " ++ show t2 ++ ")"
  show (TApp (TConst (TCPi k)) (TAbs u t))
                    = "~" ++ show u ++ ":" ++ show k ++ "{" ++ show t ++ "}"
  show (TApp t1 t2) = show t1 ++ "(" ++ show t2 ++ ")"
  show (TObj sigs)  = "Obj" ++ show (Set.toList sigs)

instance Show TypeConst where
  show TCInteger = "Integer"
  show TCString  = "String"
  show TCArrow   = "Arrow"
  show (TCPi k)  = "Pi[" ++ show k ++ "]"

instance Show Unique where
  show u = "<" ++ show (hashUnique u) ++ ">"

instance Show SigExpr where
  show (SigExpr vn te) = "(" ++ vn ++ ":" ++ show te ++ ")"

instance Show Sig where
  show (Sig vn t) = "(" ++ vn ++ ":" ++ show t ++ ")"

instance Show Expr where
  show (EConst ec)        = show ec
  show (EVar vn)          = vn
  show (EAbs vn te e)     = "^" ++ vn ++ ":" ++ show te ++ "{" ++ show e ++ "}"
  show (EApp e1 e2)       = show e1 ++ "(" ++ show e2 ++ ")"
  show (ETypeAbs tvn k e) = "~" ++ tvn ++ ":" ++ show k ++ "{" ++ show e ++ "}"
  show (ETypeApp e te)    = show e ++ "[" ++ show te ++ "]"

instance Show ExprConst where
  show (ECInteger i) = show i
  show (ECString s)  = show s

instance Show Value where
  show (VLiteral ec)     = show ec
  show (VClosure _ vn e) = "#<closure " ++ vn ++ ">"

-- Exceptions

unknownTypeVariable :: TypeVarName -> IO a
unknownTypeVariable tvn = fail ("unknown type variable: " ++ tvn)

unboundTypeVariable :: TypeVarName -> IO a
unboundTypeVariable tvn = fail ("unbound type variable: " ++ tvn)

unknownVariable :: VarName -> IO a
unknownVariable vn = fail ("unknown variable: " ++ vn)

unboundVariable :: VarName -> IO a
unboundVariable vn = fail ("unbound variable: " ++ vn)

unexpectedKind :: Kind -> Kind -> IO a
unexpectedKind k1 k2 = fail message
  where
    message = "unexpected " ++ show k1 ++ ", expected " ++ show k2

unexpectedType :: Type -> Type -> IO a
unexpectedType t1 t2 = fail message
  where
    message = "unexpected " ++ show t1 ++ ", expected " ++ show t2

isNotAFunctionKind :: Kind -> IO a
isNotAFunctionKind k = fail (show k ++ " is not a function kind")

isNotABaseKind :: Kind -> IO a
isNotABaseKind k = fail (show k ++ " is not a base kind")

isNotAFunctionType :: Type -> IO a
isNotAFunctionType t = fail (show t ++ " is not a function type")

isNotAPolymorphicType :: Type -> IO a
isNotAPolymorphicType t = fail (show t ++ " is not a polymorphic type")

isNotAFunction :: Value -> IO a
isNotAFunction v = fail (show v ++ " is not a function")

-- Samples

tctxt :: TypeContext
tctxt = [("Integer", KBase)
        ,("String", KBase)
        ,("Arrow", KArrow KBase (KArrow KBase KBase))
        ]

tenv :: TypeEnv
tenv = [("Integer", TConst TCInteger)
       ,("String", TConst TCString)
       ,("Arrow", TConst TCArrow)
       ]

ctxt :: Context
ctxt = []

env :: Env
env = []

execTypeExpr :: TypeContext -> TypeEnv -> TypeExpr -> IO Type
execTypeExpr tctxt tenv te = checkTypeExpr tctxt te >> evalTypeExpr tenv te

checkTE :: TypeExpr -> IO Kind
checkTE = checkTypeExpr tctxt

evalTE :: TypeExpr -> IO Type
evalTE = evalTypeExpr tenv

execTE :: TypeExpr -> IO Type
execTE = execTypeExpr tctxt tenv

execExpr :: TypeContext -> TypeEnv -> Context -> Env -> Expr -> IO Value
execExpr tctxt tenv ctxt env e = checkExpr tctxt tenv ctxt e >> evalExpr env e

checkE :: Expr -> IO Type
checkE = checkExpr tctxt tenv ctxt

evalE :: Expr -> IO Value
evalE = evalExpr env

execE :: Expr -> IO Value
execE = execExpr tctxt tenv ctxt env

teId :: TypeExpr
teId = TEApp (TEPi KBase) $ TEAbs "A" KBase
     $ TEApp (TEApp (TEVar "Arrow") (TEVar "A")) (TEVar "A")

eId :: Expr
eId = ETypeAbs "A" KBase
    $ EAbs "x" (TEVar "A") $ EVar "x"

teConst :: TypeExpr
teConst = TEApp (TEPi KBase) $ TEAbs "A" KBase
        $ TEApp (TEPi KBase) $ TEAbs "B" KBase
        $ TEApp (TEApp (TEVar "Arrow") (TEVar "A"))
                (TEApp (TEApp (TEVar "Arrow") (TEVar "B"))
                       (TEVar "A"))

eConst :: Expr
eConst = ETypeAbs "A" KBase
       $ ETypeAbs "B" KBase
       $ EAbs "x" (TEVar "A") $ EAbs "y" (TEVar "B") $ EVar "x"

teMonad :: TypeExpr
teMonad = TEAbs "M" (KArrow KBase KBase)
        $ TEObj [ SigExpr "fmap" teFMap
                , SigExpr "pure" tePure
                , SigExpr "join" teJoin
                ]

teFMap :: TypeExpr
teFMap = TEApp (TEPi KBase) $ TEAbs "A" KBase
       $ TEApp (TEPi KBase) $ TEAbs "B" KBase
       $ TEApp (TEApp (TEVar "Arrow")
                      (TEApp (TEApp (TEVar "Arrow") (TEVar "A"))
                             (TEVar "B")))
               (TEApp (TEApp (TEVar "Arrow")
                             (TEApp (TEVar "M") (TEVar "A")))
                      (TEApp (TEVar "M") (TEVar "B")))

tePure :: TypeExpr
tePure = TEApp (TEPi KBase) $ TEAbs "A" KBase
       $ TEApp (TEApp (TEVar "Arrow") (TEVar "A"))
               (TEApp (TEVar "M") (TEVar "A"))

teJoin :: TypeExpr
teJoin = TEApp (TEPi KBase) $ TEAbs "A" KBase
       $ TEApp (TEApp (TEVar "Arrow")
                      (TEApp (TEVar "M") (TEApp (TEVar "M") (TEVar "A"))))
               (TEApp (TEVar "M") (TEVar "A"))