Use different functions for table entries than for resolved calls.

The real method body goes to a method used for statically resolved
calls. That method can take the real `this` type as argument.

For table entries, which need to use `(ref any)` as receiver type,
we generate a separate function. That bridge function casts down
the receiver then performs a tail call to the real function.

With this strategy, we do not need to cast down `this` values
every time we use them. Since most calls as statically resolved,
this also means we don't need any cast at all. For dynamically
resolved calls, the overhead should be small thanks to the tail
call.

This change causes a 30% code size regression for fastLink, but
only 6% in fullLink. We should be able to optimize it further by
only emitting bridges for methods that actually appear in at
least one dispatch table.

Author: sjrd

Diff for: wasm/src/main/scala/ir2wasm/HelperFunctions.scala

@@ -549,7 +549,7 @@ object HelperFunctions {
 
         val objectClassInfo = ctx.getClassInfo(IRNames.ObjectClass)
         instrs ++= objectClassInfo.tableEntries.map { methodName =>
-          ctx.refFuncWithDeclaration(objectClassInfo.resolvedMethodInfos(methodName).wasmName)
+          ctx.refFuncWithDeclaration(objectClassInfo.resolvedMethodInfos(methodName).tableEntryName)
         }
         instrs += STRUCT_NEW(WasmTypeName.WasmStructTypeName.ObjectVTable)
         instrs += LOCAL_TEE(arrayTypeDataLocal)

Diff for: wasm/src/main/scala/ir2wasm/WasmBuilder.scala

@@ -337,7 +337,7 @@ class WasmBuilder(coreSpec: CoreSpec) {
         val proxyId = ctx.getReflectiveProxyId(proxyInfo.methodName)
         List(
           I32_CONST(proxyId),
-          REF_FUNC(proxyInfo.wasmName),
+          REF_FUNC(proxyInfo.tableEntryName),
           STRUCT_NEW(WasmStructTypeName.reflectiveProxy)
         )
       } :+ ARRAY_NEW_FIXED(WasmArrayTypeName.reflectiveProxies, reflectiveProxies.size)
@@ -417,7 +417,7 @@ class WasmBuilder(coreSpec: CoreSpec) {
         classInfo.resolvedMethodInfos.valuesIterator.filter(_.methodName.isReflectiveProxy).toList
       val typeDataFieldValues = genTypeDataFieldValues(clazz, reflectiveProxies)
       val vtableElems = classInfo.tableEntries.map { methodName =>
-        REF_FUNC(classInfo.resolvedMethodInfos(methodName).wasmName)
+        REF_FUNC(classInfo.resolvedMethodInfos(methodName).tableEntryName)
       }
       val globalVTable =
         genTypeDataGlobal(typeRef, vtableTypeName, typeDataFieldValues, vtableElems)
@@ -1068,29 +1068,26 @@ class WasmBuilder(coreSpec: CoreSpec) {
   private def genFunction(
       clazz: LinkedClass,
       method: IRTrees.MethodDef
-  )(implicit ctx: WasmContext): WasmFunction = {
-    val functionName = Names.WasmFunctionName(
-      method.flags.namespace,
-      clazz.name.name,
-      method.name.name
-    )
+  )(implicit ctx: WasmContext): Unit = {
+    val namespace = method.flags.namespace
+    val className = clazz.className
+    val methodName = method.methodName
+
+    val functionName = Names.WasmFunctionName(namespace, className, methodName)
+
+    val isHijackedClass = ctx.getClassInfo(className).kind == ClassKind.HijackedClass
 
-    // Receiver type for non-constructor methods needs to be `(ref any)` because params are invariant
-    // Otherwise, vtable can't be a subtype of the supertype's subtype
-    // Constructor can use the exact type because it won't be registered to vtables.
     val receiverTyp =
-      if (method.flags.namespace.isStatic)
+      if (namespace.isStatic)
         None
-      else if (clazz.kind == ClassKind.HijackedClass)
-        Some(transformType(IRTypes.BoxedClassToPrimType(clazz.name.name)))
-      else if (method.flags.namespace.isConstructor)
-        Some(WasmRefType.nullable(WasmTypeName.WasmStructTypeName.forClass(clazz.name.name)))
+      else if (isHijackedClass)
+        Some(transformType(IRTypes.BoxedClassToPrimType(className)))
       else
-        Some(WasmRefType.any)
+        Some(transformType(IRTypes.ClassType(className)))
 
     // Prepare for function context, set receiver and parameters
     implicit val fctx = WasmFunctionContext(
-      Some(clazz.className),
+      Some(className),
       functionName,
       receiverTyp,
       method.args,
@@ -1101,7 +1098,48 @@ class WasmBuilder(coreSpec: CoreSpec) {
     val body = method.body.getOrElse(throw new Exception("abstract method cannot be transformed"))
     WasmExpressionBuilder.generateIRBody(body, method.resultType)
 
-    fctx.buildAndAddToContext(useFunctionTypeInMainRecType = true)
+    fctx.buildAndAddToContext()
+
+    if (namespace == IRTrees.MemberNamespace.Public && !isHijackedClass) {
+      /* Also generate the bridge that is stored in the table entries. In table
+       * entries, the receiver type is always `(ref any)`.
+       *
+       * TODO: generate this only when the method is actually referred to from
+       * at least one table.
+       */
+
+      implicit val fctx = WasmFunctionContext(
+        Some(className),
+        WasmFunctionName.forTableEntry(className, methodName),
+        Some(WasmRefType.any),
+        method.args,
+        method.resultType
+      )
+
+      import fctx.instrs
+
+      val receiverLocal :: paramLocals = fctx.paramIndices: @unchecked
+
+      // Load and cast down the receiver
+      instrs += LOCAL_GET(receiverLocal)
+      receiverTyp match {
+        case Some(Types.WasmRefType(_, WasmHeapType.Any)) =>
+          () // no cast necessary
+        case Some(receiverTyp: Types.WasmRefType) =>
+          instrs += REF_CAST(receiverTyp)
+        case _ =>
+          throw new AssertionError(s"Unexpected receiver type $receiverTyp")
+      }
+
+      // Load the other parameters
+      for (paramLocal <- paramLocals)
+        instrs += LOCAL_GET(paramLocal)
+
+      // Call the statically resolved method
+      instrs += RETURN_CALL(functionName)
+
+      fctx.buildAndAddToContext(useFunctionTypeInMainRecType = true)
+    }
   }
 
   private def transformField(

Diff for: wasm/src/main/scala/ir2wasm/WasmExpressionBuilder.scala

@@ -1596,21 +1596,6 @@ private class WasmExpressionBuilder private (
       case _                      => t.tpe
     }
 
-    /* If the receiver is a Class/ModuleClass, its wasm type will be declared
-     * as `(ref any)`, and therefore we must cast it down.
-     */
-    fixedTpe match {
-      case IRTypes.ClassType(className) if className != IRNames.ObjectClass =>
-        val info = ctx.getClassInfo(className)
-        if (info.kind.isClass) {
-          instrs += REF_CAST(Types.WasmRefType(WasmStructTypeName.forClass(className)))
-        } else if (info.isInterface) {
-          instrs += REF_CAST(Types.WasmRefType(Types.WasmHeapType.ObjectType))
-        }
-      case _ =>
-        ()
-    }
-
     fixedTpe
   }
 

Diff for: wasm/src/main/scala/wasm4s/Instructions.scala

@@ -253,6 +253,7 @@ object WasmInstr {
       with StackPolymorphicInstr
   case object RETURN extends WasmSimpleInstr("return", 0x0F) with StackPolymorphicInstr
   case class CALL(i: WasmFunctionName) extends WasmFuncInstr("call", 0x10, i)
+  case class RETURN_CALL(i: WasmFunctionName) extends WasmFuncInstr("return_call", 0x12, i)
   case class THROW(i: WasmTagName) extends WasmTagInstr("throw", 0x08, i) with StackPolymorphicInstr
   case object THROW_REF extends WasmSimpleInstr("throw_ref", 0x0A) with StackPolymorphicInstr
   case class TRY_TABLE(i: BlockType, cs: List[CatchClause], label: Option[WasmLabelName] = None)

Diff for: wasm/src/main/scala/wasm4s/Names.scala

@@ -118,6 +118,9 @@ object Names {
       }
     }
 
+    def forTableEntry(clazz: IRNames.ClassName, method: IRNames.MethodName): WasmFunctionName =
+      new WasmFunctionName("t#" + clazz.nameString, method.nameString)
+
     def forExport(exportedName: String): WasmFunctionName =
       new WasmFunctionName("export", exportedName)
 

Diff for: wasm/src/main/scala/wasm4s/WasmContext.scala

@@ -576,7 +576,7 @@ class WasmContext(val module: WasmModule) extends TypeDefinableWasmContext {
         instrs += WasmInstr.I32_CONST(idx)
 
         for (method <- iface.tableEntries)
-          instrs += refFuncWithDeclaration(resolvedMethodInfos(method).wasmName)
+          instrs += refFuncWithDeclaration(resolvedMethodInfos(method).tableEntryName)
         instrs += WasmInstr.STRUCT_NEW(WasmTypeName.WasmStructTypeName.forITable(iface.name))
         instrs += WasmInstr.ARRAY_SET(WasmTypeName.WasmArrayTypeName.itables)
       }
@@ -596,7 +596,7 @@ class WasmContext(val module: WasmModule) extends TypeDefinableWasmContext {
         instrs += I32_CONST(getItableIdx(interfaceName))
 
         for (method <- interfaceInfo.tableEntries)
-          instrs += refFuncWithDeclaration(resolvedMethodInfos(method).wasmName)
+          instrs += refFuncWithDeclaration(resolvedMethodInfos(method).tableEntryName)
         instrs += STRUCT_NEW(WasmStructTypeName.forITable(interfaceName))
         instrs += ARRAY_SET(WasmArrayTypeName.itables)
       }
@@ -873,7 +873,7 @@ object WasmContext {
       val ownerClass: IRNames.ClassName,
       val methodName: IRNames.MethodName
   ) {
-    val wasmName = WasmFunctionName(IRTrees.MemberNamespace.Public, ownerClass, methodName)
+    val tableEntryName = WasmFunctionName.forTableEntry(ownerClass, methodName)
 
     private var effectivelyFinal: Boolean = true