Set correct type for generic struct in environment

3 files changed, 35 insertions, 3 deletions

diff --git a/src/test/java/org/shedlang/compiler/ast/typed/TypedStructDefinitionNodeMatcher.java b/src/test/java/org/shedlang/compiler/ast/typed/TypedStructDefinitionNodeMatcher.java
index 6dd41fe8..54b6cac8 100644
--- a/src/test/java/org/shedlang/compiler/ast/typed/TypedStructDefinitionNodeMatcher.java
+++ b/src/test/java/org/shedlang/compiler/ast/typed/TypedStructDefinitionNodeMatcher.java
@@ -41,7 +41,7 @@ public class TypedStructDefinitionNodeMatcher extends MatcherBuilder<
         )));
     }
 
-    public TypedStructDefinitionNodeMatcher withTypeConstructor(Matcher<TypeConstructor<?>> type) {
+    public TypedStructDefinitionNodeMatcher withTypeConstructor(Matcher<? super TypeConstructor<?>> type) {
         return with(has("typeOrConstructor", x -> x.typeOrConstructor(), instanceOf(
             TypeOrConstructor.Constructor.class,
             has("value", x -> (TypeConstructor<StructType>) x.value(), type)
diff --git a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerStructDefinitionTests.java b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerStructDefinitionTests.java
index 230fd811..497a3869 100644
--- a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerStructDefinitionTests.java
+++ b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerStructDefinitionTests.java
@@ -78,7 +78,7 @@ public class TypeCheckerStructDefinitionTests {
     }
 
     @Test
-    public void typeIsAddedToEnvironment() {
+    public void whenDefinitionHasNoTypeParamsThenTypeIsAddedToEnvironment() {
         var untyped = UntypedStructDefinitionNode.builder()
             .withName("ExampleStruct")
             .build();
@@ -96,6 +96,34 @@ public class TypeCheckerStructDefinitionTests {
     }
 
     @Test
+    public void whenDefinitionHasTypeParamsThenTypeConstructorIsAddedToEnvironment() {
+        var untyped = UntypedStructDefinitionNode.builder()
+            .withName("ExampleStruct")
+            .withTypeLevelParams(List.of(
+                Untyped.typeLevelParam("T1"),
+                Untyped.typeLevelParam("T2")
+            ))
+            .build();
+        var context = TypeCheckerContext.EMPTY;
+
+        var typed = typeCheckNamespaceStatement(untyped, context);
+
+        assertThat(
+            typed.context().compileTimeLookup("ExampleStruct"),
+            isOptionalOf(allOf(
+                has("type", x -> x.type(), isOptionalOf(isTypeLevelValueType(isTypeConstructor(
+                    isSequence(
+                        isStructTypeLevelParam(NAMESPACE_NAME, "ExampleStruct", "T1"),
+                        isStructTypeLevelParam(NAMESPACE_NAME, "ExampleStruct", "T2")
+                    ),
+                    isStructType(NAMESPACE_NAME, "ExampleStruct")
+                )))),
+                has("variableType", x -> x.variableType(), equalTo(VariableType.CONSTANT)
+            )))
+        );
+    }
+
+    @Test
     public void canTypeCheckStructDefinitionWithNoFields() {
         var untyped = UntypedStructDefinitionNode.builder()
             .withName("ExampleStruct")
diff --git a/src/test/java/org/shedlang/compiler/types/TypeMatchers.java b/src/test/java/org/shedlang/compiler/types/TypeMatchers.java
index cc47bad4..7f0e9faa 100644
--- a/src/test/java/org/shedlang/compiler/types/TypeMatchers.java
+++ b/src/test/java/org/shedlang/compiler/types/TypeMatchers.java
@@ -18,7 +18,7 @@ public class TypeMatchers {
         );
     }
 
-    public static Matcher<TypeConstructor<?>> isTypeConstructor(
+    public static Matcher<TypeLevelValue> isTypeConstructor(
         Matcher<? super Iterable<? extends TypeParam>> params,
         Matcher<? super Type> genericType
     ) {
@@ -42,4 +42,8 @@ public class TypeMatchers {
             has("name", x -> x.name(), equalTo(paramName))
         );
     }
+
+    public static Matcher<Type> isTypeLevelValueType(Matcher<TypeLevelValue> value) {
+        return instanceOf(TypeLevelValueType.class, has("value", x -> x.value(), value));
+    }
 }