Move method declaration to function declaration phase

1 files changed, 87 insertions, 27 deletions

diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
index 3e8e5ddf..b8a05746 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -1782,18 +1782,46 @@ public class TypeChecker {
         ));
     }
 
-    record TypeCheckStructDefinitionState(
-        UntypedStructDefinitionNode untypedNode,
+    private record TypeCheckStructDefinitionState(
         TypeOrConstructor<SimpleStructType> typeOrConstructor,
-        List<TypedStructFieldDefinitionNode> typedFieldNodes,
-        HashMap<String, TypeCheckFunctionDefinitionState> typeCheckMethodStates
+        Optional<List<TypedStructFieldDefinitionNode>> typedFieldNodes,
+        Optional<HashMap<String, TypeCheckFunctionDefinitionState>> typeCheckMethodStates
     ) {
+        private static TypeCheckStructDefinitionState forTypeOrConstructor(
+            TypeOrConstructor<SimpleStructType> typeOrConstructor
+        ) {
+            return new TypeCheckStructDefinitionState(
+                typeOrConstructor,
+                Optional.empty(),
+                Optional.empty()
+            );
+        }
+
+        private TypeCheckStructDefinitionState withTypedFieldNodes(
+            List<TypedStructFieldDefinitionNode> typedFieldNodes
+        ) {
+            return new TypeCheckStructDefinitionState(
+                this.typeOrConstructor,
+                Optional.of(typedFieldNodes),
+                this.typeCheckMethodStates
+            );
+        }
+
+        private TypeCheckStructDefinitionState withTypeCheckMethodStates(
+            HashMap<String, TypeCheckFunctionDefinitionState> typeCheckMethodStates
+        ) {
+            return new TypeCheckStructDefinitionState(
+                this.typeOrConstructor,
+                this.typedFieldNodes,
+                Optional.of(typeCheckMethodStates)
+            );
+        }
     }
 
     private static TypeCheckSteps<TypedStructDefinitionNode> typeCheckStepsStructDefinition(
         UntypedStructDefinitionNode untyped
     ) {
-        return new TypeCheckSteps<TypedStructDefinitionNode>() {
+        return new TypeCheckSteps<>() {
             @Override
             public TypeCheckerContext declareTypes(TypeCheckerContext context) {
                 var namespaceName = context.namespaceName()
@@ -1807,7 +1835,10 @@ public class TypeChecker {
                     new Scope.Struct(type)
                 );
 
-                context = context.addLexicalState(untyped.name(), typeOrConstructor);
+                var state = TypeCheckStructDefinitionState.forTypeOrConstructor(
+                    typeOrConstructor
+                );
+                context = context.addLexicalState(untyped.name(), state);
 
                 return context.declareVariableType(
                     untyped.name(),
@@ -1818,12 +1849,12 @@ public class TypeChecker {
 
             @Override
             public TypeCheckerContext defineTypes(TypeCheckerContext context) {
-                var typeOrConstructor = (TypeOrConstructor<SimpleStructType>) context.getLexicalState(
-                    TypeOrConstructor.class,
+                var state = context.getLexicalState(
+                    TypeCheckStructDefinitionState.class,
                     untyped.name()
                 );
 
-                var environment = environmentForTypeOrConstructor(typeOrConstructor);
+                var environment = environmentForTypeOrConstructor(state.typeOrConstructor);
 
                 context = context.enterBlock(environment.build());
 
@@ -1840,15 +1871,40 @@ public class TypeChecker {
                     fieldInfosBuilder = fieldInfosBuilder.addField(untypedField.name(), typedFieldNode.type().value());
                 }
 
+                context = context.exit();
+
+                context = context.addStructTypeInfo(
+                    state.typeOrConstructor.toType(),
+                    new StructTypeInfo(
+                        fieldInfosBuilder.build(),
+                        untyped.fieldsArePositional(),
+                        List.of()
+                    )
+                );
+
+                context = context.addLexicalState(
+                    untyped.name(),
+                    state.withTypedFieldNodes(typedFieldNodes)
+                );
+
+                return context;
+            }
+
+            @Override
+            public TypeCheckerContext declareFunctions(TypeCheckerContext context) {
+                var state = context.getLexicalState(
+                    TypeCheckStructDefinitionState.class,
+                    untyped.name()
+                );
+
+                var environment = environmentForTypeOrConstructor(state.typeOrConstructor);
+
+                context = context.enterBlock(environment.build());
+
                 var methodInfos = new ArrayList<MethodInfo>();
                 var typeCheckMethodStates = new HashMap<String, TypeCheckFunctionDefinitionState>();
                 for (var untypedMethodNode : untyped.methods()) {
-                    // TODO: we're declaring functions in the "define types" step,
-                    // instead of the "declare functions" step. This is probably
-                    // fine since those steps can (currently, at least) happen
-                    // in any other -- maybe they should just be merged, given
-                    // that otherwise this is invalid?
-                    var thisType = switch (typeOrConstructor) {
+                    var thisType = switch (state.typeOrConstructor) {
                         case TypeOrConstructor.Constructor<SimpleStructType> constructor ->
                             Types.constructedType(constructor.value(), constructor.value().params());
                         case TypeOrConstructor.Type<SimpleStructType> type ->
@@ -1882,21 +1938,23 @@ public class TypeChecker {
 
                 context = context.exit();
 
+                var structTypeInfo = context.structTypeInfo(
+                    state.typeOrConstructor.toType(),
+                    untyped.source()
+                );
                 context = context.addStructTypeInfo(
-                    typeOrConstructor.toType(),
+                    state.typeOrConstructor.toType(),
                     new StructTypeInfo(
-                        fieldInfosBuilder.build(),
-                        untyped.fieldsArePositional(),
+                        structTypeInfo.fields(),
+                        structTypeInfo.fieldsArePositional(),
                         methodInfos
                     )
                 );
 
-                context = context.addLexicalState(untyped.name(), new TypeCheckStructDefinitionState(
-                    untyped,
-                    typeOrConstructor,
-                    typedFieldNodes,
-                    typeCheckMethodStates
-                ));
+                context = context.addLexicalState(
+                    untyped.name(),
+                    state.withTypeCheckMethodStates(typeCheckMethodStates)
+                );
 
                 return context;
             }
@@ -1906,8 +1964,10 @@ public class TypeChecker {
                 var state = context.getLexicalState(TypeCheckStructDefinitionState.class, untyped.name());
 
                 var typedMethodNodes = new ArrayList<TypedFunctionDefinitionNode>();
-                for (var untypedMethodNode : state.untypedNode().methods()) {
-                    var typeCheckMethodState = state.typeCheckMethodStates.get(untypedMethodNode.name());
+                for (var untypedMethodNode : untyped.methods()) {
+                    var typeCheckMethodState = state.typeCheckMethodStates
+                        .orElseThrow()
+                        .get(untypedMethodNode.name());
 
                     var result = defineFunctionDefinition(
                         untypedMethodNode,
@@ -1922,7 +1982,7 @@ public class TypeChecker {
 
                 var typedNode = new TypedStructDefinitionNode(
                     state.typeOrConstructor,
-                    state.typedFieldNodes,
+                    state.typedFieldNodes.orElseThrow(),
                     untyped.fieldsArePositional(),
                     typedMethodNodes,
                     untyped.source()