Make type substitution carry what is known about a type

CHANGELOG.md

@@ -28,6 +28,18 @@
     An instance of `I` for type `X` may only be declared in the package declaring `I` or the one declaring
     `X`, and only once, so `I` for `X` means the same thing throughout a program regardless of imports.
 
+- A type class bound is now usable from inside a closure, so a bounded generic can hand work to one:
+
+        interface Producer
+            function produce() returns int
+
+        function indexLater<T: Indexable>(T x) returns Producer
+            return () -> T.toIndex(x)
+
+    Substituting a type variable now carries the instance chosen for it along with the type, rather than the
+    type alone, so lifting a body into a class of its own no longer loses it. Jass only for now: Lua reaches
+    such a class through its interface and still reports the bound as unresolvable there.
+
 - Added new pseudo-natives for debugging memory leaks:
 
         // returns the maximum type id, can be usd to

de.peeeq.wurstscript/src/main/java/de/peeeq/wurstscript/intermediatelang/interpreter/ProgramState.java

@@ -461,59 +461,6 @@ public ImType case_ImTupleType(ImTupleType tt) {
     }
 
 
-    // Helper method to substitute type variables
-    private ImType substituteTypeVars(ImType type, Map<ImTypeVar, ImType> substitutions) {
-        return type.match(new ImType.Matcher<ImType>() {
-            @Override
-            public ImType case_ImTypeVarRef(ImTypeVarRef typeVarRef) {
-                ImType concrete = substitutions.get(typeVarRef.getTypeVariable());
-                return concrete != null ? concrete : typeVarRef;
-            }
-
-            @Override
-            public ImType case_ImClassType(ImClassType classType) {
-                // Recursively substitute in type arguments
-                ImTypeArguments newArgs = JassIm.ImTypeArguments();
-                for (ImTypeArgument arg : classType.getTypeArguments()) {
-                    ImType substituted = substituteTypeVars(arg.getType(), substitutions);
-                    newArgs.add(JassIm.ImTypeArgument(substituted, typeClassBindingFor(arg)));
-                }
-                return JassIm.ImClassType(classType.getClassDef(), newArgs);
-            }
-
-            // For other types, return as-is
-            @Override
-            public ImType case_ImSimpleType(ImSimpleType t) {
-                return t;
-            }
-
-            @Override
-            public ImType case_ImArrayType(ImArrayType t) {
-                return t;
-            }
-
-            @Override
-            public ImType case_ImTupleType(ImTupleType t) {
-                return t;
-            }
-
-            @Override
-            public ImType case_ImVoid(ImVoid t) {
-                return t;
-            }
-
-            @Override
-            public ImType case_ImAnyType(ImAnyType t) {
-                return t;
-            }
-
-            @Override
-            public ImType case_ImArrayTypeMulti(ImArrayTypeMulti t) {
-                return t;
-            }
-        });
-    }
-
     public void pushStackframe(ImCompiletimeExpr f, WPos trace) {
         WLogger.trace(() -> "pushStackframe compiletime expr " + f);
         stackFrames.push(new ILStackFrame(f, trace));

de.peeeq.wurstscript/src/main/java/de/peeeq/wurstscript/translation/imtojass/ImAttrType.java

@@ -35,20 +35,7 @@ public static ImType getType(ImFunctionCall e) {
     }
 
     public static ImType substituteType(ImType type, List<ImTypeArgument> generics, List<ImTypeVar> typeVars) {
-        return type.match(new TypeRewriteMatcher() {
-
-            @Override
-            public ImType case_ImTypeVarRef(ImTypeVarRef t) {
-                int index = typeVars.indexOf(t.getTypeVariable());
-                if (index < 0) {
-                    return t;
-                } else if (index >= generics.size()) {
-                    throw new RuntimeException("Could not find replacement for " + t + " when replacing " + typeVars + " with " + generics);
-                }
-                return generics.get(index).getType();
-            }
-
-        });
+        return TypeSubst.of(typeVars, generics).apply(type);
     }
 
 

de.peeeq.wurstscript/src/main/java/de/peeeq/wurstscript/translation/imtojass/TypeRewriteMatcher.java

@@ -57,9 +57,17 @@ public ImType case_ImArrayType(ImArrayType t) {
     public ImType case_ImClassType(ImClassType t) {
         ImTypeArguments args = JassIm.ImTypeArguments();
         for (ImTypeArgument ta : t.getTypeArguments()) {
-            ImTypeArgument imTypeArgument = JassIm.ImTypeArgument(ta.getType().match(this), ta.getTypeClassBinding());
-            args.add(imTypeArgument);
+            args.add(rewriteTypeArgument(ta));
         }
         return JassIm.ImClassType(t.getClassDef(), args);
     }
+
+    /**
+     * Rewrites one type argument. Only the type is rewritten by default, since a plain rewrite has
+     * nothing to say about what is known of the argument. Subclasses which do -- substitution, where
+     * the replacing argument carries its own binding -- override this.
+     */
+    protected ImTypeArgument rewriteTypeArgument(ImTypeArgument ta) {
+        return JassIm.ImTypeArgument(ta.getType().match(this), ta.getTypeClassBinding());
+    }
 }

de.peeeq.wurstscript/src/main/java/de/peeeq/wurstscript/translation/imtojass/TypeRewriter.java

@@ -36,6 +36,16 @@ private static void rewrite(Element e, Function<ImType, ImType> rewriteFunc) {
             case ImTypeIdOfObj t -> t.setClazz((ImClassType) rewriteFunc.apply(t.getClazz()));
             case ImTypeIdOfClass t -> t.setClazz((ImClassType) rewriteFunc.apply(t.getClazz()));
             case ImMethod m -> m.setMethodClass((ImClassType) rewriteFunc.apply(m.getMethodClass()));
+            case ImTypeVarDispatch d -> {
+                // The dispatched type variable is a reference to a variable like any other, but it
+                // is held directly instead of as a type, so a walk over types alone passes it by.
+                // A closure capturing the type parameter it dispatches on has to rewrite it too,
+                // otherwise the lifted body still names a variable of the function it left.
+                ImType rewritten = rewriteFunc.apply(JassIm.ImTypeVarRef(d.getTypeVariable()));
+                if (rewritten instanceof ImTypeVarRef ref) {
+                    d.setTypeVariable(ref.getTypeVariable());
+                }
+            }
             case ImClass c -> {
                 List<ImClassType> newSuperClasses = new ArrayList<>();
                 for (ImClassType tt : c.getSuperClasses()) {

de.peeeq.wurstscript/src/main/java/de/peeeq/wurstscript/translation/imtojass/TypeSubst.java

@@ -0,0 +1,169 @@
+package de.peeeq.wurstscript.translation.imtojass;
+
+import de.peeeq.wurstscript.jassIm.ImFunction;
+import de.peeeq.wurstscript.jassIm.ImMethod;
+import de.peeeq.wurstscript.jassIm.ImType;
+import de.peeeq.wurstscript.jassIm.ImTypeArgument;
+import de.peeeq.wurstscript.jassIm.ImTypeClassFunc;
+import de.peeeq.wurstscript.jassIm.ImTypeVar;
+import de.peeeq.wurstscript.jassIm.ImTypeVarRef;
+import de.peeeq.wurstscript.jassIm.JassIm;
+import io.vavr.control.Either;
+import org.eclipse.jdt.annotation.Nullable;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
+
+/**
+ * Replaces type variables by the type arguments bound to them.
+ * <p>
+ * Substitution used to be expressed as two parallel lists -- the type variables of a function or
+ * class next to the type arguments given at a use of it -- paired up positionally at each point of
+ * use. That form lost information: a {@link ImTypeArgument} is a type <i>plus</i> what is known
+ * about it, but pairing the lists unwrapped the argument and returned the bare {@link ImType}, so
+ * the type class binding never survived. Callers which needed it had to re-attach it by hand, and
+ * the ones which did not know to do so silently produced an argument with no binding.
+ * <p>
+ * Keeping the argument whole fixes that: substituting into an argument position yields the argument
+ * that was bound, so its binding travels with the type it belongs to.
+ * <p>
+ * Lookup is by identity. IM nodes do not override {@code equals}, and a type variable stands for one
+ * particular declaration, so two variables which merely share a name are different variables.
+ */
+public final class TypeSubst {
+
+    private static final TypeSubst EMPTY =
+            new TypeSubst(Collections.emptyMap(), Collections.emptyList(), Collections.emptyList(), Collections.emptyList());
+
+    private final Map<ImTypeVar, ImTypeArgument> bindings;
+    /** Variables this substitution covers but for which no argument was supplied. */
+    private final List<ImTypeVar> unbound;
+    /** Kept so that a missing argument is reported the same way it was before. */
+    private final List<ImTypeVar> typeVars;
+    private final List<ImTypeArgument> typeArguments;
+
+    private TypeSubst(Map<ImTypeVar, ImTypeArgument> bindings, List<ImTypeVar> unbound,
+                      List<ImTypeVar> typeVars, List<ImTypeArgument> typeArguments) {
+        this.bindings = bindings;
+        this.unbound = unbound;
+        this.typeVars = typeVars;
+        this.typeArguments = typeArguments;
+    }
+
+    public static TypeSubst empty() {
+        return EMPTY;
+    }
+
+    /**
+     * Binds {@code typeVars} to {@code typeArguments} by position.
+     * <p>
+     * Fewer arguments than variables is allowed: the surplus variables stay unbound, and only a type
+     * which actually mentions one of them fails. Callers rely on that, because a use may legitimately
+     * leave the arguments off when it does not name the variable.
+     */
+    public static TypeSubst of(List<ImTypeVar> typeVars, List<ImTypeArgument> typeArguments) {
+        if (typeVars.isEmpty()) {
+            return EMPTY;
+        }
+        // LinkedHashMap rather than IdentityHashMap: the keys already compare by identity, and this
+        // keeps iteration order stable, which matters wherever substitution feeds emitted names.
+        Map<ImTypeVar, ImTypeArgument> bindings = new LinkedHashMap<>();
+        List<ImTypeVar> unbound = new ArrayList<>();
+        for (int i = 0; i < typeVars.size(); i++) {
+            ImTypeVar typeVar = typeVars.get(i);
+            if (i < typeArguments.size()) {
+                // A variable repeated in the list keeps its first argument, as positional lookup did.
+                bindings.putIfAbsent(typeVar, typeArguments.get(i));
+            } else if (!bindings.containsKey(typeVar)) {
+                unbound.add(typeVar);
+            }
+        }
+        return new TypeSubst(bindings, unbound, typeVars, typeArguments);
+    }
+
+    public boolean isEmpty() {
+        return bindings.isEmpty() && unbound.isEmpty();
+    }
+
+    /** Applies this substitution to a type. Any binding on the replacing argument is not part of a type. */
+    public ImType apply(ImType type) {
+        if (isEmpty()) {
+            // Nothing to replace, so hand back the type itself. Worth doing: this runs for the
+            // return type of every call, and rebuilding a class or tuple type only to get an equal
+            // one back was pure allocation.
+            return type;
+        }
+        return type.match(new TypeRewriteMatcher() {
+            @Override
+            public ImType case_ImTypeVarRef(ImTypeVarRef t) {
+                ImTypeArgument replacement = resolve(t);
+                return replacement == null ? t : replacement.getType();
+            }
+
+            @Override
+            protected ImTypeArgument rewriteTypeArgument(ImTypeArgument argument) {
+                return TypeSubst.this.apply(argument);
+            }
+        });
+    }
+
+    /**
+     * Applies this substitution to a type argument.
+     * <p>
+     * When the argument is exactly a type variable this substitution binds, the bound argument
+     * replaces it whole, so the binding recorded at the use site reaches the body being substituted
+     * into. A binding already present on the argument is more specific and is kept.
+     */
+    public ImTypeArgument apply(ImTypeArgument argument) {
+        if (isEmpty()) {
+            return argument;
+        }
+        if (argument.getType() instanceof ImTypeVarRef ref) {
+            ImTypeArgument replacement = resolve(ref);
+            if (replacement != null) {
+                return JassIm.ImTypeArgument(replacement.getType(), binding(argument, replacement));
+            }
+        }
+        return JassIm.ImTypeArgument(apply(argument.getType()), argument.getTypeClassBinding());
+    }
+
+    /**
+     * Shares the map rather than copying it. A binding is only ever replaced wholesale through
+     * {@code setTypeClassBinding}, never added to in place, so two arguments naming the same
+     * instances can name the same map.
+     */
+    private static Map<ImTypeClassFunc, Either<ImMethod, ImFunction>> binding(ImTypeArgument argument,
+                                                                              ImTypeArgument replacement) {
+        return argument.getTypeClassBinding().isEmpty()
+                ? replacement.getTypeClassBinding()
+                : argument.getTypeClassBinding();
+    }
+
+    private @Nullable ImTypeArgument resolve(ImTypeVarRef ref) {
+        ImTypeVar typeVar = ref.getTypeVariable();
+        ImTypeArgument bound = bindings.get(typeVar);
+        if (bound != null) {
+            return bound;
+        }
+        if (unbound.contains(typeVar)) {
+            throw new RuntimeException("Could not find replacement for " + ref
+                    + " when replacing " + typeVars + " with " + typeArguments);
+        }
+        return null;
+    }
+
+    @Override
+    public String toString() {
+        StringBuilder sb = new StringBuilder("[");
+        bindings.forEach((typeVar, argument) -> {
+            if (sb.length() > 1) {
+                sb.append(", ");
+            }
+            sb.append(typeVar.getName()).append(" -> ").append(argument.getType());
+        });
+        return sb.append("]").toString();
+    }
+}

de.peeeq.wurstscript/src/test/java/tests/wurstscript/tests/TypeClassTests.java

@@ -67,6 +67,60 @@ public void dispatchRuntime() {
         );
     }
 
+    /**
+     * A closure lifts its body into a class of its own, capturing the enclosing type variables. The
+     * requirement is dispatched from inside that class, so the binding has to survive being carried
+     * across into it.
+     */
+    @Test
+    public void dispatchInsideClosure() {
+        testAssertOkLines(true,
+            "package test",
+            "native testSuccess()",
+            "interface ToIndex<T:>",
+            "    function toIndex(T x) returns int",
+            "implements ToIndex<int>",
+            "    function toIndex(int x) returns int",
+            "        return x * 2",
+            "interface Producer",
+            "    function produce() returns int",
+            "function foo<Q: ToIndex>(Q x) returns int",
+            "    Producer p = () -> Q.toIndex(x)",
+            "    return p.produce()",
+            "init",
+            "    if foo(21) == 42",
+            "        testSuccess()"
+        );
+    }
+
+    /**
+     * The same closure is still rejected for Lua, and this pins that it is rejected clearly rather
+     * than mistranslated. Lua keeps generics erased and specialises only what it can reach through a
+     * concrete type; a closure is reached through its interface, so the specialised class exists but
+     * nothing calls it. Making that work is a change to Lua's erasure, not to substitution. Should
+     * it be made, this test fails and becomes the success case above.
+     */
+    @Test
+    public void dispatchInsideClosureIsRejectedForLua() {
+        test().testLua(true).executeProg().expectError("could not be resolved for the Lua target").lines(
+            "package test",
+            "native testSuccess()",
+            "interface ToIndex<T:>",
+            "    function toIndex(T x) returns int",
+            "implements ToIndex<int>",
+            "    function toIndex(int x) returns int",
+            "        return x * 2",
+            "interface Producer",
+            "    function produce() returns int",
+            "function foo<Q: ToIndex>(Q x) returns int",
+            "    Producer p = () -> Q.toIndex(x)",
+            "    return p.produce()",
+            "init",
+            "    if foo(21) == 42",
+            "        testSuccess()"
+        );
+    }
+
     /** Each type argument picks its own instance, so one generic serves several types. */
     @Test
     public void twoInstancesOfOneClass() {