Moving local fun callee generation to CallReceiver

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

@@ -2277,15 +2277,6 @@ public class ExpressionCodegen extends JetVisitor<StackValue, StackValue> implem
     ) {
         CallableDescriptor descriptor = resolvedCall.getResultingDescriptor();
 
-        //generate callee for local function
-        Type calleeType = callableMethod.getGenerateCalleeType();
-        if (calleeType != null) {
-            StackValue value = findLocalOrCapturedValue(descriptor.getOriginal());
-            assert value != null : "Local fun should be found in locals or in captured params: " + resolvedCall;
-            value.put(calleeType, v);
-        }
-
-        //generate this (absent for local fun) and receiver
         if (!(descriptor instanceof ConstructorDescriptor)) { // otherwise already
             receiver = StackValue.receiver(resolvedCall, receiver, this, callableMethod);
             receiver.put(receiver.type, v);

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

@@ -330,12 +330,16 @@ public abstract class StackValue {
             ExpressionCodegen codegen,
             @Nullable CallableMethod callableMethod
     ) {
-        if (resolvedCall.getThisObject().exists() || resolvedCall.getReceiverArgument().exists()) {
+        if (resolvedCall.getThisObject().exists() || resolvedCall.getReceiverArgument().exists() || isLocalFunCall(callableMethod)) {
             return new CallReceiver(resolvedCall, receiver, codegen, callableMethod, true);
         }
         return receiver;
     }
 
+    private static boolean isLocalFunCall(@Nullable CallableMethod callableMethod) {
+        return callableMethod != null && callableMethod.getGenerateCalleeType() != null;
+    }
+
     public static StackValue receiverWithoutReceiverArgument(StackValue receiverWithParameter) {
         if (receiverWithParameter instanceof CallReceiver) {
             CallReceiver callReceiver = (CallReceiver) receiverWithParameter;
@@ -1139,7 +1143,11 @@ public abstract class StackValue {
                 else {
                     return codegen.typeMapper.mapType(descriptor.getExpectedThisObject().getType());
                 }
-            } else {
+            }
+            else if (isLocalFunCall(callableMethod)) {
+                return callableMethod.getGenerateCalleeType();
+            }
+            else {
                 return Type.VOID_TYPE;
             }
         }
@@ -1162,6 +1170,13 @@ public abstract class StackValue {
                     genReceiver(v, thisObject, type, null, 0);
                 }
 
+                depth = 1;
+            } else if (isLocalFunCall(callableMethod)) {
+                assert receiver == none() || receiverArgument.exists(): "Receiver should be present only for local extension function: " + callableMethod;
+                StackValue value = codegen.findLocalOrCapturedValue(descriptor.getOriginal());
+                assert value != null : "Local fun should be found in locals or in captured params: " + resolvedCall;
+                value.put(callableMethod.getGenerateCalleeType(), v);
+
                 depth = 1;
             }
             else {

compiler/testData/codegen/box/functions/localFunctions/kt4514.kt

@@ -0,0 +1,9 @@
+fun box(): String {
+    fun String.f() = this
+    val vf: String.() -> String = { this }
+
+    val localExt = "O".f() + "K"?.f()
+    if (localExt != "OK") return "localExt $localExt"
+
+    return "O".vf() + "K"?.vf()
+}

compiler/testData/codegen/box/functions/localFunctions/kt4784.kt

@@ -0,0 +1,20 @@
+open class T(var value: Int) {}
+
+fun plusAssign(): T {
+
+    fun T.plusAssign(s: Int) {
+        value += s
+    }
+
+    var t  = T(1)
+    t += 1
+
+    return t
+}
+
+fun box(): String {
+    val result = plusAssign().value
+    if (result != 2) return "fail 1: $result"
+
+    return "OK"
+}

compiler/testData/codegen/box/functions/localFunctions/localExtensionOnNullableParameter.kt

@@ -0,0 +1,21 @@
+open class T(var value: Int) {}
+
+fun localExtensionOnNullableParameter(): T {
+
+    fun T.local(s: Int) {
+        value += s
+    }
+
+    var t: T? = T(1)
+    t?.local(2)
+
+    return t!!
+}
+
+
+fun box(): String {
+    val result = localExtensionOnNullableParameter().value
+    if (result != 3) return "fail 2: $result"
+
+    return "OK"
+}

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

@@ -2961,6 +2961,11 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
                 doTest("compiler/testData/codegen/box/functions/localFunctions/kt4119_2.kt");
             }
             
+            @TestMetadata("kt4514.kt")
+            public void testKt4514() throws Exception {
+                doTest("compiler/testData/codegen/box/functions/localFunctions/kt4514.kt");
+            }
+            
             @TestMetadata("kt4777.kt")
             public void testKt4777() throws Exception {
                 doTest("compiler/testData/codegen/box/functions/localFunctions/kt4777.kt");
@@ -2971,6 +2976,16 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
                 doTest("compiler/testData/codegen/box/functions/localFunctions/kt4783.kt");
             }
             
+            @TestMetadata("kt4784.kt")
+            public void testKt4784() throws Exception {
+                doTest("compiler/testData/codegen/box/functions/localFunctions/kt4784.kt");
+            }
+            
+            @TestMetadata("localExtensionOnNullableParameter.kt")
+            public void testLocalExtensionOnNullableParameter() throws Exception {
+                doTest("compiler/testData/codegen/box/functions/localFunctions/localExtensionOnNullableParameter.kt");
+            }
+            
             @TestMetadata("localFunctionInConstructor.kt")
             public void testLocalFunctionInConstructor() throws Exception {
                 doTest("compiler/testData/codegen/box/functions/localFunctions/localFunctionInConstructor.kt");