Define struct type info earlier

2 files changed, 23 insertions, 5 deletions

diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeCheckSteps.java b/src/main/java/org/shedlang/compiler/typechecker/TypeCheckSteps.java
index be94fb24..db9789e3 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeCheckSteps.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeCheckSteps.java
@@ -5,6 +5,10 @@ public interface TypeCheckSteps<T> {
         return context;
     }
 
+    default TypeCheckerContext defineTypes(TypeCheckerContext context) {
+        return context;
+    }
+
     default TypeCheckerContext declareFunctions(TypeCheckerContext context) {
         return context;
     }
@@ -12,6 +16,6 @@ public interface TypeCheckSteps<T> {
     TypeCheckerResult<T> finish(TypeCheckerContext context);
 
     default TypeCheckerResult<T> typeCheckAllPhases(TypeCheckerContext context) {
-        return finish(declareFunctions(declareTypes(context)));
+        return finish(declareFunctions(defineTypes(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 bd5b43fd..54e9d9b9 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -676,6 +676,10 @@ public class TypeChecker {
         }
 
         for (var typeCheckSteps : typeCheckStepSets) {
+            context = typeCheckSteps.defineTypes(context);
+        }
+
+        for (var typeCheckSteps : typeCheckStepSets) {
             context = typeCheckSteps.declareFunctions(context);
         }
 
@@ -836,7 +840,7 @@ public class TypeChecker {
             }
 
             @Override
-            public TypeCheckerResult<TypedStructDefinitionNode> finish(TypeCheckerContext context) {
+            public TypeCheckerContext defineTypes(TypeCheckerContext context) {
                 var typedFieldNodes = new ArrayList<TypedStructFieldDefinitionNode>();
                 var fieldTypes = OrderedPMap.<String, Type>empty();
 
@@ -852,14 +856,24 @@ public class TypeChecker {
 
                 var type = context.getLexicalState(StructType.class, untyped.name());
 
-                // TODO: this needs to happen before type-checking function bodies
                 context = context.addStructTypeInfo(type, new StructTypeInfo(fieldTypes));
 
-                return context.toResult(new TypedStructDefinitionNode(
+                var typedNode = new TypedStructDefinitionNode(
                     untyped.name(),
                     typedFieldNodes,
                     untyped.source()
-                ));
+                );
+
+                context = context.addLexicalState(untyped.name(), typedNode);
+
+                return context;
+            }
+
+            @Override
+            public TypeCheckerResult<TypedStructDefinitionNode> finish(TypeCheckerContext context) {
+                var typedNode = context.getLexicalState(TypedStructDefinitionNode.class, untyped.name());
+
+                return context.toResult(typedNode);
             }
         };
     }