module Typecheck where

import Prelude hiding (lookup)
import Parser.AbsSchmim
import Data.Map.Lazy
import Control.Monad.Reader

type Context = Map Ident Type
type Typechecker a = ReaderT Context Maybe a

assert :: MonadPlus m => Bool -> m ()
assert True = return ()
assert False = mzero

typeof :: Exp -> Typechecker Type
typeof ETr = return TBool
typeof EFl = return TBool
typeof (ENum _) = return TInt
typeof (EVar x) = do
    mt <- asks $ lookup x
    case mt of 
        (Just t) -> return t
        Nothing -> lift Nothing
typeof (EAdd e1 e2) = do
    t1 <- typeof e1
    t2 <- typeof e2
    assert (t1 == TInt)
    assert (t2 == TInt)
    return TInt
typeof (ESub e1 e2) = do
    t1 <- typeof e1
    t2 <- typeof e2
    assert (t1 == TInt)
    assert (t2 == TInt)
    return TInt
typeof (EMul e1 e2) = do
    t1 <- typeof e1
    t2 <- typeof e2
    assert (t1 == TInt)
    assert (t2 == TInt)
    return TInt
typeof (EDiv e1 e2) = do
    t1 <- typeof e1
    t2 <- typeof e2
    assert (t1 == TInt)
    assert (t2 == TInt)
    return TInt
typeof (EAnd e1 e2) = do
    t1 <- typeof e1
    t2 <- typeof e2
    assert (t1 == TBool)
    assert (t2 == TBool)
    return TBool
typeof (EOr e1 e2) = do
    t1 <- typeof e1
    t2 <- typeof e2
    assert (t1 == TBool)
    assert (t2 == TBool)
    return TBool
typeof (ENot e) = do
    t <- typeof e
    assert (t == TBool)
    return TBool
typeof (EEq e1 e2) = do
    t1 <- typeof e1
    t2 <- typeof e2
    assert (t1 == TInt)
    assert (t2 == TInt)
    return TBool
typeof (ELeq e1 e2) = do
    t1 <- typeof e1
    t2 <- typeof e2
    assert (t1 == TInt)
    assert (t2 == TInt)
    return TBool
typeof (EGeq e1 e2) = do
    t1 <- typeof e1
    t2 <- typeof e2
    assert (t1 == TInt)
    assert (t2 == TInt)
    return TBool
typeof (ELess e1 e2) = do
    t1 <- typeof e1
    t2 <- typeof e2
    assert (t1 == TInt)
    assert (t2 == TInt)
    return TBool
typeof (EGrt e1 e2) = do
    t1 <- typeof e1
    t2 <- typeof e2
    assert (t1 == TInt)
    assert (t2 == TInt)
    return TBool
typeof (EIfte e1 e2 e3) = do
    t1 <- typeof e1
    t2 <- typeof e2
    t3 <- typeof e3
    assert (t1 == TBool)
    assert (t2 == t3)
    return t2
typeof (EAbs x t e) = do
    t' <- local (insert x t) $ typeof e
    return $ TFun t t'
typeof (EApp e []) = typeof e
typeof (EApp e1 [e2]) = do
    t1 <- typeof e1
    t2 <- typeof e2
    case t1 of
        (TFun t t')
            | t == t2 -> return t'
        _ -> lift Nothing
typeof (EApp e1 es) = do
    t1 <- typeof e1
    ts <- forM es typeof
    applyTypes t1 ts
    where
        applyTypes t [] = return t
        applyTypes (TFun t t') (t1:ts) = do
            assert $ t == t1
            applyTypes t' ts
typeof (EVrnt l e t@(TVrnt ls)) = do
    t1 <- typeof e
    (Labeled l' t2) <- lift $ lookupLabel l ls
    assert $ t1 == t2
    return t
    where
        lookupLabel l [] = Nothing
        lookupLabel l (lbl@(Labeled l1 t):ls)
            | l == l1 = Just lbl
            | otherwise = lookupLabel l ls
typeof (EVrnt l e _) = lift Nothing
typeof (ETpl es) = do
    ts <- forM es typeof
    return $ TTpl ts
typeof (EProj e n) = do
    t <- typeof e
    case t of
        (TTpl ts)
            | length ts >= fromIntegral n -> return $ ts !! (fromIntegral (n - 1))
        _ -> lift Nothing
typeof (EMtch e ms) = do
    t <- typeof e
    case t of
        (TVrnt ls) -> do
            assert $ (length ms == length ls)
            ts <- forM (zip ms ls) getMatchingType
            assert $ allSame ts
            return $ head ts
        _ -> lift Nothing
    where
        getMatchingType (Matching l1 x e, Labeled l2 t) = do
            assert $ l1 == l2
            local (insert x t) $ typeof e
        allSame [] = True
        allSame [_] = True
        allSame (x:y:zs)
            | x == y = allSame (y:zs)
            | otherwise = False
typeof (ENil t) = do
    case t of
        (TList _) -> return t
        _ -> lift Nothing
typeof (ECons e1 e2) = do
    t1 <- typeof e1
    t2 <- typeof e2
    case t2 of
        (TList t) -> do
            assert $ t == t1
            return t2
        _ -> lift Nothing
typeof (EHead e) = do
    t <- typeof e
    case t of
        (TList t') -> return t'
        _ -> lift Nothing
typeof (ETail e) = do
    t <- typeof e
    case t of
        (TList _) -> return t
        _ -> lift Nothing
typeof (EIsnil e) = do
    t <- typeof e
    case t of
        (TList _) -> return TBool
        _ -> lift Nothing
typeof (ELet f ps e1 e2) = do
    t1 <- local (insertAll ps) $ typeof e1
    local (insert f (prependParamTypes t1 ps)) $ typeof e2
    where
        insertAll [] env = env
        insertAll ((Param x t):ps) env = insertAll ps $ insert x t env
        prependParamTypes t ps = Prelude.foldr (\(Param _ t2) t1 -> TFun t2 t1) t ps
typeof (EFix f) = do
    t <- typeof f
    case t of
        (TFun t1 t2)
            | t1 == t2 -> return t1
        _ -> lift Nothing
typeof (ELetrec f t ps e1 e2) = do
    typeof (EApp (EAbs f t e2) [(EFix (EAbs f t e1'))])
    where
        e1' = Prelude.foldr (\(Param x t') e -> (EAbs x t' e)) e1 ps