Support references to local extensions in codegen

compiler/backend/src/org/jetbrains/jet/codegen/ExpressionCodegen.java

@@ -2507,6 +2507,9 @@ public class ExpressionCodegen extends JetVisitor<StackValue, StackValue> implem
                 Call call = CallMaker.makeCall(fakeExpression, NO_RECEIVER, null, fakeExpression, fakeArguments);
                 Callable callable = codegen.resolveToCallable(codegen.accessibleFunctionDescriptor(referencedFunction), false);
                 StackValue receiver = generateCallee(codegen, callable);
+                if (extensionReceiver.exists()) {
+                    receiver = StackValue.composed(receiver, receiverParameterStackValue(signature));
+                }
                 result = codegen.invokeFunctionWithCalleeOnStack(call, receiver, fakeResolvedCall, callable);
             }
 
@@ -2574,13 +2577,15 @@ public class ExpressionCodegen extends JetVisitor<StackValue, StackValue> implem
             if (receiver == null) return NO_RECEIVER;
 
             JetExpression receiverExpression = JetPsiFactory.createExpression(state.getProject(), "callableReferenceFakeReceiver");
-
-            Type firstParameterType = signature.getAsmMethod().getArgumentTypes()[0];
-            // 0 is this (the closure class), 1 is the method's first parameter
-            codegen.tempVariables.put(receiverExpression, StackValue.local(1, firstParameterType));
-
+            codegen.tempVariables.put(receiverExpression, receiverParameterStackValue(signature));
             return new ExpressionReceiver(receiverExpression, receiver.getType());
         }
+
+        @NotNull
+        private static StackValue.Local receiverParameterStackValue(@NotNull JvmMethodSignature signature) {
+            // 0 is this (the callable reference class), 1 is the invoke() method's first parameter
+            return StackValue.local(1, signature.getAsmMethod().getArgumentTypes()[0]);
+        }
     }
 
     @Override

compiler/testData/codegen/box/callableReference/local/extension.kt

@@ -0,0 +1,6 @@
+class A
+
+fun box(): String {
+    fun A.foo() = "OK"
+    return A().(A::foo)()
+}

compiler/testData/codegen/box/callableReference/local/extensionToLocalClass.kt

@@ -0,0 +1,5 @@
+fun box(): String {
+    class A
+    fun A.foo() = "OK"
+    return (::A)().(A::foo)()
+}

compiler/testData/codegen/box/callableReference/local/extensionToPrimitive.kt

@@ -0,0 +1,4 @@
+fun box(): String {
+    fun Int.is42With(that: Int) = this + 2 * that == 42
+    return if (16.(Int::is42With)(13)) "OK" else "Fail"
+}

compiler/testData/codegen/box/callableReference/local/extensionWithClosure.kt

@@ -0,0 +1,11 @@
+class A
+
+fun box(): String {
+    var result = "Fail"
+
+    fun A.ext() { result = "OK" }
+
+    val f = A::ext
+    A().f()
+    return result
+}

compiler/tests/org/jetbrains/jet/codegen/generated/BlackBoxCodegenTestGenerated.java

@@ -767,6 +767,26 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
                 doTest("compiler/testData/codegen/box/callableReference/local/enumExtendsTrait.kt");
             }
             
+            @TestMetadata("extension.kt")
+            public void testExtension() throws Exception {
+                doTest("compiler/testData/codegen/box/callableReference/local/extension.kt");
+            }
+            
+            @TestMetadata("extensionToLocalClass.kt")
+            public void testExtensionToLocalClass() throws Exception {
+                doTest("compiler/testData/codegen/box/callableReference/local/extensionToLocalClass.kt");
+            }
+            
+            @TestMetadata("extensionToPrimitive.kt")
+            public void testExtensionToPrimitive() throws Exception {
+                doTest("compiler/testData/codegen/box/callableReference/local/extensionToPrimitive.kt");
+            }
+            
+            @TestMetadata("extensionWithClosure.kt")
+            public void testExtensionWithClosure() throws Exception {
+                doTest("compiler/testData/codegen/box/callableReference/local/extensionWithClosure.kt");
+            }
+            
             @TestMetadata("genericMember.kt")
             public void testGenericMember() throws Exception {
                 doTest("compiler/testData/codegen/box/callableReference/local/genericMember.kt");