From 298f03e167d4cf3ef5b5088f355c8d8e4ee21467 Mon Sep 17 00:00:00 2001
From: Michael Williamson <mike@zwobble.org>
Date: Fri, 12 Jan 2024 17:51:59 +0000
Subject: Handle type level params when substituting in function type

---
 .../shedlang/compiler/types/TypeSubstitution.java  |  4 +-
 .../compiler/types/TypeSubstitutionTests.java      | 44 ++++++++++++++++++++++
 2 files changed, 45 insertions(+), 3 deletions(-)

(limited to 'src')

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
@@ -63,6 +63,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(
@@ -99,6 +123,26 @@ public class TypeSubstitutionTests {
             .build()));
     }
 
+    @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(
-- 
cgit v1.2.3