JS: fix coroutine transformation of callable references to local functions. See KT-16164

compiler/testData/codegen/box/coroutines/localCallableRef.kt

@@ -0,0 +1,24 @@
+// WITH_RUNTIME
+// WITH_COROUTINES
+import kotlin.coroutines.experimental.*
+import kotlin.coroutines.experimental.intrinsics.*
+
+suspend fun suspendWithValue(result: () -> String): String = suspendCoroutineOrReturn { x ->
+    x.resume(result())
+    COROUTINE_SUSPENDED
+}
+
+fun builder(c: suspend () -> Unit) {
+    c.startCoroutine(EmptyContinuation)
+}
+
+fun box(): String {
+    var result = ""
+
+    builder {
+        fun ok() = "OK"
+        result = suspendWithValue(::ok)
+    }
+
+    return result
+}

compiler/testData/codegen/light-analysis/coroutines/localCallableRef.txt

@@ -0,0 +1,31 @@
+@kotlin.Metadata
+public final class CoroutineUtilKt {
+    public final static @org.jetbrains.annotations.NotNull method handleExceptionContinuation(@org.jetbrains.annotations.NotNull p0: kotlin.jvm.functions.Function1): kotlin.coroutines.experimental.Continuation
+    public final static @org.jetbrains.annotations.NotNull method handleResultContinuation(@org.jetbrains.annotations.NotNull p0: kotlin.jvm.functions.Function1): kotlin.coroutines.experimental.Continuation
+}
+
+@kotlin.Metadata
+public class EmptyContinuation {
+    public final static field Companion: EmptyContinuation.Companion
+    private final @org.jetbrains.annotations.NotNull field context: kotlin.coroutines.experimental.CoroutineContext
+    inner class EmptyContinuation/Companion
+    public @synthetic.kotlin.jvm.GeneratedByJvmOverloads method <init>(): void
+    public method <init>(@org.jetbrains.annotations.NotNull p0: kotlin.coroutines.experimental.CoroutineContext): void
+    public synthetic method <init>(p0: kotlin.coroutines.experimental.CoroutineContext, p1: int, p2: kotlin.jvm.internal.DefaultConstructorMarker): void
+    public @org.jetbrains.annotations.NotNull method getContext(): kotlin.coroutines.experimental.CoroutineContext
+    public method resume(@org.jetbrains.annotations.Nullable p0: java.lang.Object): void
+    public method resumeWithException(@org.jetbrains.annotations.NotNull p0: java.lang.Throwable): void
+}
+
+@kotlin.Metadata
+public final static class EmptyContinuation/Companion {
+    inner class EmptyContinuation/Companion
+    private method <init>(): void
+}
+
+@kotlin.Metadata
+public final class LocalCallableRefKt {
+    public final static @org.jetbrains.annotations.NotNull method box(): java.lang.String
+    public final static method builder(@org.jetbrains.annotations.NotNull p0: kotlin.jvm.functions.Function1): void
+    public final static @org.jetbrains.annotations.Nullable method suspendWithValue(@org.jetbrains.annotations.NotNull p0: kotlin.jvm.functions.Function0, @org.jetbrains.annotations.NotNull p1: kotlin.coroutines.experimental.Continuation): java.lang.Object
+}

compiler/tests-ir-jvm/tests/org/jetbrains/kotlin/codegen/ir/IrBlackBoxCodegenTestGenerated.java

@@ -4886,6 +4886,12 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
             doTest(fileName);
         }
 
+        @TestMetadata("localCallableRef.kt")
+        public void testLocalCallableRef() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/localCallableRef.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("localDelegate.kt")
         public void testLocalDelegate() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/localDelegate.kt");

compiler/tests/org/jetbrains/kotlin/codegen/BlackBoxCodegenTestGenerated.java

@@ -4886,6 +4886,12 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
             doTest(fileName);
         }
 
+        @TestMetadata("localCallableRef.kt")
+        public void testLocalCallableRef() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/localCallableRef.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("localDelegate.kt")
         public void testLocalDelegate() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/localDelegate.kt");

js/js.inliner/src/org/jetbrains/kotlin/js/coroutine/CoroutinePasses.kt

@@ -18,8 +18,9 @@ package org.jetbrains.kotlin.js.coroutine
 
 import org.jetbrains.kotlin.js.backend.ast.*
 import org.jetbrains.kotlin.js.backend.ast.metadata.*
+import org.jetbrains.kotlin.js.inline.util.collectFreeVariables
+import org.jetbrains.kotlin.js.inline.util.replaceNames
 import org.jetbrains.kotlin.js.translate.utils.JsAstUtils
-import org.jetbrains.kotlin.utils.singletonOrEmptyList
 
 fun JsNode.collectNodesToSplit(breakContinueTargets: Map<JsContinue, JsStatement>): Set<JsNode> {
     val root = this
@@ -189,7 +190,7 @@ private fun CoroutineBlock.collectTargetBlocks(): Set<CoroutineBlock> {
     val targetBlocks = mutableSetOf<CoroutineBlock>()
     jsBlock.accept(object : RecursiveJsVisitor() {
         override fun visitDebugger(x: JsDebugger) {
-            targetBlocks += x.targetExceptionBlock.singletonOrEmptyList() + x.targetBlock.singletonOrEmptyList()
+            targetBlocks += listOfNotNull(x.targetExceptionBlock) + listOfNotNull(x.targetBlock)
         }
     })
     return targetBlocks
@@ -239,7 +240,23 @@ fun JsBlock.replaceSpecialReferences(context: CoroutineTransformationContext) {
 fun JsBlock.replaceLocalVariables(scope: JsScope, context: CoroutineTransformationContext, localVariables: Set<JsName>) {
     replaceSpecialReferences(context)
     val visitor = object : JsVisitorWithContextImpl() {
-        override fun visit(x: JsFunction, ctx: JsContext<*>) = false
+        override fun visit(x: JsFunction, ctx: JsContext<*>): Boolean = false
+
+        override fun endVisit(x: JsFunction, ctx: JsContext<in JsNode>) {
+            val freeVars = x.collectFreeVariables().intersect(localVariables)
+            if (freeVars.isNotEmpty()) {
+                val wrapperFunction = JsFunction(x.scope.parent, JsBlock(), "")
+                val wrapperInvocation = JsInvocation(wrapperFunction)
+                wrapperFunction.body.statements += JsReturn(x)
+                val nameMap = freeVars.associate { it to wrapperFunction.scope.declareTemporaryName(it.ident) }
+                for (freeVar in freeVars) {
+                    wrapperFunction.parameters += JsParameter(nameMap[freeVar]!!)
+                    wrapperInvocation.arguments += JsNameRef(scope.getFieldName(freeVar), JsLiteral.THIS)
+                }
+                x.body = replaceNames(x.body, nameMap.mapValues { it.value.makeRef() })
+                ctx.replaceMe(wrapperInvocation)
+            }
+        }
 
         override fun endVisit(x: JsNameRef, ctx: JsContext<in JsNode>) {
             if (x.qualifier == null && x.name in localVariables) {
@@ -266,8 +283,6 @@ fun JsBlock.replaceLocalVariables(scope: JsScope, context: CoroutineTransformati
             else {
                 ctx.removeMe()
             }
-
-            super.endVisit(x, ctx)
         }
     }
     visitor.accept(this)

js/js.tests/test/org/jetbrains/kotlin/js/test/semantics/JsCodegenBoxTestGenerated.java

@@ -5601,6 +5601,12 @@ public class JsCodegenBoxTestGenerated extends AbstractJsCodegenBoxTest {
             doTest(fileName);
         }
 
+        @TestMetadata("localCallableRef.kt")
+        public void testLocalCallableRef() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/localCallableRef.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("localDelegate.kt")
         public void testLocalDelegate() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/localDelegate.kt");