JS: don't treat return, break and continue as split points when possible

Author: Alexey Andreev

js/js.dart-ast/src/com/google/dart/compiler/backend/js/ast/JsFor.java

@@ -107,7 +107,18 @@ public void traverse(JsVisitorWithContext v, JsContext ctx) {
             if (initExpression != null) {
                 initExpression = v.accept(initExpression);
             } else if (initVars != null) {
-                initVars = v.acceptStatement(initVars);
+                JsStatement newInitVars = v.<JsStatement>acceptStatement(initVars);
+                if (newInitVars instanceof JsVars) {
+                    initVars = (JsVars) newInitVars;
+                }
+                else {
+                    initVars = null;
+                    if (newInitVars instanceof JsExpressionStatement) {
+                        initExpression = ((JsExpressionStatement) newInitVars).getExpression();
+                    } else if (newInitVars != null) {
+                        ctx.addPrevious(newInitVars);
+                    }
+                }
             }
 
             if (condition != null) {

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

@@ -51,7 +51,7 @@ class CoroutineBodyTransformer(
 
     fun preProcess(node: JsNode) {
         breakContinueTargetStatements += node.collectBreakContinueTargets()
-        nodesToSplit = node.collectNodesToSplit()
+        nodesToSplit = node.collectNodesToSplit(breakContinueTargetStatements)
     }
 
     fun postProcess(): List<CoroutineBlock> {

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

@@ -21,11 +21,13 @@ import com.google.dart.compiler.backend.js.ast.metadata.*
 import org.jetbrains.kotlin.js.translate.utils.JsAstUtils
 import org.jetbrains.kotlin.utils.singletonOrEmptyList
 
-fun JsNode.collectNodesToSplit(): Set<JsNode> {
+fun JsNode.collectNodesToSplit(breakContinueTargets: Map<JsContinue, JsStatement>): Set<JsNode> {
+    val root = this
     val nodes = mutableSetOf<JsNode>()
 
     val visitor = object : RecursiveJsVisitor() {
         var childrenInSet = false
+        var finallyLevel = 0
 
         override fun visitInvocation(invocation: JsInvocation) {
             super.visitInvocation(invocation)
@@ -37,23 +39,41 @@ fun JsNode.collectNodesToSplit(): Set<JsNode> {
 
         override fun visitReturn(x: JsReturn) {
             super.visitReturn(x)
-            nodes += x
-            childrenInSet = true
+
+            if (root in nodes || finallyLevel > 0) {
+                nodes += x
+                childrenInSet = true
+            }
         }
 
         override fun visitBreak(x: JsBreak) {
             super.visitBreak(x)
 
-            // It's a simplification
-            // TODO: don't split break and continue statements when possible
-            nodes += x
-            childrenInSet = true
+            val breakTarget = breakContinueTargets[x]!!
+            if (breakTarget in nodes) {
+                nodes += x
+                childrenInSet = true
+            }
         }
 
         override fun visitContinue(x: JsContinue) {
             super.visitContinue(x)
-            nodes += x
-            childrenInSet = true
+
+            val continueTarget = breakContinueTargets[x]!!
+            if (continueTarget in nodes) {
+                nodes += x
+                childrenInSet = true
+            }
+        }
+
+        override fun visitTry(x: JsTry) {
+            if (x.finallyBlock != null) {
+                finallyLevel++
+            }
+            super.visitTry(x)
+            if (x.finallyBlock != null) {
+                finallyLevel--
+            }
         }
 
         override fun visitElement(node: JsNode) {
@@ -71,8 +91,12 @@ fun JsNode.collectNodesToSplit(): Set<JsNode> {
         }
     }
 
-    visitor.accept(this)
-    visitor.accept(this)
+    while (true) {
+        val countBefore = nodes.size
+        visitor.accept(this)
+        val countAfter = nodes.size
+        if (countAfter == countBefore) break
+    }
 
     return nodes
 }