JS: inline all calls in a statement in one pass

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

@@ -5,6 +5,9 @@
 package com.google.dart.compiler.backend.js.ast;
 
 import org.jetbrains.annotations.Nullable;
+import sun.reflect.generics.reflectiveObjects.NotImplementedException;
+
+import java.util.List;
 
 /**
  * The context in which a JsNode visitation occurs. This represents the set of
@@ -13,9 +16,19 @@ import org.jetbrains.annotations.Nullable;
  */
 public abstract class JsContext<T extends JsNode> {
 
-  public abstract <R extends T> void insertAfter(R node);
+  public <R extends T> void addPrevious(R node) {
+    throw new NotImplementedException();
+  }
 
-  public abstract <R extends T> void insertBefore(R node);
+  public <R extends T> void addPrevious(List<R> nodes) {
+    for (R node : nodes) {
+      addPrevious(node);
+    }
+  }
+
+  public <R extends T> void addNext(R node) {
+    throw new NotImplementedException();
+  }
 
   public abstract void removeMe();
 

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

@@ -26,8 +26,7 @@ import com.intellij.util.SmartList;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
-import java.util.List;
+import java.util.*;
-import java.util.Stack;
 
 /**
  * A visitor for iterating through and modifying an AST.
@@ -37,47 +36,71 @@ public class JsVisitorWithContextImpl extends JsVisitorWithContext {
     private final Stack<JsContext<JsStatement>> statementContexts = new Stack<JsContext<JsStatement>>();
 
     public class ListContext<T extends JsNode> extends JsContext<T> {
-        private List<T> collection;
+        private List<T> nodes;
         private int index;
 
+        // Those are reset in every iteration of traverse()
+        private final List<T> previous = new SmartList<T>();
+        private final List<T> next = new SmartList<T>();
+        private boolean removed = false;
+
         @Override
-        public <R extends T> void insertAfter(R node) {
+        public <R extends T> void addPrevious(R node) {
-            //noinspection unchecked
+            previous.add(node);
-            collection.add(index + 1, (T) node);
         }
 
         @Override
-        public <R extends T> void insertBefore(R node) {
+        public <R extends T> void addNext(R node) {
-            //noinspection unchecked
+            next.add(node);
-            collection.add(index++, (T) node);
         }
 
         @Override
         public void removeMe() {
-            collection.remove(index--);
+            removed = true;
         }
 
         @Override
         public <R extends T> void replaceMe(R node) {
-            checkReplacement(collection.get(index), node);
+            checkReplacement(nodes.get(index), node);
-            collection.set(index, node);
+            nodes.set(index, node);
+            removed = false;
         }
         
         @Nullable
         @Override
         public T getCurrentNode() {
-            if (index < collection.size()) {
+            if (!removed && index < nodes.size()) {
-                return collection.get(index);
+                return nodes.get(index);
             }
 
             return null;
         }
 
-        protected void traverse(List<T> collection) {
+        protected void traverse(List<T> nodes) {
-            this.collection = collection;
+            assert previous.isEmpty(): "addPrevious() was called before traverse()";
-            for (index = 0; index < collection.size(); ++index) {
+            assert next.isEmpty(): "addNext() was called before traverse()";
-                T node = collection.get(index);
+            this.nodes = nodes;
-                doTraverse(node, this);
+
+            for (index = 0; index < nodes.size(); index++) {
+                removed = false;
+                previous.clear();
+                next.clear();
+                doTraverse(getCurrentNode(), this);
+
+                if (!previous.isEmpty()) {
+                    nodes.addAll(index, previous);
+                    index += previous.size();
+                }
+
+                if (removed) {
+                    nodes.remove(index);
+                    index--;
+                }
+
+                if (!next.isEmpty()) {
+                    nodes.addAll(index + 1, next);
+                    index += next.size();
+                }
             }
         }
     }
@@ -88,16 +111,6 @@ public class JsVisitorWithContextImpl extends JsVisitorWithContext {
     private class NodeContext<T extends JsNode> extends JsContext<T> {
         protected T node;
 
-        @Override
-        public <R extends T> void insertAfter(R node) {
-            throw new UnsupportedOperationException();
-        }
-
-        @Override
-        public <R extends T> void insertBefore(R node) {
-            throw new UnsupportedOperationException();
-        }
-
         @Override
         public void removeMe() {
             throw new UnsupportedOperationException();

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

@@ -48,36 +48,6 @@ public class JsInliner extends JsVisitorWithContextImpl {
     private final Stack<JsFunction> namedFunctionsStack = new Stack<JsFunction>();
     private final LinkedList<JsCallInfo> inlineCallInfos = new LinkedList<JsCallInfo>();
 
-    /**
-     * A statement can contain more, than one inlineable sub-expressions.
-     * When inline call is expanded, current statement is shifted forward,
-     * but still has same statement context with same index on stack.
-     *
-     * The shifting is intentional, because there could be function literals,
-     * that need to be inlined, after expansion.
-     *
-     * After shifting following inline expansion in the same statement could be
-     * incorrect, because wrong statement index is used.
-     *
-     * To prevent this, after every shift this flag is set to true,
-     * so that visitor wont go deeper until statement is visited.
-     *
-     * Example:
-     *  inline fun f(g: () -> Int): Int { val a = g(); return a }
-     *  inline fun Int.abs(): Int = if (this < 0) -this else this
-     *
-     *  val g = { 10 }
-     *  >> val h = f(g).abs()    // last statement context index
-     *
-     *  val g = { 10 }           // after inline
-     *  >> val f$result          // statement index was not changed
-     *  val a = g()
-     *  f$result = a
-     *  val h = f$result.abs()   // current expression still here; incorrect to inline abs(),
-     *                           //  because statement context on stack point to different statement
-     */
-    private boolean lastStatementWasShifted = false;
-
     public static JsProgram process(@NotNull TranslationContext context) {
         JsProgram program = context.program();
         IdentityHashMap<JsName, JsFunction> functions = collectNamedFunctions(program);
@@ -150,7 +120,7 @@ public class JsInliner extends JsVisitorWithContextImpl {
             inline(call, context);
         }
 
-        return !lastStatementWasShifted;
+        return true;
     }
 
     @Override
@@ -175,36 +145,21 @@ public class JsInliner extends JsVisitorWithContextImpl {
         JsStatement inlineableBody = inlineableResult.getInlineableBody();
         JsExpression resultExpression = inlineableResult.getResultExpression();
         JsContext<JsStatement> statementContext = inliningContext.getStatementContext();
+        // body of inline function can contain call to lambdas that need to be inlined
         accept(inlineableBody);
+        statementContext.addPrevious(flattenStatement(inlineableBody));
 
         /**
          * Assumes, that resultExpression == null, when result is not needed.
          * @see FunctionInlineMutator.isResultNeeded()
          */
         if (resultExpression == null) {
-            statementContext.removeCurrentStatement();
+            statementContext.removeMe();
-        } else {
+            return;
-            context.replaceMe(resultExpression);
         }
 
-        /** @see #lastStatementWasShifted */
+        resultExpression = accept(resultExpression);
-        statementContext.shiftCurrentStatementForward();
+        context.replaceMe(resultExpression);
-    }
-
-    /**
-     * Prevents JsInliner from traversing sub-expressions,
-     * when current statement was shifted forward.
-     */
-    @Override
-    protected <T extends JsNode> void doTraverse(T node, JsContext ctx) {
-        if (node instanceof JsStatement) {
-            /** @see #lastStatementWasShifted */
-            lastStatementWasShifted = false;
-        }
-
-        if (!lastStatementWasShifted) {
-            super.doTraverse(node, ctx);
-        }
     }
 
     @NotNull

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

@@ -36,7 +36,7 @@ class NamingContext(
         if (renamingApplied) throw RuntimeException("RenamingContext has been applied already")
 
         val result = replaceNames(target, renamings)
-        statementContext.insertAllBefore(declarations)
+        statementContext.addPrevious(declarations)
         renamingApplied = true
 
         return result

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

@@ -43,16 +43,17 @@ class ReturnReplacingVisitor(private val resultRef: JsNameRef?, private val brea
     override fun endVisit(x: JsReturn?, ctx: JsContext<*>?) {
         if (x == null || ctx == null) return
 
-        if (breakLabel != null) {
+        ctx.removeMe()
-            ctx.insertAfter(JsBreak(breakLabel))
-        }
 
         val returnReplacement = getReturnReplacement(x.getExpression())
         if (returnReplacement != null) {
-            ctx.insertBefore(JsExpressionStatement(returnReplacement))
+            ctx.addNext(JsExpressionStatement(returnReplacement))
+        }
+
+        if (breakLabel != null) {
+            ctx.addNext(JsBreak(breakLabel))
         }
 
-        ctx.removeMe()
     }
 
     private fun getReturnReplacement(returnExpression: JsExpression?): JsExpression? {

js/js.translator/testData/inline/cases/expressionBodyWithLambdaCall.kt

@@ -0,0 +1,30 @@
+package foo
+
+// CHECK_CONTAINS_NO_CALLS: test
+
+// A copy of stdlib run function.
+// Copied to not to depend on run implementation.
+// It's important, that the body is just `return fn()`.
+inline fun evaluate<T>(fn: ()->T): T = fn()
+
+fun test(n: Int): Int {
+    return evaluate {
+        var i = n
+        var sum = 0
+
+        while (i > 0) {
+            sum += i
+            i--
+        }
+
+        sum
+    }
+}
+
+fun box(): String {
+    assertEquals(6, test(3))
+    assertEquals(0, test(0))
+    assertEquals(0, test(-1))
+
+    return "OK"
+}