1 files changed, 71 insertions, 21 deletions

diff --git a/src/Typecheck.hs b/src/Typecheck.hs
index c86187d..2ed189d 100644
--- a/src/Typecheck.hs
+++ b/src/Typecheck.hs
@@ -1,31 +1,46 @@
 {-# LANGUAGE TupleSections #-}
 module Typecheck
 (inferExpr,
- checkExpr)
+ checkExpr,
+ Environment)
   where
 import Syntax
-import Control.Monad
+--import Control.Monad
+import Control.Monad.Except
 type Environment = [(Idnt,Typ)]
 
-assertInt TInt = Just ()
-assertInt _ = Nothing
-assertRec (TRecord ts) = Just ts
-assertRec _ = Nothing
-assertVar (TVariant ts) = Just ts
-assertVar _ = Nothing
-assertPtr (TPtr t) = Just t
-assertPtr _ = Nothing
+type Error = Except String
 
+assertInt :: Typ -> Expr -> Error ()
+assertInt TInt _ = return ()
+assertInt t e = throwError $ "Expected type of expression "++ show e ++ " to be (TInt), instead got" ++ show t 
+  
+assertRec :: Typ -> Expr -> Error TypeDict
+assertRec (TRecord ts) _ = return ts
+assertRec t e = throwError $ "Expected type of expression "++ show e ++ " to be (TRecord _), instead got" ++ show t 
 
-inferExpr :: Expr -> Environment -> Maybe Typ
-inferExpr (Var name) env = lookup name env
-inferExpr (IntLit _) _ = Just TInt
+assertVar :: Typ -> Expr -> Error TypeDict
+assertVar (TVariant ts) _ = return ts
+assertVar t e = throwError $ "Expected type of expression "++ show e ++ " to be (TVariant _), instead got" ++ show t 
+
+assertPtr :: Typ -> Expr -> Error Typ
+assertPtr (TPtr t) _ = return t
+assertPtr t e = throwError $ "Expected type of expression "++ show e ++ " to be (TPtr _), instead got" ++ show t 
+
+envlookup :: Idnt -> TypeDict -> Expr -> Error Typ
+envlookup name env expr = case lookup name env of
+  Just t -> return t
+  Nothing -> throwError $ "Variable "++ name ++ "not defined in expr" ++ show expr
+
+inferExpr :: Expr -> Environment -> Error Typ
+inferExpr e@(Var name) env = envlookup name env e
+inferExpr (IntLit _) _ = return TInt
 inferExpr (Record ts) env = do
-  tts <- sequence $ map (\(v,e) -> (v,) <$> inferExpr e env) ts
+  tts <- mapM (\(v,e) -> (v,) <$> inferExpr e env) ts
   return $ TRecord tts
-inferExpr (Variant t v e) env = do
-  ts <- assertVar t
-  tt <- lookup v ts
+inferExpr ee@(Variant t v e) env = do
+  ts <- assertVar t e
+  tt <- envlookup v ts ee
   checkExpr e env tt
   return t
 inferExpr (Add e1 e2) env = do
@@ -41,11 +56,46 @@ inferExpr (Neg e) env = do
   return TInt
 inferExpr (Deref e1) env = do
   t <- inferExpr e1 env
-  tt <- assertPtr t
+  tt <- assertPtr t e1
   return tt
+inferExpr c@(Match e cases) env = do
+  varType <- inferExpr e env
+  variants <- assertVar varType e
+  -- TODO ensure variants and cases cover the same names
+  caseTypes <- mapM (\(v,e) ->(v,) <$> inferExpr e env) cases
+  ret <- allSame caseTypes
+  return $ snd ret
+  where
+    allSame :: [(Idnt,Typ)] -> Error (Idnt,Typ)
+    allSame (x:[]) = return x
+    allSame [] = throwError $ "Empty case "++ show c
+    allSame ((xx,x):xs) = do
+      (v,t) <- allSame xs
+      if t==x then return (v,t) else
+        throwError $ "Different types in branches " ++ show v ++ "," ++ show xx ++ "of expression" ++ show c
+inferExpr (Proj name e) env = do
+  t <- inferExpr e env
+  tt <- assertRec t e
+  case lookup name tt of
+    Just typ -> return typ
+    Nothing -> throwError $ "Record " ++show e ++ "doesn't have field" ++ name
 
-checkExpr :: Expr -> Environment -> Typ -> Maybe ()
+checkExpr :: Expr -> Environment -> Typ -> Error ()
+checkExpr e@(Add e1 e2) env t = do
+  assertInt t e
+  checkExpr e1 env TInt
+  checkExpr e2 env TInt
+checkExpr e@(Mult e1 e2) env t = do
+  assertInt t e
+  checkExpr e1 env TInt
+  checkExpr e2 env TInt
+checkExpr e@(Neg ee) env t = do
+  assertInt t e
+  checkExpr ee env TInt
+checkExpr (Deref e1) env t = do
+ checkExpr e1 env (TPtr t)
 checkExpr e env t = do
   t' <- inferExpr e env
-  guard (t==t')
-  return ()
+  if t == t' then return ()
+    else throwError $ "Expected type of expression " ++
+         show e ++ " to be"++ show t  ++" instead got" ++ show t'