Type check methods fully

1 files changed, 43 insertions, 11 deletions

diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
index 375e2120..ddd5a341 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -903,7 +903,9 @@ public class TypeChecker {
 
             @Override
             public TypeCheckerResult<TypedFunctionDefinitionNode> finish(TypeCheckerContext context) {
-                return defineFunctionDefinition(untyped, context);
+                var state = context.getLexicalState(TypeCheckFunctionDefinitionState.class, untyped.name());
+
+                return defineFunctionDefinition(untyped, state, context);
             }
         };
     }
@@ -986,9 +988,9 @@ public class TypeChecker {
 
     private static TypeCheckerResult<TypedFunctionDefinitionNode> defineFunctionDefinition(
         UntypedFunctionDefinitionNode untyped,
+        TypeCheckFunctionDefinitionState state,
         TypeCheckerContext context
     ) {
-        var state = context.getLexicalState(TypeCheckFunctionDefinitionState.class, untyped.name());
         var typedSignature = state.typedSignatureNode;
         var type = state.type;
 
@@ -1643,6 +1645,14 @@ public class TypeChecker {
         ));
     }
 
+    record TypeCheckStructDefinitionState(
+        UntypedStructDefinitionNode untypedNode,
+        TypeOrConstructor<SimpleStructType> typeOrConstructor,
+        List<TypedStructFieldDefinitionNode> typedFieldNodes,
+        HashMap<String, TypeCheckFunctionDefinitionState> typeCheckMethodStates
+    ) {
+    }
+
     private static TypeCheckSteps<TypedStructDefinitionNode> typeCheckStepsStructDefinition(
         UntypedStructDefinitionNode untyped
     ) {
@@ -1693,8 +1703,8 @@ public class TypeChecker {
                     fieldInfosBuilder = fieldInfosBuilder.addField(untypedField.name(), typedFieldNode.type().value());
                 }
 
-                var typedMethodNodes = new ArrayList<TypedFunctionDefinitionNode>();
                 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
@@ -1709,6 +1719,8 @@ public class TypeChecker {
                     var typeCheckFunctionState = declareFunctionResult.value();
                     context = declareFunctionResult.context();
 
+                    typeCheckMethodStates.put(untypedMethodNode.name(), typeCheckFunctionState);
+
                     methodInfos.add(new MethodInfo(
                         untypedMethodNode.name(),
                         declareFunctionResult.value().type()
@@ -1726,22 +1738,42 @@ public class TypeChecker {
                     )
                 );
 
-                var typedNode = new TypedStructDefinitionNode(
+                context = context.addLexicalState(untyped.name(), new TypeCheckStructDefinitionState(
+                    untyped,
                     typeOrConstructor,
                     typedFieldNodes,
-                    untyped.fieldsArePositional(),
-                    List.of(),
-                    untyped.source()
-                );
-
-                context = context.addLexicalState(untyped.name(), typedNode);
+                    typeCheckMethodStates
+                ));
 
                 return context;
             }
 
             @Override
             public TypeCheckerResult<TypedStructDefinitionNode> finish(TypeCheckerContext context) {
-                var typedNode = context.getLexicalState(TypedStructDefinitionNode.class, untyped.name());
+                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());
+
+                    var result = defineFunctionDefinition(
+                        untypedMethodNode,
+                        typeCheckMethodState,
+                        context
+                    );
+                    var typedMethodNode = result.value();
+                    context = result.context();
+
+                    typedMethodNodes.add(typedMethodNode);
+                }
+
+                var typedNode = new TypedStructDefinitionNode(
+                    state.typeOrConstructor,
+                    state.typedFieldNodes,
+                    untyped.fieldsArePositional(),
+                    typedMethodNodes,
+                    untyped.source()
+                );
 
                 return context.toResult(typedNode);
             }