unification??

gistfile1.hs

{-# OPTIONS_GHC -Wall
                -fno-warn-missing-signatures
                -fno-warn-hi-shadowing
                -fno-warn-unused-do-bind
                -fno-warn-name-shadowing #-}

module Poly where

--import Control.Monad.Identity
import Control.Monad.Error
import Control.Monad.State
import Data.Unique
import qualified Data.Map as M


type TypeError = String
type ThrowsTyErr = ErrorT TypeError IO


type TC = StateT TCState ThrowsTyErr

debug :: String -> TC ()
debug = liftIO . putStrLn

-- Env holds let bindings. we must use `retrieve` for getting from env.
-- retrieve copies generic types so that we can instantiate types
-- in each call of generic functions.
--type Env = [(String, Ty)]
type Env = M.Map String Ty

retrieve :: Env -> VarEnv -> NonGenerics -> String -> TC Ty
-- FIXME: search var in non-generic var list
retrieve env varenv nongenerics name = do
    case M.lookup name env of
        Nothing -> throwError "unbound var -- retrieve"
        Just ty -> liftM fst (liftIO $ freshTy nongenerics M.empty (prune varenv ty))

type CopyEnv = M.Map TypeVar TypeVar

freshTy :: NonGenerics -> CopyEnv -> Ty -> IO (Ty, CopyEnv)
freshTy nongen copyEnv (TyArr a b) = do
    (param, copyEnv') <- freshTy nongen copyEnv a
    (result, copyEnv'') <- freshTy nongen copyEnv' b
    return $ (TyArr param result, copyEnv'')
freshTy nongen copyEnv (TyList a) = do
    (a', copyEnv') <- freshTy nongen copyEnv a
    return (TyList a', copyEnv')
freshTy nongen copyEnv t@(TyVar v) =
    if v `elem` nongen
        then return (t, copyEnv)
        else case M.lookup v copyEnv of
                 Just t' -> return (TyVar t', copyEnv)
                 Nothing -> do
                     newVar@(TyVar v') <- newTyVar
                     let copyEnv' = M.insert v v' copyEnv
                     return (newVar, copyEnv')
freshTy _ env t = return (t, env)


--type NonGenerics = [(TypeVar, Ty)]
type NonGenerics = [TypeVar]

newNonGenerics :: NonGenerics
newNonGenerics = []

--type VarEnv = [(TypeVar, Ty)]
type VarEnv = M.Map TypeVar Ty

prune :: VarEnv -> Ty -> Ty
prune env (TyVar var) =
    case M.lookup var env of
        Nothing -> TyVar var
        Just v -> prune env v
prune env (TyArr v1 v2) =
    let v1' = prune env v1
        v2' = prune env v2
    in
        TyArr v1' v2'
prune env (TyList a) =
    TyList (prune env a)
prune _ ty = ty

--data TCState = TCState { env :: Env }
data TCState = TCState

--newTCState :: TCState
--newTCState = TCState []

runTC = runStateT

data Term = TmTrue
          | TmFalse
          | TmInt Int
          | TmUnit
          | TmIf Term Term Term
          | TmVar String
          | TmAbs String Term
          | TmApp Term Term
          | TmList
          | TmLet Declaration Term
    deriving Show

data Declaration = Def String Term
                 | Rec String Term
    deriving Show

data Ty = TyBool
        | TyInt
        | TyUnit
        | TyArr Ty Ty
        | TyList Ty
        -- | TyVar (Either TypeVar Ty)
        | TyVar TypeVar
    deriving (Eq)

instance Show Ty where
    show TyInt = "Int"
    show TyBool = "Bool"
    show TyUnit = "Unit"
    show (TyArr a b) = "(" ++ show a ++ " -> " ++ show b ++ ")"
    show (TyList l) = "[" ++ show l ++ "]"
    show (TyVar v) = show v

type TypeVar = Unique

instance Show Unique where
    show u = "t" ++ (show $ hashUnique u)

newTyVar :: IO Ty
newTyVar = liftM TyVar (liftIO newUnique)

initEnv :: IO Env
initEnv = do
    var <- newTyVar
    return $ M.fromList
           [ ("first", TyArr (TyList var) var)
           , ("rest", TyArr (TyList var) (TyList var))
           , ("null", TyArr (TyList var) TyBool)
           , ("cons", TyArr var (TyArr (TyList var) (TyList var)))
           , ("empty", TyList var)
           ]

occursInType :: Ty -> Ty -> Bool
occursInType (TyVar var1) (TyVar var2) = var1 == var2
occursInType ty@TyVar{} (TyArr a b) = occursInType ty a || occursInType ty b
occursInType ty@TyVar{} (TyList t) = occursInType ty t
occursInType _ _ = False

unify :: VarEnv -> Ty -> Ty -> TC VarEnv
unify varenv t1 t2 =
    let t1' = prune varenv t1
        t2' = prune varenv t2
    in do
        debug $ "t1': " ++ show t1' ++ " t2': " ++ show t2'
        case (t1', t2') of
            (TyVar v, _) -> if occursInType t1' t2'
                                then throwError "thpe clash -- unify"
                                else do
                                    debug $ "unifying variable " ++ show t1' ++ " with " ++ show t2'
                                    return $ M.insert v t2' varenv
            (_, TyVar _) -> unify varenv t2' t1'
            (TyArr a1 b1, TyArr a2 b2) -> do
                env <- unify varenv a1 a2
                unify env b1 b2
            (TyList a, TyList b) -> unify varenv a b
            _ -> throwError "error -- unify"


analyzeExp :: Env -> VarEnv -> NonGenerics -> Term -> TC Ty
analyzeExp _ _ _ (TmTrue) = return TyBool
analyzeExp _ _ _ (TmFalse) = return TyBool
analyzeExp _ _ _ (TmUnit) = return TyUnit
analyzeExp _ _ _ (TmInt _) = return TyInt

analyzeExp env varenv nongenerics (TmVar name) = do
    debug $ "analyze var: " ++ name
    r <- retrieve env varenv nongenerics name
    return $ prune varenv r
analyzeExp env varenv nongenerics (TmIf guard ifE thenE) = do
    debug "analyze if"
    guardTy <- analyzeExp env varenv nongenerics guard
    debug $ "unifying guard to Bool"
    varenv' <- unify varenv guardTy TyBool

    ifEty <- analyzeExp env varenv' nongenerics ifE
    thenEty <- analyzeExp env varenv' nongenerics thenE

    debug $ "unifying ifE and thenE"
    varenv'' <- unify varenv' ifEty thenEty

    return $ prune varenv'' ifEty
analyzeExp env varenv nongenerics (TmAbs param body) = do
    debug "analyze abs"

    binder@(TyVar v) <- liftIO newTyVar

    let env' = M.insert param binder env
    --extendEnv param binder
    let nongenerics' = v:nongenerics    
    --extendNonGenerics binder

    bodyTy <- analyzeExp env' varenv nongenerics' body
    debug "finished analyzing body"
    return $ prune varenv (TyArr binder bodyTy)
analyzeExp env varenv nongenerics (TmApp fun param) = do
    debug "analyze app"
    funty <- analyzeExp env varenv nongenerics fun
    debug $ "finished analyzing fun: " ++ show funty

    argty <- analyzeExp env varenv nongenerics param
    debug $ "finished analyzing param: " ++ show argty

    res@(TyVar v) <- liftIO newTyVar
    debug $ "unifying fun: (" ++ show funty ++ ") (" ++ show (TyArr argty res) ++ ")"
    varenv' <- unify varenv funty (TyArr argty res)

    debug $ show $ M.toList varenv'
    --debug $ show $ M.toList env

    --case M.lookup v varenv' of
        --Just t -> return t
        --Nothing -> return res
    return $ prune varenv' res
analyzeExp env varenv nongenerics TmList = do
    var <- liftIO newTyVar
    return $ prune varenv (TyList var)
analyzeExp env varenv nongenerics (TmLet decl scope) = do
    env' <- analyzeDecl env varenv nongenerics decl
    analyzeExp env' varenv nongenerics scope

analyzeDecl :: Env -> VarEnv -> NonGenerics -> Declaration -> TC Env
analyzeDecl env varenv nongenerics (Def name term) = do
    ty <- analyzeExp env varenv nongenerics term
    return $ M.insert name ty env
analyzeDecl env varenv nongenerics (Rec name term) = do
    newVar <- liftIO newTyVar
    let env' = M.insert name newVar env
    ty <- analyzeExp env' varenv nongenerics term
    return $ M.insert name ty env'



main :: IO ()
main = do
    env <- initEnv
    let varenv = M.empty
    let nongenerics = []

    --let term = TmApp (TmVar "first") (TmApp (TmApp (TmVar "cons") TmTrue) TmList)
    --let term = TmApp (TmApp (TmVar "cons") TmTrue) TmList
    let term = TmLet (Def "id1" (TmAbs "a" (TmVar "a"))) (TmApp (TmVar "id1") TmTrue)
    ty <- runErrorT $ flip evalStateT TCState (analyzeExp env varenv nongenerics term)
    putStrLn $ show ty
    return ()