JS inline: replace returns with result variable assignment and break

js/js.inliner/src/org/jetbrains/k2js/inline/FunctionInlineMutator.java

@@ -17,19 +17,27 @@
 package org.jetbrains.k2js.inline;
 
 import com.google.dart.compiler.backend.js.ast.*;
+import com.intellij.util.containers.ContainerUtil;
 import org.jetbrains.annotations.NotNull;
 
 import static org.jetbrains.k2js.inline.InlinePackage.aliasArgumentsIfNeeded;
+import static org.jetbrains.k2js.inline.InlinePackage.renameLocals;
 
 import java.util.*;
 
 class FunctionInlineMutator {
+    private static final String BREAK_LABEL = "break_inlined";
+    private static final String RESULT_LABEL = "result_inlined";
+
     private JsBlock body;
+    private final boolean isResultNeeded;
     private final JsFunction invokedFunction;
     private final List<JsExpression> arguments;
     private final List<JsParameter> parameters;
     private final RenamingContext<JsBlock> renamingContext;
     private final InsertionPoint<JsStatement> insertionPoint;
+    private JsNameRef resultExpr = null;
+    private JsLabel breakLabel = null;
 
     public static InlineableResult getInlineableCallReplacement(
             @NotNull JsInvocation call,
@@ -39,13 +47,24 @@ class FunctionInlineMutator {
         mutator.process();
 
         JsStatement inlineableBody = mutator.body;
-        return new InlineableResult(inlineableBody, call);
+        if (mutator.breakLabel != null) {
+            mutator.breakLabel.setStatement(inlineableBody);
+            inlineableBody = mutator.breakLabel;
+        }
+
+        JsExpression resultExpression = null;
+        if (mutator.isResultNeeded) {
+            resultExpression = mutator.resultExpr;
+        }
+
+        return new InlineableResult(inlineableBody, resultExpression);
     }
 
     private FunctionInlineMutator(@NotNull JsInvocation call, @NotNull InliningContext inliningContext) {
         FunctionContext functionContext = inliningContext.getFunctionContext();
         invokedFunction = functionContext.getFunctionDefinition(call);
         body = invokedFunction.getBody().deepCopy();
+        isResultNeeded = inliningContext.isResultNeeded(call);
         arguments = call.getArguments();
         parameters = invokedFunction.getParameters();
         renamingContext = inliningContext.getRenamingContext();
@@ -56,6 +75,7 @@ class FunctionInlineMutator {
         aliasArgumentsIfNeeded(renamingContext, arguments, parameters);
         renameLocals(renamingContext, invokedFunction);
         applyRenaming();
+        replaceReturns();
     }
 
     private void applyRenaming() {
@@ -64,4 +84,37 @@ class FunctionInlineMutator {
         Collection<JsVars> declarations = renamingResult.getDeclarations();
         insertionPoint.insertAllBefore(declarations);
     }
+
+    private void replaceReturns() {
+        int returnCount = ReturnCounter.countReturns(body);
+        if (returnCount == 0) {
+            // TODO return Unit (KT-5647)
+            resultExpr = JsLiteral.UNDEFINED;
+        } else {
+            doReplaceReturns(returnCount);
+        }
+    }
+
+    private void doReplaceReturns(int returnCount) {
+        JsReturn returnOnTop = ContainerUtil.findInstance(body.getStatements(), JsReturn.class);
+        boolean hasReturnOnTopLevel = returnOnTop != null;
+
+        if (isResultNeeded) {
+            JsName resultName = renamingContext.getFreshName(getResultLabel());
+            renamingContext.newVar(resultName, null);
+            resultExpr = resultName.makeRef();
+        }
+
+        boolean needBreakLabel = !(returnCount == 1 && hasReturnOnTopLevel);
+        JsNameRef breakLabelRef = null;
+
+        if (needBreakLabel) {
+            JsName breakName = renamingContext.getFreshName(getBreakLabel());
+            breakLabelRef = breakName.makeRef();
+            breakLabel = new JsLabel(breakName);
+        }
+
+        assert resultExpr == null || resultExpr instanceof JsNameRef;
+        ReplaceReturnVisitor.replaceReturn(body, (JsNameRef) resultExpr, breakLabelRef);
+    }
 }

js/js.inliner/src/org/jetbrains/k2js/inline/ReplaceReturnVisitor.java

@@ -0,0 +1,78 @@
+/*
+ * Copyright 2010-2014 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.jetbrains.k2js.inline;
+
+import com.google.dart.compiler.backend.js.ast.*;
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+import org.jetbrains.k2js.translate.utils.JsAstUtils;
+
+import static org.jetbrains.k2js.inline.util.UtilPackage.canHaveSideEffect;
+
+class ReplaceReturnVisitor extends JsVisitorWithContextImpl {
+    private final JsNameRef resultRef;
+
+    @Nullable
+    private final JsNameRef breakLabel;
+
+    @NotNull
+    public static JsNode replaceReturn(JsStatement statement, @Nullable JsNameRef resultRef, @Nullable JsNameRef breakLabel) {
+        return new ReplaceReturnVisitor(resultRef, breakLabel).accept(statement);
+    }
+
+    private ReplaceReturnVisitor(@Nullable JsNameRef resultRef, @Nullable JsNameRef breakLabel) {
+        this.resultRef = resultRef;
+        this.breakLabel = breakLabel;
+    }
+
+    @Override
+    public void endVisit(JsReturn x, JsContext ctx) {
+        if (breakLabel != null) {
+            JsBreak breakFunction = new JsBreak(breakLabel);
+            ctx.insertAfter(breakFunction);
+        }
+
+        JsExpression returnExpression = x.getExpression();
+        if (returnExpression == null) {
+            ctx.removeMe();
+            return;
+        }
+
+        if (resultRef != null) {
+            JsExpression resultAssignment = JsAstUtils.assignment(resultRef, returnExpression);
+            ctx.replaceMe(new JsExpressionStatement(resultAssignment));
+            return;
+        }
+
+        if (canHaveSideEffect(returnExpression)) {
+            JsStatement replacement = new JsExpressionStatement(returnExpression);
+            ctx.replaceMe(replacement);
+        } else {
+            ctx.removeMe();
+        }
+    }
+
+    @Override
+    public boolean visit(JsObjectLiteral x, JsContext ctx) {
+        return false;
+    }
+
+    @Override
+    public boolean visit(JsFunction x, JsContext ctx) {
+        return false;
+    }
+}

js/js.inliner/src/org/jetbrains/k2js/inline/ReturnCounter.java

@@ -0,0 +1,44 @@
+/*
+ * Copyright 2010-2014 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.jetbrains.k2js.inline;
+
+import com.google.dart.compiler.backend.js.ast.JsNode;
+import com.google.dart.compiler.backend.js.ast.JsReturn;
+import com.google.dart.compiler.backend.js.ast.RecursiveJsVisitor;
+
+public class ReturnCounter extends RecursiveJsVisitor {
+
+    private int count = 0;
+
+    public static int countReturns(JsNode root) {
+        ReturnCounter counter = new ReturnCounter();
+        root.accept(counter);
+
+        return counter.getCount();
+    }
+
+    ReturnCounter() {}
+
+    private int getCount() {
+        return count;
+    }
+
+    @Override
+    public void visitReturn(JsReturn x) {
+        count++;
+    }
+}

js/js.inliner/src/org/jetbrains/k2js/inline/util/sideEffectUtils.kt

@@ -0,0 +1,63 @@
+/*
+ * Copyright 2010-2014 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.jetbrains.k2js.inline.util
+
+import com.google.dart.compiler.backend.js.ast.JsExpression
+import com.google.dart.compiler.backend.js.ast.RecursiveJsVisitor
+import com.google.dart.compiler.backend.js.ast.JsNode
+import com.google.dart.compiler.backend.js.ast.JsBinaryOperation
+import com.google.dart.compiler.backend.js.ast.JsPostfixOperation
+import com.google.dart.compiler.backend.js.ast.JsPrefixOperation
+import com.google.dart.compiler.backend.js.ast.JsNew
+import com.google.dart.compiler.backend.js.ast.JsInvocation
+
+public fun canHaveSideEffect(expression: JsExpression): Boolean {
+    val visitor = SideEffectExpessionVisitor()
+    visitor.accept(expression)
+    return visitor.canHaveSideEffect
+}
+
+private class SideEffectExpessionVisitor() : RecursiveJsVisitor() {
+    public var canHaveSideEffect: Boolean = false
+        private set
+
+    override fun visitElement(node: JsNode?) {
+        if (!canHaveSideEffect) {
+            super<RecursiveJsVisitor>.visitElement(node)
+        }
+    }
+
+    override fun visitBinaryExpression(x: JsBinaryOperation?) {
+        canHaveSideEffect = true
+    }
+
+    override fun visitPostfixOperation(x: JsPostfixOperation?) {
+        canHaveSideEffect = true
+    }
+
+    override fun visitPrefixOperation(x: JsPrefixOperation?) {
+        canHaveSideEffect = true
+    }
+
+    override fun visitNew(x: JsNew?) {
+        canHaveSideEffect = true
+    }
+
+    override fun visitInvocation(invocation: JsInvocation?) {
+        canHaveSideEffect = true
+    }
+}