Add notion of VariableType

9 files changed, 94 insertions, 68 deletions

diff --git a/src/main/java/org/shedlang/compiler/ShedCompiler.java b/src/main/java/org/shedlang/compiler/ShedCompiler.java
index 45411d3c..fcfb127c 100644
--- a/src/main/java/org/shedlang/compiler/ShedCompiler.java
+++ b/src/main/java/org/shedlang/compiler/ShedCompiler.java
@@ -42,7 +42,7 @@ public class ShedCompiler {
             .addNamespaceFields(ShedBuiltins.NAMESPACE_NAME, P.orderedMap(
                 Map.entry(ShedBuiltins.PRINT_NAME, ShedBuiltins.PRINT_TYPE)
             ))
-            .enterNamespaceWithDefinedVariables(P.map(
+            .enterGlobal(P.map(
                 Map.entry("Bool", Types.metaType(Types.BOOL)),
                 Map.entry("Int", Types.metaType(Types.INT)),
                 Map.entry("Io", Types.effectType(Effects.IO)),
diff --git a/src/main/java/org/shedlang/compiler/typechecker/DeclaredVariable.java b/src/main/java/org/shedlang/compiler/typechecker/DeclaredVariable.java
new file mode 100644
index 00000000..08d45542
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/typechecker/DeclaredVariable.java
@@ -0,0 +1,10 @@
+package org.shedlang.compiler.typechecker;
+
+import org.shedlang.compiler.sources.Source;
+
+public record DeclaredVariable(
+    String name,
+    VariableType variableType,
+    Source source
+) {
+}
diff --git a/src/main/java/org/shedlang/compiler/typechecker/DeclaredVariables.java b/src/main/java/org/shedlang/compiler/typechecker/DeclaredVariables.java
index faa8f4e2..2e7e236e 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/DeclaredVariables.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/DeclaredVariables.java
@@ -6,9 +6,9 @@ import org.shedlang.compiler.ast.untyped.*;
 import java.util.List;
 
 class DeclaredVariables {
-    static PVector<UntypedVariableDeclarationNode> forFunctionStatements(
+    static PVector<DeclaredVariable> forFunctionStatements(
         List<UntypedFunctionStatementNode> statements,
-        PVector<UntypedVariableDeclarationNode> variables
+        PVector<DeclaredVariable> variables
     ) {
         for (var statement : statements) {
             variables = DeclaredVariables.forFunctionStatement(statement, variables);
@@ -17,9 +17,9 @@ class DeclaredVariables {
         return variables;
     }
 
-    static PVector<UntypedVariableDeclarationNode> forFunctionStatement(
+    static PVector<DeclaredVariable> forFunctionStatement(
         UntypedFunctionStatementNode statement,
-        PVector<UntypedVariableDeclarationNode> variables
+        PVector<DeclaredVariable> variables
     ) {
         return switch (statement) {
             case UntypedAssignNode assign ->
@@ -42,9 +42,9 @@ class DeclaredVariables {
         };
     }
 
-    public static PVector<UntypedVariableDeclarationNode> forNamespaceStatements(
+    public static PVector<DeclaredVariable> forNamespaceStatements(
         List<UntypedNamespaceStatementNode> statements,
-        PVector<UntypedVariableDeclarationNode> variables
+        PVector<DeclaredVariable> variables
     ) {
         for (var statement : statements) {
             variables = DeclaredVariables.forNamespaceStatement(statement, variables);
@@ -53,40 +53,49 @@ class DeclaredVariables {
         return variables;
     }
 
-    static PVector<UntypedVariableDeclarationNode> forNamespaceStatement(
+    static PVector<DeclaredVariable> forNamespaceStatement(
         UntypedNamespaceStatementNode statement,
-        PVector<UntypedVariableDeclarationNode> variables
+        PVector<DeclaredVariable> variables
     ) {
         return switch (statement) {
             case UntypedFunctionDefinitionNode functionDefinition ->
-                variables.plus(functionDefinition);
+                variables.plus(new DeclaredVariable(
+                    functionDefinition.name(),
+                    VariableType.CONSTANT,
+                    functionDefinition.source()
+                ));
 
             case UntypedValNode val ->
                 forVal(val, variables);
         };
     }
 
-    static PVector<UntypedVariableDeclarationNode> forVal(
+    static PVector<DeclaredVariable> forVal(
         UntypedValNode statement,
-        PVector<UntypedVariableDeclarationNode> variables
+        PVector<DeclaredVariable> variables
     ) {
-        return forTarget(statement.target(), variables);
+        return forTarget(statement.target(), VariableType.VAL, variables);
     }
 
-    static PVector<UntypedVariableDeclarationNode> forVar(
+    static PVector<DeclaredVariable> forVar(
         UntypedVarNode statement,
-        PVector<UntypedVariableDeclarationNode> variables
+        PVector<DeclaredVariable> variables
     ) {
-        return forTarget(statement.target(), variables);
+        return forTarget(statement.target(), VariableType.VAR, variables);
     }
 
-    private static PVector<UntypedVariableDeclarationNode> forTarget(
+    private static PVector<DeclaredVariable> forTarget(
         UntypedTargetNode target,
-        PVector<UntypedVariableDeclarationNode> variables
+        VariableType variableType,
+        PVector<DeclaredVariable> variables
     ) {
         return switch (target) {
             case UntypedTargetNode.Variable variable ->
-                variables.plus(variable);
+                variables.plus(new DeclaredVariable(
+                    variable.name(),
+                    variableType,
+                    variable.source()
+                ));
         };
     }
 }
diff --git a/src/main/java/org/shedlang/compiler/typechecker/DependencyChecker.java b/src/main/java/org/shedlang/compiler/typechecker/DependencyChecker.java
index b1d7c1ac..72eb5088 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/DependencyChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/DependencyChecker.java
@@ -4,7 +4,6 @@ import org.pcollections.Empty;
 import org.pcollections.PMap;
 import org.pcollections.PVector;
 import org.shedlang.compiler.ast.typed.FreeVariable;
-import org.shedlang.compiler.ast.untyped.UntypedVariableDeclarationNode;
 import org.shedlang.compiler.typechecker.errors.ReferenceBeforeDefinitionError;
 import org.shedlang.compiler.util.Iterables;
 
@@ -16,7 +15,7 @@ class DependencyChecker {
     static final DependencyChecker EMPTY = new DependencyChecker(Empty.map(), Empty.vector());
 
     private record ImmediateDependencies(
-        Collection<UntypedVariableDeclarationNode> variableDeclarations,
+        Collection<DeclaredVariable> variableDeclarations,
         Set<FreeVariable> dependencies
     ) {}
 
@@ -40,7 +39,7 @@ class DependencyChecker {
     }
 
     DependencyChecker addImmediateDependencies(
-        Collection<UntypedVariableDeclarationNode> variableDeclarations,
+        Collection<DeclaredVariable> variableDeclarations,
         Set<FreeVariable> dependencies
     ) {
         return new DependencyChecker(
diff --git a/src/main/java/org/shedlang/compiler/typechecker/EnvironmentBuilder.java b/src/main/java/org/shedlang/compiler/typechecker/EnvironmentBuilder.java
index 284b4d31..dca5abb5 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/EnvironmentBuilder.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/EnvironmentBuilder.java
@@ -2,7 +2,6 @@ package org.shedlang.compiler.typechecker;
 
 import org.pcollections.Empty;
 import org.pcollections.PMap;
-import org.shedlang.compiler.ast.untyped.UntypedVariableDeclarationNode;
 import org.shedlang.compiler.sources.Source;
 import org.shedlang.compiler.typechecker.errors.DuplicateVariableName;
 
@@ -17,21 +16,23 @@ public class EnvironmentBuilder {
         this.environment = environment;
     }
 
-    public EnvironmentBuilder addUndefined(Iterable<UntypedVariableDeclarationNode> nodes) {
+    public EnvironmentBuilder addUndefined(
+        Iterable<DeclaredVariable> variables
+    ) {
         var builder = this;
 
-        for (var node : nodes) {
-            builder = builder.addUndefined(node.name(), node.source());
+        for (var variable : variables) {
+            builder = builder.addUndefined(variable);
         }
 
         return builder;
     }
 
-    public EnvironmentBuilder addUndefined(String name, Source source) {
-        if (environment.containsKey(name)) {
-            throw new DuplicateVariableName(name, source);
+    public EnvironmentBuilder addUndefined(DeclaredVariable variable) {
+        if (environment.containsKey(variable.name())) {
+            throw new DuplicateVariableName(variable.name(), variable.source());
         }
-        return new EnvironmentBuilder(environment.plus(name, Variable.undefined()));
+        return new EnvironmentBuilder(environment.plus(variable.name(), Variable.undefined(variable.variableType())));
     }
 
     public EnvironmentBuilder add(String name, Variable variable, Source source) {
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
index 499f41a3..0828e59a 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -134,7 +134,7 @@ public class TypeChecker {
         var targetResult = typeCheckTarget(
             untyped.target(),
             expression.type(),
-            TargetType.ASSIGN,
+            Optional.empty(),
             context
         );
         var target = targetResult.value();
@@ -448,7 +448,7 @@ public class TypeChecker {
         for (var param : typedParams.all()) {
             bodyEnvironment = bodyEnvironment.add(
                 param.name(),
-                Variable.defined(param.type().value()),
+                Variable.defined(param.type().value(), VariableType.PARAM),
                 param.source()
             );
         }
@@ -812,7 +812,7 @@ public class TypeChecker {
         var targetResult = typeCheckTarget(
             untyped.target(),
             expression.type(),
-            TargetType.DEFINE,
+            Optional.of(VariableType.VAL),
             context
         );
         var target = targetResult.value();
@@ -832,7 +832,7 @@ public class TypeChecker {
         var targetResult = typeCheckTarget(
             untyped.target(),
             expression.type(),
-            TargetType.DEFINE,
+            Optional.of(VariableType.VAR),
             context
         );
         var target = targetResult.value();
@@ -844,40 +844,36 @@ public class TypeChecker {
     private static TypeCheckerResult<TypedTargetNode> typeCheckTarget(
         UntypedTargetNode target,
         Type expressionType,
-        TargetType targetType,
+        Optional<VariableType> variableType,
         TypeCheckerContext context
     ) {
         return switch (target) {
             case UntypedTargetNode.Variable variable -> {
-                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());
-                    }
+                if (variableType.isEmpty()) {
+                    var targetTypeResult = runtimeLookup(
+                        variable.name(),
+                        variable.source(),
+                        context
+                    );
+                    var targetType = targetTypeResult.value();
+                    context = targetTypeResult.context();
+
+                    checkIsSubType(
+                        expressionType,
+                        targetType,
+                        target.source(),
+                        context
+                    );
+                } else {
+                    context = context.define(
+                        variable.name(),
+                        expressionType,
+                        target.source()
+                    );
                 }
 
                 yield context.toResult(new TypedTargetNode.Variable(variable.name(), target.source()));
             }
         };
     }
-
-    private enum TargetType {
-        ASSIGN,
-        DEFINE,
-    }
 }
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java b/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
index a4cfedce..1f50d24b 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
@@ -137,9 +137,9 @@ public class TypeCheckerContext {
         );
     }
 
-    public TypeCheckerContext enterNamespaceWithDefinedVariables(PMap<String, Type> variables) {
+    public TypeCheckerContext enterGlobal(PMap<String, Type> variables) {
         var environment = variables.entrySet().stream()
-            .map(x -> Map.entry(x.getKey(), Variable.defined(x.getValue())))
+            .map(x -> Map.entry(x.getKey(), Variable.defined(x.getValue(), VariableType.VAL)))
             .collect(PCollectors.toMap());
 
         return enterNamespace(environment);
@@ -275,7 +275,10 @@ public class TypeCheckerContext {
 
         return this.updateLexicalContext(
             lexicalContext -> new LexicalContext(
-                lexicalContext.environment.plus(name, Variable.defined(type)),
+                lexicalContext.environment.plus(name, new Variable(
+                    Optional.of(type),
+                    variable.variableType()
+                )),
                 lexicalContext.effect,
                 lexicalContext.returnType,
                 lexicalContext.boundVariables,
diff --git a/src/main/java/org/shedlang/compiler/typechecker/Variable.java b/src/main/java/org/shedlang/compiler/typechecker/Variable.java
index 6c6be749..5c820159 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/Variable.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/Variable.java
@@ -4,13 +4,13 @@ import org.shedlang.compiler.types.Type;
 
 import java.util.Optional;
 
-public record Variable(Optional<Type> type) {
-    public static Variable defined(Type type) {
-        return new Variable(Optional.of(type));
+public record Variable(Optional<Type> type, VariableType variableType) {
+    public static Variable defined(Type type, VariableType variableType) {
+        return new Variable(Optional.of(type), variableType);
     }
 
-    public static Variable undefined() {
-        return new Variable(Optional.empty());
+    public static Variable undefined(VariableType variableType) {
+        return new Variable(Optional.empty(), variableType);
     }
 
     public boolean isDefined() {
diff --git a/src/main/java/org/shedlang/compiler/typechecker/VariableType.java b/src/main/java/org/shedlang/compiler/typechecker/VariableType.java
new file mode 100644
index 00000000..7b6b646e
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/typechecker/VariableType.java
@@ -0,0 +1,8 @@
+package org.shedlang.compiler.typechecker;
+
+public enum VariableType {
+    CONSTANT,
+    PARAM,
+    VAL,
+    VAR,
+}