Merge pull request #118 from sjrd/better-receiver-type-for-static-dispatch

Better receiver types for static dispatch.

Author: tanishiking

wasm/src/main/scala/ir2wasm/HelperFunctions.scala

@@ -550,7 +550,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)

wasm/src/main/scala/ir2wasm/TypeTransformer.scala

@@ -30,17 +30,6 @@ object TypeTransformer {
       case _                   => List(transformType(t))
     }
 
-  def transformTypeRef(t: IRTypes.TypeRef)(implicit ctx: ReadOnlyWasmContext): Types.WasmType =
-    t match {
-      case arrayTypeRef: IRTypes.ArrayTypeRef =>
-        Types.WasmRefType.nullable(
-          Names.WasmTypeName.WasmStructTypeName.forArrayClass(arrayTypeRef)
-        )
-      case IRTypes.ClassRef(className) =>
-        transformClassByName(className)
-      case IRTypes.PrimRef(tpe) => transformPrimType(tpe)
-    }
-
   /** Transforms a value type to a unique Wasm type.
     *
     * This method cannot be used for `void` and `nothing`, since they have no corresponding Wasm
@@ -54,28 +43,23 @@ object TypeTransformer {
         Types.WasmRefType.nullable(
           Names.WasmTypeName.WasmStructTypeName.forArrayClass(tpe.arrayTypeRef)
         )
-      case IRTypes.ClassType(className) => transformClassByName(className)
+      case IRTypes.ClassType(className) => transformClassType(className)
       case IRTypes.RecordType(fields)   => ???
       case IRTypes.StringType | IRTypes.UndefType =>
         Types.WasmRefType.any
       case p: IRTypes.PrimTypeWithRef => transformPrimType(p)
     }
 
-  private def transformClassByName(
+  def transformClassType(
       className: IRNames.ClassName
-  )(implicit ctx: ReadOnlyWasmContext): Types.WasmType = {
-    className match {
-      case _ =>
-        val info = ctx.getClassInfo(className)
-        if (info.isAncestorOfHijackedClass)
-          Types.WasmRefType.anyref
-        else if (info.isInterface)
-          Types.WasmRefType.nullable(Types.WasmHeapType.ObjectType)
-        else
-          Types.WasmRefType.nullable(
-            Names.WasmTypeName.WasmStructTypeName.forClass(className)
-          )
-    }
+  )(implicit ctx: ReadOnlyWasmContext): Types.WasmRefType = {
+    val info = ctx.getClassInfo(className)
+    if (info.isAncestorOfHijackedClass)
+      Types.WasmRefType.anyref
+    else if (info.isInterface)
+      Types.WasmRefType.nullable(Types.WasmHeapType.ObjectType)
+    else
+      Types.WasmRefType.nullable(WasmTypeName.WasmStructTypeName.forClass(className))
   }
 
   private def transformPrimType(

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)
@@ -489,29 +489,35 @@ class WasmBuilder(coreSpec: CoreSpec) {
       ctor.name.name
     )
 
+    val resultTyp = WasmRefType(typeName)
+
     implicit val fctx = WasmFunctionContext(
       WasmFunctionName.loadModule(clazz.className),
       Nil,
-      List(WasmRefType.nullable(typeName))
+      List(resultTyp)
     )
 
+    val instanceLocal = fctx.addLocal("instance", resultTyp)
+
     import fctx.instrs
 
-    // global.get $module_name
-    // ref.if_null
-    //   ref.null $module_type
-    //   call $module_init ;; should set to global
-    // end
-    // global.get $module_name
-    instrs += GLOBAL_GET(globalInstanceName) // [rt]
-    instrs += REF_IS_NULL // [rt] -> [i32] (bool)
-    instrs += IF(BlockType.ValueType())
-    instrs += CALL(WasmFunctionName.newDefault(clazz.name.name))
-    instrs += GLOBAL_SET(globalInstanceName)
-    instrs += GLOBAL_GET(globalInstanceName)
-    instrs += CALL(ctorName)
-    instrs += END
-    instrs += GLOBAL_GET(globalInstanceName) // [rt]
+    fctx.block(resultTyp) { nonNullLabel =>
+      // load global, return if not null
+      instrs += GLOBAL_GET(globalInstanceName)
+      instrs += BR_ON_NON_NULL(nonNullLabel)
+
+      // create an instance and call its constructor
+      instrs += CALL(WasmFunctionName.newDefault(clazz.name.name))
+      instrs += LOCAL_TEE(instanceLocal)
+      instrs += CALL(ctorName)
+
+      // store it in the global
+      instrs += LOCAL_GET(instanceLocal)
+      instrs += GLOBAL_SET(globalInstanceName)
+
+      // return it
+      instrs += LOCAL_GET(instanceLocal)
+    }
 
     fctx.buildAndAddToContext()
   }
@@ -1076,31 +1082,28 @@ class WasmBuilder(coreSpec: CoreSpec) {
   private def genFunction(
       clazz: LinkedClass,
       method: IRTrees.MethodDef
-  )(implicit ctx: WasmContext): WasmFunction = {
+  )(implicit ctx: WasmContext): Unit = {
     implicit val pos = method.pos
 
-    val functionName = Names.WasmFunctionName(
-      method.flags.namespace,
-      clazz.name.name,
-      method.name.name
-    )
+    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(transformClassType(className).toNonNullable)
 
     // Prepare for function context, set receiver and parameters
     implicit val fctx = WasmFunctionContext(
-      Some(clazz.className),
+      Some(className),
       functionName,
       receiverTyp,
       method.args,
@@ -1111,7 +1114,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(

wasm/src/main/scala/ir2wasm/WasmExpressionBuilder.scala

@@ -1699,21 +1699,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
   }
 
@@ -1908,8 +1893,7 @@ private class WasmExpressionBuilder private (
      * if the given class is an ancestor of hijacked classes (which in practice
      * is only the case for j.l.Object).
      */
-    val instanceTyp =
-      Types.WasmRefType.nullable(WasmStructTypeName.forClass(n.className))
+    val instanceTyp = Types.WasmRefType(WasmStructTypeName.forClass(n.className))
     val localInstance = fctx.addSyntheticLocal(instanceTyp)
 
     fctx.markPosition(n)
@@ -1941,7 +1925,7 @@ private class WasmExpressionBuilder private (
     val primLocal = fctx.addSyntheticLocal(primTyp)
 
     val boxClassType = IRTypes.ClassType(boxClassName)
-    val boxTyp = TypeTransformer.transformType(boxClassType)(ctx)
+    val boxTyp = TypeTransformer.transformClassType(boxClassName)(ctx).toNonNullable
     val instanceLocal = fctx.addSyntheticLocal(boxTyp)
 
     /* The generated code is as follows. Before the codegen, the stack contains
@@ -1986,12 +1970,12 @@ private class WasmExpressionBuilder private (
 
   private def genWrapAsThrowable(tree: IRTrees.WrapAsThrowable): IRTypes.Type = {
     val throwableClassType = IRTypes.ClassType(IRNames.ThrowableClass)
-    val throwableTyp = TypeTransformer.transformType(throwableClassType)(ctx)
+    val nonNullThrowableTyp = Types.WasmRefType(Types.WasmHeapType.ThrowableType)
 
-    val jsExceptionClassType = IRTypes.ClassType(SpecialNames.JSExceptionClass)
-    val jsExceptionTyp = TypeTransformer.transformType(jsExceptionClassType)(ctx)
+    val jsExceptionTyp =
+      TypeTransformer.transformClassType(SpecialNames.JSExceptionClass)(ctx).toNonNullable
 
-    fctx.block(throwableTyp) { doneLabel =>
+    fctx.block(nonNullThrowableTyp) { doneLabel =>
       genTree(tree.expr, IRTypes.AnyType)
 
       fctx.markPosition(tree)
@@ -2000,7 +1984,7 @@ private class WasmExpressionBuilder private (
       instrs += BR_ON_CAST(
         doneLabel,
         Types.WasmRefType.anyref,
-        Types.WasmRefType(Types.WasmHeapType.ThrowableType)
+        nonNullThrowableTyp
       )
 
       // otherwise, wrap in a new JavaScriptException

wasm/src/main/scala/wasm4s/Instructions.scala

@@ -258,6 +258,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)

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)
 

wasm/src/main/scala/wasm4s/Types.scala

@@ -23,7 +23,10 @@ object Types {
   case object WasmInt8 extends WasmPackedType("i8", 0x78)
   case object WasmInt16 extends WasmPackedType("i16", 0x77)
 
-  final case class WasmRefType(nullable: Boolean, heapType: WasmHeapType) extends WasmType
+  final case class WasmRefType(nullable: Boolean, heapType: WasmHeapType) extends WasmType {
+    def toNullable: WasmRefType = WasmRefType(true, heapType)
+    def toNonNullable: WasmRefType = WasmRefType(false, heapType)
+  }
 
   object WasmRefType {
 

wasm/src/main/scala/wasm4s/WasmContext.scala

@@ -578,7 +578,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)
       }
@@ -598,7 +598,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)
       }
@@ -879,7 +879,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