Statically-typed values with dynamically-kinded types

DKT.hs

{-# language FlexibleContexts #-}
{-# language TypeOperators #-}

module DKT where
  
import Control.Monad (guard)
import Control.Monad.Error.Class (throwError)
import Control.Monad.Trans (lift)
import Control.Monad.Trans.State
import Control.Monad.Trans.Writer
  
-- In languages like Haskell and PureScript, types have kinds and kinds
-- are checked at compile time by a method such as unification. The kind 
-- system is like a simple type system, but "one level up".
--
-- However, this approach disallows certain types, such as self application:
-- 
--   newtype Omega f = Omega (f f)
--
-- (or else what kind should f have?) But if we instantiate f to something
-- like Proxy, or Const x, we get something reasonable. Should this be allowed?
--
-- In a dynamically-typed language like the untyped lambda calculus, we could write a 
-- value-level term like Omega: \x. x x, and apply it to some values, because their
-- types would be known and checked only at runtime.
--
-- What would happen if we tried to use this idea one level up, for tracking kinds?
-- The type system would still be statically checked, but we could make types themselves
-- "dynamically-kinded". What would this mean?
--
-- Just like a dynamically-typed language can make assertions about the type of a value
-- at runtime, a dynamically-kinded language can make assertions about the kind of a type
-- at compile time (compile time being the "runtime" for type level computations :D)
-- Instead of assigning every type a kind early on in the typechecking process,
-- we can provide a way to evaluate the kind of a type, and defer this evaluation until
-- it is really needed.
  
-- Here is a simple language of base types, constructors, applications and 
-- function types.
data Ty
  = TyInt
  | TyCon String
  | TyVar String
  | Ty :$ Ty
  | Ty :-> Ty
  deriving (Show, Read, Eq)
  
-- Kinds are a bit more complicated. They can either be fully-evaluated or not.
-- An example of a kind which is not fully evaluated is the kind of a 
-- partially-applied type constructor (KindCtor).
data Kind
  = KindHard HardKind
  | KindStar
  | KindCtor Ctor
  deriving (Show, Read, Eq)

data HardKind
  = KindVar String
  | KindFresh Int
  deriving (Show, Read, Eq)

-- As we evaluate kinds, we might get stuck trying to check the kind of a 
-- subexpression. For example, what is the kind of the type Omega a? It depends on a,
-- but we require the type a to be applicable to itself. For an abstract type, this
-- can't be known up front, so we defer this decision by emitting a constraint.
--
-- Our constraints have the form "a type of kind K1, when applied to type T, returns
-- a type of kind K2". In the example above, we'd evaluate the kind of a, and assert that the result
-- could be applied to type a, returning a type of kind *.
--
-- Such constraints could be stored in an environment, to be checked as necessary when
-- the appropriate types get instantiated.
type Constraint = (HardKind, Ty, Kind)

-- Data constructors have type arguments and a type which appears in the body.
data Ctor = Ctor
  { ctorArgs :: [String]
  , ctorBody :: Ty
  } deriving (Show, Read, Eq)
  
type Env = [(String, Ctor)]

-- Some example type constructors and an environment containing them
identity = TyCon "Identity"
omega = TyCon "Omega"
phantom = TyCon "Phantom"

env :: Env
env =
  [ ("Identity", Ctor ["a"] (TyVar "a"))
  , ("Omega", Ctor ["f"] (TyVar "f" :$ TyVar "f"))
  , ("Phantom", Ctor ["a"] TyInt)
  ]

-- The type checking monad provides a gensym, and a way to emit generated constraints.
type Check = StateT Int (WriterT [Constraint] (Either String))
  
runCheck :: Check a -> Either String (a, [Constraint])  
runCheck = runWriterT . flip evalStateT 0

-- Now for the main event. Evaluating the kind of a type is broken up into two
-- mutually-recursive functions, eval and apply.
evalKind :: Env -> Ty -> Check Kind
evalKind env ty = eval ty where
  eval TyInt = pure KindStar
  eval (TyCon ctor) = do
    ctor@(Ctor args body) <- maybe (throwError "unknown ctor") pure (lookup ctor env)
    case args of
      [] -> do
        k <- eval body
        unifiesWithKindStar k
        pure KindStar
      _ -> pure (KindCtor ctor)
  eval (TyVar i) = pure (KindHard (KindVar i))
  eval (t1 :$ t2) = do
    k <- eval t1
    apply k t2
  eval (t1 :-> t2) = do
    k1 <- eval t1
    unifiesWithKindStar k1
    k2 <- eval t2
    unifiesWithKindStar k2
    pure KindStar
    
  apply (KindHard hk) ty = do
    gs <- get
    modify (1 +)
    lift $ tell [(hk, ty, KindHard (KindFresh gs))]
    pure (KindHard (KindFresh gs))
  apply KindStar _ = throwError "Cannot apply a ground type to an argument"
  apply (KindCtor (Ctor [] _)) _ = throwError "Cannot apply a constructor which takes no arguments"
  apply (KindCtor (Ctor [x] body)) ty = do
    k <- eval (replace x ty body)
    unifiesWithKindStar k
    pure KindStar
  -- The partial application case - we substitute the first argument into
  -- the body, and return a KindCtor to represent the unevaluated result.
  apply (KindCtor (Ctor (x : xs) body)) ty =
    pure (KindCtor (Ctor xs (replace x ty body)))
  
  replace arg ty (TyVar arg') | arg == arg' = ty
  replace arg ty (t1 :$ t2) = replace arg ty t1 :$ replace arg ty t2
  replace arg ty (t1 :-> t2) = replace arg ty t1 :-> replace arg ty t2
  replace arg ty other = other

  -- I'm too lazy to write a proper unification algorithm and solve
  -- for the substitution right now, but this function exists to check that
  -- a kind can be unified with *.
  unifiesWithKindStar :: Kind -> Check ()
  unifiesWithKindStar KindStar = pure ()
  unifiesWithKindStar KindHard{} = pure ()
  unifiesWithKindStar k = throwError $ "Expected kind *, got " ++ show k