wasm gc: cast receivers on devirtualized calls to actual argument type, avoid this cast when possible

core/src/main/java/org/teavm/backend/wasm/gc/PreciseTypeInference.java

@@ -15,23 +15,23 @@
  */
 package org.teavm.backend.wasm.gc;
 
-import java.util.ArrayList;
-import java.util.Collections;
-import java.util.List;
 import org.teavm.model.ClassHierarchy;
-import org.teavm.model.ClassReader;
 import org.teavm.model.MethodReference;
 import org.teavm.model.Program;
 import org.teavm.model.ValueType;
 import org.teavm.model.analysis.BaseTypeInference;
+import org.teavm.model.instructions.InvocationType;
 
 public class PreciseTypeInference extends BaseTypeInference<ValueType> {
     public static final ValueType OBJECT_TYPE = ValueType.object("java.lang.Object");
     private ClassHierarchy hierarchy;
+    private WasmGCMethodReturnTypes returnTypes;
 
-    public PreciseTypeInference(Program program, MethodReference reference, ClassHierarchy hierarchy) {
+    public PreciseTypeInference(Program program, MethodReference reference, ClassHierarchy hierarchy,
+            WasmGCMethodReturnTypes returnTypes) {
         super(program, reference);
         this.hierarchy = hierarchy;
+        this.returnTypes = returnTypes;
     }
 
     @Override
@@ -87,7 +87,7 @@ public class PreciseTypeInference extends BaseTypeInference<ValueType> {
                     } else if (hierarchy.isSuperType(q, p, false)) {
                         return b;
                     }
-                    return ValueType.object(findCommonSuperclass(first, second));
+                    return ValueType.object(WasmGCUtil.findCommonSuperclass(hierarchy, first, second));
                 } else {
                     return OBJECT_TYPE;
                 }
@@ -96,34 +96,16 @@ public class PreciseTypeInference extends BaseTypeInference<ValueType> {
             }
         }
     }
-
-    private String findCommonSuperclass(ClassReader a, ClassReader b) {
-        var firstPath = findPathToRoot(a);
-        Collections.reverse(firstPath);
-        var secondPath = findPathToRoot(b);
-        Collections.reverse(secondPath);
-        if (firstPath.get(0) != secondPath.get(0)) {
-            return "java.lang.Object";
-        }
-        var max = Math.max(firstPath.size(), secondPath.size());
-        var index = 1;
-        while (index < max && firstPath.get(index) == secondPath.get(index)) {
-            ++index;
-        }
-        return firstPath.get(index).getName();
-    }
-
-    private List<ClassReader> findPathToRoot(ClassReader cls) {
-        var path = new ArrayList<ClassReader>();
-        while (cls != null) {
-            path.add(cls);
-            cls = cls.getParent() != null ? hierarchy.getClassSource().get(cls.getParent()) : null;
-        }
-        return path;
-    }
-
     @Override
     protected ValueType elementType(ValueType valueType) {
         return ((ValueType.Array) valueType).getItemType();
     }
+
+    @Override
+    protected ValueType methodReturnType(InvocationType invocationType, MethodReference methodRef) {
+        if (invocationType == InvocationType.SPECIAL) {
+            return returnTypes.returnTypeOf(methodRef);
+        }
+        return super.methodReturnType(invocationType, methodRef);
+    }
 }

core/src/main/java/org/teavm/backend/wasm/gc/WasmGCMethodReturnTypes.java

@@ -0,0 +1,64 @@
+/*
+ *  Copyright 2024 Alexey Andreev.
+ *
+ *  Licensed under the Apache License, Version 2.0 (the "License");
+ *  you may not use this file except in compliance with the License.
+ *  You may obtain a copy of the License at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *  Unless required by applicable law or agreed to in writing, software
+ *  distributed under the License is distributed on an "AS IS" BASIS,
+ *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *  See the License for the specific language governing permissions and
+ *  limitations under the License.
+ */
+package org.teavm.backend.wasm.gc;
+
+import java.util.HashMap;
+import java.util.Map;
+import org.teavm.dependency.DependencyInfo;
+import org.teavm.model.ClassHierarchy;
+import org.teavm.model.MethodReference;
+import org.teavm.model.ValueType;
+
+public class WasmGCMethodReturnTypes {
+    private DependencyInfo dependencyInfo;
+    private ClassHierarchy hierarchy;
+    private Map<MethodReference, ValueType> cache = new HashMap<>();
+
+    public WasmGCMethodReturnTypes(DependencyInfo dependencyInfo, ClassHierarchy hierarchy) {
+        this.dependencyInfo = dependencyInfo;
+        this.hierarchy = hierarchy;
+    }
+
+    public ValueType returnTypeOf(MethodReference reference) {
+        return cache.computeIfAbsent(reference, this::calculateReturnType);
+    }
+
+    private ValueType calculateReturnType(MethodReference reference) {
+        if (!(reference.getReturnType() instanceof ValueType.Object)) {
+            return reference.getReturnType();
+        }
+        var method = dependencyInfo.getMethod(reference);
+        if (method == null) {
+            return reference.getReturnType();
+        }
+        var types = method.getResult().getTypes();
+        if (types.length == 0) {
+            return reference.getReturnType();
+        }
+        var type = hierarchy.getClassSource().get(types[0]);
+        if (type == null) {
+            return reference.getReturnType();
+        }
+        for (var i = 1; i < types.length; ++i) {
+            var otherType = hierarchy.getClassSource().get(types[i]);
+            if (otherType == null) {
+                return reference.getReturnType();
+            }
+            type = hierarchy.getClassSource().get(WasmGCUtil.findCommonSuperclass(hierarchy, type, otherType));
+        }
+        return ValueType.object(type.getName());
+    }
+}

core/src/main/java/org/teavm/backend/wasm/gc/WasmGCUtil.java

@@ -0,0 +1,54 @@
+/*
+ *  Copyright 2024 Alexey Andreev.
+ *
+ *  Licensed under the Apache License, Version 2.0 (the "License");
+ *  you may not use this file except in compliance with the License.
+ *  You may obtain a copy of the License at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *  Unless required by applicable law or agreed to in writing, software
+ *  distributed under the License is distributed on an "AS IS" BASIS,
+ *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *  See the License for the specific language governing permissions and
+ *  limitations under the License.
+ */
+package org.teavm.backend.wasm.gc;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import org.teavm.model.ClassHierarchy;
+import org.teavm.model.ClassReader;
+
+public final class WasmGCUtil {
+    private WasmGCUtil() {
+    }
+
+
+    public static String findCommonSuperclass(ClassHierarchy hierarchy, ClassReader a, ClassReader b) {
+        var firstPath = findPathToRoot(hierarchy, a);
+        Collections.reverse(firstPath);
+        var secondPath = findPathToRoot(hierarchy, b);
+        Collections.reverse(secondPath);
+        if (firstPath.get(0) != secondPath.get(0)) {
+            return "java.lang.Object";
+        }
+        var max = Math.max(firstPath.size(), secondPath.size());
+        var index = 1;
+        while (index < max && firstPath.get(index) == secondPath.get(index)) {
+            ++index;
+        }
+        return firstPath.get(index).getName();
+    }
+
+    private static List<ClassReader> findPathToRoot(ClassHierarchy hierarchy, ClassReader cls) {
+        var path = new ArrayList<ClassReader>();
+        while (cls != null) {
+            path.add(cls);
+            cls = cls.getParent() != null ? hierarchy.getClassSource().get(cls.getParent()) : null;
+        }
+        return path;
+    }
+
+}

core/src/main/java/org/teavm/backend/wasm/gc/WasmGCVariableCategoryProvider.java

@@ -23,9 +23,11 @@ import org.teavm.model.util.VariableCategoryProvider;
 public class WasmGCVariableCategoryProvider implements VariableCategoryProvider {
     private ClassHierarchy hierarchy;
     private PreciseTypeInference inference;
+    private WasmGCMethodReturnTypes returnTypes;
 
-    public WasmGCVariableCategoryProvider(ClassHierarchy hierarchy) {
+    public WasmGCVariableCategoryProvider(ClassHierarchy hierarchy, WasmGCMethodReturnTypes returnTypes) {
         this.hierarchy = hierarchy;
+        this.returnTypes = returnTypes;
     }
 
     public PreciseTypeInference getTypeInference() {
@@ -34,7 +36,7 @@ public class WasmGCVariableCategoryProvider implements VariableCategoryProvider
 
     @Override
     public Object[] getCategories(Program program, MethodReference method) {
-        inference = new PreciseTypeInference(program, method, hierarchy);
+        inference = new PreciseTypeInference(program, method, hierarchy, returnTypes);
         var result = new Object[program.variableCount()];
         for (int i = 0; i < program.variableCount(); ++i) {
             var type = inference.typeOf(program.variableAt(i));

core/src/main/java/org/teavm/backend/wasm/generate/common/methods/BaseWasmGenerationVisitor.java

@@ -846,6 +846,10 @@ public abstract class BaseWasmGenerationVisitor implements StatementVisitor, Exp
                 accept(argument);
                 call.getArguments().add(result);
             }
+            if (expr.getType() == InvocationType.SPECIAL) {
+                var firstArg = call.getArguments().get(0);
+                call.getArguments().set(0, mapFirstArgumentForCall(firstArg, function, expr.getMethod()));
+            }
             if (callSiteId != null) {
                 callSiteId.addToLastArg(call.getArguments());
             }
@@ -903,13 +907,17 @@ public abstract class BaseWasmGenerationVisitor implements StatementVisitor, Exp
         }
     }
 
+    protected WasmExpression mapFirstArgumentForCall(WasmExpression argument, WasmFunction function,
+            MethodReference method) {
+        return argument;
+    }
+
     protected abstract WasmExpression generateVirtualCall(
             WasmLocal instance,
             MethodReference method,
             List<WasmExpression> arguments
     );
 
-
     private boolean needsCallSiteId() {
         return isManaged();
     }

core/src/main/java/org/teavm/backend/wasm/generate/gc/WasmGCDeclarationsGenerator.java

@@ -19,6 +19,7 @@ import java.util.List;
 import java.util.function.Predicate;
 import org.teavm.backend.wasm.BaseWasmFunctionRepository;
 import org.teavm.backend.wasm.WasmFunctionTypes;
+import org.teavm.backend.wasm.gc.WasmGCMethodReturnTypes;
 import org.teavm.backend.wasm.generate.WasmNameProvider;
 import org.teavm.backend.wasm.generate.gc.classes.WasmGCClassGenerator;
 import org.teavm.backend.wasm.generate.gc.classes.WasmGCClassInfoProvider;
@@ -45,6 +46,7 @@ public class WasmGCDeclarationsGenerator {
     public final WasmFunctionTypes functionTypes;
     private final WasmGCClassGenerator classGenerator;
     private final WasmGCMethodGenerator methodGenerator;
+    private final WasmGCMethodReturnTypes returnTypes;
 
     public WasmGCDeclarationsGenerator(
             WasmModule module,
@@ -60,6 +62,7 @@ public class WasmGCDeclarationsGenerator {
         var virtualTables = createVirtualTableProvider(classes, virtualMethods);
         functionTypes = new WasmFunctionTypes(module);
         var names = new WasmNameProvider();
+        returnTypes = new WasmGCMethodReturnTypes(dependencyInfo, hierarchy);
         methodGenerator = new WasmGCMethodGenerator(
                 module,
                 hierarchy,
@@ -69,6 +72,7 @@ public class WasmGCDeclarationsGenerator {
                 functionTypes,
                 names,
                 diagnostics,
+                returnTypes,
                 customGenerators
         );
         var tags = new TagRegistry(classes, hierarchy);

core/src/main/java/org/teavm/backend/wasm/generate/gc/classes/WasmGCStandardClasses.java

@@ -34,7 +34,7 @@ public class WasmGCStandardClasses {
 
     public WasmGCClassInfo stringClass() {
         if (stringClassInfo == null) {
-            stringClassInfo = classGenerator.getClassInfo("java.lang.Class");
+            stringClassInfo = classGenerator.getClassInfo("java.lang.String");
         }
         return stringClassInfo;
     }

core/src/main/java/org/teavm/backend/wasm/generate/gc/methods/WasmGCGenerationVisitor.java

@@ -357,6 +357,31 @@ public class WasmGCGenerationVisitor extends BaseWasmGenerationVisitor {
         result = invocation(expr, null, false);
     }
 
+    @Override
+    protected WasmExpression mapFirstArgumentForCall(WasmExpression argument, WasmFunction function,
+            MethodReference method) {
+        argument.acceptVisitor(typeInference);
+        var actualType = typeInference.getResult();
+        var expectedType = function.getType().getParameterTypes().get(0);
+        if (actualType == expectedType || !(actualType instanceof WasmType.CompositeReference)
+                || !(expectedType instanceof WasmType.CompositeReference)) {
+            return argument;
+        }
+        var actualComposite = ((WasmType.CompositeReference) actualType).composite;
+        var expectedComposite = ((WasmType.CompositeReference) expectedType).composite;
+        if (!(actualComposite instanceof WasmStructure) || !(expectedComposite instanceof WasmStructure)) {
+            return argument;
+        }
+
+        var actualStruct = (WasmStructure) actualComposite;
+        var expectedStruct = (WasmStructure) expectedComposite;
+        if (!actualStruct.isSupertypeOf(expectedStruct)) {
+            return argument;
+        }
+
+        return new WasmCast(argument, expectedComposite.getReference());
+    }
+
     @Override
     public void visit(QualificationExpr expr) {
         if (expr.getQualified() == null) {

core/src/main/java/org/teavm/backend/wasm/generate/gc/methods/WasmGCMethodGenerator.java

@@ -25,6 +25,7 @@ import org.teavm.ast.decompilation.Decompiler;
 import org.teavm.backend.lowlevel.generate.NameProvider;
 import org.teavm.backend.wasm.BaseWasmFunctionRepository;
 import org.teavm.backend.wasm.WasmFunctionTypes;
+import org.teavm.backend.wasm.gc.WasmGCMethodReturnTypes;
 import org.teavm.backend.wasm.gc.WasmGCVariableCategoryProvider;
 import org.teavm.backend.wasm.generate.gc.classes.WasmGCClassInfoProvider;
 import org.teavm.backend.wasm.generate.gc.classes.WasmGCStandardClasses;
@@ -76,6 +77,7 @@ public class WasmGCMethodGenerator implements BaseWasmFunctionRepository {
     private WasmGCClassInfoProvider classInfoProvider;
     private WasmGCStandardClasses standardClasses;
     private WasmGCStringProvider strings;
+    private WasmGCMethodReturnTypes returnTypes;
 
     public WasmGCMethodGenerator(
             WasmModule module,
@@ -86,6 +88,7 @@ public class WasmGCMethodGenerator implements BaseWasmFunctionRepository {
             WasmFunctionTypes functionTypes,
             NameProvider names,
             Diagnostics diagnostics,
+            WasmGCMethodReturnTypes returnTypes,
             WasmGCCustomGeneratorProvider customGenerators
     ) {
         this.module = module;
@@ -96,6 +99,7 @@ public class WasmGCMethodGenerator implements BaseWasmFunctionRepository {
         this.functionTypes = functionTypes;
         this.names = names;
         this.diagnostics = diagnostics;
+        this.returnTypes = returnTypes;
         this.customGenerators = customGenerators;
     }
 
@@ -139,7 +143,7 @@ public class WasmGCMethodGenerator implements BaseWasmFunctionRepository {
     }
 
     private WasmFunction createStaticFunction(MethodReference methodReference) {
-        var returnType = typeMapper.mapType(methodReference.getReturnType());
+        var returnType = typeMapper.mapType(returnTypes.returnTypeOf(methodReference));
         var parameterTypes = new WasmType[methodReference.parameterCount()];
         for (var i = 0; i < parameterTypes.length; ++i) {
             parameterTypes[i] = typeMapper.mapType(methodReference.parameterType(i));
@@ -166,7 +170,7 @@ public class WasmGCMethodGenerator implements BaseWasmFunctionRepository {
     }
 
     private WasmFunction createInstanceFunction(MethodReference methodReference) {
-        var returnType = typeMapper.mapType(methodReference.getReturnType());
+        var returnType = typeMapper.mapType(returnTypes.returnTypeOf(methodReference));
         var parameterTypes = new WasmType[methodReference.parameterCount() + 1];
         parameterTypes[0] = typeMapper.mapType(ValueType.object(methodReference.getClassName()));
         for (var i = 0; i < methodReference.parameterCount(); ++i) {
@@ -206,7 +210,7 @@ public class WasmGCMethodGenerator implements BaseWasmFunctionRepository {
 
     private void generateRegularMethodBody(MethodHolder method, WasmFunction function) {
         var decompiler = getDecompiler();
-        var categoryProvider = new WasmGCVariableCategoryProvider(hierarchy);
+        var categoryProvider = new WasmGCVariableCategoryProvider(hierarchy, returnTypes);
         var allocator = new RegisterAllocator(categoryProvider);
         allocator.allocateRegisters(method.getReference(), method.getProgram(), friendlyToDebugger);
         var ast = decompiler.decompileRegular(method);

core/src/main/java/org/teavm/backend/wasm/model/WasmStructure.java

@@ -38,6 +38,16 @@ public class WasmStructure extends WasmCompositeType {
         this.supertype = supertype;
     }
 
+    public boolean isSupertypeOf(WasmStructure subtype) {
+        while (subtype != null) {
+            if (subtype == this) {
+                return true;
+            }
+            subtype = subtype.getSupertype();
+        }
+        return false;
+    }
+
     @Override
     public void acceptVisitor(WasmCompositeTypeVisitor visitor) {
         visitor.visit(this);

core/src/main/java/org/teavm/backend/wasm/render/WasmBinaryRenderingVisitor.java

@@ -1030,7 +1030,7 @@ class WasmBinaryRenderingVisitor implements WasmExpressionVisitor {
         expression.getValue().acceptVisitor(this);
         writer.writeByte(0xfb);
         writer.writeByte(23);
-        writer.writeType(expression.getTargetType(), module);
+        writer.writeHeapType(expression.getTargetType(), module);
         popLocation();
     }
 

core/src/main/java/org/teavm/model/analysis/BaseTypeInference.java

@@ -43,6 +43,7 @@ import org.teavm.model.instructions.FloatConstantInstruction;
 import org.teavm.model.instructions.GetElementInstruction;
 import org.teavm.model.instructions.GetFieldInstruction;
 import org.teavm.model.instructions.IntegerConstantInstruction;
+import org.teavm.model.instructions.InvocationType;
 import org.teavm.model.instructions.InvokeInstruction;
 import org.teavm.model.instructions.IsInstanceInstruction;
 import org.teavm.model.instructions.LongConstantInstruction;
@@ -169,7 +170,7 @@ public abstract class BaseTypeInference<T> {
 
     protected abstract T elementType(T t);
 
-    protected T methodReturnType(MethodReference methodRef) {
+    protected T methodReturnType(InvocationType invocationType, MethodReference methodRef) {
         return mapType(methodRef.getReturnType());
     }
 
@@ -293,7 +294,7 @@ public abstract class BaseTypeInference<T> {
 
         @Override
         public void visit(InvokeInstruction insn) {
-            type(insn.getReceiver(), methodReturnType(insn.getMethod()));
+            type(insn.getReceiver(), methodReturnType(insn.getType(), insn.getMethod()));
         }
 
         @Override

jso/impl/src/main/java/org/teavm/jso/impl/JSTypeInference.java

@@ -21,6 +21,7 @@ import org.teavm.model.MethodReference;
 import org.teavm.model.Program;
 import org.teavm.model.ValueType;
 import org.teavm.model.analysis.BaseTypeInference;
+import org.teavm.model.instructions.InvocationType;
 
 class JSTypeInference extends BaseTypeInference<JSType> {
     private JSTypeHelper typeHelper;
@@ -85,7 +86,7 @@ class JSTypeInference extends BaseTypeInference<JSType> {
     }
 
     @Override
-    protected JSType methodReturnType(MethodReference methodRef) {
+    protected JSType methodReturnType(InvocationType invocationType, MethodReference methodRef) {
         if (!methodRef.getReturnType().isObject(Object.class)) {
             return mapType(methodRef.getReturnType());
         }