From 33314b25c0b24aa490422ad5275b4814b94dbc4b Mon Sep 17 00:00:00 2001
From: Svyatoslav Kuzmich <svyatoslav.kuzmich@jetbrains.com>
Date: Wed, 12 May 2021 21:27:50 +0300
Subject: [PATCH] [Wasm] Add function interface replacer lowering

Replace function interface super types with runtime interfaces.
Needed for interface calls machinery.
---
 .../kotlin/backend/wasm/WasmLoweringPhases.kt |  8 +++
 .../lower/WasmFunctionInterfaceReplacer.kt    | 54 +++++++++++++++++++
 2 files changed, 62 insertions(+)
 create mode 100644 compiler/ir/backend.wasm/src/org/jetbrains/kotlin/backend/wasm/lower/WasmFunctionInterfaceReplacer.kt

diff --git a/compiler/ir/backend.wasm/src/org/jetbrains/kotlin/backend/wasm/WasmLoweringPhases.kt b/compiler/ir/backend.wasm/src/org/jetbrains/kotlin/backend/wasm/WasmLoweringPhases.kt
index 81c2cccc9e9..158733cacc3 100644
--- a/compiler/ir/backend.wasm/src/org/jetbrains/kotlin/backend/wasm/WasmLoweringPhases.kt
+++ b/compiler/ir/backend.wasm/src/org/jetbrains/kotlin/backend/wasm/WasmLoweringPhases.kt
@@ -335,6 +335,13 @@ private val wasmNullSpecializationLowering = makeWasmModulePhase(
     description = "Specialize assigning Nothing? values to other types."
 )
 
+private val wasmFunctionInterfaceReplacer = makeWasmModulePhase(
+    ::WasmFunctionInterfaceReplacer,
+    name = "WasmFunctionInterfaceReplacer",
+    description = "Replace function interface with concrete runtime interfaces"
+)
+
+
 private val staticMembersLoweringPhase = makeWasmModulePhase(
     ::StaticMembersLowering,
     name = "StaticMembersLowering",
@@ -513,6 +520,7 @@ val wasmPhases = NamedCompilerPhase(
 
             virtualDispatchReceiverExtractionPhase then
             staticMembersLoweringPhase then
+            wasmFunctionInterfaceReplacer then
             wasmNullSpecializationLowering then
             validateIrAfterLowering
 )
diff --git a/compiler/ir/backend.wasm/src/org/jetbrains/kotlin/backend/wasm/lower/WasmFunctionInterfaceReplacer.kt b/compiler/ir/backend.wasm/src/org/jetbrains/kotlin/backend/wasm/lower/WasmFunctionInterfaceReplacer.kt
new file mode 100644
index 00000000000..5172ffa2530
--- /dev/null
+++ b/compiler/ir/backend.wasm/src/org/jetbrains/kotlin/backend/wasm/lower/WasmFunctionInterfaceReplacer.kt
@@ -0,0 +1,54 @@
+/*
+ * Copyright 2010-2019 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.backend.wasm.lower
+
+import org.jetbrains.kotlin.backend.common.FileLoweringPass
+import org.jetbrains.kotlin.backend.common.IrElementTransformerVoidWithContext
+import org.jetbrains.kotlin.backend.wasm.WasmBackendContext
+import org.jetbrains.kotlin.ir.IrStatement
+import org.jetbrains.kotlin.ir.declarations.IrClass
+import org.jetbrains.kotlin.ir.declarations.IrFile
+import org.jetbrains.kotlin.ir.types.IrSimpleType
+import org.jetbrains.kotlin.ir.types.IrType
+import org.jetbrains.kotlin.ir.types.typeWithArguments
+import org.jetbrains.kotlin.ir.util.isFunction
+import org.jetbrains.kotlin.ir.util.render
+import org.jetbrains.kotlin.ir.visitors.transformChildrenVoid
+
+/**
+ * Replace kotlin.Function{N} etc super types with  kotlin.wasm.internal.Function{N} super types.
+ * This is a workaround for having a concrete IR for these interfaces.
+ */
+class WasmFunctionInterfaceReplacer(val context: WasmBackendContext) : FileLoweringPass {
+
+    private fun replaceType(type: IrType): IrType {
+        if (type.isFunction()) {
+            require(type is IrSimpleType)
+            val klass: IrClass = type.classifier.owner as IrClass
+
+            // No need to replace just "kotlin.Function"
+            if (klass.name.identifier == "Function") {
+                return type
+            }
+
+            val arity = klass.typeParameters.size - 1
+            return context.wasmSymbols.functionN(arity).typeWithArguments(type.arguments)
+        }
+
+        return type
+    }
+
+    override fun lower(irFile: IrFile) {
+        irFile.transformChildrenVoid(object : IrElementTransformerVoidWithContext() {
+            override fun visitClassNew(declaration: IrClass): IrStatement {
+                if (declaration.superTypes.any { it.isFunction() }) {
+                    declaration.superTypes = declaration.superTypes.map { replaceType(it) }
+                }
+                return super.visitClassNew(declaration)
+            }
+        })
+    }
+}
\ No newline at end of file