[Wasm] Support passing lambdas from Kotlin to JS

This commit is contained in:
Svyatoslav Kuzmich
2021-10-25 14:23:08 +03:00
parent 766377a041
commit f7d3cb269f
5 changed files with 549 additions and 11 deletions
@@ -504,6 +504,9 @@ val wasmPhases = NamedCompilerPhase(
excludeDeclarationsFromCodegenPhase then
expectDeclarationsRemovingPhase then
jsInteropFunctionsLowering then
jsInteropFunctionCallsLowering then
// TODO: Need some helpers from stdlib
// arrayConstructorPhase then
wrapInlineDeclarationsWithReifiedTypeParametersPhase then
@@ -532,9 +535,6 @@ val wasmPhases = NamedCompilerPhase(
delegateToPrimaryConstructorLoweringPhase then
// Common prefix ends
jsInteropFunctionsLowering then
jsInteropFunctionCallsLowering then
enumEntryInstancesLoweringPhase then
enumEntryInstancesBodyLoweringPhase then
enumClassCreateInitializerLoweringPhase then
@@ -208,6 +208,8 @@ class WasmSymbols(
private val wasmDataRefClass = getIrClass(FqName("kotlin.wasm.internal.reftypes.dataref"))
val wasmDataRefType by lazy { wasmDataRefClass.defaultType }
private val externalInterfaceClass = getIrClass(FqName("kotlin.wasm.internal.ExternalInterfaceType"))
val externalInterfaceType by lazy { externalInterfaceClass.defaultType }
inner class JsInteropAdapters {
val kotlinToJsStringAdapter = getInternalFunction("kotlinToJsStringAdapter")
@@ -228,6 +230,9 @@ class WasmSymbols(
private val jsNameClass = getIrClass(FqName("kotlin.js.JsName"))
val jsNameConstructor by lazy { jsNameClass.constructors.single() }
private val jsFunClass = getIrClass(FqName("kotlin.JsFun"))
val jsFunConstructor by lazy { jsFunClass.constructors.single() }
private fun findClass(memberScope: MemberScope, name: Name): ClassDescriptor =
memberScope.getContributedClassifier(name, NoLookupLocation.FROM_BACKEND) as ClassDescriptor
@@ -14,20 +14,23 @@ import org.jetbrains.kotlin.backend.wasm.WasmBackendContext
import org.jetbrains.kotlin.backend.wasm.ir2wasm.isBuiltInWasmRefType
import org.jetbrains.kotlin.backend.wasm.ir2wasm.isExported
import org.jetbrains.kotlin.backend.wasm.ir2wasm.isExternalType
import org.jetbrains.kotlin.ir.UNDEFINED_OFFSET
import org.jetbrains.kotlin.ir.backend.js.utils.getJsNameOrKotlinName
import org.jetbrains.kotlin.ir.builders.*
import org.jetbrains.kotlin.ir.builders.declarations.addValueParameter
import org.jetbrains.kotlin.ir.builders.declarations.buildFun
import org.jetbrains.kotlin.ir.declarations.IrClass
import org.jetbrains.kotlin.ir.declarations.IrDeclaration
import org.jetbrains.kotlin.ir.declarations.IrDeclarationParent
import org.jetbrains.kotlin.ir.declarations.IrSimpleFunction
import org.jetbrains.kotlin.ir.expressions.IrBody
import org.jetbrains.kotlin.ir.expressions.IrCall
import org.jetbrains.kotlin.ir.expressions.IrExpression
import org.jetbrains.kotlin.ir.types.IrType
import org.jetbrains.kotlin.ir.types.makeNullable
import org.jetbrains.kotlin.ir.expressions.*
import org.jetbrains.kotlin.ir.expressions.impl.IrFunctionExpressionImpl
import org.jetbrains.kotlin.ir.types.*
import org.jetbrains.kotlin.ir.util.*
import org.jetbrains.kotlin.ir.visitors.IrElementTransformerVoid
import org.jetbrains.kotlin.ir.visitors.transformChildrenVoid
import org.jetbrains.kotlin.name.Name
import kotlin.math.absoluteValue
/**
* Create wrappers for external and @JsExport functions when type adaptation is needed
@@ -37,6 +40,17 @@ class JsInteropFunctionsLowering(val context: WasmBackendContext) : DeclarationT
val symbols = context.wasmSymbols
val adapters = symbols.jsInteropAdapters
// Used to for export lambdas
object KOTLIN_WASM_CLOSURE_FOR_JS_CLOSURE : IrStatementOriginImpl("KOTLIN_WASM_CLOSURE_FOR_JS_CLOSURE")
private val closureCallExports = mutableMapOf<IrSimpleType, IrSimpleFunction>()
private val kotlinClosureToJsConverters = mutableMapOf<IrSimpleType, IrSimpleFunction>()
private val jsClosureCallers = mutableMapOf<IrSimpleType, IrSimpleFunction>()
private val jsToKotlinClosures = mutableMapOf<IrSimpleType, IrSimpleFunction>()
val additionalDeclarations = mutableListOf<IrDeclaration>()
lateinit var currentParent: IrDeclarationParent
override fun transformFlat(declaration: IrDeclaration): List<IrDeclaration>? {
if (declaration.isFakeOverride) return null
if (declaration !is IrSimpleFunction) return null
@@ -46,10 +60,14 @@ class JsInteropFunctionsLowering(val context: WasmBackendContext) : DeclarationT
check(!(isExported && isExternal)) { "Exported external declarations are not supported: ${declaration.fqNameWhenAvailable}" }
check(declaration.parent !is IrClass) { "Interop members are not supported: ${declaration.fqNameWhenAvailable}" }
return if (isExternal)
additionalDeclarations.clear()
currentParent = declaration.parent
val newDeclarations = if (isExternal)
transformExternalFunction(declaration)
else
transformExportFunction(declaration)
return (newDeclarations ?: listOf(declaration)) + additionalDeclarations
}
/**
@@ -156,9 +174,9 @@ class JsInteropFunctionsLowering(val context: WasmBackendContext) : DeclarationT
irCall(functionToCall).let { call ->
for ((index, valueParameter) in function.valueParameters.withIndex()) {
val get = irGet(valueParameter)
call.putValueArgument(index, valueParametersAdapters[index]?.adapt(get, builder) ?: get)
call.putValueArgument(index, valueParametersAdapters[index].adaptIfNeeded(get, builder))
}
resultAdapter?.adapt(call, builder) ?: call
resultAdapter.adaptIfNeeded(call, builder)
}
)
}
@@ -191,6 +209,42 @@ class JsInteropFunctionsLowering(val context: WasmBackendContext) : DeclarationT
if (isBuiltInWasmRefType(this))
return null
if (this is IrSimpleType && this.isFunction()) {
val functionTypeInfo = FunctionTypeInfo(this, toJs = true)
// Kotlin's closures are objects that implement FunctionN interface.
// JavaScript can receive opaque reference to them but cannot call them directly.
// Thus, we export helper "caller" method that JavaScript will use to call kotlin closures:
//
// @JsExport
// fun __callFunction_<signatureHash>(f: dataref, p1: JsType1, p2: JsType2, ...): JsTypeRes {
// return adapt(
// cast<FunctionN>(f).invoke(
// adapt(p1),
// adapt(p2),
// ...
// )
// )
// }
//
closureCallExports.getOrPut(this) {
createKotlinClosureCaller(functionTypeInfo)
}
// Converter functions creates new JavaScript closures that delegate to Kotlin closures
// using above-mentioned "caller" export:
//
// @JsFun("""(f) => {
// (p1, p2, ...) => <wasm-exports>.__callFunction_<signatureHash>(f, p1, p2, ...)
// }""")
// external fun __convertKotlinClosureToJsClosure_<signatureHash>(f: dataref): ExternalRef
//
val kotlinToJsClosureConvertor = kotlinClosureToJsConverters.getOrPut(this) {
createKotlinToJsClosureConvertor(functionTypeInfo)
}
return FunctionBasedAdapter(kotlinToJsClosureConvertor)
}
return SendKotlinObjectToJsAdapter(this)
}
@@ -220,15 +274,244 @@ class JsInteropFunctionsLowering(val context: WasmBackendContext) : DeclarationT
if (isBuiltInWasmRefType(this))
return null
if (this is IrSimpleType && this.isFunction()) {
val functionTypeInfo = FunctionTypeInfo(this, toJs = false)
// JavaScript's closures are external references that cannot be called directly in WebAssembly.
// Thus, we import helper "caller" method that WebAssembly will use to call JS closures:
//
// @JsFun("(f, p0, p1, ...) => f(p0, p1, ...)")
// external fun __callJsClosure_<signatureHash>(f: ExternalRef, p0: JsType1, p1: JsType2, ...): JsResType
//
val jsClosureCaller = jsClosureCallers.getOrPut(this) {
createJsClosureCaller(functionTypeInfo)
}
// Converter functions creates new Kotlin closure that delegate to JS closure
// using above-mentioned "caller" import:
//
// fun __convertJsClosureToKotlinClosure_<signatureHash>(f: ExternalRef) : FunctionN<KotlinType1, ..., KotlinResType> =
// { p0: KotlinType1, p1: KotlinType2, ... ->
// adapt(__callJsClosure_<signatureHash>(f, adapt(p0), adapt(p1), ..))
// }
//
val jsToKotlinClosure = jsToKotlinClosures.getOrPut(this) {
createJsToKotlinClosureConverter(functionTypeInfo, jsClosureCaller)
}
return FunctionBasedAdapter(jsToKotlinClosure)
}
return ReceivingKotlinObjectFromJsAdapter(this)
}
private fun createKotlinClosureCaller(info: FunctionTypeInfo): IrSimpleFunction {
val result = context.irFactory.buildFun {
name = Name.identifier("__callFunction_${info.hashString}")
returnType = info.adaptedResultType
}
result.parent = currentParent
result.addValueParameter {
name = Name.identifier("f")
type = context.wasmSymbols.wasmDataRefType
}
var count = 0
info.adaptedParameterTypes.forEach { type ->
result.addValueParameter {
this.name = Name.identifier("p" + count++.toString())
this.type = type
}
}
val builder = context.createIrBuilder(result.symbol)
result.body = builder.irBlockBody {
val invokeFun = info.functionType.classOrNull!!.owner.functions.single { it.name == Name.identifier("invoke") }
val callInvoke = irCall(invokeFun.symbol, info.originalResultType).also { call ->
call.dispatchReceiver =
ReceivingKotlinObjectFromJsAdapter(invokeFun.dispatchReceiverParameter!!.type)
.adapt(irGet(result.valueParameters[0]), builder)
for (i in info.adaptedParameterTypes.indices) {
call.putValueArgument(i, info.parametersAdapters[i].adaptIfNeeded(irGet(result.valueParameters[i + 1]), builder))
}
}
+irReturn(info.resultAdapter.adaptIfNeeded(callInvoke, builder))
}
result.annotations += builder.irCallConstructor(context.wasmSymbols.jsExportConstructor, typeArguments = emptyList())
additionalDeclarations += result
return result
}
private fun createKotlinToJsClosureConvertor(info: FunctionTypeInfo): IrSimpleFunction {
val result = context.irFactory.buildFun {
name = Name.identifier("__convertKotlinClosureToJsClosure_${info.hashString}")
returnType = context.wasmSymbols.externalInterfaceType
}
result.parent = currentParent
result.addValueParameter {
name = Name.identifier("f")
type = context.wasmSymbols.wasmDataRefType
}
val builder = context.createIrBuilder(result.symbol)
// TODO: Cache created JS closures
val arity = info.parametersAdapters.size
val jsCode = buildString {
append("(f) => (")
appendParameterList(arity)
append(") => wasmInstance.exports.__callFunction_")
append(info.hashString)
append("(f, ")
appendParameterList(arity)
append(")")
}
result.annotations += builder.irCallConstructor(context.wasmSymbols.jsFunConstructor, typeArguments = emptyList()).also {
it.putValueArgument(0, builder.irString(jsCode))
}
additionalDeclarations += result
return result
}
private fun createJsToKotlinClosureConverter(
info: FunctionTypeInfo,
jsClosureCaller: IrSimpleFunction,
): IrSimpleFunction {
val functionType = info.functionType
val result = context.irFactory.buildFun {
name = Name.identifier("__convertJsClosureToKotlinClosure_${info.hashString}")
returnType = functionType
}
result.parent = currentParent
result.addValueParameter {
name = Name.identifier("f")
type = context.wasmSymbols.externalInterfaceType
}
val builder = context.createIrBuilder(result.symbol)
val backendContext = context
result.body = builder.irBlockBody {
val lambda = context.irFactory.buildFun {
name = Name.identifier("kotlinClosure_${info.hashString}")
returnType = info.originalResultType
}
lambda.parent = result
val lambdaBuilder = backendContext.createIrBuilder(lambda.symbol)
info.originalParameterTypes.forEachIndexed { index, irType ->
lambda.addValueParameter {
name = Name.identifier("p$index")
type = irType
}
}
lambda.body = lambdaBuilder.irBlockBody {
val jsClosureCallerCall = irCall(jsClosureCaller)
jsClosureCallerCall.putValueArgument(0, irGet(result.valueParameters[0]))
for ((adapterIndex, paramAdapter) in info.parametersAdapters.withIndex()) {
jsClosureCallerCall.putValueArgument(
adapterIndex + 1,
paramAdapter.adaptIfNeeded(
irGet(lambda.valueParameters[adapterIndex]),
lambdaBuilder
)
)
}
+irReturn(info.resultAdapter.adaptIfNeeded(jsClosureCallerCall, lambdaBuilder))
}
+irReturn(
IrFunctionExpressionImpl(
UNDEFINED_OFFSET,
UNDEFINED_OFFSET,
functionType,
lambda,
origin = KOTLIN_WASM_CLOSURE_FOR_JS_CLOSURE,
)
)
}
additionalDeclarations += result
return result
}
private fun createJsClosureCaller(info: FunctionTypeInfo): IrSimpleFunction {
val result = context.irFactory.buildFun {
name = Name.identifier("__callJsClosure_${info.hashString}")
returnType = info.adaptedResultType
}
result.parent = currentParent
result.addValueParameter {
name = Name.identifier("f")
type = symbols.externalInterfaceType
}
val arity = info.adaptedParameterTypes.size
repeat(arity) { paramIndex ->
result.addValueParameter {
name = Name.identifier("p$paramIndex")
type = info.adaptedParameterTypes[paramIndex]
}
}
val builder = context.createIrBuilder(result.symbol)
val jsFun = buildString {
append("(f, ")
appendParameterList(arity)
append(") => f(")
appendParameterList(arity)
append(")")
}
result.annotations += builder.irCallConstructor(context.wasmSymbols.jsFunConstructor, typeArguments = emptyList()).also {
it.putValueArgument(0, builder.irString(jsFun))
}
additionalDeclarations += result
return result
}
inner class FunctionTypeInfo(val functionType: IrSimpleType, toJs: Boolean) {
init {
require(functionType.arguments.all { it is IrTypeProjection }) {
"Star projection is not supported in function type interop ${functionType.render()}"
}
}
val hashString: String =
functionType.hashCode().absoluteValue.toString(Character.MAX_RADIX)
val originalParameterTypes: List<IrType> =
functionType.arguments.dropLast(1).map { (it as IrTypeProjection).type }
val originalResultType: IrType =
(functionType.arguments.last() as IrTypeProjection).type
val parametersAdapters: List<InteropTypeAdapter?> =
originalParameterTypes.map { parameterType ->
if (toJs)
parameterType.jsToKotlinAdapterIfNeeded(isReturn = false)
else
parameterType.kotlinToJsAdapterIfNeeded(isReturn = false)
}
val resultAdapter: InteropTypeAdapter? =
if (toJs)
originalResultType.kotlinToJsAdapterIfNeeded(isReturn = true)
else
originalResultType.jsToKotlinAdapterIfNeeded(isReturn = true)
val adaptedParameterTypes: List<IrType> =
originalParameterTypes.zip(parametersAdapters).map { (parameterType, adapter) ->
(if (toJs) adapter?.fromType else adapter?.toType) ?: parameterType
}
val adaptedResultType: IrType =
(if (toJs) resultAdapter?.toType else resultAdapter?.fromType) ?: originalResultType
}
interface InteropTypeAdapter {
val fromType: IrType
val toType: IrType
fun adapt(expression: IrExpression, builder: IrBuilderWithScope): IrExpression
}
fun InteropTypeAdapter?.adaptIfNeeded(expression: IrExpression, builder: IrBuilderWithScope): IrExpression =
this?.adapt(expression, builder) ?: expression
/**
* Adapter implemented as a single function call
*/
@@ -288,6 +571,14 @@ class JsInteropFunctionsLowering(val context: WasmBackendContext) : DeclarationT
}
}
private fun StringBuilder.appendParameterList(size: Int) =
repeat(size) {
append("p")
append(it)
if (it + 1 < size)
append(", ")
}
/**
* Redirect calls to external and @JsExport functions to created wrappers
*/
@@ -0,0 +1,237 @@
// FILE: externals.js
function apply7(f) {
return f("a")("b")("c")("d")("e")("f")("g")
}
function extenalWithLambda(x, dc) {
x(true, 1, 2, '3'.charCodeAt(0), 4, 5n, 6.1, 7.1, "S", create123Array(), dc);
}
function create123Array() { return [1, 2, 3]; }
function is123Array(ei) {
return ei[0] === 1 && ei[1] === 2 && ei[2] === 3;
}
function externalWithLambdas2(
boolean, // () -> Boolean,
byte, // () -> Byte,
short, // () -> Short,
char, // () -> Char,
int, // () -> Int,
long, // () -> Long,
float, // () -> Float,
double, // () -> Double,
string, // () -> String,
ei, // () -> EI,
dc, // () -> DC,
dcGetY, // (DC) -> Int
) {
let result = 0
function test(boolean) {
if (boolean !== true) throw "Fail" + result
result++
}
test(boolean() === true)
test(byte() === 100)
test(short() === 200)
test(char() === "я".charCodeAt(0))
test(int() === 300)
test(long() === 400n)
test(float() === 500.5)
test(double() === 600.5)
test(string() === "700")
test(is123Array(ei()))
test(dcGetY(dc()) == 800)
return result
}
function createJsLambda() {
return (
boolean,
byte,
short,
char,
int,
long,
float,
double,
string,
ei,
dc,
dcGetY
) => {
let result = 0;
function test(x) {
if (x !== true) throw "Fail" + result;
result++;
}
test(boolean === true);
test(byte === 100);
test(short === 200);
test(char === "я".charCodeAt(0));
test(int === 300);
test(long === 400n);
test(float === 500.5);
test(double === 600.5);
test(string === "700");
test(is123Array(ei));
test(dcGetY(dc) == 800);
return result
};
}
// MODULE: main
// FILE: externals.kt
external fun createJsLambda(): (Boolean, Byte, Short, Char, Int, Long, Float, Double, String, EI, DC, (DC) -> Int) -> Int
external fun apply7(f: (String) -> ((String) -> ((String) -> ((String) -> ((String) -> ((String) -> ((String) -> String))))))): String
external interface EI
external fun is123Array(x: EI): Boolean
external fun create123Array(): EI
data class DC(val x: Int, val y: Int)
external fun extenalWithLambda(
x: (Boolean, Byte, Short, Char, Int, Long, Float, Double, String, EI, DC) -> Unit,
dc: DC,
)
external fun externalWithLambdas2(
boolean: () -> Boolean,
byte: () -> Byte,
short: () -> Short,
char: () -> Char,
int: () -> Int,
long: () -> Long,
float: () -> Float,
double: () -> Double,
string: () -> String,
ei: () -> EI,
dc: () -> DC,
dcGetY: (DC) -> Int,
): Int
@JsExport
fun exportedF() : (Int, Int, Int) -> (String, String) -> String =
{ i1: Int, i2: Int, i3: Int ->
{ s1: String, s2: String ->
"$s1${i1 + i2 + i3}$s2"
}
}
typealias SS2S = (String, String) -> String
@JsExport
fun complexHigherOrder(x: (SS2S, SS2S) -> SS2S): (SS2S, SS2S) -> SS2S =
{ ss2s1: SS2S, ss2s2: SS2S -> x(ss2s1, ss2s2) }
fun complexHigherOrerTest() {
val x = { ss2s1: SS2S, ss2s2: SS2S ->
{ s1: String, s2: String -> ss2s1(s1, s2) + "_" + ss2s2(s1, s2) }
}
val y = complexHigherOrder(x)
val res = y({s1, s2 -> s1 + "+" + s2}, {s1, s2 -> s1 + "-" + s2})("abc", "xyz")
if (res != "abc+xyz_abc-xyz") {
error("Fail complexHigherOrderTest $res")
}
}
fun box(): String {
val res = apply7 { s1 -> { s2 -> { s3 -> { s4 -> { s5 -> { s6 -> { s7 -> s1 + s2 + s3 + s4 + s5 + s6 + s7 } } } } } } }
if (res != "abcdefg") return "Fail"
var extenalWithLambdasCount = 0
fun test(x: Boolean) {
if (!x) error("Fail")
extenalWithLambdasCount++
}
extenalWithLambda({ bool: Boolean,
byte: Byte,
short: Short,
char: Char,
int: Int,
long: Long,
float: Float,
double: Double,
string: String,
ei: EI,
dc: DC ->
test(bool == true)
test(byte == 1.toByte())
test(short == 2.toShort())
test(char == '3')
test(int == 4)
test(long == 5L)
test(float == 6.1f)
test(double == 7.1)
test(string == "S")
test(is123Array(ei))
test(dc.x == 100 && dc.y == 200)
}, DC(100, 200))
if (extenalWithLambdasCount != 11) return "Fail 1"
val externalWithLambdas2Count = externalWithLambdas2(
boolean = { true },
byte = { 100.toByte() },
short = { 200.toShort() },
char = { 'я' },
int = { 300 },
long = { 400L },
float = { 500.5f },
double = { 600.5 },
string = { "700" },
ei = { create123Array() },
dc = { DC(800, 800) },
dcGetY = { it.y }
)
if (externalWithLambdas2Count != 11) return "Fail externalWithLambdas2"
val jsLambda = createJsLambda()
val jsLambdaCount = jsLambda(
true,
100.toByte(),
200.toShort(),
'я',
300,
400L,
500.5f,
600.5,
"700",
create123Array(),
DC(800, 800),
{ it.y }
)
if (jsLambdaCount != 11)
return "Fail 3"
complexHigherOrerTest()
return "OK"
}
// FILE: functionTypes__after.js
const exportedFres = main.exportedF()(1, 20, 300)("<", ">");
if (exportedFres !== "<321>")
throw "exportedF fail " + exportedFres;
(function testComplexHighOrder() {
const x = (ss2s1, ss2s2) => (s1, s2) => ss2s1(s1, s2) + "_" + ss2s2(s1, s2);
const y = main.complexHigherOrder(x);
const res = y((s1, s2) => s1 + "+" + s2, (s1, s2) => s1 + "-" + s2)("abc", "xyz");
if (res !== "abc+xyz_abc-xyz") {
throw "Fail complexHigherOrderTest " + res;
}
})();
@@ -40,6 +40,11 @@ public class IrCodegenWasmJsInteropWasmTestGenerated extends AbstractIrCodegenWa
runTest("compiler/testData/codegen/boxWasmJsInterop/externals.kt");
}
@TestMetadata("functionTypes.kt")
public void testFunctionTypes() throws Exception {
runTest("compiler/testData/codegen/boxWasmJsInterop/functionTypes.kt");
}
@TestMetadata("jsExport.kt")
public void testJsExport() throws Exception {
runTest("compiler/testData/codegen/boxWasmJsInterop/jsExport.kt");