Add type-checking for while nodes

4 files changed, 110 insertions, 2 deletions

diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
index 3eff92bf..651a4f79 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -559,8 +559,8 @@ public class TypeChecker {
             case UntypedVarNode var ->
                 typeCheckVar(var, context);
 
-            case UntypedWhileNode untypedWhileNode ->
-                throw new UnsupportedOperationException("TODO");
+            case UntypedWhileNode while_ ->
+                typeCheckWhile(while_, context);
         };
     }
 
@@ -885,4 +885,25 @@ public class TypeChecker {
             }
         };
     }
+
+    private static TypeCheckerResult<TypedFunctionStatementNode> typeCheckWhile(
+        UntypedWhileNode untyped,
+        TypeCheckerContext context
+    ) {
+        var typedConditionResult = typeCheckExpression(untyped.condition(), context);
+        var typedCondition = typedConditionResult.value();
+        context = typedConditionResult.context();
+
+        checkExpressionType(typedCondition, Types.BOOL, context);
+
+        var typedBodyResult = typeCheckBlockNested(untyped.body(), context);
+        var typedBody = typedBodyResult.value();
+        context = typedBodyResult.context();
+
+        return context.toResult(new TypedWhileNode(
+            typedCondition,
+            typedBody,
+            untyped.source()
+        ));
+    }
 }
diff --git a/src/test/java/org/shedlang/compiler/ast/typed/TypedNodeMatchers.java b/src/test/java/org/shedlang/compiler/ast/typed/TypedNodeMatchers.java
index 9ee8c8fa..5b841f4f 100644
--- a/src/test/java/org/shedlang/compiler/ast/typed/TypedNodeMatchers.java
+++ b/src/test/java/org/shedlang/compiler/ast/typed/TypedNodeMatchers.java
@@ -259,4 +259,15 @@ public class TypedNodeMatchers {
             has("expression", x -> x.expression(), expression)
         );
     }
+
+    public static Matcher<TypedFunctionStatementNode> isTypedWhileNode(
+        Matcher<? super TypedExpressionNode> condition,
+        Matcher<? super TypedBlockNode> body
+    ) {
+        return instanceOf(
+            TypedWhileNode.class,
+            has("condition", x -> x.condition(), condition),
+            has("body", x -> x.body(), body)
+        );
+    }
 }
diff --git a/src/test/java/org/shedlang/compiler/ast/untyped/Untyped.java b/src/test/java/org/shedlang/compiler/ast/untyped/Untyped.java
index 1fe48649..753bec1a 100644
--- a/src/test/java/org/shedlang/compiler/ast/untyped/Untyped.java
+++ b/src/test/java/org/shedlang/compiler/ast/untyped/Untyped.java
@@ -171,4 +171,11 @@ public class Untyped {
     ) {
         return new UntypedVarNode(target, expression, Sources.NULL);
     }
+
+    public static UntypedWhileNode while_(
+        UntypedExpressionNode condition,
+        UntypedBlockNode body
+    ) {
+        return new UntypedWhileNode(condition, body, Sources.NULL);
+    }
 }
diff --git a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerWhileTests.java b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerWhileTests.java
new file mode 100644
index 00000000..209be5e1
--- /dev/null
+++ b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerWhileTests.java
@@ -0,0 +1,69 @@
+package org.shedlang.compiler.typechecker;
+
+import org.junit.jupiter.api.Test;
+import org.shedlang.compiler.ast.untyped.Untyped;
+import org.shedlang.compiler.typechecker.errors.UnexpectedTypeError;
+import org.shedlang.compiler.typechecker.typesets.TypeSets;
+import org.shedlang.compiler.types.Types;
+
+import java.util.List;
+
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.shedlang.compiler.ast.typed.TypedNodeMatchers.*;
+import static org.zwobble.precisely.AssertThat.assertThat;
+import static org.zwobble.precisely.Matchers.*;
+
+public class TypeCheckerWhileTests {
+    @Test
+    public void conditionIsTypeChecked() {
+        var untyped = Untyped.while_(
+            Untyped.boolLiteral(true),
+            Untyped.block(List.of())
+        );
+        var context = TypeCheckerContext.EMPTY;
+
+        var typed = TypeChecker.typeCheckFunctionStatement(untyped, context);
+
+        assertThat(typed.value(), isTypedWhileNode(
+            isTypedBoolLiteralNode(true),
+            anything()
+        ));
+    }
+
+    @Test
+    public void whenConditionIsNotBoolThenErrorIsThrown() {
+        var untyped = Untyped.while_(
+            Untyped.intLiteral(0),
+            Untyped.block(List.of())
+        );
+        var context = TypeCheckerContext.EMPTY;
+
+        var error = assertThrows(
+            UnexpectedTypeError.class,
+            () -> TypeChecker.typeCheckFunctionStatement(untyped, context)
+        );
+
+        assertThat(error.expected(), equalTo(TypeSets.type(Types.BOOL)));
+        assertThat(error.actual(), equalTo(Types.INT));
+    }
+
+    @Test
+    public void bodyIsTypeChecked() {
+        var untyped = Untyped.while_(
+            Untyped.boolLiteral(true),
+            Untyped.block(List.of(
+                Untyped.expressionStatement(Untyped.intLiteral(42))
+            ))
+        );
+        var context = TypeCheckerContext.EMPTY;
+
+        var typed = TypeChecker.typeCheckFunctionStatement(untyped, context);
+
+        assertThat(typed.value(), isTypedWhileNode(
+            anything(),
+            isTypedBlockNode(isSequence(
+                isTypedExpressionStatementNode(isTypedIntLiteralNode(42))
+            ))
+        ));
+    }
+}