Handle type level params when substituting in function type

2 files changed, 45 insertions, 3 deletions

diff --git a/src/main/java/org/shedlang/compiler/types/TypeSubstitution.java b/src/main/java/org/shedlang/compiler/types/TypeSubstitution.java
index cc7ddbc9..4b29895c 100644
--- a/src/main/java/org/shedlang/compiler/types/TypeSubstitution.java
+++ b/src/main/java/org/shedlang/compiler/types/TypeSubstitution.java
@@ -1,7 +1,6 @@
 package org.shedlang.compiler.types;
 
 import java.util.List;
-import java.util.Optional;
 import java.util.function.BiFunction;
 
 public class TypeSubstitution {
@@ -107,8 +106,7 @@ public class TypeSubstitution {
         TypeSubstitutionMap map
     ) {
         return new FunctionType(
-            // TODO: don't clear type level params
-            Optional.empty(),
+            functionType.typeLevelParams(),
             functionType.positionalParams().stream()
                 .map(param -> substituteType(param, map))
                 .toList(),
diff --git a/src/test/java/org/shedlang/compiler/types/TypeSubstitutionTests.java b/src/test/java/org/shedlang/compiler/types/TypeSubstitutionTests.java
index 064012f9..6eb7105e 100644
--- a/src/test/java/org/shedlang/compiler/types/TypeSubstitutionTests.java
+++ b/src/test/java/org/shedlang/compiler/types/TypeSubstitutionTests.java
@@ -64,6 +64,30 @@ public class TypeSubstitutionTests {
     }
 
     @Test
+    public void substitutingFunctionPreservesTypeLevelParams() {
+        var typeSubstitutionMap = TypeSubstitutionMap.of(
+            Map.entry(TypeParam.forStructType(STRUCT_TYPE, "T"), Types.INT)
+        );
+
+        var result = TypeSubstitution.substituteType(
+            FunctionType.builder()
+                .withTypeLevelParams(List.of(
+                    TypeParam.forFunction(NAMESPACE_NAME, "f", "T")
+                ))
+                .build(),
+            typeSubstitutionMap
+        );
+
+        assertThat(result, equalTo(
+            FunctionType.builder()
+                .withTypeLevelParams(List.of(
+                    TypeParam.forFunction(NAMESPACE_NAME, "f", "T")
+                ))
+                .build())
+        );
+    }
+
+    @Test
     public void substitutingFunctionSubstitutesPositionalArgs() {
         var typeSubstitutionMap = TypeSubstitutionMap.of(
             Map.entry(TypeParam.forStructType(STRUCT_TYPE, "T"), Types.INT)
@@ -100,6 +124,26 @@ public class TypeSubstitutionTests {
     }
 
     @Test
+    public void substitutingFunctionPreservesEffect() {
+        var typeSubstitutionMap = TypeSubstitutionMap.of(
+            Map.entry(TypeParam.forStructType(STRUCT_TYPE, "T"), Types.INT)
+        );
+
+        var result = TypeSubstitution.substituteType(
+            FunctionType.builder()
+                .withEffect(Effects.IO)
+                .build(),
+            typeSubstitutionMap
+        );
+
+        assertThat(result, equalTo(
+            FunctionType.builder()
+                .withEffect(Effects.IO)
+                .build())
+        );
+    }
+
+    @Test
     public void substitutingFunctionSubstitutesReturnType() {
         var typeSubstitutionMap = TypeSubstitutionMap.of(
             Map.entry(TypeParam.forStructType(STRUCT_TYPE, "T"), Types.INT)