Store native type info with methods

4 files changed, 95 insertions, 10 deletions

diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
index 59f961ab..5d08633d 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -1765,14 +1765,17 @@ public class TypeChecker {
                     TypeCheckNativeTypeDeclarationState.class,
                     untyped.name()
                 );
+                var typeOrConstructor = state.typeOrConstructor;
+
                 var typedMethodNodes = new ArrayList<TypedNativeFunctionDeclarationNode>();
+                var methodInfos = new ArrayList<MethodInfo>();
 
                 for (var untypedMethodNode : untyped.methods()) {
                     context = context.enterBlock(P.map(
                         Map.entry(
                             "This",
                             Variable.defined(
-                                Types.metaType(state.typeOrConstructor.toType()),
+                                Types.metaType(typeOrConstructor.toType()),
                                 VariableType.VAL
                             )
                         )
@@ -1793,12 +1796,26 @@ public class TypeChecker {
                         untypedMethodNode.source()
                     );
                     typedMethodNodes.add(typedMethodNode);
+
+                    var methodType = functionTypeToMethodType(typedSignatureNode.value());
+                    var methodInfo = new MethodInfo(
+                        untypedMethodNode.name(),
+                        methodType
+                    );
+                    methodInfos.add(methodInfo);
                 }
 
                 var newState = new TypeCheckNativeTypeDeclarationState(
-                    state.typeOrConstructor,
+                    typeOrConstructor,
                     typedMethodNodes
                 );
+
+                var nativeTypeInfo = new NativeTypeInfo(methodInfos);
+                context = context.addNativeTypeInfo(
+                    typeOrConstructor.toType(),
+                    nativeTypeInfo
+                );
+
                 return context.addLexicalState(untyped.name(), newState);
             }
 
@@ -2075,14 +2092,7 @@ public class TypeChecker {
                     typeCheckMethodStates.put(untypedMethodNode.name(), typeCheckFunctionState);
 
                     var functionType = declareFunctionResult.value().type();
-                    // TODO: check first parameter type?
-                    var methodType = new FunctionType(
-                        functionType.typeLevelParams(),
-                        functionType.positionalParams().stream().skip(1).toList(),
-                        functionType.namedParams(),
-                        functionType.effect(),
-                        functionType.returnType()
-                    );
+                    var methodType = functionTypeToMethodType(functionType);
                     methodInfos.add(new MethodInfo(
                         untypedMethodNode.name(),
                         methodType
@@ -2146,6 +2156,17 @@ public class TypeChecker {
         };
     }
 
+    private static FunctionType functionTypeToMethodType(FunctionType functionType) {
+        // TODO: check first parameter type?
+        return new FunctionType(
+            functionType.typeLevelParams(),
+            functionType.positionalParams().stream().skip(1).toList(),
+            functionType.namedParams(),
+            functionType.effect(),
+            functionType.returnType()
+        );
+    }
+
     private static TypeCheckSteps<? extends TypedNamespaceStatementNode> typeCheckStepsSumTypeDefinition(
         UntypedSumTypeDefinitionNode untyped
     ) {
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java b/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
index 359fa410..e9addbc2 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
@@ -477,6 +477,16 @@ public class TypeCheckerContext {
             .orElseThrow(() -> new MissingNamespaceTypeInfoError(type, source));
     }
 
+    public TypeCheckerContext addNativeTypeInfo(
+        SimpleNativeType type,
+        NativeTypeInfo nativeTypeInfo
+    ) {
+        return new TypeCheckerContext(
+            this.typeInfos.addNativeTypeInfo(type, nativeTypeInfo),
+            this.stack
+        );
+    }
+
     public TypeCheckerContext addStructTypeInfo(
         SimpleStructType structType,
         StructTypeInfo structTypeInfo
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeInfos.java b/src/main/java/org/shedlang/compiler/typechecker/TypeInfos.java
index 5905d48b..1bb6e800 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeInfos.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeInfos.java
@@ -12,6 +12,7 @@ import java.util.Set;
 public class TypeInfos {
     public static final TypeInfos EMPTY = new TypeInfos(
         Empty.map(),
+        Empty.map(),
         P.orderedPMap(),
         Empty.map(),
         Empty.map(),
@@ -19,6 +20,7 @@ public class TypeInfos {
     );
 
     private final PMap<NamespaceName, NamespaceTypeInfo> namespaceTypeInfos;
+    private final PMap<SimpleNativeType, NativeTypeInfo> nativeTypeInfos;
     private final OrderedPMap<SimpleStructType, StructTypeInfo> structTypeInfos;
     private final PMap<SimpleSumType, SumTypeInfo> sumTypeInfos;
     private final PMap<StructType, OrderedPSet<SumType>> structSupertypes;
@@ -26,12 +28,14 @@ public class TypeInfos {
 
     public TypeInfos(
         PMap<NamespaceName, NamespaceTypeInfo> namespaceTypeInfos,
+        PMap<SimpleNativeType, NativeTypeInfo> nativeTypeInfos,
         OrderedPMap<SimpleStructType, StructTypeInfo> structTypeInfos,
         PMap<SimpleSumType, SumTypeInfo> sumTypeInfos,
         PMap<StructType, OrderedPSet<SumType>> structSupertypes,
         PMap<UserDefinedEffect, UserDefinedEffectInfo> userDefinedEffectInfos
     ) {
         this.namespaceTypeInfos = namespaceTypeInfos;
+        this.nativeTypeInfos = nativeTypeInfos;
         this.structTypeInfos = structTypeInfos;
         this.sumTypeInfos = sumTypeInfos;
         this.structSupertypes = structSupertypes;
@@ -44,6 +48,7 @@ public class TypeInfos {
     ) {
         return new TypeInfos(
             this.namespaceTypeInfos.plus(namespaceName, namespaceTypeInfo),
+            this.nativeTypeInfos,
             this.structTypeInfos,
             this.sumTypeInfos,
             this.structSupertypes,
@@ -55,12 +60,31 @@ public class TypeInfos {
         return Optional.ofNullable(namespaceTypeInfos.get(namespaceName));
     }
 
+    public TypeInfos addNativeTypeInfo(
+        SimpleNativeType type,
+        NativeTypeInfo nativeTypeInfo
+    ) {
+        return new TypeInfos(
+            this.namespaceTypeInfos,
+            this.nativeTypeInfos.plus(type, nativeTypeInfo),
+            this.structTypeInfos,
+            this.sumTypeInfos,
+            this.structSupertypes,
+            this.userDefinedEffectInfos
+        );
+    }
+
+    public Optional<NativeTypeInfo> nativeTypeInfo(SimpleNativeType type) {
+        return Optional.ofNullable(nativeTypeInfos.get(type));
+    }
+
     public TypeInfos addStructTypeInfo(
         SimpleStructType type,
         StructTypeInfo structTypeInfo
     ) {
         return new TypeInfos(
             this.namespaceTypeInfos,
+            this.nativeTypeInfos,
             this.structTypeInfos.plus(type, structTypeInfo),
             this.sumTypeInfos,
             this.structSupertypes,
@@ -120,6 +144,7 @@ public class TypeInfos {
 
         return new TypeInfos(
             this.namespaceTypeInfos,
+            this.nativeTypeInfos,
             this.structTypeInfos,
             this.sumTypeInfos.plus(sumType, sumTypeInfo),
             structSupertypes,
@@ -188,6 +213,7 @@ public class TypeInfos {
     public TypeInfos addUserDefinedEffectInfo(UserDefinedEffect effect, UserDefinedEffectInfo effectInfo) {
         return new TypeInfos(
             this.namespaceTypeInfos,
+            this.nativeTypeInfos,
             this.structTypeInfos,
             this.sumTypeInfos,
             this.structSupertypes,
diff --git a/src/main/java/org/shedlang/compiler/types/NativeTypeInfo.java b/src/main/java/org/shedlang/compiler/types/NativeTypeInfo.java
new file mode 100644
index 00000000..c87266d8
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/types/NativeTypeInfo.java
@@ -0,0 +1,28 @@
+package org.shedlang.compiler.types;
+
+import org.pcollections.PVector;
+import org.shedlang.compiler.util.P;
+
+import java.util.List;
+
+public record NativeTypeInfo(
+    List<MethodInfo> methods
+) {
+    public static Builder builder() {
+        return new Builder(P.vector());
+    }
+
+    public record Builder(
+        PVector<MethodInfo> methods
+    ) {
+        public NativeTypeInfo build() {
+            return new NativeTypeInfo(methods);
+        }
+
+        public Builder addMethod(String name, FunctionType type) {
+            return new Builder(
+                this.methods.plus(new MethodInfo(name, type))
+            );
+        }
+    }
+}