kind_inference.hs

import Control.Monad (foldM)
import Control.Monad.State (StateT, modify, get, put, lift, evalStateT)
import Data.Map (Map)
import qualified Data.Map as Map

data KindExpr
  = TCon String
  | TVar String
  | TAp  KindExpr KindExpr
  deriving (Show)
           

data Kind
  = Star
  | KVar String
  | KAp Kind Kind
  deriving (Eq, Show)

type Sub = Map String Kind

type Env = Map String Kind

type KindInfer = StateT Int (Either String)

emptySub :: Sub
emptySub = Map.empty

-- TODO: temp hack:
getConstructorKind :: String -> KindInfer Kind
getConstructorKind "[]" = return (KAp Star Star)
getConstructorKind _    = return Star

infer :: ([String], [KindExpr]) -> Either String Sub
infer problem = runKinds (inferKinds problem)

runKinds :: KindInfer a -> Either String a
runKinds f = evalStateT f 0

newKindVar :: KindInfer Kind
newKindVar = do
  n <- get
  let var = KVar $ "k" ++ show n
  put (n + 1)
  return var

inferKinds :: ([String], [KindExpr]) -> KindInfer Sub
inferKinds (variables, exprs) = do
  varKinds <- mapM (const newKindVar) variables
  let env = Map.fromList $ zip variables varKinds
  sub <- foldM (inferExpr env) emptySub exprs
  return $ Map.map (apply sub) env

inferExpr :: Env -> Sub -> KindExpr -> KindInfer Sub
inferExpr env sub expr = do
  (s1, k) <- inferKind env sub expr
  let sub1 = compose sub s1
  sub2 <- unifyKinds (apply sub1 k) Star
  return $ compose sub1 sub2

inferKind :: Env -> Sub -> KindExpr -> KindInfer (Sub, Kind)
inferKind env sub expr = case expr of
  TVar v -> case Map.lookup v env of
    Nothing -> error $ "undefined type var: " ++ v
    Just k  -> return (emptySub, k)
  TCon c -> do
    k <- getConstructorKind c
    return (emptySub, k)
  TAp t1 t2 -> do
    (s1, k1) <- inferKind env sub t1
    let sub1 = compose sub s1
    (s2, k2) <- inferKind env sub1 t2
    let sub2 = compose sub1 s2

    retKind <- newKindVar
    let lhsKind = (KAp k2 retKind)
    sub3 <- unifyKinds (apply sub2 lhsKind) (apply sub2 k1)
    let sub' = compose sub2 sub3
    return (sub', apply sub' retKind)

unifyKinds :: Kind -> Kind -> KindInfer Sub
unifyKinds k1 k2 = case (k1, k2) of
  (Star, Star)           ->
    return emptySub
  (KVar v1, KVar v2)     ->
    if v1 == v2
    then return emptySub
    else return $ Map.singleton v1 k2
  (KVar v1, _)           ->
    varBind v1 k2
  (_, KVar v2)           ->
    varBind v2 k1
  (KAp l1 r1, KAp l2 r2) -> do
    sub1 <- unifyKinds l1 l2
    sub2 <- unifyKinds (apply sub1 r1) (apply sub1 r2)
    return $ compose sub1 sub2
  (_, _)                 ->
    lift $ Left $ "kinds do not unify: " ++ show k1 ++ ", " ++ show k2

varBind :: String -> Kind -> KindInfer Sub
varBind var k
  | (KVar var) == k   = return emptySub
  | containsVar var k = lift $ Left $ "infinite kind"
  | otherwise         = return $ Map.singleton var k

containsVar :: String -> Kind -> Bool
containsVar var k = case k of
  Star    -> False
  KVar s  -> s == var
  KAp l r -> containsVar var l || containsVar var r

compose :: Sub -> Sub -> Sub
compose older newer =
  Map.union (Map.map (apply newer) older) newer


apply :: Sub -> Kind -> Kind
apply sub k = case k of
  Star    -> Star
  KVar v1 -> case Map.lookup v1 sub of
    Nothing -> k
    Just k' -> k'
  KAp l r ->
    KAp (apply sub l) (apply sub r)

{-
(Note: I want to find a way to generalize KindExpr so that this code only depends on a typeclass)
-}

-- Note that some 'constants' can have kinds other than *:
data HigherKinds3 a = HigherKinds3 (a [])

-- Infinite kind:
-- data HigherKinds4 a b = HigherKinds4 (a Int) (b b)

data HigherKinds a b = HigherKinds (a Int) (b a)
-- Kind:
-- HigherKinds :: (* -> *) -> ((* -> *) -> *) -> *


{-

a: k1, b: k2

Examine `(a Int)`
Int: *
a: k1 = k3 -> k4
k3 = *
a: k3 -> k4
k4 = *
a: * -> *

Examine `(b a)`
a: * -> *
b: k2 = k5 -> k6
k6 = *
k5 = * -> *
b: (* -> *) -> *

Does the fact that (a Int) appears without more types being applied mean that it must be of kind *?
It looks like yes, because this fails:

    data HigherKinds2 a = HigherKinds2 (a Int) a


How should the problem be expressed?

{variables :: [String], exprs :: [KindExpr]}

KindExpr = Type String | Application KindExpr KindExpr
Or just use the Type type
TYpe = TVar String | TCon String | TAp Type Type
Kind = Star | KVar String | KAp Kind Kind

Algo:
-}