From d94641005c2b35800ee4499d2c17ba06579dabe6 Mon Sep 17 00:00:00 2001
From: Michael Williamson <mike@zwobble.org>
Date: Mon, 25 Dec 2023 11:47:18 +0000
Subject: Add type substitution for function types

---
 .../compiler/types/TypeSubstitutionTests.java      | 54 ++++++++++++++++++++++
 1 file changed, 54 insertions(+)

(limited to 'src/test/java')

diff --git a/src/test/java/org/shedlang/compiler/types/TypeSubstitutionTests.java b/src/test/java/org/shedlang/compiler/types/TypeSubstitutionTests.java
index b7eaa0aa..cd676d71 100644
--- a/src/test/java/org/shedlang/compiler/types/TypeSubstitutionTests.java
+++ b/src/test/java/org/shedlang/compiler/types/TypeSubstitutionTests.java
@@ -35,6 +35,60 @@ public class TypeSubstitutionTests {
         assertThat(result, equalTo(Types.BOTTOM));
     }
 
+    @Test
+    public void substitutingFunctionSubstitutesPositionalArgs() {
+        var typeSubstitutionMap = TypeSubstitutionMap.of(
+            Map.entry(TypeParam.forStructType(STRUCT_TYPE, "T"), Types.INT)
+        );
+
+        var result = TypeSubstitution.substitute(
+            FunctionType.builder()
+                .addPositionalParam(TypeParam.forStructType(STRUCT_TYPE, "T"))
+                .build(),
+            typeSubstitutionMap
+        );
+
+        assertThat(result, equalTo(FunctionType.builder()
+            .addPositionalParam(Types.INT)
+            .build()));
+    }
+
+    @Test
+    public void substitutingFunctionSubstitutesNamedArgs() {
+        var typeSubstitutionMap = TypeSubstitutionMap.of(
+            Map.entry(TypeParam.forStructType(STRUCT_TYPE, "T"), Types.INT)
+        );
+
+        var result = TypeSubstitution.substitute(
+            FunctionType.builder()
+                .addNamedParam("x", TypeParam.forStructType(STRUCT_TYPE, "T"))
+                .build(),
+            typeSubstitutionMap
+        );
+
+        assertThat(result, equalTo(FunctionType.builder()
+            .addNamedParam("x", Types.INT)
+            .build()));
+    }
+
+    @Test
+    public void substitutingFunctionSubstitutesReturnType() {
+        var typeSubstitutionMap = TypeSubstitutionMap.of(
+            Map.entry(TypeParam.forStructType(STRUCT_TYPE, "T"), Types.INT)
+        );
+
+        var result = TypeSubstitution.substitute(
+            FunctionType.builder()
+                .withReturnType(TypeParam.forStructType(STRUCT_TYPE, "T"))
+                .build(),
+            typeSubstitutionMap
+        );
+
+        assertThat(result, equalTo(FunctionType.builder()
+            .withReturnType(Types.INT)
+            .build()));
+    }
+
     @Test
     public void substitutingIntDoesNothing() {
         var typeSubstitutionMap = TypeSubstitutionMap.of(
-- 
cgit v1.2.3