Add simple effect subsitution

2 files changed, 27 insertions, 2 deletions

diff --git a/src/main/java/org/shedlang/compiler/types/TypeSubstitution.java b/src/main/java/org/shedlang/compiler/types/TypeSubstitution.java
index de055e1d..fc0a355b 100644
--- a/src/main/java/org/shedlang/compiler/types/TypeSubstitution.java
+++ b/src/main/java/org/shedlang/compiler/types/TypeSubstitution.java
@@ -6,6 +6,25 @@ import java.util.function.BiFunction;
 public class TypeSubstitution {
     private TypeSubstitution() {}
 
+    public static Effect substituteEffect(Effect effect, TypeSubstitutionMap map) {
+        return switch (effect) {
+            case EffectParam effectParam ->
+                map.substituteEffect(effectParam);
+
+            case EffectUnion effectUnion ->
+                effect;
+
+            case IoEffect ioEffect ->
+                effect;
+
+            case NoneEffect noneEffect ->
+                effect;
+
+            case UserDefinedEffect userDefinedEffect ->
+                effect;
+        };
+    }
+
     public static Type substituteType(Type type, TypeSubstitutionMap map) {
         return switch (type) {
             case BoolType boolType ->
@@ -20,7 +39,8 @@ public class TypeSubstitution {
             case ConstructedSumType constructedType ->
                 substituteConstructedSumType(constructedType, map);
 
-            case FunctionType functionType -> substituteFunctionType(functionType, map);
+            case FunctionType functionType ->
+                substituteFunctionType(functionType, map);
 
             case IntType intType ->
                 type;
@@ -114,7 +134,7 @@ public class TypeSubstitution {
             functionType.namedParams().stream()
                 .map(param -> new NamedParam(param.name(), substituteType(param.type(), map)))
                 .toList(),
-            functionType.effect(),
+            substituteEffect(functionType.effect(), map),
             substituteType(functionType.returnType(), map)
         );
     }
diff --git a/src/main/java/org/shedlang/compiler/types/TypeSubstitutionMap.java b/src/main/java/org/shedlang/compiler/types/TypeSubstitutionMap.java
index 02c602af..731a7b85 100644
--- a/src/main/java/org/shedlang/compiler/types/TypeSubstitutionMap.java
+++ b/src/main/java/org/shedlang/compiler/types/TypeSubstitutionMap.java
@@ -37,6 +37,11 @@ public class TypeSubstitutionMap {
         return (Type) this.map.getOrDefault(typeParam, typeParam);
     }
 
+    public Effect substituteEffect(EffectParam effectParam) {
+        // TODO: Validate map? Check here?
+        return (Effect) this.map.getOrDefault(effectParam, effectParam);
+    }
+
     @Override
     public boolean equals(Object o) {
         if (this == o) return true;