Create method information for methods in struct

5 files changed, 81 insertions, 14 deletions

diff --git a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedStructDefinitionNode.java b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedStructDefinitionNode.java
index af335ee5..666a9cef 100644
--- a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedStructDefinitionNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedStructDefinitionNode.java
@@ -93,5 +93,16 @@ public record UntypedStructDefinitionNode(
 
             );
         }
+
+        public Builder addMethod(UntypedFunctionDefinitionNode method) {
+            return new Builder(
+                this.name,
+                this.visibility,
+                this.typeLevelParams,
+                this.fields,
+                this.fieldsArePositional,
+                this.methods.plus(method)
+            );
+        }
     }
 }
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
index 1d3a9535..375e2120 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -888,7 +888,17 @@ public class TypeChecker {
         return new TypeCheckSteps<TypedFunctionDefinitionNode>() {
             @Override
             public TypeCheckerContext declareFunctions(TypeCheckerContext context) {
-                return declareFunctionDefinition(untyped, context);
+                var declareResult = declareFunctionDefinition(untyped, Optional.empty(), context);
+                var state = declareResult.value();
+                context = declareResult.context();
+
+                context = context.declareVariableType(
+                    untyped.name(),
+                    state.type(),
+                    untyped.source()
+                );
+
+                return context.addLexicalState(untyped.name(), state);
             }
 
             @Override
@@ -905,8 +915,9 @@ public class TypeChecker {
     ) {
     }
 
-    private static TypeCheckerContext declareFunctionDefinition(
+    private static TypeCheckerResult<TypeCheckFunctionDefinitionState> declareFunctionDefinition(
         UntypedFunctionDefinitionNode untyped,
+        Optional<Type> thisType,
         TypeCheckerContext context
     ) {
         var namespaceName = context.namespaceName()
@@ -927,7 +938,16 @@ public class TypeChecker {
             );
 
         var positionalParams = typedSignature.params().positional().stream()
-            .map(typedParamNodeToParam)
+            .map(typedParamNode -> {
+                // TODO: check that type is not present
+                return thisType
+                    .map(type -> new PartialFunctionSignature.Param(
+                        typedParamNode.name(),
+                        type,
+                        typedParamNode.source()
+                    ))
+                    .orElseGet(() -> typedParamNodeToParam.apply(typedParamNode));
+            })
             .toList();
 
         var namedParams = typedSignature.params().named().stream()
@@ -961,9 +981,7 @@ public class TypeChecker {
 
         var state = new TypeCheckFunctionDefinitionState(typedSignature, partialSignature, type);
 
-        return context
-            .declareVariableType(untyped.name(), type, untyped.source())
-            .addLexicalState(untyped.name(), state);
+        return context.toResult(state);
     }
 
     private static TypeCheckerResult<TypedFunctionDefinitionNode> defineFunctionDefinition(
@@ -1675,13 +1693,36 @@ public class TypeChecker {
                     fieldInfosBuilder = fieldInfosBuilder.addField(untypedField.name(), typedFieldNode.type().value());
                 }
 
+                var typedMethodNodes = new ArrayList<TypedFunctionDefinitionNode>();
+                var methodInfos = new ArrayList<MethodInfo>();
+                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 declareFunctionResult = declareFunctionDefinition(
+                        untypedMethodNode,
+                        Optional.of(typeOrConstructor.toType()),
+                        context
+                    );
+                    var typeCheckFunctionState = declareFunctionResult.value();
+                    context = declareFunctionResult.context();
+
+                    methodInfos.add(new MethodInfo(
+                        untypedMethodNode.name(),
+                        declareFunctionResult.value().type()
+                    ));
+                }
+
                 context = context.exit();
 
                 context = context.addStructTypeInfo(
                     typeOrConstructor.toType(),
                     new StructTypeInfo(
                         fieldInfosBuilder.build(),
-                        untyped.fieldsArePositional()
+                        untyped.fieldsArePositional(),
+                        methodInfos
                     )
                 );
 
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeInfos.java b/src/main/java/org/shedlang/compiler/typechecker/TypeInfos.java
index 503e97a8..5905d48b 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeInfos.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeInfos.java
@@ -5,6 +5,7 @@ import org.shedlang.compiler.types.*;
 import org.shedlang.compiler.util.P;
 import org.shedlang.compiler.util.PCollectors;
 
+import java.util.List;
 import java.util.Optional;
 import java.util.Set;
 
@@ -96,7 +97,9 @@ public class TypeInfos {
 
             return new StructTypeInfo(
                 fieldsInfoBuilder.build(),
-                typeInfo.fieldsArePositional()
+                typeInfo.fieldsArePositional(),
+                // TODO: handle methods
+                List.of()
             );
         });
     }
diff --git a/src/main/java/org/shedlang/compiler/types/MethodInfo.java b/src/main/java/org/shedlang/compiler/types/MethodInfo.java
new file mode 100644
index 00000000..8c5548e7
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/types/MethodInfo.java
@@ -0,0 +1,4 @@
+package org.shedlang.compiler.types;
+
+public record MethodInfo(String name, FunctionType type) {
+}
diff --git a/src/main/java/org/shedlang/compiler/types/StructTypeInfo.java b/src/main/java/org/shedlang/compiler/types/StructTypeInfo.java
index a34a9c10..b8a7df38 100644
--- a/src/main/java/org/shedlang/compiler/types/StructTypeInfo.java
+++ b/src/main/java/org/shedlang/compiler/types/StructTypeInfo.java
@@ -1,10 +1,15 @@
 package org.shedlang.compiler.types;
 
+import org.pcollections.PVector;
+import org.shedlang.compiler.util.P;
+
+import java.util.List;
 import java.util.Optional;
 
 public record StructTypeInfo(
     FieldsInfo<Type> fields,
-    boolean fieldsArePositional
+    boolean fieldsArePositional,
+    List<MethodInfo> methods
 ) {
     public int fieldCount() {
         return fields().fieldCount();
@@ -19,28 +24,31 @@ public record StructTypeInfo(
     }
 
     public static Builder builder() {
-        return new Builder(FieldsInfo.builder(), false);
+        return new Builder(FieldsInfo.builder(), false, P.vector());
     }
 
     public record Builder(
         FieldsInfo.Builder<Type> fields,
-        boolean fieldsArePositional
+        boolean fieldsArePositional,
+        PVector<MethodInfo> methods
     ) {
         public StructTypeInfo build() {
-            return new StructTypeInfo(fields.build(), fieldsArePositional);
+            return new StructTypeInfo(fields.build(), fieldsArePositional, methods);
         }
 
         public Builder addField(String name, Type type) {
             return new Builder(
                 this.fields.addField(name, type),
-                fieldsArePositional
+                fieldsArePositional,
+                this.methods
             );
         }
 
         public Builder withFieldsArePositional(boolean fieldsArePositional) {
             return new Builder(
                 this.fields,
-                fieldsArePositional
+                fieldsArePositional,
+                this.methods
             );
         }
     }