JS: do not create variable for result, if inline function has one top level return

js/js.inliner/src/org/jetbrains/kotlin/js/inline/FunctionInlineMutator.kt

@@ -23,6 +23,7 @@ import org.jetbrains.kotlin.js.inline.clean.removeDefaultInitializers
 import org.jetbrains.kotlin.js.inline.context.InliningContext
 import org.jetbrains.kotlin.js.inline.context.NamingContext
 import org.jetbrains.kotlin.js.inline.util.*
+import org.jetbrains.kotlin.js.inline.util.rewriters.ReturnReplacingVisitor
 import org.jetbrains.kotlin.js.translate.utils.JsAstUtils.newVar
 import org.jetbrains.kotlin.js.translate.utils.ast.*
 
@@ -55,17 +56,11 @@ class FunctionInlineMutator private (private val call: JsInvocation, private val
         aliasArgumentsIfNeeded(namingContext, arguments, parameters)
         renameLocalNames(namingContext, invokedFunction)
         removeStatementsAfterTopReturn()
+        processReturns()
 
-        if (isResultNeeded && canBeExpression(body)) {
-            resultExpr = asExpression(body)
-            body.getStatements().clear()
-
-            /** JsExpression can be immutable, so need to reassign  */
-            resultExpr = namingContext.applyRenameTo(resultExpr!!) as JsExpression
-        }
-        else {
-            processReturns()
-            namingContext.applyRenameTo(body)
+        namingContext.applyRenameTo(body)
+        resultExpr = resultExpr?.let {
+            namingContext.applyRenameTo(it) as JsExpression
         }
     }
 
@@ -103,33 +98,29 @@ class FunctionInlineMutator private (private val call: JsInvocation, private val
         if (returnCount == 0) {
             // TODO return Unit (KT-5647)
             resultExpr = JsLiteral.UNDEFINED
+            return
         }
-        else {
-            doReplaceReturns(returnCount)
+
+        if (returnCount == 1) {
+            val statements = body.getStatements()
+            val lastTopLevelStatement = statements[statements.lastIndex]
+
+            if (lastTopLevelStatement is JsReturn) {
+                resultExpr = lastTopLevelStatement.getExpression()
+                statements.remove(statements.lastIndex)
+                return
+            }
         }
+
+        doReplaceReturns()
     }
 
-    private fun doReplaceReturns(returnCount: Int) {
-        val returnOnTop = ContainerUtil.findInstance(body.getStatements(), javaClass<JsReturn>())
-        val hasReturnOnTopLevel = returnOnTop != null
-
-        val needBreakLabel = !(returnCount == 1 && hasReturnOnTopLevel)
-        var breakLabelRef: JsNameRef? = null
-
-        if (needBreakLabel) {
-            val breakName = namingContext.getFreshName(getBreakLabel())
-            breakLabelRef = breakName.makeRef()
-            breakLabel = JsLabel(breakName)
-        }
-
+    private fun doReplaceReturns() {
         val resultReference = getResultReference()
         if (resultReference != null) {
             resultExpr = resultReference
         }
-
         assert(resultExpr == null || resultExpr is JsNameRef)
-        replaceReturns(body, resultExpr as? JsNameRef, breakLabelRef)
-    }
 
         val breakName = namingContext.getFreshName(getBreakLabel())
         breakLabel = JsLabel(breakName)
@@ -235,12 +226,7 @@ class FunctionInlineMutator private (private val call: JsInvocation, private val
                 inlineableBody = breakLabel
             }
 
-            var resultExpression: JsExpression? = null
-            if (mutator.isResultNeeded) {
-                resultExpression = mutator.resultExpr
-            }
-
-            return InlineableResult(inlineableBody, resultExpression)
+            return InlineableResult(inlineableBody, mutator.resultExpr)
         }
 
         platformStatic
@@ -270,13 +256,5 @@ class FunctionInlineMutator private (private val call: JsInvocation, private val
             val statements = body.getStatements()
             return statements.size() == 1 && statements.get(0) is JsReturn
         }
-
-        private fun asExpression(body: JsBlock): JsExpression {
-            assert(canBeExpression(body))
-
-            val statements = body.getStatements()
-            val returnStatement = statements.get(0) as JsReturn
-            return returnStatement.getExpression()
-        }
     }
 }

js/js.inliner/src/org/jetbrains/kotlin/js/inline/JsInliner.java

@@ -184,8 +184,8 @@ public class JsInliner extends JsVisitorWithContextImpl {
         JsExpression resultExpression = inlineableResult.getResultExpression();
         JsContext<JsStatement> statementContext = inliningContext.getStatementContext();
         // body of inline function can contain call to lambdas that need to be inlined
-        JsStatement statement = accept(inlineableBody);
-        assert inlineableBody == statement;
+        JsStatement inlineableBodyWithLambdasInlined = accept(inlineableBody);
+        assert inlineableBody == inlineableBodyWithLambdasInlined;
 
         statementContext.addPrevious(flattenStatement(inlineableBody));
 
@@ -199,7 +199,17 @@ public class JsInliner extends JsVisitorWithContextImpl {
         }
 
         resultExpression = accept(resultExpression);
-        context.replaceMe(resultExpression);
+        JsStatement currentStatement = statementContext.getCurrentNode();
+
+        if (currentStatement instanceof JsExpressionStatement &&
+            ((JsExpressionStatement) currentStatement).getExpression() == call &&
+            (resultExpression == null || !canHaveSideEffect(resultExpression))
+        ) {
+            statementContext.removeMe();
+        }
+        else {
+            context.replaceMe(resultExpression);
+        }
     }
 
     @NotNull

js/js.inliner/src/org/jetbrains/kotlin/js/inline/context/NamingContext.kt

@@ -30,16 +30,15 @@ class NamingContext(
 ) {
     private val renamings = IdentityHashMap<JsName, JsExpression>()
     private val declarations = ArrayList<JsVars>()
-    private var renamingApplied = false
+    private var addedDeclarations = false
 
     public fun applyRenameTo(target: JsNode): JsNode {
-        if (renamingApplied) throw RuntimeException("RenamingContext has been applied already")
+        if (!addedDeclarations) {
+            statementContext.addPrevious(declarations)
+            addedDeclarations = true
+        }
 
-        val result = replaceNames(target, renamings)
-        statementContext.addPrevious(declarations)
-        renamingApplied = true
-
-        return result
+        return replaceNames(target, renamings)
     }
 
     public fun replaceName(name: JsName, replacement: JsExpression) {

js/js.inliner/src/org/jetbrains/kotlin/js/inline/util/rewriteUtils.kt

@@ -30,10 +30,6 @@ public fun <T : JsNode> replaceNames(node: T, replaceMap: IdentityHashMap<JsName
     return NameReplacingVisitor(replaceMap).accept(node)!!
 }
 
-public fun replaceReturns(scope: JsNode, resultRef: JsNameRef?, breakLabel: JsNameRef?): JsNode {
-    return ReturnReplacingVisitor(resultRef, breakLabel).accept(scope)!!
-}
-
 public fun replaceThisReference<T : JsNode>(node: T, replacement: JsExpression) {
     ThisReplacingVisitor(replacement).accept(node)
 }

js/js.inliner/src/org/jetbrains/kotlin/js/inline/util/rewriters/ReturnReplacingVisitor.kt

@@ -28,7 +28,7 @@ import com.google.dart.compiler.backend.js.ast.JsVisitorWithContextImpl
 import org.jetbrains.kotlin.js.inline.util.canHaveSideEffect
 import org.jetbrains.kotlin.js.translate.utils.JsAstUtils
 
-class ReturnReplacingVisitor(private val resultRef: JsNameRef?, private val breakLabel: JsNameRef?) : JsVisitorWithContextImpl() {
+public class ReturnReplacingVisitor(private val resultRef: JsNameRef?, private val breakLabel: JsNameRef?) : JsVisitorWithContextImpl() {
 
     /**
      * Prevents replacing returns in object literal

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

@@ -35,6 +35,12 @@ public class InlineSizeReductionTestGenerated extends AbstractInlineSizeReductio
         JetTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("js/js.translator/testData/inlineSizeReduction/cases"), Pattern.compile("^(.+)\\.kt$"), true);
     }
 
+    @TestMetadata("oneTopLevelReturn.kt")
+    public void testOneTopLevelReturn() throws Exception {
+        String fileName = JetTestUtils.navigationMetadata("js/js.translator/testData/inlineSizeReduction/cases/oneTopLevelReturn.kt");
+        doTest(fileName);
+    }
+
     @TestMetadata("simpleReturnFunction.kt")
     public void testSimpleReturnFunction() throws Exception {
         String fileName = JetTestUtils.navigationMetadata("js/js.translator/testData/inlineSizeReduction/cases/simpleReturnFunction.kt");

js/js.translator/testData/inlineSizeReduction/cases/oneTopLevelReturn.kt

@@ -0,0 +1,39 @@
+package foo
+
+// CHECK_CONTAINS_NO_CALLS: test1
+// CHECK_CONTAINS_NO_CALLS: test2
+// CHECK_VARS_COUNT: function=test1 count=0
+// CHECK_VARS_COUNT: function=test2 count=1
+
+var log = ""
+
+inline fun run1(fn: ()->Int): Int {
+    log += "1;"
+    return 1 + fn()
+}
+
+inline fun run2(fn: ()->Int): Int {
+    log += "2;"
+    return 2 + run1(fn)
+}
+
+inline fun run3(fn: ()->Int): Int {
+    log += "3;"
+    return 3 + run2(fn)
+}
+
+fun test1(x: Int): Int = run3 { x }
+
+fun test2(x: Int): Int {
+    val result = 1 + run3 { x }
+    return result
+}
+
+fun box(): String {
+    assertEquals(7, test1(1))
+    assertEquals("3;2;1;", log)
+
+    assertEquals(8, test2(1))
+
+    return "OK"
+}