From ca4926527ed9f7bb2ada88a71228e4a90e0d2539 Mon Sep 17 00:00:00 2001
From: janmasrovira <janmasrovira@gmail.com>
Date: Thu, 16 Feb 2023 12:54:53 +0100
Subject: [PATCH] Add type annotation to case expression (#1849)

---
 src/Juvix/Compiler/Internal/Extra.hs                            | 1 +
 src/Juvix/Compiler/Internal/Language.hs                         | 2 ++
 src/Juvix/Compiler/Internal/Translation/FromAbstract.hs         | 1 +
 .../Translation/FromInternal/Analysis/ArityChecking/Checker.hs  | 1 +
 .../Translation/FromInternal/Analysis/TypeChecking/Checker.hs   | 1 +
 5 files changed, 6 insertions(+)

diff --git a/src/Juvix/Compiler/Internal/Extra.hs b/src/Juvix/Compiler/Internal/Extra.hs
index 71f9acdf2..fbf36a70b 100644
--- a/src/Juvix/Compiler/Internal/Extra.hs
+++ b/src/Juvix/Compiler/Internal/Extra.hs
@@ -250,6 +250,7 @@ instance HasExpressions Case where
   leafExpressions f l = do
     _caseBranches :: NonEmpty CaseBranch <- traverse (leafExpressions f) (l ^. caseBranches)
     _caseExpression <- leafExpressions f (l ^. caseExpression)
+    _caseExpressionType <- traverse (leafExpressions f) (l ^. caseExpressionType)
     pure Case {..}
     where
       _caseParens = l ^. caseParens
diff --git a/src/Juvix/Compiler/Internal/Language.hs b/src/Juvix/Compiler/Internal/Language.hs
index b453e94bb..d19b02edd 100644
--- a/src/Juvix/Compiler/Internal/Language.hs
+++ b/src/Juvix/Compiler/Internal/Language.hs
@@ -155,6 +155,8 @@ instance Hashable CaseBranch
 
 data Case = Case
   { _caseExpression :: Expression,
+    -- | The typechecker fills this field
+    _caseExpressionType :: Maybe Expression,
     _caseBranches :: NonEmpty CaseBranch,
     _caseParens :: Bool
   }
diff --git a/src/Juvix/Compiler/Internal/Translation/FromAbstract.hs b/src/Juvix/Compiler/Internal/Translation/FromAbstract.hs
index 8edd0aec7..c6d940fc8 100644
--- a/src/Juvix/Compiler/Internal/Translation/FromAbstract.hs
+++ b/src/Juvix/Compiler/Internal/Translation/FromAbstract.hs
@@ -346,6 +346,7 @@ goCase c = do
   _caseExpression <- goExpression (c ^. Abstract.caseExpression)
   _caseBranches <- mapM goCaseBranch (c ^. Abstract.caseBranches)
   let _caseParens = c ^. Abstract.caseParens
+      _caseExpressionType :: Maybe Expression = Nothing
   return Case {..}
 
 goCaseBranch :: Members '[NameIdGen] r => Abstract.CaseBranch -> Sem r CaseBranch
diff --git a/src/Juvix/Compiler/Internal/Translation/FromInternal/Analysis/ArityChecking/Checker.hs b/src/Juvix/Compiler/Internal/Translation/FromInternal/Analysis/ArityChecking/Checker.hs
index 214a3101d..50a4ef122 100644
--- a/src/Juvix/Compiler/Internal/Translation/FromInternal/Analysis/ArityChecking/Checker.hs
+++ b/src/Juvix/Compiler/Internal/Translation/FromInternal/Analysis/ArityChecking/Checker.hs
@@ -356,6 +356,7 @@ checkCase ari l = do
   _caseBranches <- mapM checkCaseBranch (l ^. caseBranches)
   _caseExpression <- checkExpression ArityUnit (l ^. caseExpression)
   let _caseParens = l ^. caseParens
+      _caseExpressionType :: Maybe Expression = Nothing
   return Case {..}
   where
     checkCaseBranch :: CaseBranch -> Sem r CaseBranch
diff --git a/src/Juvix/Compiler/Internal/Translation/FromInternal/Analysis/TypeChecking/Checker.hs b/src/Juvix/Compiler/Internal/Translation/FromInternal/Analysis/TypeChecking/Checker.hs
index 0a6865bcc..f794dc0a1 100644
--- a/src/Juvix/Compiler/Internal/Translation/FromInternal/Analysis/TypeChecking/Checker.hs
+++ b/src/Juvix/Compiler/Internal/Translation/FromInternal/Analysis/TypeChecking/Checker.hs
@@ -555,6 +555,7 @@ inferExpression' hint e = case e of
         Just hi -> return hi
       typedCaseExpression <- inferExpression' Nothing (c ^. caseExpression)
       let _caseExpression = typedCaseExpression ^. typedExpression
+          _caseExpressionType = Just (typedCaseExpression ^. typedType)
           goBranch :: CaseBranch -> Sem r CaseBranch
           goBranch b =
             traverseOf