Handle type level params when substituting in function type

1 files changed, 44 insertions, 0 deletions

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)