From afbc700db24f88e30c342413e4db60e285da3492 Mon Sep 17 00:00:00 2001
From: Samuel Hammersberg <samuel.hammersberg@gmail.com>
Date: Mon, 20 Feb 2023 16:43:54 +0100
Subject: [PATCH] Fixed the type checker accidentally chucking cases in some
 cases.

---
 src/TypeChecker.hs | 19 ++++++++++++++-----
 1 file changed, 14 insertions(+), 5 deletions(-)

diff --git a/src/TypeChecker.hs b/src/TypeChecker.hs
index 9cb9c39..bb31a2f 100644
--- a/src/TypeChecker.hs
+++ b/src/TypeChecker.hs
@@ -154,6 +154,20 @@ check cxt exp typ = case exp of
         unless (typeEq t1 typ) $ throwError "Wrong lamda type!"
         pure $ T.EAbs t1 (x, t) e'
 
+    ECase e cs t -> do
+        (e',t1) <- infer cxt e
+        unless (typeEq t t1) $
+            throwError "Inferred type and type annotation doesn't match"
+        case traverse (\(CaseMatch c e) -> do
+            -- //TODO check c as well
+            e' <- check cxt e t
+            unless (typeEq t t1) $
+                throwError "Inferred type and type annotation doesn't match"
+            pure (t1, T.Case c e')
+            ) cs of
+                Right cs -> pure $ T.ECase t1 e' cs
+                Left e   -> throwError e
+
     ELet b e -> do
         let cxt' = insertBind b cxt
         b' <- checkBind cxt' b
@@ -165,11 +179,6 @@ check cxt exp typ = case exp of
             throwError "Inferred type and type annotation doesn't match"
         check cxt e t
 
-    ECase e _ t -> do
-        unless (typeEq t typ) $
-            throwError "Inferred type and type annotation doesn't match"
-        check cxt e t
-
 -- | Check if types are equivalent. Doesn't handle coercion or polymorphism.
 typeEq :: Type -> Type -> Bool
 typeEq (TFun t t1) (TFun q q1) = typeEq t q && typeEq t1 q1