From 29293af45b6d9e33eddf5ef5bf53ca5c1d39ab3d Mon Sep 17 00:00:00 2001
From: Troels Henriksen <athas@sigkill.dk>
Date: Sun, 13 Oct 2024 19:17:39 +0200
Subject: [PATCH] Branch reasons.

---
 src/Language/Futhark/TypeChecker/Constraints.hs | 10 ++++++++++
 src/Language/Futhark/TypeChecker/Terms2.hs      |  6 +++---
 2 files changed, 13 insertions(+), 3 deletions(-)

diff --git a/src/Language/Futhark/TypeChecker/Constraints.hs b/src/Language/Futhark/TypeChecker/Constraints.hs
index 6ed3b64a84..c8ceb92763 100644
--- a/src/Language/Futhark/TypeChecker/Constraints.hs
+++ b/src/Language/Futhark/TypeChecker/Constraints.hs
@@ -69,6 +69,7 @@ data Reason
     ReasonAscription Loc Type Type
   | ReasonRetType Loc Type Type
   | ReasonApply Loc (Maybe (QualName VName)) Exp Type Type
+  | ReasonBranches Loc Type Type
   deriving (Eq, Show)
 
 instance Located Reason where
@@ -77,6 +78,7 @@ instance Located Reason where
   locOf (ReasonAscription l _ _) = l
   locOf (ReasonRetType l _ _) = l
   locOf (ReasonApply l _ _ _ _) = l
+  locOf (ReasonBranches l _ _) = l
 
 data Ct
   = CtEq Reason Type Type
@@ -359,6 +361,14 @@ cannotUnify reason notes bcs t1 t2 = do
                 <+> "to"
                 <+> dquotes (align $ shorten $ group $ pretty e)
                 <> " (invalid type)."
+    ReasonBranches loc former latter -> do
+      former' <- enrichType former
+      latter' <- enrichType latter
+      typeError loc notes . stack $
+        [ "Branches differ in type.",
+          "Former:" <+> pretty former',
+          "Latter:" <+> pretty latter'
+        ]
 
 -- Precondition: 'v' is currently flexible.
 subTyVar :: Reason -> BreadCrumbs -> VName -> Type -> SolveM ()
diff --git a/src/Language/Futhark/TypeChecker/Terms2.hs b/src/Language/Futhark/TypeChecker/Terms2.hs
index 89b5320a67..41594c4282 100644
--- a/src/Language/Futhark/TypeChecker/Terms2.hs
+++ b/src/Language/Futhark/TypeChecker/Terms2.hs
@@ -810,7 +810,7 @@ checkCases mt rest_cs =
     (c, Just cs) -> do
       (c', c_t) <- checkCase mt c
       (cs', cs_t) <- checkCases mt cs
-      ctEq (Reason (locOf c)) c_t cs_t
+      ctEq (ReasonBranches (locOf c) c_t cs_t) c_t cs_t
       pure (NE.cons c' cs', c_t)
 
 -- | An unmatched pattern. Used in in the generation of
@@ -1203,8 +1203,8 @@ checkExp (AppExp (If e1 e2 e3 loc) _) = do
   if_t <- newType loc SizeLifted "if_t" NoUniqueness
 
   ctEq (Reason (locOf e1')) e1_t (Scalar (Prim Bool))
-  ctEq (Reason (locOf loc)) e2_t if_t
-  ctEq (Reason (locOf loc)) e3_t if_t
+  ctEq (ReasonBranches (locOf loc) e2_t e3_t) e2_t if_t
+  ctEq (ReasonBranches (locOf loc) e2_t e3_t) e3_t if_t
 
   if_t' <- asStructType if_t
   pure $ AppExp (If e1' e2' e3' loc) (Info $ AppRes if_t' [])