JS: optimize variable representation in coroutines

Don't convert local variables to fields of coroutine object
when variable is both used and defined in a single block.

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

@@ -43,11 +43,14 @@ class CoroutineFunctionTransformer(private val function: JsFunction, name: Strin
         val globalCatchBlockIndex = coroutineBlocks.indexOf(context.globalCatchBlock)
 
         coroutineBlocks.forEach { it.jsBlock.collectAdditionalLocalVariables() }
-        coroutineBlocks.forEach { it.jsBlock.replaceLocalVariables(context, localVariables) }
+
+        val survivingLocalVars = coroutineBlocks.collectVariablesSurvivingBetweenBlocks(
+                localVariables, function.parameters.map { it.name }.toSet())
+        coroutineBlocks.forEach { it.jsBlock.replaceLocalVariables(context, survivingLocalVars) }
 
         val additionalStatements = mutableListOf<JsStatement>()
         generateDoResume(coroutineBlocks, context, additionalStatements)
-        generateContinuationConstructor(context, additionalStatements, globalCatchBlockIndex)
+        generateContinuationConstructor(context, additionalStatements, globalCatchBlockIndex, survivingLocalVars)
 
         generateCoroutineInstantiation(context)
 
@@ -57,7 +60,8 @@ class CoroutineFunctionTransformer(private val function: JsFunction, name: Strin
     private fun generateContinuationConstructor(
             context: CoroutineTransformationContext,
             statements: MutableList<JsStatement>,
-            globalCatchBlockIndex: Int
+            globalCatchBlockIndex: Int,
+            survivingLocalVars: Set<JsName>
     ) {
         val psiElement = context.metadata.psiElement
 
@@ -93,7 +97,7 @@ class CoroutineFunctionTransformer(private val function: JsFunction, name: Strin
             if (context.metadata.hasReceiver) {
                 assignToField(context.receiverFieldName, context.receiverFieldName.makeRef(), psiElement)
             }
-            for (localVariable in localVariables) {
+            for (localVariable in survivingLocalVars) {
                 val value = if (localVariable !in parameterNames) Namer.getUndefinedExpression() else localVariable.makeRef()
                 assignToField(context.getFieldName(localVariable), value, psiElement)
             }

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

@@ -21,6 +21,7 @@ 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.js.translate.utils.splitToRanges
 
 fun JsNode.collectNodesToSplit(breakContinueTargets: Map<JsContinue, JsStatement>): Set<JsNode> {
     val root = this
@@ -239,6 +240,79 @@ fun JsBlock.replaceSpecialReferences(context: CoroutineTransformationContext) {
     visitor.accept(this)
 }
 
+fun List<CoroutineBlock>.collectVariablesSurvivingBetweenBlocks(localVariables: Set<JsName>, parameters: Set<JsName>): Set<JsName> {
+    val varDefinedIn = localVariables.associate { it to mutableSetOf<Int>() }
+    val varDeclaredIn = localVariables.associate { it to mutableSetOf<Int>() }
+    val varUsedIn = localVariables.associate { it to mutableSetOf<Int>() }
+
+    for ((blockIndex, block) in withIndex()) {
+        for (statement in block.statements) {
+            statement.accept(object : RecursiveJsVisitor() {
+                override fun visitNameRef(nameRef: JsNameRef) {
+                    super.visitNameRef(nameRef)
+                    varUsedIn[nameRef.name]?.add(blockIndex)
+                }
+
+                override fun visit(x: JsVars.JsVar) {
+                    varDeclaredIn[x.name]?.add(blockIndex)
+                    if (x.initExpression != null) {
+                        varDefinedIn[x.name]?.add(blockIndex)
+                    }
+                    super.visit(x)
+
+                }
+
+                override fun visitBinaryExpression(x: JsBinaryOperation) {
+                    val lhs = x.arg1
+                    if (x.operator.isAssignment && lhs is JsNameRef) {
+                        varDefinedIn[lhs.name]?.add(blockIndex)?.let {
+                            accept(x.arg2)
+                            return
+                        }
+                    }
+                    super.visitBinaryExpression(x)
+                }
+
+                override fun visitFunction(x: JsFunction) {
+                    x.name?.let {
+                        varDefinedIn[it]?.add(blockIndex)
+                    }
+                }
+
+                override fun visitLabel(x: JsLabel) {
+                    accept(x.statement)
+                }
+
+                override fun visitBreak(x: JsBreak) {}
+
+                override fun visitContinue(x: JsContinue) {}
+            })
+        }
+    }
+
+    fun JsName.isLocalInBlock(): Boolean {
+        val def = varDefinedIn[this]!!
+        val use = varUsedIn[this]!!
+        val decl = varDeclaredIn[this]!!
+        if (def.size == 1 && use.size == 1) {
+            val singleDef = def.single()
+            val singleUse = use.single()
+            return singleDef == singleUse && decl.isNotEmpty()
+        }
+        return use.isEmpty()
+    }
+
+    return localVariables.filterNot { localVar ->
+        if (localVar in parameters) {
+            varUsedIn[localVar]!!.isEmpty() && varDefinedIn[localVar]!!.isEmpty() && varDeclaredIn[localVar]!!.isEmpty()
+        }
+        else {
+            localVar.isLocalInBlock()
+        }
+    }.toSet()
+
+}
+
 fun JsBlock.replaceLocalVariables(context: CoroutineTransformationContext, localVariables: Set<JsName>) {
     replaceSpecialReferences(context)
 
@@ -269,22 +343,36 @@ fun JsBlock.replaceLocalVariables(context: CoroutineTransformationContext, local
         }
 
         override fun endVisit(x: JsVars, ctx: JsContext<in JsStatement>) {
-            val assignments = x.vars.mapNotNull {
+            if (x.vars.none { it.name in localVariables }) return
-                val fieldName = context.getFieldName(it.name)
+
-                val initExpression = it.initExpression
+            val statements = mutableListOf<JsStatement>()
-                if (initExpression != null) {
+            for ((range, shouldReplace) in x.vars.splitToRanges { it.name in localVariables }) {
-                    JsAstUtils.assignment(JsNameRef(fieldName, JsThisRef()), it.initExpression)
+                if (shouldReplace) {
+                    val assignments = x.vars.mapNotNull {
+                        val fieldName = context.getFieldName(it.name)
+                        val initExpression = it.initExpression
+                        if (initExpression != null) {
+                            JsAstUtils.assignment(JsNameRef(fieldName, JsThisRef()), it.initExpression)
+                        }
+                        else {
+                            null
+                        }
+                    }
+                    if (assignments.isNotEmpty()) {
+                        statements += JsExpressionStatement(JsAstUtils.newSequence(assignments))
+                    }
                 }
                 else {
-                    null
+                    statements += JsVars(*range.toTypedArray())
                 }
             }
 
-            if (assignments.isNotEmpty()) {
+            if (statements.size == 1) {
-                ctx.replaceMe(JsExpressionStatement(JsAstUtils.newSequence(assignments)))
+                ctx.replaceMe(statements[0])
             }
             else {
                 ctx.removeMe()
+                ctx.addPrevious(statements)
             }
         }
     }

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

@@ -912,6 +912,21 @@ public class BoxJsTestGenerated extends AbstractBoxJsTest {
         }
     }
 
+    @TestMetadata("js/js.translator/testData/box/coroutines")
+    @TestDataPath("$PROJECT_ROOT")
+    @RunWith(JUnit3RunnerWithInners.class)
+    public static class Coroutines extends AbstractBoxJsTest {
+        public void testAllFilesPresentInCoroutines() throws Exception {
+            KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("js/js.translator/testData/box/coroutines"), Pattern.compile("^([^_](.+))\\.kt$"), TargetBackend.JS, true);
+        }
+
+        @TestMetadata("localVarOptimization.kt")
+        public void testLocalVarOptimization() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("js/js.translator/testData/box/coroutines/localVarOptimization.kt");
+            doTest(fileName);
+        }
+    }
+
     @TestMetadata("js/js.translator/testData/box/crossModuleRef")
     @TestDataPath("$PROJECT_ROOT")
     @RunWith(JUnit3RunnerWithInners.class)

js/js.tests/test/org/jetbrains/kotlin/js/test/utils/DirectiveTestUtils.java

@@ -299,6 +299,27 @@ public class DirectiveTestUtils {
         }
     };
 
+    private static final DirectiveHandler DECLARES_VARIABLE = new DirectiveHandler("DECLARES_VARIABLE") {
+        @Override
+        void processEntry(@NotNull JsNode ast, @NotNull ArgumentsHelper arguments) throws Exception {
+            String functionName = arguments.getNamedArgument("function");
+            String varName = arguments.getNamedArgument("name");
+            JsFunction function = AstSearchUtil.getFunction(ast, functionName);
+            boolean[] varDeclared = new boolean[1];
+            function.accept(new RecursiveJsVisitor() {
+                @Override
+                public void visit(@NotNull JsVars.JsVar x) {
+                    super.visit(x);
+                    if (x.getName().getIdent().equals(varName)) {
+                        varDeclared[0] = true;
+                    }
+                }
+            });
+
+            assertTrue("Function " + functionName + " does not declare variable " + varName, varDeclared[0]);
+        }
+    };
+
     private static final List<DirectiveHandler> DIRECTIVE_HANDLERS = Arrays.asList(
             FUNCTION_CONTAINS_NO_CALLS,
             FUNCTION_NOT_CALLED,
@@ -319,7 +340,8 @@ public class DirectiveTestUtils {
             NOT_REFERENCED,
             HAS_INLINE_METADATA,
             HAS_NO_INLINE_METADATA,
-            HAS_NO_CAPTURED_VARS
+            HAS_NO_CAPTURED_VARS,
+            DECLARES_VARIABLE
     );
 
     public static void processDirectives(@NotNull JsNode ast, @NotNull String sourceCode) throws Exception {

js/js.translator/testData/box/coroutines/localVarOptimization.kt

@@ -0,0 +1,49 @@
+// EXPECTED_REACHABLE_NODES: 1124
+// DECLARES_VARIABLE: function=doResume name=k
+// PROPERTY_READ_COUNT: name=local$o count=1
+// PROPERTY_WRITE_COUNT: name=local$o count=2
+import kotlin.coroutines.experimental.*
+
+var next: () -> Unit = {}
+var complete = false
+var log = ""
+
+suspend fun foo(x: String): String = suspendCoroutine { continuation ->
+    log += "[$x]"
+    next = { continuation.resume(x) }
+}
+
+fun build(x: suspend () -> Unit) {
+    next = {
+        x.startCoroutine(object : Continuation<Unit> {
+            override val context = EmptyCoroutineContext
+
+            override fun resume(x: Unit) {
+                complete = true
+            }
+
+            override fun resumeWithException(x: Throwable) {
+                complete = true
+            }
+        })
+    }
+}
+
+fun box(): String {
+    build {
+        val o = foo("O")
+        log += "-"
+        val k = foo("K")
+        log += ":"
+        log += "{$o$k}"
+    }
+
+    while (!complete) {
+        next()
+        log += "#"
+    }
+
+    if (log != "[O]#-[K]#:{OK}#") return "fail: $log"
+
+    return "OK"
+}

js/js.translator/testData/lineNumbers/coroutine.kt

@@ -14,4 +14,4 @@ suspend fun bar(): Unit {
     println(a + b)
 }
 
-// LINES: 4 4 4 5 5 5 5 6 4 4 4 15 9 9 9 9 15 9 9 9 9 * 15 10 10 11 11 11 11 11 * 11 12 12 13 13 13 13 13 13 14 14
+// LINES: 4 4 4 5 5 5 5 6 4 4 4 15 9 9 9 9 15 9 9 9 * 15 10 10 11 11 11 11 11 * 11 12 12 13 13 13 13 13 13 13 14 14