Add return type to Case constructors

src/ClosureConversion.hs

@@ -84,10 +84,11 @@ convertAppliedTerm term args =
     LambdaLifted.App fun arg -> do
       arg' <- convertTerm arg
       convertAppliedTerm fun $ arg' : args
-    LambdaLifted.Case scrutinee branches defaultBranch ->
+    LambdaLifted.Case scrutinee type_ branches defaultBranch ->
       applyArgs args $
         ClosureConverted.Case
           <$> convertTerm scrutinee
+          <*> convertTerm type_
           <*> convertBranches branches
           <*> mapM convertTerm defaultBranch
 

src/ClosureConverted/Domain.hs

@@ -32,7 +32,7 @@ data Head
 type Type = Value
 
 data Branches where
-  Branches :: Environment v -> Syntax.Branches v -> Maybe (Syntax.Term v) -> Branches
+  Branches :: Type -> Environment v -> Syntax.Branches v -> Maybe (Syntax.Term v) -> Branches
 
 data Closure where
   Closure :: Environment v -> Scope Syntax.Term v -> Closure

src/ClosureConverted/Evaluation.hs

@@ -71,9 +71,10 @@ evaluate env term =
       fun' <- evaluate env fun
       args' <- mapM (evaluate env) args
       apply env fun' args'
-    Syntax.Case scrutinee branches defaultBranch -> do
+    Syntax.Case scrutinee type_ branches defaultBranch -> do
       scrutineeValue <- evaluate env scrutinee
-      case_ scrutineeValue $ Domain.Branches env branches defaultBranch
+      type' <- lazy $ evaluate env type_
+      case_ scrutineeValue $ Domain.Branches (Domain.Lazy type') env branches defaultBranch
 
 chooseConstructorBranch
   :: Domain.Environment v
@@ -190,7 +191,7 @@ applyTelescope env tele args k =
       pure Nothing
 
 case_ :: Domain.Value -> Domain.Branches -> M Domain.Value
-case_ scrutinee branches@(Domain.Branches env branches' defaultBranch) =
+case_ scrutinee branches@(Domain.Branches _ env branches' defaultBranch) =
   case (scrutinee, branches') of
     (Domain.Con constr _params args, Syntax.ConstructorBranches _ constructorBranches) ->
       chooseConstructorBranch env constr args constructorBranches defaultBranch

src/ClosureConverted/Readback.hs

@@ -62,12 +62,13 @@ readbackGlobal env global spine =
       readbackGroupedSpine env global' groupedSpine
 
 readbackGroupedElimination :: Domain.Environment v -> Syntax.Term v -> Domain.GroupedElimination -> M (Syntax.Term v)
-readbackGroupedElimination env eliminee elimination =
+readbackGroupedElimination env eliminee elimination = do
   case elimination of
     Domain.GroupedApps args -> do
       args' <- mapM (readback env) args
       ClosureConversion.applyArgs args' $ pure eliminee
-    Domain.GroupedCase (Domain.Branches env' branches defaultBranch) -> do
+    Domain.GroupedCase (Domain.Branches type_ env' branches defaultBranch) -> do
+      type' <- readback env type_
       branches' <- case branches of
         Syntax.ConstructorBranches constructorTypeName constructorBranches ->
           Syntax.ConstructorBranches constructorTypeName <$> OrderedHashMap.forMUnordered constructorBranches (readbackConstructorBranch env env')
@@ -82,7 +83,7 @@ readbackGroupedElimination env eliminee elimination =
       defaultBranch' <- forM defaultBranch \branch -> do
         branch' <- Evaluation.evaluate env' branch
         readback env branch'
-      pure $ Syntax.Case eliminee branches' defaultBranch'
+      pure $ Syntax.Case eliminee type' branches' defaultBranch'
 
 readbackGroupedSpine :: (Foldable f) => Domain.Environment v -> Syntax.Term v -> f Domain.GroupedElimination -> M (Syntax.Term v)
 readbackGroupedSpine =

src/ClosureConverted/Syntax.hs

@@ -57,7 +57,7 @@ data Term v
   | Pi !Name !(Type v) !(Scope Type v)
   | Closure !Name.Lifted [Term v]
   | ApplyClosure !(Term v) [Term v]
-  | Case !(Term v) (Branches v) !(Maybe (Term v))
+  | Case !(Term v) !(Type v) (Branches v) !(Maybe (Term v))
   deriving (Eq, Show, Generic, Hashable)
 
 type Type = Term

src/ClosureConverted/TypeOf.hs

@@ -10,7 +10,6 @@ import qualified ClosureConverted.Domain as Domain
 import qualified ClosureConverted.Evaluation as Evaluation
 import qualified ClosureConverted.Readback as Readback
 import qualified ClosureConverted.Syntax as Syntax
-import qualified Data.OrderedHashMap as OrderedHashMap
 import qualified Environment
 import qualified Literal
 import Monad
@@ -68,7 +67,7 @@ typeOf context value =
   case value of
     Domain.Neutral head spine -> do
       headType <- typeOfHead context head
-      typeOfSpineApplication context headType spine
+      typeOfSpineApplication headType spine
     Domain.Con con params args -> do
       conType <- fetch $ Query.ClosureConvertedConstructorType con
       conType' <-
@@ -113,38 +112,19 @@ typeOfHead context head =
 
 typeOfSpineApplication
   :: (Foldable f)
-  => Context v
-  -> Domain.Type
+  => Domain.Type
   -> f Domain.Elimination
   -> M Domain.Type
 typeOfSpineApplication =
-  foldlM . typeOfElimination
+  foldlM typeOfElimination
 
-typeOfElimination :: Context v -> Domain.Type -> Domain.Elimination -> M Domain.Type
-typeOfElimination context type_ elimination =
+typeOfElimination :: Domain.Type -> Domain.Elimination -> M Domain.Type
+typeOfElimination headType elimination =
   case elimination of
     Domain.App arg ->
-      typeOfApplication type_ arg
-    Domain.Case (Domain.Branches env branches defaultBranch) ->
-      case defaultBranch of
-        Just term -> do
-          value' <- Evaluation.evaluate env term
-          typeOf context value'
-        Nothing ->
-          case branches of
-            Syntax.ConstructorBranches _ constructorBranches ->
-              case OrderedHashMap.elems constructorBranches of
-                branchTele : _ ->
-                  typeOfTelescope context env branchTele
-                [] ->
-                  panic "TODO closure converted type of branchless case"
-            Syntax.LiteralBranches literalBranches ->
-              case OrderedHashMap.elems literalBranches of
-                body : _ -> do
-                  body' <- Evaluation.evaluate env body
-                  typeOf context body'
-                [] ->
-                  panic "TODO closure converted type of branchless case"
+      typeOfApplication headType arg
+    Domain.Case (Domain.Branches type_ _ _ _) ->
+      pure type_
 
 typeOfApplication
   :: Domain.Type

src/ClosureConvertedToAssembly.hs

@@ -1038,7 +1038,7 @@ storeTerm env term returnLocation returnType =
       panic "st c" -- TODO
     Syntax.ApplyClosure {} ->
       panic $ "st ac " <> show term -- TODO
-    Syntax.Case scrutinee branches maybeDefaultBranch -> do
+    Syntax.Case scrutinee _type branches maybeDefaultBranch -> do
       let defaultBranch =
             fromMaybe
               (Syntax.Apply (Name.Lifted Builtin.UnknownName 0) [Syntax.Global $ Name.Lifted Builtin.UnitName 0])

src/Core/Domain.hs

@@ -62,7 +62,7 @@ data Closure where
   Closure :: Environment v -> Scope Syntax.Term v -> Closure
 
 data Branches where
-  Branches :: Environment v -> Syntax.Branches v -> Maybe (Syntax.Term v) -> Branches
+  Branches :: Type -> Environment v -> Syntax.Branches v -> Maybe (Syntax.Term v) -> Branches
 
 var :: Var -> Value
 var v = Neutral (Var v) mempty

src/Core/Domain/Showable.hs

@@ -46,7 +46,7 @@ data Closure where
 deriving instance Show Closure
 
 data Branches where
-  Branches :: Environment v -> Syntax.Branches v -> Maybe (Syntax.Term v) -> Branches
+  Branches :: Type -> Environment v -> Syntax.Branches v -> Maybe (Syntax.Term v) -> Branches
 
 deriving instance Show Branches
 
@@ -89,9 +89,10 @@ closureTo (Domain.Closure env term) =
   flip Closure term <$> environmentTo env
 
 branchesTo :: Domain.Branches -> M Branches
-branchesTo (Domain.Branches env branches defaultBranch) = do
+branchesTo (Domain.Branches type_ env branches defaultBranch) = do
+  type' <- to type_
   env' <- environmentTo env
-  pure $ Branches env' branches defaultBranch
+  pure $ Branches type' env' branches defaultBranch
 
 environmentTo :: Domain.Environment v -> M (Environment v)
 environmentTo env = do

src/Core/Evaluation.hs

@@ -104,9 +104,10 @@ evaluate env term =
       fun' <- evaluate env fun
       arg' <- evaluate env arg
       apply fun' plicity arg'
-    Syntax.Case scrutinee branches defaultBranch -> do
+    Syntax.Case scrutinee type_ branches defaultBranch -> do
       scrutineeValue <- evaluate env scrutinee
-      case_ scrutineeValue $ Domain.Branches env branches defaultBranch
+      type' <- lazyEvaluate env type_
+      case_ scrutineeValue $ Domain.Branches type' env branches defaultBranch
     Syntax.Spanned _ term' ->
       evaluate env term'
 
@@ -232,7 +233,7 @@ apply fun plicity arg =
       panic "applying non-function"
 
 case_ :: Domain.Value -> Domain.Branches -> M Domain.Value
-case_ scrutinee branches@(Domain.Branches env branches' defaultBranch) =
+case_ scrutinee branches@(Domain.Branches _ env branches' defaultBranch) =
   case (scrutinee, branches') of
     (Domain.Con constr args, Syntax.ConstructorBranches _ constructorBranches) ->
       chooseConstructorBranch env constr (toList args) constructorBranches defaultBranch

src/Core/Pretty.hs

@@ -142,11 +142,13 @@ prettyTerm prec env term =
     Syntax.App t1 plicity t2 ->
       prettyParen (prec > appPrec) $
         prettyTerm appPrec env t1 <+> Plicity.prettyAnnotation plicity <> prettyTerm (appPrec + 1) env t2
-    Syntax.Case scrutinee branches defaultBranch ->
+    Syntax.Case scrutinee type_ branches defaultBranch ->
       prettyParen (prec > casePrec) $
         "case"
           <+> prettyTerm 0 env scrutinee
           <+> "of"
+          <+> "->"
+          <+> prettyTerm 0 env type_
           <> line
           <> indent
             2

src/Core/Readback.hs

@@ -77,7 +77,8 @@ readbackElimination env eliminee elimination =
     Domain.App plicity arg -> do
       arg' <- readback env arg
       pure $ Syntax.App eliminee plicity arg'
-    Domain.Case (Domain.Branches env' branches defaultBranch) -> do
+    Domain.Case (Domain.Branches type_ env' branches defaultBranch) -> do
+      type' <- readback env type_
       branches' <- case branches of
         Syntax.ConstructorBranches constructorTypeName constructorBranches ->
           Syntax.ConstructorBranches constructorTypeName <$> OrderedHashMap.forMUnordered constructorBranches (mapM $ readbackConstructorBranch env env')
@@ -92,7 +93,7 @@ readbackElimination env eliminee elimination =
       defaultBranch' <- forM defaultBranch \branch -> do
         branch' <- Evaluation.evaluate env' branch
         readback env branch'
-      pure $ Syntax.Case eliminee branches' defaultBranch'
+      pure $ Syntax.Case eliminee type' branches' defaultBranch'
 
 readbackSpine :: Domain.Environment v -> Syntax.Term v -> Domain.Spine -> M (Syntax.Term v)
 readbackSpine =

src/Core/Syntax.hs

@@ -39,7 +39,7 @@ data Term v
   | Fun !(Type v) !Plicity !(Type v)
   | Lam !Bindings !(Type v) !Plicity !(Scope Term v)
   | App !(Term v) !Plicity !(Term v)
-  | Case !(Term v) (Branches v) !(Maybe (Term v))
+  | Case !(Term v) !(Type v) (Branches v) !(Maybe (Term v))
   | Spanned !Span.Relative !(Term v)
   deriving (Eq, Show, Generic, Hashable)
 
@@ -160,8 +160,11 @@ dependencies term =
     Fun domain _ target -> dependencies domain <> dependencies target
     Lam _ type_ _ body -> dependencies type_ <> dependencies body
     App function _ argument -> dependencies function <> dependencies argument
-    Case scrutinee branches defaultBranch ->
-      dependencies scrutinee <> branchesDependencies branches <> foldMap dependencies defaultBranch
+    Case scrutinee type_ branches defaultBranch ->
+      dependencies scrutinee
+        <> dependencies type_
+        <> branchesDependencies branches
+        <> foldMap dependencies defaultBranch
     Spanned _ term' -> dependencies term'
 
 letsDependencies :: Lets v -> HashSet Name.Qualified

src/Core/TypeOf.hs

@@ -9,7 +9,6 @@ import qualified Core.Bindings as Bindings
 import qualified Core.Domain as Domain
 import qualified Core.Evaluation as Evaluation
 import qualified Core.Syntax as Syntax
-import qualified Data.OrderedHashMap as OrderedHashMap
 import Data.Tsil (Tsil)
 import qualified Elaboration
 import Elaboration.Context (Context)
@@ -71,30 +70,12 @@ typeOfHead context hd =
       Evaluation.evaluate Environment.empty $ Meta.entryType solution
 
 typeOfElimination :: Context v -> Domain.Type -> Domain.Elimination -> M Domain.Type
-typeOfElimination context type_ elimination =
+typeOfElimination context headType elimination =
   case elimination of
     Domain.App plicity arg -> do
-      typeOfApplication context type_ plicity arg
-    Domain.Case (Domain.Branches env branches defaultBranch) ->
-      case defaultBranch of
-        Just term -> do
-          value' <- Evaluation.evaluate env term
-          typeOf context value'
-        Nothing ->
-          case branches of
-            Syntax.ConstructorBranches _ constructorBranches ->
-              case OrderedHashMap.elems constructorBranches of
-                (_, branchTele) : _ ->
-                  typeOfTelescope context env branchTele
-                [] ->
-                  panic "TODO type of branchless case"
-            Syntax.LiteralBranches literalBranches ->
-              case OrderedHashMap.elems literalBranches of
-                (_, body) : _ -> do
-                  body' <- Evaluation.evaluate env body
-                  typeOf context body'
-                [] ->
-                  panic "TODO type of branchless case"
+      typeOfApplication context headType plicity arg
+    Domain.Case (Domain.Branches type_ _ _ _) ->
+      pure type_
 
 typeOfSpineApplication :: Context v -> Domain.Type -> Domain.Spine -> M Domain.Type
 typeOfSpineApplication =

src/Core/Zonking.hs

@@ -144,10 +144,11 @@ zonk env metas postponed term =
         Right fun' -> do
           result <- Syntax.App fun' plicity <$> zonkTerm env metas postponed arg
           pure $ Right result
-    Syntax.Case scrutinee branches defaultBranch -> do
+    Syntax.Case scrutinee type_ branches defaultBranch -> do
       result <-
         Syntax.Case
           <$> zonkTerm env metas postponed scrutinee
+          <*> zonkTerm env metas postponed type_
           <*> zonkBranches env metas postponed branches
           <*> forM defaultBranch (zonkTerm env metas postponed)
       pure $ Right result

src/Elaboration/Context.hs

@@ -347,10 +347,11 @@ freeVars context value = do
       case elimination of
         Domain.App _plicity arg ->
           freeVars context arg
-        Domain.Case (Domain.Branches env branches defaultBranch) -> do
+        Domain.Case (Domain.Branches type_ env branches defaultBranch) -> do
+          typeVars <- freeVars context type_
           defaultBranchVars <- mapM (freeVars context <=< lift . Evaluation.evaluate env) defaultBranch
           brVars <- branchVars context env branches
-          pure $ fold defaultBranchVars <> brVars
+          pure $ typeVars <> fold defaultBranchVars <> brVars
 
     abstractionVars name type' closure = do
       typeVars <- freeVars context type'
@@ -655,14 +656,14 @@ forceNeutral context head_ spine
     findMatchingDefaultBranch ((eqSpine, coveredConstrs, coveredLits) : rest)
       | Just (spinePrefix, Domain.Spine Seq.Empty spineSuffix) <- Domain.matchSpinePrefix spine eqSpine =
           case spineSuffix of
-            (Domain.Branches env (Syntax.ConstructorBranches typeName cbrs) (Just defaultBranch), args) Seq.:<| spineSuffix' -> do
+            (Domain.Branches type_ env (Syntax.ConstructorBranches typeName cbrs) (Just defaultBranch), args) Seq.:<| spineSuffix' -> do
               eq <- Unification.equalSpines context spinePrefix eqSpine
               if eq && all (\c -> HashSet.member (Name.QualifiedConstructor typeName c) coveredConstrs) (OrderedHashMap.keys cbrs)
                 then pure $ Just do
                   branchValue <- Evaluation.evaluate env defaultBranch
                   Evaluation.applySpine branchValue $ Domain.Spine args spineSuffix'
                 else findMatchingDefaultBranch rest
-            (Domain.Branches env (Syntax.LiteralBranches lbrs) (Just defaultBranch), args) Seq.:<| spineSuffix' -> do
+            (Domain.Branches type_ env (Syntax.LiteralBranches lbrs) (Just defaultBranch), args) Seq.:<| spineSuffix' -> do
               eq <- Unification.equalSpines context spinePrefix eqSpine
               if eq && all (`HashSet.member` coveredLits) (OrderedHashMap.keys lbrs)
                 then pure $ Just do

src/Elaboration/Equation.hs

@@ -202,7 +202,8 @@ occursElimination context occ flexibility elimination =
     Domain.Case branches -> occursBranches context occ flexibility branches
 
 occursBranches :: Context v -> (Domain.Head -> Bool) -> Flexibility -> Domain.Branches -> M ()
-occursBranches context occ flexibility (Domain.Branches outerEnv branches defaultBranch) = do
+occursBranches context occ flexibility (Domain.Branches type_ outerEnv branches defaultBranch) = do
+  occurs context occ flexibility type_
   case branches of
     Syntax.ConstructorBranches _ constructorBranches ->
       forM_ constructorBranches $ mapM_ $ occursTele context outerEnv

src/Elaboration/Matching.hs

@@ -740,8 +740,9 @@ splitConstructor outerContext config scrutineeValue scrutineeHead scrutineeSpine
                 Just <$> check context' config Postponement.CanPostpone
 
           scrutinee <- Elaboration.readback context scrutineeValue
+          expectedType <- Elaboration.readback context config.expectedType
 
-          pure $ Syntax.Case scrutinee (Syntax.ConstructorBranches typeName $ OrderedHashMap.fromList branches) defaultBranch
+          pure $ Syntax.Case scrutinee expectedType (Syntax.ConstructorBranches typeName $ OrderedHashMap.fromList branches) defaultBranch
         ((plicity1, param) : params', Domain.Telescope.Extend _ _ plicity2 targetClosure)
           | plicity1 == plicity2 -> do
               target <- targetClosure param
@@ -858,8 +859,9 @@ splitLiteral context config scrutineeValue scrutineeHead scrutineeSpine span lit
     Just <$> check context' config Postponement.CanPostpone
 
   scrutinee <- Elaboration.readback context scrutineeValue
+  expectedType <- Elaboration.readback context config.expectedType
 
-  pure $ f $ Syntax.Case scrutinee (Syntax.LiteralBranches $ OrderedHashMap.fromList branches) defaultBranch
+  pure $ f $ Syntax.Case scrutinee expectedType (Syntax.LiteralBranches $ OrderedHashMap.fromList branches) defaultBranch
 
 findLiteralMatches
   :: Context v
@@ -911,9 +913,11 @@ splitEqualityOr context config matches k =
               Right context' -> do
                 result <- check context' config Postponement.CanPostpone
                 scrutinee <- Elaboration.readback context' scrutineeValue
+                expectedType <- Elaboration.readback context' config.expectedType
                 pure $
                   Syntax.Case
                     scrutinee
+                    expectedType
                     ( Syntax.ConstructorBranches
                         Builtin.EqualsName
                         (OrderedHashMap.fromList [(Name.unqualifyConstructor Builtin.ReflName, ([], Telescope.Empty result))])