Set correct type for generic struct in environment

6 files changed, 55 insertions, 7 deletions

diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
index 7c0dc632..54a3722c 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -1431,8 +1431,11 @@ public class TypeChecker {
 
                 context = context.addLexicalState(untyped.name(), typeOrConstructor);
 
-                // TODO: set correct type for type constructor in environment
-                return context.declareVariableType(untyped.name(), Types.metaType(type), untyped.source());
+                return context.declareVariableType(
+                    untyped.name(),
+                    typeOrConstructor.toMetaType(),
+                    untyped.source()
+                );
             }
 
             @Override
diff --git a/src/main/java/org/shedlang/compiler/types/TypeOrConstructor.java b/src/main/java/org/shedlang/compiler/types/TypeOrConstructor.java
index fea54044..6c512117 100644
--- a/src/main/java/org/shedlang/compiler/types/TypeOrConstructor.java
+++ b/src/main/java/org/shedlang/compiler/types/TypeOrConstructor.java
@@ -10,12 +10,18 @@ public sealed interface TypeOrConstructor<T extends Type> {
     }
 
     T toType();
+    org.shedlang.compiler.types.Type toMetaType();
 
     record Type<T extends org.shedlang.compiler.types.Type>(T value) implements TypeOrConstructor<T> {
         @Override
         public T toType() {
             return value;
         }
+
+        @Override
+        public org.shedlang.compiler.types.Type toMetaType() {
+            return Types.metaType(value);
+        }
     }
 
     record Constructor<T extends org.shedlang.compiler.types.Type>(TypeConstructor<T> value) implements TypeOrConstructor<T> {
@@ -23,5 +29,10 @@ public sealed interface TypeOrConstructor<T extends Type> {
         public T toType() {
             return value.genericType();
         }
+
+        @Override
+        public org.shedlang.compiler.types.Type toMetaType() {
+            return Types.typeConstructorType(value);
+        }
     }
 }
diff --git a/src/main/java/org/shedlang/compiler/types/Types.java b/src/main/java/org/shedlang/compiler/types/Types.java
index b2303558..76c02f07 100644
--- a/src/main/java/org/shedlang/compiler/types/Types.java
+++ b/src/main/java/org/shedlang/compiler/types/Types.java
@@ -1,7 +1,5 @@
 package org.shedlang.compiler.types;
 
-import org.shedlang.compiler.ast.untyped.UntypedTypeLevelParamNode;
-
 import java.util.List;
 import java.util.stream.Collectors;
 
@@ -40,6 +38,10 @@ public class Types {
         return new TypeLevelValueType(effect);
     }
 
+    public static Type typeConstructorType(TypeConstructor<?> typeConstructor) {
+        return new TypeLevelValueType(typeConstructor);
+    }
+
     public static Type commonSupertype(List<Type> subtypes) {
         var uniqueSubtypes = subtypes.stream()
             .filter(subtype -> subtype != Types.BOTTOM)
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));
+    }
 }