Define struct type info earlier

4 files changed, 66 insertions, 6 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);
             }
         };
     }
diff --git a/src/test/java/org/shedlang/compiler/ast/typed/TypedNodeMatchers.java b/src/test/java/org/shedlang/compiler/ast/typed/TypedNodeMatchers.java
index a21e3375..7bd1402d 100644
--- a/src/test/java/org/shedlang/compiler/ast/typed/TypedNodeMatchers.java
+++ b/src/test/java/org/shedlang/compiler/ast/typed/TypedNodeMatchers.java
@@ -83,6 +83,15 @@ public class TypedNodeMatchers {
         );
     }
 
+    public static Matcher<TypedExpressionNode> isTypedExpressionNode(
+        Type type
+    ) {
+        return instanceOf(
+            TypedExpressionNode.class,
+            has("type", x -> x.type(), equalTo(type))
+        );
+    }
+
     public static Matcher<TypedFunctionStatementNode> isTypedExpressionStatementNode(
         Matcher<TypedExpressionNode> expression
     ) {
diff --git a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerNamespaceTests.java b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerNamespaceTests.java
index 5a081076..226769bf 100644
--- a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerNamespaceTests.java
+++ b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerNamespaceTests.java
@@ -324,7 +324,7 @@ public class TypeCheckerNamespaceTests {
     }
 
     @Test
-    public void functionCanUseTypeDefinedLater() {
+    public void functionSignatureCanUseTypeDefinedLater() {
         var untyped = Untyped.namespace(NamespaceName.of("Example", "Project"), List.of(
             UntypedFunctionDefinitionNode.builder()
                 .withName("printUser")
@@ -362,6 +362,39 @@ public class TypeCheckerNamespaceTests {
     }
 
     @Test
+    public void functionBodyCanAccessFieldOnTypeDefinedLater() {
+        var untyped = Untyped.namespace(NamespaceName.of("Example", "Project"), List.of(
+            UntypedFunctionDefinitionNode.builder()
+                .withName("printUser")
+                .addPositionalParam("user", Untyped.typeLevelReference("User"))
+                .withReturnType(Untyped.typeLevelReference("Unit"))
+                .addStatement(Untyped.expressionStatement(Untyped.fieldAccess(Untyped.reference("user"), "name")))
+                .build(),
+            UntypedStructDefinitionNode.builder()
+                .withName("User")
+                .addField(Untyped.structFieldDefinition("name", Untyped.typeLevelReference("String")))
+                .build()
+        ));
+        var context = TypeCheckerContext.EMPTY
+            .enterGlobal(P.map(Map.entry("String", Types.metaType(Types.STRING))))
+            .enterGlobal(P.map(Map.entry("Unit", Types.metaType(Types.UNIT))));
+
+        var result = TypeChecker.typeCheckNamespace(untyped, context);
+
+        assertThat(
+            result.value(),
+            isTypedNamespaceNode()
+                .withBody(isSequence(
+                    isTypedFunctionDefinitionNode()
+                        .withBody(isTypedBlockNode(isSequence(
+                            isTypedExpressionStatementNode(isTypedExpressionNode(Types.STRING))
+                        ))),
+                    isTypedStructDefinitionNode()
+                ))
+        );
+    }
+
+    @Test
     public void structsCanHaveMutuallyRecursiveFields() {
         var namespaceName = NamespaceName.of("Example", "Project");
         var untyped = Untyped.namespace(namespaceName, List.of(