Use type params in struct fields

3 files changed, 75 insertions, 6 deletions

diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
index 54a3722c..964a0000 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -3,6 +3,7 @@ package org.shedlang.compiler.typechecker;
 import org.shedlang.compiler.ast.typed.*;
 import org.shedlang.compiler.ast.untyped.*;
 import org.shedlang.compiler.sources.Source;
+import org.shedlang.compiler.sources.Sources;
 import org.shedlang.compiler.typechecker.errors.NoSuchFieldError;
 import org.shedlang.compiler.typechecker.errors.*;
 import org.shedlang.compiler.typechecker.typesets.TypeSetTypeLevelValue;
@@ -1440,6 +1441,26 @@ public class TypeChecker {
 
             @Override
             public TypeCheckerContext defineTypes(TypeCheckerContext context) {
+                var typeOrConstructor = (TypeOrConstructor<StructType>) context.getLexicalState(
+                    TypeOrConstructor.class,
+                    untyped.name()
+                );
+
+                var environment = EnvironmentBuilder.EMPTY;
+
+                if (typeOrConstructor instanceof TypeOrConstructor.Constructor<StructType> constructor) {
+                    for (var param : constructor.value().params()) {
+                        // TODO: proper source
+                        environment = environment.add(
+                            param.name(),
+                            Variable.defined(Types.metaType(param), VariableType.CONSTANT),
+                            Sources.NULL
+                        );
+                    }
+                }
+
+                context = context.enterBlock(environment.build());
+
                 var typedFieldNodes = new ArrayList<TypedStructFieldDefinitionNode>();
                 var fieldInfosBuilder = FieldsInfo.builder();
 
@@ -1453,10 +1474,7 @@ public class TypeChecker {
                     fieldInfosBuilder = fieldInfosBuilder.addField(untypedField.name(), typedFieldNode.type().value());
                 }
 
-                var typeOrConstructor = (TypeOrConstructor<StructType>) context.getLexicalState(
-                    TypeOrConstructor.class,
-                    untyped.name()
-                );
+                context = context.exit();
 
                 context = context.addStructTypeInfo(
                     typeOrConstructor.toType(),
diff --git a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerStructDefinitionTests.java b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerStructDefinitionTests.java
index 497a3869..76d399db 100644
--- a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerStructDefinitionTests.java
+++ b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerStructDefinitionTests.java
@@ -196,4 +196,54 @@ public class TypeCheckerStructDefinitionTests {
             ))
         );
     }
+
+    @Test
+    public void typeParamsCanBeUsedInFieldTypes() {
+        var untyped = UntypedStructDefinitionNode.builder()
+            .withName("ExampleStruct")
+            .withTypeLevelParams(List.of(
+                Untyped.typeLevelParam("T1"),
+                Untyped.typeLevelParam("T2")
+            ))
+            .addField(Untyped.structFieldDefinition("x1", Untyped.typeLevelReference("T1")))
+            .addField(Untyped.structFieldDefinition("x2", Untyped.typeLevelReference("T2")))
+            .build();
+        var context = TypeCheckerContext.EMPTY
+            .enterGlobal(P.map());
+
+        var typed = typeCheckNamespaceStatement(untyped, context);
+
+        var structTypeInfo = typed.context().toTypeInfos()
+            .structTypeInfo(Types.struct(NAMESPACE_NAME, "ExampleStruct"))
+            .orElseThrow();
+        assertThat(
+            structTypeInfo,
+            has("fields", x -> x.fields(), isSequence(
+                allOf(
+                    has("name", x -> x.name(), equalTo("x1")),
+                    has("type", x -> x.type(), isStructTypeLevelParam(NAMESPACE_NAME, "ExampleStruct", "T1"))
+                ),
+                allOf(
+                    has("name", x -> x.name(), equalTo("x2")),
+                    has("type", x -> x.type(), isStructTypeLevelParam(NAMESPACE_NAME, "ExampleStruct", "T2"))
+                )
+            ))
+        );
+    }
+
+    @Test
+    public void typeParamsCannotBeUsedAfterStructDefinition() {
+        var untyped = UntypedStructDefinitionNode.builder()
+            .withName("ExampleStruct")
+            .withTypeLevelParams(List.of(
+                Untyped.typeLevelParam("T1")
+            ))
+            .build();
+        var context = TypeCheckerContext.EMPTY
+            .enterGlobal(P.map());
+
+        var typed = typeCheckNamespaceStatement(untyped, context);
+
+        assertThat(typed.context().compileTimeLookup("T1"), isOptionalEmpty());
+    }
 }
diff --git a/src/test/java/org/shedlang/compiler/types/TypeMatchers.java b/src/test/java/org/shedlang/compiler/types/TypeMatchers.java
index 7f0e9faa..b52bbdff 100644
--- a/src/test/java/org/shedlang/compiler/types/TypeMatchers.java
+++ b/src/test/java/org/shedlang/compiler/types/TypeMatchers.java
@@ -29,12 +29,13 @@ public class TypeMatchers {
         );
     }
 
-    public static Matcher<TypeParam> isStructTypeLevelParam(
+    public static Matcher<Type> isStructTypeLevelParam(
         NamespaceName namespaceName,
         String structName,
         String paramName
     ) {
-        return allOf(
+        return instanceOf(
+            TypeParam.class,
             has("scope", x -> x.scope(), instanceOf(
                 Scope.Struct.class,
                 has("type", x -> x.type(), isStructType(namespaceName, structName))