Create type constructor for structs with type params

11 files changed, 102 insertions, 13 deletions

diff --git a/src/main/java/org/shedlang/compiler/ast/typed/TypedStructDefinitionNode.java b/src/main/java/org/shedlang/compiler/ast/typed/TypedStructDefinitionNode.java
index eee3fcb8..bec48751 100644
--- a/src/main/java/org/shedlang/compiler/ast/typed/TypedStructDefinitionNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/typed/TypedStructDefinitionNode.java
@@ -5,25 +5,27 @@ import org.shedlang.compiler.sources.Source;
 import org.shedlang.compiler.types.NamespaceName;
 import org.shedlang.compiler.types.StructType;
 import org.shedlang.compiler.sources.Sources;
+import org.shedlang.compiler.types.TypeOrConstructor;
 import org.shedlang.compiler.types.Types;
 import org.shedlang.compiler.util.P;
 
 import java.util.List;
 
 public record TypedStructDefinitionNode(
-    StructType type,
+    TypeOrConstructor<StructType> typeOrConstructor,
     List<TypedStructFieldDefinitionNode> fields,
     Source source
 ) implements TypedNamespaceStatementNode {
     public static Builder builder() {
+        var namespaceName = NamespaceName.of("Default", "Test", "Namespace");
         return new Builder(
-            Types.struct(NamespaceName.of("Default", "Test", "Namespace"), "X"),
+            TypeOrConstructor.type(Types.struct(namespaceName, "X")),
             P.vector()
         );
     }
 
     public record Builder(
-        StructType type,
+        TypeOrConstructor<StructType> type,
         PVector<TypedStructFieldDefinitionNode> fields
     ) {
         public TypedStructDefinitionNode build() {
@@ -36,7 +38,7 @@ public record TypedStructDefinitionNode(
 
         public Builder withType(StructType type) {
             return new Builder(
-                type,
+                TypeOrConstructor.type(type),
                 this.fields
             );
         }
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 148bd066..c1611322 100644
--- a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedStructDefinitionNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedStructDefinitionNode.java
@@ -45,6 +45,15 @@ public record UntypedStructDefinitionNode(
             );
         }
 
+        public Builder withTypeLevelParams(List<UntypedTypeLevelParamNode> typeLevelParams) {
+            return new Builder(
+                this.name,
+                this.visibility,
+                Optional.of(typeLevelParams),
+                this.fields
+            );
+        }
+
         public Builder addField(UntypedStructFieldDefinitionNode field) {
             return new Builder(
                 this.name,
diff --git a/src/main/java/org/shedlang/compiler/shid/ShidCodeGen.java b/src/main/java/org/shedlang/compiler/shid/ShidCodeGen.java
index 89bf275c..8c42982b 100644
--- a/src/main/java/org/shedlang/compiler/shid/ShidCodeGen.java
+++ b/src/main/java/org/shedlang/compiler/shid/ShidCodeGen.java
@@ -540,15 +540,17 @@ public class ShidCodeGen {
         TypedStructDefinitionNode node,
         Context context
     ) {
+        var structType = node.typeOrConstructor().toType();
+
         var defineStructInstruction = new ShidDefineStruct(
-            node.type(),
+            structType,
             node.fields().stream()
                 .map(field -> new ShidDefineStruct.Field(field.name(), field.source()))
                 .toList(),
             node.source()
         );
 
-        var localName = context.shedVariableToLocal(node.type().name());
+        var localName = context.shedVariableToLocal(structType.name());
         var storeInstruction = new ShidLocalStore(localName, node.source());
 
         return context
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
index b789e888..7c0dc632 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -1416,8 +1416,22 @@ public class TypeChecker {
 
                 var type = new StructType(namespaceName, untyped.name());
 
-                context = context.addLexicalState(untyped.name(), type);
-
+                var typeOrConstructor = untyped.typeLevelParams()
+                    .map(typeLevelParamNodes -> TypeOrConstructor.constructor(new TypeConstructor<>(
+                        type.name(),
+                        typeLevelParamNodes.stream()
+                            .map(typeLevelParamNode -> new TypeParam(
+                                new Scope.Struct(type),
+                                typeLevelParamNode.name()
+                            ))
+                            .toList(),
+                        type
+                    )))
+                    .orElseGet(() -> TypeOrConstructor.type(type));
+
+                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());
             }
 
@@ -1436,12 +1450,18 @@ public class TypeChecker {
                     fieldInfosBuilder = fieldInfosBuilder.addField(untypedField.name(), typedFieldNode.type().value());
                 }
 
-                var type = context.getLexicalState(StructType.class, untyped.name());
+                var typeOrConstructor = (TypeOrConstructor<StructType>) context.getLexicalState(
+                    TypeOrConstructor.class,
+                    untyped.name()
+                );
 
-                context = context.addStructTypeInfo(type, new StructTypeInfo(fieldInfosBuilder.build()));
+                context = context.addStructTypeInfo(
+                    typeOrConstructor.toType(),
+                    new StructTypeInfo(fieldInfosBuilder.build())
+                );
 
                 var typedNode = new TypedStructDefinitionNode(
-                    type,
+                    typeOrConstructor,
                     typedFieldNodes,
                     untyped.source()
                 );
diff --git a/src/main/java/org/shedlang/compiler/types/Scope.java b/src/main/java/org/shedlang/compiler/types/Scope.java
new file mode 100644
index 00000000..0ba3c776
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/types/Scope.java
@@ -0,0 +1,5 @@
+package org.shedlang.compiler.types;
+
+public sealed interface Scope {
+    record Struct(StructType type) implements Scope {}
+}
diff --git a/src/main/java/org/shedlang/compiler/types/Type.java b/src/main/java/org/shedlang/compiler/types/Type.java
index a1de4068..27d133ee 100644
--- a/src/main/java/org/shedlang/compiler/types/Type.java
+++ b/src/main/java/org/shedlang/compiler/types/Type.java
@@ -6,5 +6,5 @@ package org.shedlang.compiler.types;
  * found through auxiliary data.
  */
 public sealed interface Type extends TypeLevelValue
-    permits BoolType, BottomType, FunctionType, IntType, NamespaceType, StringType, StructType, SumType, TopType, TupleType, TypeLevelValueType, UnitType {
+    permits BoolType, BottomType, FunctionType, IntType, NamespaceType, StringType, StructType, SumType, TopType, TupleType, TypeLevelValueType, TypeParam, UnitType {
 }
diff --git a/src/main/java/org/shedlang/compiler/types/TypeConstructor.java b/src/main/java/org/shedlang/compiler/types/TypeConstructor.java
new file mode 100644
index 00000000..78aa0b61
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/types/TypeConstructor.java
@@ -0,0 +1,14 @@
+package org.shedlang.compiler.types;
+
+import java.util.List;
+
+public record TypeConstructor<T extends Type>(
+    String name,
+    List<TypeParam> params,
+    T genericType
+) implements TypeLevelValue {
+    @Override
+    public String describe() {
+        return name;
+    }
+}
diff --git a/src/main/java/org/shedlang/compiler/types/TypeLevelValue.java b/src/main/java/org/shedlang/compiler/types/TypeLevelValue.java
index f10232bb..a002c88a 100644
--- a/src/main/java/org/shedlang/compiler/types/TypeLevelValue.java
+++ b/src/main/java/org/shedlang/compiler/types/TypeLevelValue.java
@@ -1,5 +1,5 @@
 package org.shedlang.compiler.types;
 
-public sealed interface TypeLevelValue permits Effect, Type {
+public sealed interface TypeLevelValue permits Effect, Type, TypeConstructor {
     String describe();
 }
diff --git a/src/main/java/org/shedlang/compiler/types/TypeOrConstructor.java b/src/main/java/org/shedlang/compiler/types/TypeOrConstructor.java
new file mode 100644
index 00000000..fea54044
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/types/TypeOrConstructor.java
@@ -0,0 +1,27 @@
+package org.shedlang.compiler.types;
+
+public sealed interface TypeOrConstructor<T extends Type> {
+    static <T extends org.shedlang.compiler.types.Type> TypeOrConstructor<T> type(T value) {
+        return new Type<>(value);
+    }
+
+    static <T extends org.shedlang.compiler.types.Type> TypeOrConstructor<T> constructor(TypeConstructor<T> value) {
+        return new Constructor<>(value);
+    }
+
+    T toType();
+
+    record Type<T extends org.shedlang.compiler.types.Type>(T value) implements TypeOrConstructor<T> {
+        @Override
+        public T toType() {
+            return value;
+        }
+    }
+
+    record Constructor<T extends org.shedlang.compiler.types.Type>(TypeConstructor<T> value) implements TypeOrConstructor<T> {
+        @Override
+        public T toType() {
+            return value.genericType();
+        }
+    }
+}
diff --git a/src/main/java/org/shedlang/compiler/types/TypeParam.java b/src/main/java/org/shedlang/compiler/types/TypeParam.java
new file mode 100644
index 00000000..240d701d
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/types/TypeParam.java
@@ -0,0 +1,8 @@
+package org.shedlang.compiler.types;
+
+public record TypeParam(Scope scope, String name) implements Type {
+    @Override
+    public String describe() {
+        return name;
+    }
+}
diff --git a/src/main/java/org/shedlang/compiler/types/Types.java b/src/main/java/org/shedlang/compiler/types/Types.java
index efc9f000..b2303558 100644
--- a/src/main/java/org/shedlang/compiler/types/Types.java
+++ b/src/main/java/org/shedlang/compiler/types/Types.java
@@ -1,5 +1,7 @@
 package org.shedlang.compiler.types;
 
+import org.shedlang.compiler.ast.untyped.UntypedTypeLevelParamNode;
+
 import java.util.List;
 import java.util.stream.Collectors;