virtual extension functions support and tests.

translator/src/org/jetbrains/k2js/translate/expression/FunctionTranslator.java

@@ -10,13 +10,13 @@ import org.jetbrains.k2js.translate.context.TemporaryVariable;
 import org.jetbrains.k2js.translate.context.TranslationContext;
 import org.jetbrains.k2js.translate.general.AbstractTranslator;
 import org.jetbrains.k2js.translate.general.Translation;
+import org.jetbrains.k2js.translate.utils.DescriptorUtils;
 
 import java.util.ArrayList;
 import java.util.List;
 
 import static org.jetbrains.k2js.translate.utils.BindingUtils.getDescriptorForElement;
 import static org.jetbrains.k2js.translate.utils.BindingUtils.getFunctionDescriptor;
-import static org.jetbrains.k2js.translate.utils.DescriptorUtils.isExtensionFunction;
 import static org.jetbrains.k2js.translate.utils.TranslationUtils.functionWithScope;
 
 
@@ -29,6 +29,8 @@ public final class FunctionTranslator extends AbstractTranslator {
     private final JetDeclarationWithBody functionDeclaration;
     @NotNull
     private final JsFunction functionObject;
+    @NotNull
+    private final TranslationContext functionBodyContext;
 
     @NotNull
     public static FunctionTranslator newInstance(@NotNull JetDeclarationWithBody function,
@@ -41,6 +43,7 @@ public final class FunctionTranslator extends AbstractTranslator {
         super(context);
         this.functionDeclaration = functionDeclaration;
         this.functionObject = createFunctionObject();
+        this.functionBodyContext = functionBodyContext();
     }
 
     @NotNull
@@ -61,12 +64,18 @@ public final class FunctionTranslator extends AbstractTranslator {
 
     @NotNull
     private JsFunction generateFunctionObject() {
-        TranslationContext innerContext = functionBodyContext();
-        functionObject.setParameters(translateParameters(functionDeclaration.getValueParameters(), innerContext));
-        functionObject.setBody(translateBody(innerContext));
+        functionObject.setParameters(translateParameters(functionDeclaration.getValueParameters(), functionBodyContext));
+        functionObject.setBody(translateBody(functionBodyContext));
+        restoreContext();
         return functionObject;
     }
 
+    private void restoreContext() {
+        if (isExtensionFunction()) {
+            functionBodyContext.removeAliasForThis();
+        }
+    }
+
     private JsFunction createFunctionObject() {
         if (isDeclaration()) {
             return functionWithScope
@@ -156,13 +165,20 @@ public final class FunctionTranslator extends AbstractTranslator {
     }
 
     private void mayBeAddThisParameterForExtensionFunction(@NotNull List<JsParameter> jsParameters) {
-        if (!(functionDeclaration instanceof JetNamedFunction)) return;
-
-        FunctionDescriptor functionDescriptor = getFunctionDescriptor(context().bindingContext(), functionDeclaration);
-        if (isExtensionFunction(functionDescriptor)) {
-            JsName receiver = context().jsScope().declareName("receiver");
+        boolean isExtensionFunction = isExtensionFunction();
+        if (isExtensionFunction) {
+            JsName receiver = functionBodyContext.jsScope().declareName("receiver");
             context().aliaser().setAliasForThis(receiver);
             jsParameters.add(new JsParameter(receiver));
         }
     }
+
+    private boolean isExtensionFunction() {
+        if (!(functionDeclaration instanceof JetNamedFunction)) {
+            return false;
+        }
+
+        FunctionDescriptor functionDescriptor = getFunctionDescriptor(context().bindingContext(), functionDeclaration);
+        return DescriptorUtils.isExtensionFunction(functionDescriptor);
+    }
 }

translator/src/org/jetbrains/k2js/translate/reference/CallTranslator.java

@@ -161,7 +161,8 @@ public final class CallTranslator extends AbstractTranslator {
     private JsExpression qualifiedMethodReference(@NotNull DeclarationDescriptor descriptor) {
         JsExpression methodReference = ReferenceTranslator.translateReference(descriptor, context());
         if (isExtensionFunction()) {
-            AstUtil.setQualifier(methodReference, context().getQualifierForDescriptor(functionDescriptor));
+            AstUtil.setQualifier(methodReference,
+                    TranslationUtils.getExtensionFunctionImplicitReceiver(context(), functionDescriptor));
         } else if (receiver != null) {
             AstUtil.setQualifier(methodReference, receiver);
         }

translator/src/org/jetbrains/k2js/translate/utils/TranslationUtils.java

@@ -205,6 +205,20 @@ public final class TranslationUtils {
                 return TranslationUtils.getThisQualifier(context);
             }
         }
+        return getImplicitReceiverBasedOnOwner(context, referencedDescriptor);
+    }
+
+    @NotNull
+    public static JsExpression getExtensionFunctionImplicitReceiver(@NotNull TranslationContext context,
+                                                                    @NotNull FunctionDescriptor descriptor) {
+        JsExpression receiverBasedOnOwner = getImplicitReceiverBasedOnOwner(context, descriptor);
+        assert receiverBasedOnOwner != null : "Extension function must have receiver based on owner.";
+        return receiverBasedOnOwner;
+    }
+
+    @Nullable
+    private static JsExpression getImplicitReceiverBasedOnOwner(@NotNull TranslationContext context,
+                                                                @NotNull DeclarationDescriptor referencedDescriptor) {
         if (isOwnedByClass(referencedDescriptor)) {
             return TranslationUtils.getThisQualifier(context);
         }

translator/test/org/jetbrains/k2js/test/ExtensionFunctionTest.java

@@ -32,4 +32,14 @@ public final class ExtensionFunctionTest extends TranslationTest {
     public void extensionUsedInsideClass() throws Exception {
         testFooBoxIsTrue("extensionUsedInsideClass.kt");
     }
+
+    @Test
+    public void virtualExtension() throws Exception {
+        testFooBoxIsTrue("virtualExtension.kt");
+    }
+
+    @Test
+    public void virtualExtensionOverride() throws Exception {
+        testFooBoxIsTrue("virtualExtensionOverride.kt");
+    }
 }

translator/testFiles/extensionFunction/cases/virtualExtension.kt

@@ -0,0 +1,15 @@
+namespace foo
+
+class A(var a : Int) {
+
+    fun Int.modify() : Int {
+        return this * 3;
+    }
+
+    fun eval() = a.modify();
+}
+
+fun box(): Boolean {
+    val a = A(4)
+    return (a.eval() == 12)
+}

translator/testFiles/extensionFunction/cases/virtualExtensionOverride.kt

@@ -0,0 +1,20 @@
+namespace foo
+
+open class A(var a : Int) {
+
+    open fun Int.modify() : Int {
+        return this * 3;
+    }
+
+    fun eval() = a.modify();
+}
+
+class B(a : Int) : A(a) {
+    override fun Int.modify() : Int {
+        return this - 2;
+    }
+}
+
+fun box(): Boolean {
+    return (A(4).eval() == 12) && (A(2).eval() == 6) && (B(3).eval() == 1)
+}