JS inline: support inline extension functions

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

@@ -20,22 +20,26 @@ 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.needToAlias;
 import static org.jetbrains.k2js.inline.InlinePackage.aliasArgumentsIfNeeded;
 import static org.jetbrains.k2js.inline.InlinePackage.renameLocals;
+import static org.jetbrains.k2js.inline.InlinePackage.replaceThisReference;
+import static org.jetbrains.k2js.inline.InlinePackage.hasThisReference;
 
 import java.util.*;
 
 class FunctionInlineMutator {
     private static final String BREAK_LABEL = "break_inlined";
     private static final String RESULT_LABEL = "result_inlined";
+    private static final String THIS_ALIAS = "self";
 
-    private JsBlock body;
-    private final boolean isResultNeeded;
+    private final JsInvocation call;
+    private final InliningContext inliningContext;
     private final JsFunction invokedFunction;
-    private final List<JsExpression> arguments;
-    private final List<JsParameter> parameters;
+    private final boolean isResultNeeded;
     private final RenamingContext<JsBlock> renamingContext;
     private final InsertionPoint<JsStatement> insertionPoint;
+    private JsBlock body;
     private JsNameRef resultExpr = null;
     private JsLabel breakLabel = null;
 
@@ -61,23 +65,44 @@ class FunctionInlineMutator {
     }
 
     private FunctionInlineMutator(@NotNull JsInvocation call, @NotNull InliningContext inliningContext) {
+        this.inliningContext = inliningContext;
+        this.call = call;
+
         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();
         insertionPoint = inliningContext.getStatementContext().getInsertionPoint();
     }
 
     private void process() {
-        aliasArgumentsIfNeeded(renamingContext, arguments, parameters);
+        replaceThis();
+        aliasArgumentsIfNeeded(renamingContext, getArguments(), getParameters());
         renameLocals(renamingContext, invokedFunction);
         applyRenaming();
         replaceReturns();
     }
 
+    private void replaceThis() {
+        if (!hasThisReference(body)) {
+            return;
+        }
+
+        JsExpression thisReplacement = inliningContext.getThisReplacement(call);
+        if (thisReplacement == JsLiteral.THIS) {
+            return;
+        }
+
+        if (needToAlias(thisReplacement)) {
+            JsName thisName = renamingContext.getFreshName(THIS_ALIAS);
+            renamingContext.newVar(thisName, thisReplacement);
+            thisReplacement = thisName.makeRef();
+        }
+
+        replaceThisReference(body, thisReplacement);
+    }
+
     private void applyRenaming() {
         RenamingResult<JsBlock> renamingResult = renamingContext.applyRename(body);
         body = renamingResult.getRenamed();
@@ -117,4 +142,14 @@ class FunctionInlineMutator {
         assert resultExpr == null || resultExpr instanceof JsNameRef;
         ReplaceReturnVisitor.replaceReturn(body, (JsNameRef) resultExpr, breakLabelRef);
     }
+
+    @NotNull
+    private List<JsExpression> getArguments() {
+        return inliningContext.getArguments(call);
+    }
+
+    @NotNull
+    private List<JsParameter> getParameters() {
+        return invokedFunction.getParameters();
+    }
 }

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

@@ -19,6 +19,8 @@ package org.jetbrains.k2js.inline;
 import com.google.dart.compiler.backend.js.ast.*;
 import org.jetbrains.annotations.NotNull;
 
+import java.util.List;
+
 interface InliningContext {
     @NotNull
     <T extends JsNode> RenamingContext<T> getRenamingContext();
@@ -29,5 +31,11 @@ interface InliningContext {
     @NotNull
     FunctionContext getFunctionContext();
 
+    @NotNull
+    JsExpression getThisReplacement(JsInvocation call);
+
+    @NotNull
+    List<JsExpression> getArguments(JsInvocation call);
+
     boolean isResultNeeded(JsInvocation call);
 }

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

@@ -72,4 +72,42 @@ class InvocationUtil {
 
         return (JsFunction) returnExpr;
     }
+
+    /**
+     * Tests if invocation is JavaScript call function
+     *
+     * @return true  if invocation is something like `x.call(thisReplacement)`
+     *         false otherwise
+     */
+    public static boolean isCallInvocation(JsInvocation invocation) {
+        JsExpression qualifier = invocation.getQualifier();
+        if (!(qualifier instanceof JsNameRef)) {
+            return false;
+        }
+
+        JsNameRef qualifierNameRef = (JsNameRef) qualifier;
+        return qualifierNameRef.getIdent().equals("call")
+               && invocation.getArguments().size() == 1;
+    }
+
+    public static boolean hasReceiver(@NotNull JsInvocation invocation) {
+        return getReceiverImpl(invocation) != null;
+    }
+
+    @NotNull
+    public static JsExpression getReceiver(@NotNull JsInvocation invocation) {
+        JsExpression receiver = getReceiverImpl(invocation);
+        assert receiver != null;
+        return receiver;
+    }
+
+    @Nullable
+    private static JsExpression getReceiverImpl(@NotNull JsInvocation invocation) {
+        JsExpression qualifier = invocation.getQualifier();
+        if (!(qualifier instanceof JsNameRef)) {
+            return null;
+        }
+
+        return  ((JsNameRef) qualifier).getQualifier();
+    }
 }

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

@@ -22,6 +22,7 @@ import org.jetbrains.annotations.Nullable;
 
 import java.util.IdentityHashMap;
 import java.util.Stack;
+import java.util.List;
 
 import static org.jetbrains.k2js.inline.FunctionInlineMutator.getInlineableCallReplacement;
 
@@ -70,6 +71,31 @@ public class JsInliner extends JsVisitorWithContextImpl {
             return functionContext;
         }
 
+        @NotNull
+        @Override
+        public JsExpression getThisReplacement(JsInvocation call) {
+            if (InvocationUtil.isCallInvocation(call)) {
+                return call.getArguments().get(0);
+            }
+
+            if (InvocationUtil.hasReceiver(call)) {
+                return InvocationUtil.getReceiver(call);
+            }
+
+            return JsLiteral.THIS;
+        }
+
+        @NotNull
+        @Override
+        public List<JsExpression> getArguments(JsInvocation call) {
+            List<JsExpression> arguments = call.getArguments();
+            if (InvocationUtil.isCallInvocation(call)) {
+                return arguments.subList(1, arguments.size());
+            }
+
+            return arguments;
+        }
+
         @Override
         public boolean isResultNeeded(JsInvocation call) {
             JsStatement currentStatement = getStatementContext().getCurrentStatement();

js/js.inliner/src/org/jetbrains/k2js/inline/ThisReplacing.kt

@@ -0,0 +1,60 @@
+/*
+ * 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.*
+
+fun replaceThisReference<T : JsNode>(node: T, replacement: JsExpression) {
+    val visitor = ThisReplacingVisitor(replacement)
+    visitor.accept(node)
+}
+
+fun hasThisReference(node: JsNode): Boolean {
+    val visitor = ContainsThisVisitor()
+    visitor.accept(node)
+    return visitor.containsThis
+}
+
+private class ThisReplacingVisitor(private val thisReplacement: JsExpression) : JsVisitorWithContextImpl() {
+    override fun endVisit(x: JsLiteral.JsThisRef?, ctx: JsContext?) {
+        ctx?.replaceMe(thisReplacement)
+    }
+
+    override fun visit(x: JsFunction?, ctx: JsContext?) = false
+
+    override fun visit(x: JsObjectLiteral?, ctx: JsContext?) = false
+}
+
+private class ContainsThisVisitor(): RecursiveJsVisitor() {
+    public var containsThis: Boolean = false
+        private set
+
+    override fun visitElement(node: JsNode?) {
+        if (!containsThis) {
+            super<RecursiveJsVisitor>.visitElement(node)
+        }
+    }
+
+    override fun visitThis(x: JsLiteral.JsThisRef?) {
+        containsThis = true
+    }
+
+
+    override fun visitFunction(x: JsFunction?) { }
+
+    override fun visitObjectLiteral(x: JsObjectLiteral?) { }
+}