Add specialised type substitution functions

2 files changed, 46 insertions, 23 deletions

diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java b/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
index f1797ef3..239e8d4d 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
@@ -488,14 +488,13 @@ public class TypeCheckerContext {
         var variants = typeInfo.variants().stream()
             .map(variant -> switch (variant) {
                 case SumVariant.ConstructedStruct constructedStruct ->
-                    // TODO: overload to avoid cast?
-                    SumVariant.constructedStruct((ConstructedType<StructType>) TypeSubstitution.substitute(
+                    SumVariant.constructedStruct(TypeSubstitution.substituteConstructedType(
                         constructedStruct.type(),
                         typeSubstitutionMap
                     ));
 
                 case SumVariant.Struct struct ->
-                    SumVariant.struct((StructType) TypeSubstitution.substitute(
+                    SumVariant.struct(TypeSubstitution.substituteStructType(
                         struct.type(),
                         typeSubstitutionMap
                     ));
diff --git a/src/main/java/org/shedlang/compiler/types/TypeSubstitution.java b/src/main/java/org/shedlang/compiler/types/TypeSubstitution.java
index 5d2b7077..e03ff225 100644
--- a/src/main/java/org/shedlang/compiler/types/TypeSubstitution.java
+++ b/src/main/java/org/shedlang/compiler/types/TypeSubstitution.java
@@ -5,16 +5,14 @@ public class TypeSubstitution {
 
     public static Type substitute(Type type, TypeSubstitutionMap map) {
         return switch (type) {
-            case BoolType boolType -> type;
+            case BoolType boolType ->
+                type;
 
-            case BottomType bottomType -> type;
+            case BottomType bottomType ->
+                type;
 
-            case ConstructedType<?> constructedType -> new ConstructedType<>(
-                constructedType.constructor(),
-                constructedType.args().stream()
-                    .map(arg -> substitute(arg, map))
-                    .toList()
-            );
+            case ConstructedType<?> constructedType ->
+                substituteConstructedType(constructedType, map);
 
             case FunctionType functionType -> new FunctionType(
                 functionType.positionalParams().stream()
@@ -27,34 +25,60 @@ public class TypeSubstitution {
                 substitute(functionType.returnType(), map)
             );
 
-            case IntType intType -> type;
+            case IntType intType ->
+                type;
 
-            case NamespaceType namespaceType -> type;
+            case NamespaceType namespaceType ->
+                type;
 
-            case StringType stringType -> type;
+            case StringType stringType ->
+                type;
 
             case StructType structType ->
-                type;
+                substituteStructType(structType, map);
 
             case SumType sumType -> {
                 throw new UnsupportedOperationException("TODO");
             }
 
-            case TopType topType -> type;
+            case TopType topType ->
+                type;
 
-            case TupleType tupleType -> new TupleType(
-                tupleType.elements().stream()
-                    .map(element -> substitute(element, map))
-                    .toList()
-            );
+            case TupleType tupleType ->
+                new TupleType(
+                    tupleType.elements().stream()
+                        .map(element -> substitute(element, map))
+                        .toList()
+                );
 
             case TypeLevelValueType typeLevelValueType -> {
                 throw new UnsupportedOperationException("TODO");
             }
 
-            case TypeParam typeParam -> map.substitute(typeParam);
+            case TypeParam typeParam ->
+                map.substitute(typeParam);
 
-            case UnitType unitType -> type;
+            case UnitType unitType ->
+                type;
         };
     }
+
+    public static <T extends StructuredType> ConstructedType<T> substituteConstructedType(
+        ConstructedType<T> constructedType,
+        TypeSubstitutionMap map
+    ) {
+        return new ConstructedType<>(
+            constructedType.constructor(),
+            constructedType.args().stream()
+                .map(arg -> substitute(arg, map))
+                .toList()
+        );
+    }
+
+    public static StructType substituteStructType(
+        StructType type,
+        TypeSubstitutionMap map
+    ) {
+        return type;
+    }
 }