From 43525abda3c91151bc392f5ed279d1fda0cee018 Mon Sep 17 00:00:00 2001
From: Alexey Andreev <Alexey.Andreev@jetbrains.com>
Date: Mon, 14 Nov 2016 14:40:21 +0300
Subject: [PATCH] JS: don't treat return, break and continue as split points
 when possible

---
 .../dart/compiler/backend/js/ast/JsFor.java   | 13 +++++-
 .../js/coroutine/CoroutineBodyTransformer.kt  |  2 +-
 .../kotlin/js/coroutine/CoroutinePasses.kt    | 46 ++++++++++++++-----
 3 files changed, 48 insertions(+), 13 deletions(-)

diff --git a/js/js.dart-ast/src/com/google/dart/compiler/backend/js/ast/JsFor.java b/js/js.dart-ast/src/com/google/dart/compiler/backend/js/ast/JsFor.java
index 5d6daaff112..ca4c6daccda 100644
--- a/js/js.dart-ast/src/com/google/dart/compiler/backend/js/ast/JsFor.java
+++ b/js/js.dart-ast/src/com/google/dart/compiler/backend/js/ast/JsFor.java
@@ -107,7 +107,18 @@ public class JsFor extends SourceInfoAwareJsNode implements JsStatement {
             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) {
diff --git a/js/js.inliner/src/org/jetbrains/kotlin/js/coroutine/CoroutineBodyTransformer.kt b/js/js.inliner/src/org/jetbrains/kotlin/js/coroutine/CoroutineBodyTransformer.kt
index 8b66e4a7de8..01cb0a6ce94 100644
--- a/js/js.inliner/src/org/jetbrains/kotlin/js/coroutine/CoroutineBodyTransformer.kt
+++ b/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> {
diff --git a/js/js.inliner/src/org/jetbrains/kotlin/js/coroutine/CoroutinePasses.kt b/js/js.inliner/src/org/jetbrains/kotlin/js/coroutine/CoroutinePasses.kt
index 14d198a9247..c9865c04a1e 100644
--- a/js/js.inliner/src/org/jetbrains/kotlin/js/coroutine/CoroutinePasses.kt
+++ b/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
 }