Kinds.hs

-----------------------------------------------------------------------------
--
-- Module      :  Language.PureScript.TypeChecker.Kinds
-- Copyright   :  (c) Phil Freeman 2013
-- License     :  MIT
--
-- Maintainer  :  Phil Freeman <paf31@cantab.net>
-- Stability   :  experimental
-- Portability :
--
-- |
-- This module implements the kind checker
--
-----------------------------------------------------------------------------

{-# OPTIONS_GHC -fno-warn-orphans #-}
{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE MultiParamTypeClasses #-}

module Language.PureScript.TypeChecker.Kinds (
    kindOf,
    kindsOf,
    kindsOfAll
) where

import Language.PureScript.Types
import Language.PureScript.Kinds
import Language.PureScript.Names
import Language.PureScript.TypeChecker.Monad
import Language.PureScript.Pretty

import Control.Monad.State
import Control.Monad.Error
import Control.Monad.Reader
import Control.Monad.Unify

import Control.Applicative

import qualified Data.Map as M

instance Unifiable Check Kind where
  unknown = KUnknown
  isUnknown (KUnknown u) = Just u
  isUnknown _ = Nothing
  KUnknown u1 ?= KUnknown u2 | u1 == u2 = return ()
  KUnknown u ?= k = replace u k
  k ?= KUnknown u = replace u k
  Star ?= Star = return ()
  Bang ?= Bang = return ()
  Row k1 ?= Row k2 = k1 ?= k2
  FunKind k1 k2 ?= FunKind k3 k4 = do
    k1 ?= k3
    k2 ?= k4
  k1 ?= k2 = UnifyT . lift . throwError $ "Cannot unify " ++ prettyPrintKind k1 ++ " with " ++ prettyPrintKind k2 ++ "."

-- |
-- Infer the kind of a single type
--
kindOf :: ModuleName -> Type -> Check Kind
kindOf moduleName ty = liftUnify $ starIfUnknown <$> infer ty

-- |
-- Infer the kind of a type constructor with a collection of arguments and a collection of associated data constructors
--
kindsOf :: ModuleName -> ProperName -> [String] -> [Type] -> Check Kind
kindsOf moduleName name args ts = fmap starIfUnknown . liftUnify $ do
  tyCon <- fresh
  kargs <- replicateM (length args) fresh
  let dict = (name, tyCon) : zip (map ProperName args) kargs
  bindLocalTypeVariables moduleName dict $
    solveTypes ts kargs tyCon

-- |
-- Simultaneously infer the kinds of several mutually recursive type constructors
--
kindsOfAll :: ModuleName -> [(ProperName, [String], Type)] -> [(ProperName, [String], [Type])] -> Check ([Kind], [Kind])
kindsOfAll moduleName syns tys = fmap tidyUp . liftUnify $ do
  synVars <- replicateM (length syns) fresh
  let dict = zipWith (\(name, _, _) var -> (name, var)) syns synVars
  bindLocalTypeVariables moduleName dict $ do
    tyCons <- replicateM (length tys) fresh
    let dict = zipWith (\(name, _, _) tyCon -> (name, tyCon)) tys tyCons
    bindLocalTypeVariables moduleName dict $ do
      data_ks <- zipWithM (\tyCon (_, args, ts) -> do
        kargs <- replicateM (length args) fresh
        let argDict = zip (map ProperName args) kargs
        bindLocalTypeVariables moduleName argDict $
          solveTypes ts kargs tyCon) tyCons tys
      syn_ks <- zipWithM (\synVar (_, args, ty) -> do
        kargs <- replicateM (length args) fresh
        let argDict = zip (map ProperName args) kargs
        bindLocalTypeVariables moduleName argDict $
          solveTypes [ty] kargs synVar) synVars syns
      return (syn_ks, data_ks)
  where
  tidyUp (ks1, ks2) = (map starIfUnknown ks1, map starIfUnknown ks2)

-- |
-- Solve the set of kind constraints associated with the data constructors for a type constructor
--
solveTypes :: [Type] -> [Kind] -> Kind -> UnifyT Check Kind
solveTypes ts kargs tyCon = do
  ks <- mapM infer ts
  tyCon ?= foldr FunKind Star kargs
  forM_ ks $ \k -> k ?= Star
  return tyCon

-- |
-- Default all unknown kinds to the Star kind of types
--
starIfUnknown :: Kind -> Kind
starIfUnknown (KUnknown _) = Star
starIfUnknown (FunKind k1 k2) = FunKind (starIfUnknown k1) (starIfUnknown k2)
starIfUnknown k = k

-- |
-- Infer a kind for a type
--
infer :: Type -> UnifyT Check Kind
infer (Object row) = do
  k <- infer row
  k ?= Row Star
  return Star
infer (TypeVar v) = do
  Just moduleName <- checkCurrentModule <$> get
  UnifyT . lift $ lookupTypeVariable moduleName (Qualified Nothing (ProperName v))
infer (TypeConstructor v) = do
  env <- liftCheck getEnv
  Just moduleName <- checkCurrentModule <$> get
  case M.lookup (qualify moduleName v) (types env) of
    Nothing -> UnifyT . lift . throwError $ "Unknown type constructor '" ++ show v ++ "'"
    Just (kind, _) -> return kind
infer (TypeApp t1 t2) = do
  k0 <- fresh
  k1 <- infer t1
  k2 <- infer t2
  k1 ?= FunKind k2 k0
  return k0
infer (ForAll ident ty) = do
  k <- fresh
  Just moduleName <- checkCurrentModule <$> get
  bindLocalTypeVariables moduleName [(ProperName ident, k)] $ infer ty
infer REmpty = do
  k <- fresh
  return $ Row k
infer (RCons _ ty row) = do
  k1 <- infer ty
  k2 <- infer row
  k2 ?= Row k1
  return $ Row k1
infer (ConstrainedType deps ty) = do
  mapM_ (infer . snd) deps
  k <- infer ty
  k ?= Star
  return Star
infer _ = error "Invalid argument to infer"