Combine assign, val and var statements

16 files changed, 157 insertions, 268 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
index 24bc4532..c543ee30 100644
--- a/src/main/java/org/shedlang/compiler/ast/typed/TypedAssignNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/typed/TypedAssignNode.java
@@ -6,7 +6,7 @@ public record TypedAssignNode(
     TypedTargetNode target,
     TypedExpressionNode expression,
     Source source
-) implements TypedFunctionStatementNode {
+) implements TypedFunctionStatementNode, TypedNamespaceStatementNode {
     @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 b2995bff..57a585a0 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 TypedWhileNode, TypedAssignNode, TypedBlockReturnNode, TypedExpressionStatementNode, TypedFunctionReturnNode, TypedValNode, TypedVarNode {
+public sealed interface TypedFunctionStatementNode extends TypedNode permits TypedWhileNode, TypedAssignNode, TypedBlockReturnNode, TypedExpressionStatementNode, TypedFunctionReturnNode {
     TypedFunctionStatementBehaviour behaviour();
 }
diff --git a/src/main/java/org/shedlang/compiler/ast/typed/TypedNamespaceStatementNode.java b/src/main/java/org/shedlang/compiler/ast/typed/TypedNamespaceStatementNode.java
index 112d14bc..f2186711 100644
--- a/src/main/java/org/shedlang/compiler/ast/typed/TypedNamespaceStatementNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/typed/TypedNamespaceStatementNode.java
@@ -1,4 +1,4 @@
 package org.shedlang.compiler.ast.typed;
 
-public sealed interface TypedNamespaceStatementNode extends TypedNode permits TypedFunctionDefinitionNode, TypedValNode {
+public sealed interface TypedNamespaceStatementNode extends TypedNode permits TypedAssignNode, TypedFunctionDefinitionNode {
 }
diff --git a/src/main/java/org/shedlang/compiler/ast/typed/TypedTargetNode.java b/src/main/java/org/shedlang/compiler/ast/typed/TypedTargetNode.java
index 7af7203f..4a86975b 100644
--- a/src/main/java/org/shedlang/compiler/ast/typed/TypedTargetNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/typed/TypedTargetNode.java
@@ -3,5 +3,7 @@ package org.shedlang.compiler.ast.typed;
 import org.shedlang.compiler.sources.Source;
 
 public sealed interface TypedTargetNode extends TypedNode {
-    record Variable(String name, Source source) implements TypedTargetNode {}
+    record Reassign(String name, Source source) implements TypedTargetNode {}
+    record Val(String name, Source source) implements TypedTargetNode {}
+    record Var(String name, Source source) implements TypedTargetNode {}
 }
diff --git a/src/main/java/org/shedlang/compiler/ast/typed/TypedValNode.java b/src/main/java/org/shedlang/compiler/ast/typed/TypedValNode.java
deleted file mode 100644
index c25b9d2d..00000000
--- a/src/main/java/org/shedlang/compiler/ast/typed/TypedValNode.java
+++ /dev/null
@@ -1,14 +0,0 @@
-package org.shedlang.compiler.ast.typed;
-
-import org.shedlang.compiler.sources.Source;
-
-public record TypedValNode(
-    TypedTargetNode target,
-    TypedExpressionNode expression,
-    Source source
-) implements TypedFunctionStatementNode, TypedNamespaceStatementNode {
-    @Override
-    public TypedFunctionStatementBehaviour behaviour() {
-        return TypedFunctionStatementBehaviour.CONTINUE;
-    }
-}
diff --git a/src/main/java/org/shedlang/compiler/ast/typed/TypedVarNode.java b/src/main/java/org/shedlang/compiler/ast/typed/TypedVarNode.java
deleted file mode 100644
index 3c698b52..00000000
--- a/src/main/java/org/shedlang/compiler/ast/typed/TypedVarNode.java
+++ /dev/null
@@ -1,14 +0,0 @@
-package org.shedlang.compiler.ast.typed;
-
-import org.shedlang.compiler.sources.Source;
-
-public record TypedVarNode(
-    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/untyped/UntypedAssignNode.java b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedAssignNode.java
index 9e6827f8..6b28bd89 100644
--- a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedAssignNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedAssignNode.java
@@ -6,7 +6,7 @@ public record UntypedAssignNode(
     UntypedTargetNode target,
     UntypedExpressionNode expression,
     Source source
-) implements UntypedFunctionStatementNode {
+) implements UntypedFunctionStatementNode, UntypedNamespaceStatementNode {
     @Override
     public UntypedFunctionStatementBehaviour behaviour() {
         return UntypedFunctionStatementBehaviour.CONTINUE;
diff --git a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionStatementNode.java b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionStatementNode.java
index 5fadfcaf..90b0c29b 100644
--- a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionStatementNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionStatementNode.java
@@ -1,7 +1,5 @@
 package org.shedlang.compiler.ast.untyped;
 
-import java.util.List;
-
-public sealed interface UntypedFunctionStatementNode extends UntypedNode permits UntypedAssignNode, UntypedBlockReturnNode, UntypedExpressionStatementNode, UntypedFunctionReturnNode, UntypedValNode, UntypedVarNode, UntypedWhileNode {
+public sealed interface UntypedFunctionStatementNode extends UntypedNode permits UntypedAssignNode, UntypedBlockReturnNode, UntypedExpressionStatementNode, UntypedFunctionReturnNode, UntypedWhileNode {
     UntypedFunctionStatementBehaviour behaviour();
 }
diff --git a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedNamespaceStatementNode.java b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedNamespaceStatementNode.java
index 9e4a9395..90e7142f 100644
--- a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedNamespaceStatementNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedNamespaceStatementNode.java
@@ -1,4 +1,4 @@
 package org.shedlang.compiler.ast.untyped;
 
-public sealed interface UntypedNamespaceStatementNode extends UntypedNode permits UntypedFunctionDefinitionNode, UntypedValNode {
+public sealed interface UntypedNamespaceStatementNode extends UntypedNode permits UntypedAssignNode, UntypedFunctionDefinitionNode {
 }
diff --git a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedTargetNode.java b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedTargetNode.java
index a70965bf..b7e9153c 100644
--- a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedTargetNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedTargetNode.java
@@ -3,5 +3,7 @@ package org.shedlang.compiler.ast.untyped;
 import org.shedlang.compiler.sources.Source;
 
 public sealed interface UntypedTargetNode extends UntypedNode {
-    record Variable(String name, Source source) implements UntypedTargetNode, UntypedVariableDeclarationNode {}
+    record Reassign(String name, Source source) implements UntypedTargetNode, UntypedVariableDeclarationNode {}
+    record Val(String name, Source source) implements UntypedTargetNode, UntypedVariableDeclarationNode {}
+    record Var(String name, Source source) implements UntypedTargetNode, UntypedVariableDeclarationNode {}
 }
diff --git a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedValNode.java b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedValNode.java
deleted file mode 100644
index e1353ab9..00000000
--- a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedValNode.java
+++ /dev/null
@@ -1,14 +0,0 @@
-package org.shedlang.compiler.ast.untyped;
-
-import org.shedlang.compiler.sources.Source;
-
-public record UntypedValNode(
-    UntypedTargetNode target,
-    UntypedExpressionNode expression,
-    Source source
-) implements UntypedFunctionStatementNode, UntypedNamespaceStatementNode {
-    @Override
-    public UntypedFunctionStatementBehaviour behaviour() {
-        return UntypedFunctionStatementBehaviour.CONTINUE;
-    }
-}
diff --git a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedVarNode.java b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedVarNode.java
deleted file mode 100644
index 3cd533b9..00000000
--- a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedVarNode.java
+++ /dev/null
@@ -1,14 +0,0 @@
-package org.shedlang.compiler.ast.untyped;
-
-import org.shedlang.compiler.sources.Source;
-
-public record UntypedVarNode(
-    UntypedTargetNode target,
-    UntypedExpressionNode expression,
-    Source source
-) implements UntypedFunctionStatementNode {
-    @Override
-    public UntypedFunctionStatementBehaviour behaviour() {
-        return UntypedFunctionStatementBehaviour.CONTINUE;
-    }
-}
diff --git a/src/main/java/org/shedlang/compiler/parser/Parser.java b/src/main/java/org/shedlang/compiler/parser/Parser.java
index a5a48756..c7ab639f 100644
--- a/src/main/java/org/shedlang/compiler/parser/Parser.java
+++ b/src/main/java/org/shedlang/compiler/parser/Parser.java
@@ -107,11 +107,9 @@ public class Parser {
     private static UntypedAssignNode parseAssign(TokenIterator<TokenType> tokens) {
         var start = tokens.position();
 
-        tokens.skip(TokenType.KEYWORD_ASSIGN);
-
         var target = parseTarget(tokens);
 
-        tokens.skip(TokenType.SYMBOL_COLON_EQUALS);
+        tokens.skip(TokenType.SYMBOL_EQUALS);
 
         var expression = parseTopLevelExpression(tokens);
 
@@ -303,18 +301,12 @@ public class Parser {
 
     static UntypedFunctionStatementNode parseFunctionStatement(TokenIterator<TokenType> tokens) {
         return switch (tokens.peek().tokenType()) {
-            case KEYWORD_ASSIGN ->
+            case KEYWORD_ASSIGN, KEYWORD_VAL, KEYWORD_VAR ->
                 parseAssign(tokens);
 
             case KEYWORD_RETURN ->
                 parseFunctionReturn(tokens);
 
-            case KEYWORD_VAL ->
-                parseVal(tokens);
-
-            case KEYWORD_VAR ->
-                parseVar(tokens);
-
             case KEYWORD_WHILE ->
                 parseWhile(tokens);
 
@@ -407,7 +399,7 @@ public class Parser {
                 parseFunctionDefinition(tokens);
 
             case KEYWORD_VAL ->
-                parseVal(tokens);
+                parseAssign(tokens);
 
             default ->
                 throw new UnexpectedTokenException("namespace statement", tokens.peek());
@@ -559,50 +551,43 @@ public class Parser {
     private static UntypedTargetNode parseTarget(TokenIterator<TokenType> tokens) {
         var start = tokens.position();
 
-        var name = tokens.nextValue(TokenType.IDENTIFIER);
-
-        var end = tokens.position();
-        var source = start.to(end);
-
-        return new UntypedTargetNode.Variable(name, source);
-    }
-
-    private static UntypedValNode parseVal(TokenIterator<TokenType> tokens) {
-        var start = tokens.position();
-
-        tokens.skip(TokenType.KEYWORD_VAL);
-
-        var target = parseTarget(tokens);
-
-        tokens.skip(TokenType.SYMBOL_EQUALS);
+        return switch (tokens.peek().tokenType()) {
+            case KEYWORD_ASSIGN -> {
+                tokens.skip(TokenType.KEYWORD_ASSIGN);
 
-        var expression = parseTopLevelExpression(tokens);
+                var name = tokens.nextValue(TokenType.IDENTIFIER);
 
-        tokens.skip(TokenType.SYMBOL_SEMICOLON);
+                var end = tokens.position();
+                var source = start.to(end);
 
-        var end = tokens.position();
-        var source = start.to(end);
+                yield new UntypedTargetNode.Reassign(name, source);
+            }
 
-        return new UntypedValNode(target, expression, source);
-    }
+            case KEYWORD_VAL -> {
+                tokens.skip(TokenType.KEYWORD_VAL);
 
-    private static UntypedVarNode parseVar(TokenIterator<TokenType> tokens) {
-        var start = tokens.position();
+                var name = tokens.nextValue(TokenType.IDENTIFIER);
 
-        tokens.skip(TokenType.KEYWORD_VAR);
+                var end = tokens.position();
+                var source = start.to(end);
 
-        var target = parseTarget(tokens);
+                yield new UntypedTargetNode.Val(name, source);
+            }
 
-        tokens.skip(TokenType.SYMBOL_EQUALS);
+            case KEYWORD_VAR -> {
+                tokens.skip(TokenType.KEYWORD_VAR);
 
-        var expression = parseTopLevelExpression(tokens);
+                var name = tokens.nextValue(TokenType.IDENTIFIER);
 
-        tokens.skip(TokenType.SYMBOL_SEMICOLON);
+                var end = tokens.position();
+                var source = start.to(end);
 
-        var end = tokens.position();
-        var source = start.to(end);
+                yield new UntypedTargetNode.Var(name, source);
+            }
 
-        return new UntypedVarNode(target, expression, source);
+            default ->
+                throw new UnexpectedTokenException("namespace statement", tokens.peek());
+        };
     }
 
     private static UntypedFunctionStatementNode parseWhile(
diff --git a/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java b/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java
index 1621a50c..2c9d6065 100644
--- a/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java
+++ b/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java
@@ -3,12 +3,10 @@ package org.shedlang.compiler.stackir;
 import org.pcollections.*;
 import org.shedlang.compiler.ShedBuiltins;
 import org.shedlang.compiler.ast.typed.*;
-import org.shedlang.compiler.sources.Source;
 import org.shedlang.compiler.stackir.instructions.*;
 import org.shedlang.compiler.util.P;
 
 import java.util.List;
-import java.util.function.BiFunction;
 
 public class StackCompiler {
     private StackCompiler() {
@@ -16,7 +14,7 @@ public class StackCompiler {
 
     private static Context compileAssign(TypedAssignNode node, Context context) {
         context = compileExpression(node.expression(), context);
-        context = compileTarget(node.target(), VarStoreInstruction::new, context);
+        context = compileTarget(node.target(), context);
         return context;
     }
 
@@ -239,12 +237,6 @@ public class StackCompiler {
             case TypedFunctionReturnNode functionReturn ->
                 compileFunctionReturn(functionReturn, context);
 
-            case TypedValNode val ->
-                compileVal(val, context);
-
-            case TypedVarNode var ->
-                compileVar(var, context);
-
             case TypedWhileNode while_ ->
                 compileWhile(while_, context);
         };
@@ -340,11 +332,11 @@ public class StackCompiler {
 
     public static Context compileNamespaceStatement(TypedNamespaceStatementNode node, Context context) {
         return switch (node) {
+            case TypedAssignNode assign ->
+                compileAssign(assign, context);
+
             case TypedFunctionDefinitionNode functionDefinition ->
                 compileFunctionDefinition(functionDefinition, context);
-
-            case TypedValNode val ->
-                compileVal(val, context);
         };
     }
 
@@ -374,40 +366,45 @@ public class StackCompiler {
 
     private static Context compileTarget(
         TypedTargetNode target,
-        BiFunction<String, Source, Instruction> store,
         Context context
     ) {
         return switch (target) {
-            case TypedTargetNode.Variable variable -> {
+            case TypedTargetNode.Reassign reassign -> {
                 context = context.addInstruction(
-                    store.apply(
-                        context.shedVariableToLocal(variable.name()),
-                        variable.source()
+                    new VarStoreInstruction(
+                        context.shedVariableToLocal(reassign.name()),
+                        reassign.source()
                     )
                 );
 
                 yield context;
             }
-        };
-    }
 
-    private static Context compileVal(TypedValNode val, Context context) {
-        context = compileExpression(val.expression(), context);
-        context = compileTarget(val.target(), LocalStoreInstruction::new, context);
-        return context;
-    }
+            case TypedTargetNode.Val val -> {
+                context = context.addInstruction(
+                    new LocalStoreInstruction(
+                        context.shedVariableToLocal(val.name()),
+                        val.source()
+                    )
+                );
 
-    private static Context compileVar(TypedVarNode var, Context context) {
-        context = compileExpression(var.expression(), context);
-        context = compileTarget(var.target(), VarStoreInstruction::new, context);
+                yield context;
+            }
 
-        // TODO: better to store this information in the block node with the bound variables?
-        switch (var.target()) {
-            case TypedTargetNode.Variable variable ->
-                context = context.markShedVariableAsVar(variable.name());
-        }
+            case TypedTargetNode.Var var -> {
+                context = context.addInstruction(
+                    new VarStoreInstruction(
+                        context.shedVariableToLocal(var.name()),
+                        var.source()
+                    )
+                );
 
-        return context;
+                // TODO: better to store this information in the block node with the bound variables?
+                context = context.markShedVariableAsVar(var.name());
+
+                yield context;
+            }
+        };
     }
 
     private static Context compileWhile(TypedWhileNode node, Context context) {
diff --git a/src/main/java/org/shedlang/compiler/typechecker/DeclaredVariables.java b/src/main/java/org/shedlang/compiler/typechecker/DeclaredVariables.java
index 485aa8fa..87823b46 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/DeclaredVariables.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/DeclaredVariables.java
@@ -6,6 +6,13 @@ import org.shedlang.compiler.ast.untyped.*;
 import java.util.List;
 
 class DeclaredVariables {
+    static PVector<DeclaredVariable> forAssign(
+        UntypedAssignNode statement,
+        PVector<DeclaredVariable> variables
+    ) {
+        return forTarget(statement.target(), variables);
+    }
+
     static PVector<DeclaredVariable> forFunctionStatements(
         List<UntypedFunctionStatementNode> statements,
         PVector<DeclaredVariable> variables
@@ -23,7 +30,7 @@ class DeclaredVariables {
     ) {
         return switch (statement) {
             case UntypedAssignNode assign ->
-                variables;
+                forAssign(assign, variables);
 
             case UntypedBlockReturnNode blockReturn ->
                 variables;
@@ -34,12 +41,6 @@ class DeclaredVariables {
             case UntypedFunctionReturnNode functionReturn ->
                 variables;
 
-            case UntypedValNode val ->
-                forVal(val, variables);
-
-            case UntypedVarNode var ->
-                forVar(var, variables);
-
             case UntypedWhileNode while_ ->
                 variables;
         };
@@ -61,43 +62,38 @@ class DeclaredVariables {
         PVector<DeclaredVariable> variables
     ) {
         return switch (statement) {
+            case UntypedAssignNode assign ->
+                forAssign(assign, variables);
+
             case UntypedFunctionDefinitionNode functionDefinition ->
                 variables.plus(new DeclaredVariable(
                     functionDefinition.name(),
                     VariableType.CONSTANT,
                     functionDefinition.source()
                 ));
-
-            case UntypedValNode val ->
-                forVal(val, variables);
         };
     }
 
-    static PVector<DeclaredVariable> forVal(
-        UntypedValNode statement,
-        PVector<DeclaredVariable> variables
-    ) {
-        return forTarget(statement.target(), VariableType.VAL, variables);
-    }
-
-    static PVector<DeclaredVariable> forVar(
-        UntypedVarNode statement,
-        PVector<DeclaredVariable> variables
-    ) {
-        return forTarget(statement.target(), VariableType.VAR, variables);
-    }
-
     private static PVector<DeclaredVariable> forTarget(
         UntypedTargetNode target,
-        VariableType variableType,
         PVector<DeclaredVariable> variables
     ) {
         return switch (target) {
-            case UntypedTargetNode.Variable variable ->
+            case UntypedTargetNode.Reassign reassign ->
+                variables;
+
+            case UntypedTargetNode.Var var ->
+                variables.plus(new DeclaredVariable(
+                    var.name(),
+                    VariableType.VAR,
+                    var.source()
+                ));
+
+            case UntypedTargetNode.Val val ->
                 variables.plus(new DeclaredVariable(
-                    variable.name(),
-                    variableType,
-                    variable.source()
+                    val.name(),
+                    VariableType.VAL,
+                    val.source()
                 ));
         };
     }
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
index 5302ee6c..5db00b44 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -125,6 +125,17 @@ public class TypeChecker {
         return context.toResult(new TypedArgsNode(typedPositionalArgs, typedNamedArgs, args.source()));
     }
 
+    private static TypeCheckSteps<TypedAssignNode> typeCheckStepsAssign(
+        UntypedAssignNode untyped
+    ) {
+        return new TypeCheckSteps<>() {
+            @Override
+            public TypeCheckerResult<TypedAssignNode> define(TypeCheckerContext context) {
+                return typeCheckAssign(untyped, context);
+            }
+        };
+    }
+
     private static TypeCheckerResult<TypedAssignNode> typeCheckAssign(
         UntypedAssignNode untyped,
         TypeCheckerContext context
@@ -136,7 +147,6 @@ public class TypeChecker {
         var targetResult = typeCheckTarget(
             untyped.target(),
             expression.type(),
-            Optional.empty(),
             context
         );
         var target = targetResult.value();
@@ -376,7 +386,7 @@ public class TypeChecker {
         });
     }
 
-    private static TypeCheckSteps<TypedFunctionDefinitionNode> typeCheckFunctionDefinition(
+    private static TypeCheckSteps<TypedFunctionDefinitionNode> typeCheckStepsFunctionDefinition(
         UntypedFunctionDefinitionNode untyped
     ) {
         return new TypeCheckSteps<TypedFunctionDefinitionNode>() {
@@ -553,12 +563,6 @@ public class TypeChecker {
             case UntypedFunctionReturnNode functionReturn ->
                 typeCheckFunctionReturn(functionReturn, context);
 
-            case UntypedValNode val ->
-                typeCheckVal(val).typeCheckAllPhases(context);
-
-            case UntypedVarNode var ->
-                typeCheckVar(var, context);
-
             case UntypedWhileNode while_ ->
                 typeCheckWhile(while_, context);
         };
@@ -656,15 +660,16 @@ public class TypeChecker {
                 .collect(Collectors.toSet());
 
             switch (untypedStatement) {
-                case UntypedFunctionDefinitionNode untypedFunctionDefinition -> {
-                    dependencyChecker = dependencyChecker.addDeferredDependencies(
-                        untypedFunctionDefinition.name(),
+                case UntypedAssignNode untypedAssign -> {
+                    // TODO: prevent var/assign at top-level (or just prevent capture, which we need to do anyway?)
+                    dependencyChecker = dependencyChecker.addImmediateDependencies(
+                        DeclaredVariables.forAssign(untypedAssign, P.vector()),
                         dependencies
                     );
                 }
-                case UntypedValNode untypedVal -> {
-                    dependencyChecker = dependencyChecker.addImmediateDependencies(
-                        DeclaredVariables.forVal(untypedVal, P.vector()),
+                case UntypedFunctionDefinitionNode untypedFunctionDefinition -> {
+                    dependencyChecker = dependencyChecker.addDeferredDependencies(
+                        untypedFunctionDefinition.name(),
                         dependencies
                     );
                 }
@@ -705,11 +710,11 @@ public class TypeChecker {
         UntypedNamespaceStatementNode untyped
     ) {
         return switch (untyped) {
-            case UntypedFunctionDefinitionNode functionDefinition ->
-                typeCheckFunctionDefinition(functionDefinition);
+            case UntypedAssignNode assign ->
+                typeCheckStepsAssign(assign);
 
-            case UntypedValNode val ->
-                typeCheckVal(val);
+            case UntypedFunctionDefinitionNode functionDefinition ->
+                typeCheckStepsFunctionDefinition(functionDefinition);
         };
     }
 
@@ -795,93 +800,53 @@ public class TypeChecker {
         );
     }
 
-    private static TypeCheckSteps<TypedValNode> typeCheckVal(
-        UntypedValNode untyped
-    ) {
-        return new TypeCheckSteps<>() {
-            @Override
-            public TypeCheckerResult<TypedValNode> define(TypeCheckerContext context) {
-                return defineVal(untyped, context);
-            }
-        };
-    }
-
-    private static TypeCheckerResult<TypedValNode> defineVal(
-        UntypedValNode untyped,
-        TypeCheckerContext context
-    ) {
-        var expressionResult = typeCheckExpression(untyped.expression(), context);
-        var expression = expressionResult.value();
-        context = expressionResult.context();
-
-        var targetResult = typeCheckTarget(
-            untyped.target(),
-            expression.type(),
-            Optional.of(VariableType.VAL),
-            context
-        );
-        var target = targetResult.value();
-        context = targetResult.context();
-
-        return context.toResult(new TypedValNode(target, expression, untyped.source()));
-    }
-
-    private static TypeCheckerResult<TypedVarNode> typeCheckVar(
-        UntypedVarNode untyped,
-        TypeCheckerContext context
-    ) {
-        var expressionResult = typeCheckExpression(untyped.expression(), context);
-        var expression = expressionResult.value();
-        context = expressionResult.context();
-
-        var targetResult = typeCheckTarget(
-            untyped.target(),
-            expression.type(),
-            Optional.of(VariableType.VAR),
-            context
-        );
-        var target = targetResult.value();
-        context = targetResult.context();
-
-        return context.toResult(new TypedVarNode(target, expression, untyped.source()));
-    }
-
     private static TypeCheckerResult<TypedTargetNode> typeCheckTarget(
         UntypedTargetNode target,
         Type expressionType,
-        Optional<VariableType> variableType,
         TypeCheckerContext context
     ) {
         return switch (target) {
-            case UntypedTargetNode.Variable variable -> {
-                if (variableType.isEmpty()) {
-                    var targetVariableResult = runtimeLookup(
-                        variable.name(),
-                        variable.source(),
-                        context
-                    );
-                    var targetVariable = targetVariableResult.value();
-                    context = targetVariableResult.context();
-
-                    if (targetVariable.variableType() != VariableType.VAR) {
-                        throw new CanOnlyAssignToVarError(target.source());
-                    }
-
-                    checkIsSubType(
-                        expressionType,
-                        targetVariable.type().orElseThrow(),
-                        target.source(),
-                        context
-                    );
-                } else {
-                    context = context.define(
-                        variable.name(),
-                        expressionType,
-                        target.source()
-                    );
+            case UntypedTargetNode.Reassign reassign -> {
+                var targetVariableResult = runtimeLookup(
+                    reassign.name(),
+                    reassign.source(),
+                    context
+                );
+                var targetVariable = targetVariableResult.value();
+                context = targetVariableResult.context();
+
+                if (targetVariable.variableType() != VariableType.VAR) {
+                    throw new CanOnlyAssignToVarError(target.source());
                 }
 
-                yield context.toResult(new TypedTargetNode.Variable(variable.name(), target.source()));
+                checkIsSubType(
+                    expressionType,
+                    targetVariable.type().orElseThrow(),
+                    target.source(),
+                    context
+                );
+
+                yield context.toResult(new TypedTargetNode.Reassign(reassign.name(), target.source()));
+            }
+
+            case UntypedTargetNode.Val val -> {
+                context = context.define(
+                    val.name(),
+                    expressionType,
+                    target.source()
+                );
+
+                yield context.toResult(new TypedTargetNode.Val(val.name(), target.source()));
+            }
+
+            case UntypedTargetNode.Var var -> {
+                context = context.define(
+                    var.name(),
+                    expressionType,
+                    target.source()
+                );
+
+                yield context.toResult(new TypedTargetNode.Var(var.name(), target.source()));
             }
         };
     }