Add struct type to context

4 files changed, 44 insertions, 13 deletions

diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeCheckSteps.java b/src/main/java/org/shedlang/compiler/typechecker/TypeCheckSteps.java
index 4e1e46ec..be94fb24 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeCheckSteps.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeCheckSteps.java
@@ -1,13 +1,17 @@
 package org.shedlang.compiler.typechecker;
 
 public interface TypeCheckSteps<T> {
+    default TypeCheckerContext declareTypes(TypeCheckerContext context) {
+        return context;
+    }
+
     default TypeCheckerContext declareFunctions(TypeCheckerContext context) {
         return context;
     }
 
-    TypeCheckerResult<T> define(TypeCheckerContext context);
+    TypeCheckerResult<T> finish(TypeCheckerContext context);
 
     default TypeCheckerResult<T> typeCheckAllPhases(TypeCheckerContext context) {
-        return define(declareFunctions(context));
+        return finish(declareFunctions(declareTypes(context)));
     }
 }
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
index 1681ac25..1c8bbd54 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -130,7 +130,7 @@ public class TypeChecker {
     ) {
         return new TypeCheckSteps<>() {
             @Override
-            public TypeCheckerResult<TypedAssignNode> define(TypeCheckerContext context) {
+            public TypeCheckerResult<TypedAssignNode> finish(TypeCheckerContext context) {
                 return typeCheckAssign(untyped, context);
             }
         };
@@ -402,7 +402,7 @@ public class TypeChecker {
             }
 
             @Override
-            public TypeCheckerResult<TypedFunctionDefinitionNode> define(TypeCheckerContext context) {
+            public TypeCheckerResult<TypedFunctionDefinitionNode> finish(TypeCheckerContext context) {
                 return defineFunctionDefinition(untyped, context);
             }
         };
@@ -626,7 +626,7 @@ public class TypeChecker {
             namespaceEnvironment = namespaceEnvironment.addUndefined(declaredVariables);
         }
 
-        context = context.enterNamespace(namespaceEnvironment.build());
+        context = context.enterNamespace(untyped.namespaceName(), namespaceEnvironment.build());
 
         var typedBody = new ArrayList<TypedNamespaceStatementNode>();
 
@@ -635,6 +635,10 @@ public class TypeChecker {
             .toList();
 
         for (var typeCheckSteps : typeCheckStepSets) {
+            context = typeCheckSteps.declareTypes(context);
+        }
+
+        for (var typeCheckSteps : typeCheckStepSets) {
             context = typeCheckSteps.declareFunctions(context);
         }
 
@@ -651,7 +655,7 @@ public class TypeChecker {
 
             context = context.enterFreeVariableCapture();
 
-            var typedStatementResult = typeCheckSteps.define(context);
+            var typedStatementResult = typeCheckSteps.finish(context);
             var typedStatement = typedStatementResult.value();
             context = typedStatementResult.context();
 
@@ -682,7 +686,6 @@ public class TypeChecker {
                 }
 
                 case UntypedStructDefinitionNode untypedStruct -> {
-                    throw new UnsupportedOperationException("TODO");
                 }
             }
         }
@@ -786,7 +789,14 @@ public class TypeChecker {
     ) {
         return new TypeCheckSteps<TypedStructDefinitionNode>() {
             @Override
-            public TypeCheckerResult<TypedStructDefinitionNode> define(TypeCheckerContext context) {
+            public TypeCheckerContext declareTypes(TypeCheckerContext context) {
+                // TODO: better error
+                var type = new StructType(context.namespaceName().orElseThrow(), untyped.name());
+                return context.define(untyped.name(), Types.metaType(type), untyped.source());
+            }
+
+            @Override
+            public TypeCheckerResult<TypedStructDefinitionNode> finish(TypeCheckerContext context) {
                 var fieldsResult = context.mapWithContext(untyped.fields(), (untypedField, fieldContext) -> {
                     return fieldContext.toResult(new TypedStructFieldDefinitionNode(
                         untypedField.name(),
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java b/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
index 1f50d24b..044f5041 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
@@ -21,6 +21,7 @@ public class TypeCheckerContext {
     public static final TypeCheckerContext EMPTY = new TypeCheckerContext(
         Empty.map(),
         P.stack(new LexicalContext(
+            Optional.empty(),
             Empty.map(),
             Effects.NONE,
             Optional.empty(),
@@ -56,6 +57,7 @@ public class TypeCheckerContext {
     }
 
     private static final class LexicalContext implements StackElement {
+        private final Optional<NamespaceName> namespaceName;
         private final PMap<String, Variable> environment;
         private final Effect effect;
         private final Optional<Type> returnType;
@@ -64,6 +66,7 @@ public class TypeCheckerContext {
         private final PMap<String, Object> state;
 
         public LexicalContext(
+            Optional<NamespaceName> namespaceName,
             PMap<String, Variable> environment,
             Effect effect,
             Optional<Type> returnType,
@@ -71,6 +74,7 @@ public class TypeCheckerContext {
             FreeVariables.Builder freeVariables,
             PMap<String, Object> state
         ) {
+            this.namespaceName = namespaceName;
             this.environment = environment;
             this.effect = effect;
             this.returnType = returnType;
@@ -87,6 +91,7 @@ public class TypeCheckerContext {
         @Override
         public LexicalContext withFreeVariable(String name, Source referenceSource) {
             return new LexicalContext(
+                this.namespaceName,
                 this.environment,
                 this.effect,
                 this.returnType,
@@ -142,11 +147,11 @@ public class TypeCheckerContext {
             .map(x -> Map.entry(x.getKey(), Variable.defined(x.getValue(), VariableType.VAL)))
             .collect(PCollectors.toMap());
 
-        return enterNamespace(environment);
+        return enter(Optional.empty(), environment, Effects.NONE, Optional.empty());
     }
 
-    public TypeCheckerContext enterNamespace(PMap<String, Variable> environment) {
-        return enter(environment, Effects.NONE, Optional.empty());
+    public TypeCheckerContext enterNamespace(NamespaceName namespaceName, PMap<String, Variable> environment) {
+        return enter(Optional.of(namespaceName), environment, Effects.NONE, Optional.empty());
     }
 
     public TypeCheckerContext enterFunction(
@@ -154,14 +159,15 @@ public class TypeCheckerContext {
         Effect effect,
         Type returnType
     ) {
-        return enter(environment, effect, Optional.of(returnType));
+        return enter(this.namespaceName(), environment, effect, Optional.of(returnType));
     }
 
     public TypeCheckerContext enterBlock(PMap<String, Variable> environment) {
-        return enter(environment, this.effect(), this.returnType());
+        return enter(this.namespaceName(), environment, this.effect(), this.returnType());
     }
 
     private TypeCheckerContext enter(
+        Optional<NamespaceName> newNamespaceName,
         PMap<String, Variable> newEnvironment,
         Effect newEffect,
         Optional<Type> newReturnType
@@ -169,6 +175,7 @@ public class TypeCheckerContext {
         return new TypeCheckerContext(
             this.namespaceTypeInfo,
             this.stack.plus(new LexicalContext(
+                newNamespaceName,
                 this.environment().plusAll(newEnvironment),
                 newEffect,
                 newReturnType,
@@ -235,6 +242,7 @@ public class TypeCheckerContext {
     public TypeCheckerContext addLexicalState(String name, Object state) {
         return this.updateLexicalContext(
             lexicalContext -> new LexicalContext(
+                lexicalContext.namespaceName,
                 lexicalContext.environment,
                 lexicalContext.effect,
                 lexicalContext.returnType,
@@ -275,6 +283,7 @@ public class TypeCheckerContext {
 
         return this.updateLexicalContext(
             lexicalContext -> new LexicalContext(
+                lexicalContext.namespaceName,
                 lexicalContext.environment.plus(name, new Variable(
                     Optional.of(type),
                     variable.variableType()
@@ -288,6 +297,10 @@ public class TypeCheckerContext {
         );
     }
 
+    public Optional<NamespaceName> namespaceName() {
+        return currentLexicalContext().namespaceName;
+    }
+
     private PMap<String, Variable> environment() {
         return currentLexicalContext().environment;
     }
diff --git a/src/main/java/org/shedlang/compiler/types/Types.java b/src/main/java/org/shedlang/compiler/types/Types.java
index 2f07133a..f56f0ca7 100644
--- a/src/main/java/org/shedlang/compiler/types/Types.java
+++ b/src/main/java/org/shedlang/compiler/types/Types.java
@@ -14,6 +14,10 @@ public class Types {
     public static final Type TOP = TopType.INSTANCE;
     public static final Type UNIT = UnitType.INSTANCE;
 
+    public static Type struct(NamespaceName namespaceName, String name) {
+        return new StructType(namespaceName, name);
+    }
+
     public static TupleType tuple(List<Type> elements) {
         return new TupleType(elements);
     }