Add basic type checking for when expressions

4 files changed, 80 insertions, 4 deletions

diff --git a/src/main/java/org/shedlang/compiler/ast/typed/TypedExpressionNode.java b/src/main/java/org/shedlang/compiler/ast/typed/TypedExpressionNode.java
index 93758ccf..2f479f25 100644
--- a/src/main/java/org/shedlang/compiler/ast/typed/TypedExpressionNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/typed/TypedExpressionNode.java
@@ -4,6 +4,6 @@ import org.shedlang.compiler.types.Type;
 
 public sealed interface TypedExpressionNode
     extends TypedNode
-    permits TypedBinaryOperationNode, TypedBoolLiteralNode, TypedCallNode, TypedFieldAccessNode, TypedIfNode, TypedIntLiteralNode, TypedReferenceNode, TypedStringLiteralNode, TypedTupleNode {
+    permits TypedBinaryOperationNode, TypedBoolLiteralNode, TypedCallNode, TypedFieldAccessNode, TypedIfNode, TypedIntLiteralNode, TypedReferenceNode, TypedStringLiteralNode, TypedTupleNode, TypedWhenNode {
     Type type();
 }
diff --git a/src/main/java/org/shedlang/compiler/ast/typed/TypedWhenCaseNode.java b/src/main/java/org/shedlang/compiler/ast/typed/TypedWhenCaseNode.java
new file mode 100644
index 00000000..b50003ba
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/ast/typed/TypedWhenCaseNode.java
@@ -0,0 +1,11 @@
+package org.shedlang.compiler.ast.typed;
+
+import org.shedlang.compiler.sources.Source;
+import org.shedlang.compiler.types.Type;
+
+public record TypedWhenCaseNode(
+    TypedTypeLevelExpressionNode<Type> caseType,
+    TypedBlockNode body,
+    Source source
+) implements TypedNode {
+}
diff --git a/src/main/java/org/shedlang/compiler/ast/typed/TypedWhenNode.java b/src/main/java/org/shedlang/compiler/ast/typed/TypedWhenNode.java
new file mode 100644
index 00000000..d9a4f7c9
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/ast/typed/TypedWhenNode.java
@@ -0,0 +1,14 @@
+package org.shedlang.compiler.ast.typed;
+
+import org.shedlang.compiler.sources.Source;
+import org.shedlang.compiler.types.Type;
+
+import java.util.List;
+
+public record TypedWhenNode(
+    TypedExpressionNode expression,
+    List<TypedWhenCaseNode> cases,
+    Type type,
+    Source source
+) implements TypedExpressionNode {
+}
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
index d10531ad..58396e80 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -406,7 +406,7 @@ public class TypeChecker {
                 typeCheckTuple(untypedTuple, context);
 
             case UntypedWhenNode untypedWhen ->
-                throw new UnsupportedOperationException();
+                typeCheckWhen(untypedWhen, context);
         };
     }
 
@@ -682,7 +682,7 @@ public class TypeChecker {
         var typedBody = new ArrayList<TypedNamespaceStatementNode>();
 
         var typeCheckStepSets = untyped.body().stream()
-            .map(untypedStatement -> typeCheckNamespaceStatement(untypedStatement))
+            .map(untypedStatement -> typeCheckStepsNamespaceStatement(untypedStatement))
             .toList();
 
         for (var typeCheckSteps : typeCheckStepSets) {
@@ -786,7 +786,7 @@ public class TypeChecker {
             .toResult(typed);
     }
 
-    static TypeCheckSteps<? extends TypedNamespaceStatementNode> typeCheckNamespaceStatement(
+    static TypeCheckSteps<? extends TypedNamespaceStatementNode> typeCheckStepsNamespaceStatement(
         UntypedNamespaceStatementNode untyped
     ) {
         return switch (untyped) {
@@ -1132,6 +1132,57 @@ public class TypeChecker {
         };
     }
 
+    private static TypeCheckerResult<TypedWhenNode> typeCheckWhen(
+        UntypedWhenNode untyped,
+        TypeCheckerContext context
+    ) {
+        var typedExpressionResult = typeCheckExpression(untyped.expression(), context);
+        var typedExpression = typedExpressionResult.value();
+        context = typedExpressionResult.context();
+        
+        var typedCasesResult = context.mapWithContext(
+            untyped.cases(),
+            TypeChecker::typeCheckWhenCase
+        );
+        var typedCases = typedCasesResult.value();
+        context = typedCasesResult.context();
+
+        var caseBodyTypes = typedCases.stream()
+            .map(typedCase -> typedCase.body().type())
+            .toList();
+
+        var type = Types.commonSupertype(caseBodyTypes);
+        
+        return context.toResult(new TypedWhenNode(
+            typedExpression,
+            typedCases,
+            type,
+            untyped.source()
+        ));
+    }
+
+    private static TypeCheckerResult<TypedWhenCaseNode> typeCheckWhenCase(
+        UntypedWhenCaseNode untyped,
+        TypeCheckerContext context
+    ) {
+        // TODO: check is a valid case type
+        var typedCaseType = typeCheckTypeLevelExpression(
+            TypeSets.META_TYPE,
+            untyped.caseType(),
+            context
+        );
+
+        var typedBodyResult = typeCheckBlockNested(untyped.body(), context);
+        var typedBody = typedBodyResult.value();
+        context = typedBodyResult.context();
+
+        return context.toResult(new TypedWhenCaseNode(
+            typedCaseType,
+            typedBody,
+            untyped.source()
+        ));
+    }
+
     private static TypeCheckerResult<TypedFunctionStatementNode> typeCheckWhile(
         UntypedWhileNode untyped,
         TypeCheckerContext context