implementation module check

import StdEnv
import Data.Either
import Data.List
import Data.Functor
import Data.Func
import Data.Maybe
import Data.Monoid
import Control.Applicative
import Control.Monad
import Control.Monad.Trans
import Control.Monad.RWST
import qualified Data.Map as DM
from Data.Map import instance Functor (Map k)

import ast

//Start = runRWST (infer (AST [(Function ['s','t','a','r','t'] [] (Lit (Int 42)))])
Start = runRWST (infer (TypeEnv 'DM'.newMap) t) [] {tiSupply=0,tiSubst='DM'.newMap}
where
//	t = Function ['start'] [] (Lit (Int 42))
	t =
		[Function ['id'] [] (Lit (Int 42))
		,Function ['start'] [] (App (Var ['id']) (Lit (Int 42)))
		]

check :: AST -> Either [String] AST
check (AST fs) = case sortBy (on (>) isStart) fs of
	[(Function ['start'] as _):rest]
		= case runRWST (infer (TypeEnv 'DM'.newMap) fs) [] {tiSupply=0,tiSubst='DM'.newMap} of
			Left e = Left e
			Right _ = Right (AST fs)
	_ = Left ["No start function defined"]
where
	isStart a = a=:(Function ['start'] [] _)

instance < Bool where
	< False True = True
	< _ _ = False

:: Type
	= TVar [Char]
	| TInt
	| TBool
	| TChar
	| TFun Type Type

instance == Type where
	(==) (TVar a) (TVar b) = a == b
	(==) TInt TInt = True
	(==) TBool TBool = True
	(==) TChar TChar = True
	(==) (TFun a1 a2) (TFun b1 b2) = a1 == b1 && a2 == b2
	(==) _ _ = False

instance toString Type where
	toString (TVar s) = toString s
	toString TInt = "Int"
	toString TBool = "Bool"
	toString TChar = "Char"
	toString (TFun t1 t2) = toString t1 +++ " -> " +++ toString t2

:: Scheme = Scheme [[Char]] Type
class Types a where
	ftv :: a -> [[Char]]
	apply :: Subst a -> a

instance Types Type where
	ftv (TVar n) = [n]
	ftv TInt = []
	ftv TBool = []
	ftv TChar = []
	ftv (TFun t1 t2) = union (ftv t1) (ftv t2)

	apply s (TVar n) = case 'DM'.get n s of
		Nothing = TVar n
		Just t = t
	apply s (TFun t1 t2) = TFun (apply s t1) (apply s t2)
	apply s t = t

instance Types Scheme where
	ftv (Scheme vars t) = difference (ftv t) vars
	apply s (Scheme vars t) = Scheme vars (apply (foldr 'DM'.del s vars) t)

instance Types [a] | Types a where
	ftv l = foldr union [] (map ftv l)
	apply s l = map (apply s) l

:: Subst :== 'DM'.Map [Char] Type
composeSubst s1 s2 = 'DM'.union ((apply s1) <$> s2) s1

:: TypeEnv = TypeEnv ('DM'.Map [Char] Scheme)
remove :: TypeEnv [Char] -> TypeEnv
remove (TypeEnv env) var = TypeEnv ('DM'.del var env)

instance Types TypeEnv where
	ftv (TypeEnv env) = ftv ('DM'.elems env)
	apply s (TypeEnv env) = TypeEnv (apply s <$> env)

generalize :: TypeEnv Type -> Scheme
generalize env t = Scheme (difference (ftv t) (ftv env)) t

:: TI a :== RWST TIEnv () TIState (Either [String]) a
:: TIState = {tiSupply :: Int, tiSubst :: Subst}
:: TIEnv :== [Int]

mgu :: Type Type -> TI Subst
mgu (TFun l r) (TFun l` r`) = composeSubst <$> mgu l l` <*> mgu r r`
mgu (TVar u) t = varBind u t
mgu t (TVar u) = varBind u t
mgu TInt TInt = pure 'DM'.newMap
mgu TBool TBool = pure 'DM'.newMap
mgu TChar TChar = pure 'DM'.newMap
mgu t1 t2 = liftT (Left ["cannot unify: " +++ toString t1 +++ " with " +++ toString t2])

varBind :: [Char] Type -> TI Subst
varBind u t
	| t == TVar u = pure 'DM'.newMap
	| isMember u (ftv t) = liftT (Left ["occur check fails: " +++ toString u +++ " vs. " +++ toString t])
	= pure ('DM'.singleton u t)

newTyVar :: [Char] -> TI Type
newTyVar prefix
	=   get
	>>= \t->put {t & tiSupply=t.tiSupply+1}
	>>| pure (TVar (prefix ++ fromString (toString t.tiSupply)))

instantiate :: Scheme -> TI Type
instantiate (Scheme vars t)
	= mapM (\_->newTyVar ['a']) vars
	>>= \nvars->pure (apply ('DM'.fromList (zip2 vars nvars)) t)

class infer a :: TypeEnv a -> TI (Subst, Type)

instance infer Value where
	infer _ (Int _) = pure ('DM'.newMap, TInt)
	infer _ (Bool _) = pure ('DM'.newMap, TBool)
	infer _ (Char _) = pure ('DM'.newMap, TChar)

instance infer Expression where
	infer e (Lit a) = infer e a
	infer (TypeEnv env) (Var v) = case 'DM'.get v env of
		Nothing = liftT (Left ["unbound variable: " +++ toString v])
		Just s = instantiate s >>= \t->pure ('DM'.newMap, t)
	infer env (App e1 e2)
		=              newTyVar ['a']
		>>= \tv      ->infer env e1
		>>= \(s1, t1)->infer (apply s1 env) e2
		>>= \(s2, t2)->mgu (apply s2 t1) (TFun t2 tv)
		>>= \s3->pure (composeSubst s3 (composeSubst s2 s1), apply s3 tv)
	infer env (Lambda s e)
		=        newTyVar ['l']
		>>= \tv->
			let (TypeEnv env`) = remove env s
			    env`` = TypeEnv ('DM'.union env` ('DM'.singleton s (Scheme [] tv)))
			in infer env`` e
		>>= \(s1, t1)->pure (s1, TFun (apply s1 tv) t1)

instance infer [Function] where
	infer env [] = pure ('DM'.newMap, TInt)
	infer env [Function name args body:rest]
		= infer env (foldr Lambda body args) >>= \(s1, t1)->
				let (TypeEnv env`) = remove env name
					t`             = generalize (apply s1 env) t1
					env``          = TypeEnv ('DM'.put name t` env`)
				in infer (apply s1 env``) rest >>= \(s2, t2)->pure (composeSubst s1 s2, t2)

typeInference :: ('DM'.Map [Char] Scheme) Expression -> TI Type
typeInference env e = uncurry apply <$> infer (TypeEnv env) e