Type check assignments

4 files changed, 85 insertions, 8 deletions

diff --git a/src/main/java/org/shedlang/compiler/ast/typed/TypedAssignNode.java b/src/main/java/org/shedlang/compiler/ast/typed/TypedAssignNode.java
new file mode 100644
index 00000000..24bc4532
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/ast/typed/TypedAssignNode.java
@@ -0,0 +1,14 @@
+package org.shedlang.compiler.ast.typed;
+
+import org.shedlang.compiler.sources.Source;
+
+public record TypedAssignNode(
+    TypedTargetNode target,
+    TypedExpressionNode expression,
+    Source source
+) implements TypedFunctionStatementNode {
+    @Override
+    public TypedFunctionStatementBehaviour behaviour() {
+        return TypedFunctionStatementBehaviour.CONTINUE;
+    }
+}
diff --git a/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionStatementNode.java b/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionStatementNode.java
index b1fcd49b..843d8c3e 100644
--- a/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionStatementNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionStatementNode.java
@@ -1,5 +1,5 @@
 package org.shedlang.compiler.ast.typed;
 
-public sealed interface TypedFunctionStatementNode extends TypedNode permits TypedBlockReturnNode, TypedExpressionStatementNode, TypedFunctionReturnNode, TypedValNode, TypedVarNode {
+public sealed interface TypedFunctionStatementNode extends TypedNode permits TypedAssignNode, TypedBlockReturnNode, TypedExpressionStatementNode, TypedFunctionReturnNode, TypedValNode, TypedVarNode {
     TypedFunctionStatementBehaviour behaviour();
 }
diff --git a/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java b/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java
index 2be8480d..8863253c 100644
--- a/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java
+++ b/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java
@@ -201,6 +201,9 @@ public class StackCompiler {
         Context context
     ) {
         return switch (node) {
+            case TypedAssignNode assign ->
+                throw new UnsupportedOperationException();
+
             case TypedBlockReturnNode blockReturn ->
                 compileBlockReturn(blockReturn, context);
             
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
index 8cf38fc1..ee35fa1f 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -20,11 +20,15 @@ public class TypeChecker {
     }
 
     private static void checkExpressionType(TypedExpressionNode expression, Type expectedType, TypeCheckerContext context) {
-        if (!isSubType(expression.type(), expectedType, context)) {
+        checkIsSubType(expression.type(), expectedType, expression.source(), context);
+    }
+
+    private static void checkIsSubType(Type actualType, Type expectedType, Source source, TypeCheckerContext context) {
+        if (!isSubType(actualType, expectedType, context)) {
             throw new UnexpectedTypeError(
                 TypeSets.type(expectedType),
-                expression.type(),
-                expression.source()
+                actualType,
+                source
             );
         }
     }
@@ -129,6 +133,26 @@ public class TypeChecker {
         return context.toResult(new TypedArgsNode(typedPositionalArgs, typedNamedArgs, args.source()));
     }
 
+    private static TypeCheckerResult<TypedAssignNode> typeCheckAssign(
+        UntypedAssignNode untyped,
+        TypeCheckerContext context
+    ) {
+        var expressionResult = typeCheckExpression(untyped.expression(), context);
+        var expression = expressionResult.value();
+        context = expressionResult.context();
+
+        var targetResult = typeCheckTarget(
+            untyped.target(),
+            expression.type(),
+            TargetType.ASSIGN,
+            context
+        );
+        var target = targetResult.value();
+        context = targetResult.context();
+
+        return context.toResult(new TypedAssignNode(target, expression, untyped.source()));
+    }
+
     private static TypeCheckerResult<TypedExpressionNode> typeCheckBinaryOperation(
         UntypedBinaryOperationNode untyped,
         TypeCheckerContext context
@@ -526,7 +550,7 @@ public class TypeChecker {
     ) {
         return switch (untyped) {
             case UntypedAssignNode assign ->
-                throw new UnsupportedOperationException();
+                typeCheckAssign(assign, context);
 
             case UntypedBlockReturnNode blockReturn ->
                 typeCheckBlockReturn(blockReturn, context);
@@ -795,7 +819,12 @@ public class TypeChecker {
         var expression = expressionResult.value();
         context = expressionResult.context();
 
-        var targetResult = typeCheckTarget(untyped.target(), expression.type(), context);
+        var targetResult = typeCheckTarget(
+            untyped.target(),
+            expression.type(),
+            TargetType.DEFINE,
+            context
+        );
         var target = targetResult.value();
         context = targetResult.context();
 
@@ -810,7 +839,12 @@ public class TypeChecker {
         var expression = expressionResult.value();
         context = expressionResult.context();
 
-        var targetResult = typeCheckTarget(untyped.target(), expression.type(), context);
+        var targetResult = typeCheckTarget(
+            untyped.target(),
+            expression.type(),
+            TargetType.DEFINE,
+            context
+        );
         var target = targetResult.value();
         context = targetResult.context();
 
@@ -820,14 +854,40 @@ public class TypeChecker {
     private static TypeCheckerResult<TypedTargetNode> typeCheckTarget(
         UntypedTargetNode target,
         Type expressionType,
+        TargetType targetType,
         TypeCheckerContext context
     ) {
         return switch (target) {
             case UntypedTargetNode.Variable variable -> {
-                context = context.define(variable.name(), expressionType, target.source());
+                switch (targetType) {
+                    case ASSIGN -> {
+                        var variableTypeResult = runtimeLookup(
+                            variable.name(),
+                            variable.source(),
+                            context
+                        );
+                        var variableType = variableTypeResult.value();
+                        context = variableTypeResult.context();
+
+                        checkIsSubType(
+                            expressionType,
+                            variableType,
+                            target.source(),
+                            context
+                        );
+                    }
+                    case DEFINE -> {
+                        context = context.define(variable.name(), expressionType, target.source());
+                    }
+                }
 
                 yield context.toResult(new TypedTargetNode.Variable(variable.name(), target.source()));
             }
         };
     }
+
+    private enum TargetType {
+        ASSIGN,
+        DEFINE,
+    }
 }