Fix for KT-5467: JVM codegen doesn't support capture of many extension receivers

#KT-5467 Fixed

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

@@ -2272,11 +2272,7 @@ public class ExpressionCodegen extends JetVisitor<StackValue, StackValue> implem
 
     @NotNull
     private StackValue generateReceiver(@NotNull CallableDescriptor descriptor) {
-        if (context.getCallableDescriptorWithReceiver() == descriptor) {
-            return context.getReceiverExpression(typeMapper);
-        }
-
-        return context.lookupInContext(descriptor, StackValue.local(0, OBJECT_TYPE), state, false);
+        return context.generateReceiver(descriptor, state, false);
     }
 
     // SCRIPT: generate script, move to ScriptingUtil
@@ -3495,6 +3491,7 @@ public class ExpressionCodegen extends JetVisitor<StackValue, StackValue> implem
     public StackValue visitThisExpression(@NotNull JetThisExpression expression, StackValue receiver) {
         DeclarationDescriptor descriptor = bindingContext.get(REFERENCE_TARGET, expression.getInstanceReference());
         if (descriptor instanceof ClassDescriptor) {
+            //TODO rewrite with context.lookupInContext()
             return StackValue.thisOrOuter(this, (ClassDescriptor) descriptor, false, true);
         }
         if (descriptor instanceof CallableDescriptor) {

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

@@ -1356,8 +1356,12 @@ public class ImplementationBodyCodegen extends ClassBodyCodegen {
                 DeclarationDescriptor descriptor = bindingContext.get(BindingContext.REFERENCE_TARGET, expression.getInstanceReference());
                 assert descriptor instanceof CallableDescriptor ||
                        descriptor instanceof ClassDescriptor : "'This' reference target should be class or callable descriptor but was " + descriptor;
-                if (context.getCallableDescriptorWithReceiver() != descriptor) {
-                    context.lookupInContext(descriptor, null, state, true);
+                if (descriptor instanceof ClassDescriptor) {
+                    context.lookupInContext(descriptor, StackValue.local(0, OBJECT_TYPE), state, true);
+                }
+
+                if (descriptor instanceof CallableDescriptor) {
+                    constructorContext.generateReceiver((CallableDescriptor) descriptor, state, true);
                 }
             }
         };

compiler/backend/src/org/jetbrains/jet/codegen/binding/MutableClosure.java

@@ -33,7 +33,7 @@ public final class MutableClosure implements CalculatedClosure {
     private final ResolvedCall<ConstructorDescriptor> superCall;
 
     private final ClassDescriptor enclosingClass;
-    private final CallableDescriptor enclosingReceiverDescriptor;
+    private final CallableDescriptor enclosingFunWithReceiverDescriptor;
 
     private boolean captureThis;
     private boolean captureReceiver;
@@ -48,7 +48,7 @@ public final class MutableClosure implements CalculatedClosure {
     ) {
         this.enclosingClass = enclosingClass;
         this.superCall = superCall;
-        this.enclosingReceiverDescriptor = enclosingExtensionMemberForClass(classDescriptor);
+        this.enclosingFunWithReceiverDescriptor = enclosingExtensionMemberForClass(classDescriptor);
     }
 
     @Nullable
@@ -85,18 +85,15 @@ public final class MutableClosure implements CalculatedClosure {
     @Override
     public JetType getCaptureReceiverType() {
         if (captureReceiver) {
-            //noinspection ConstantConditions
-            ReceiverParameterDescriptor parameter = enclosingReceiverDescriptor.getExtensionReceiverParameter();
-            assert parameter != null : "Receiver parameter should exist in " + enclosingReceiverDescriptor;
-            return parameter.getType();
+            return getEnclosingReceiverDescriptor().getType();
         }
 
         return null;
     }
 
     public void setCaptureReceiver() {
-        if (enclosingReceiverDescriptor == null) {
-            throw new IllegalStateException();
+        if (enclosingFunWithReceiverDescriptor == null) {
+            throw new IllegalStateException("Extension receiver parameter should exist");
         }
         this.captureReceiver = true;
     }
@@ -127,7 +124,10 @@ public final class MutableClosure implements CalculatedClosure {
         captureVariables.put(value.getDescriptor(), value);
     }
 
-    public CallableDescriptor getEnclosingReceiverDescriptor() {
-        return enclosingReceiverDescriptor;
+    @NotNull
+    public ReceiverParameterDescriptor getEnclosingReceiverDescriptor() {
+        ReceiverParameterDescriptor parameter = enclosingFunWithReceiverDescriptor.getExtensionReceiverParameter();
+        assert parameter != null : "Receiver parameter should exist in " + enclosingFunWithReceiverDescriptor;
+        return parameter;
     }
 }

compiler/backend/src/org/jetbrains/jet/codegen/context/LocalLookup.java

@@ -112,7 +112,7 @@ public interface LocalLookup {
         RECEIVER {
             @Override
             public boolean isCase(DeclarationDescriptor d) {
-                return d instanceof CallableDescriptor;
+                return d instanceof ReceiverParameterDescriptor;
             }
 
             @Override
@@ -123,9 +123,11 @@ public interface LocalLookup {
                     MutableClosure closure,
                     Type classType
             ) {
-                if (closure.getEnclosingReceiverDescriptor() != d) return null;
+                if (closure.getEnclosingReceiverDescriptor() != d) {
+                    return null;
+                }
 
-                JetType receiverType = ((CallableDescriptor) d).getExtensionReceiverParameter().getType();
+                JetType receiverType = closure.getEnclosingReceiverDescriptor().getType();
                 Type type = state.getTypeMapper().mapType(receiverType);
                 StackValue innerValue = StackValue.field(type, classType, CAPTURED_RECEIVER_FIELD, false);
                 closure.setCaptureReceiver();

compiler/backend/src/org/jetbrains/jet/codegen/context/MethodContext.java

@@ -25,14 +25,13 @@ import org.jetbrains.jet.codegen.StackValue;
 import org.jetbrains.jet.codegen.binding.MutableClosure;
 import org.jetbrains.jet.codegen.state.GenerationState;
 import org.jetbrains.jet.codegen.state.JetTypeMapper;
-import org.jetbrains.jet.lang.descriptors.CallableMemberDescriptor;
-import org.jetbrains.jet.lang.descriptors.DeclarationDescriptor;
-import org.jetbrains.jet.lang.descriptors.FunctionDescriptor;
-import org.jetbrains.jet.lang.descriptors.SimpleFunctionDescriptor;
+import org.jetbrains.jet.lang.descriptors.*;
 import org.jetbrains.jet.lang.resolve.java.AsmTypeConstants;
 import org.jetbrains.org.objectweb.asm.Label;
 import org.jetbrains.org.objectweb.asm.Type;
 
+import static org.jetbrains.jet.lang.resolve.java.AsmTypeConstants.OBJECT_TYPE;
+
 public class MethodContext extends CodegenContext<CallableMemberDescriptor> {
     private final boolean isInliningLambda;
     private Label methodStartLabel;
@@ -72,6 +71,15 @@ public class MethodContext extends CodegenContext<CallableMemberDescriptor> {
         return getParentContext().lookupInContext(d, result, state, ignoreNoOuter);
     }
 
+    @Nullable
+    public StackValue generateReceiver(@NotNull CallableDescriptor descriptor, @NotNull GenerationState state, boolean ignoreNoOuter) {
+        if (getCallableDescriptorWithReceiver() == descriptor) {
+            return getReceiverExpression(state.getTypeMapper());
+        }
+        ReceiverParameterDescriptor parameter = descriptor.getExtensionReceiverParameter();
+        return lookupInContext(parameter, StackValue.local(0, OBJECT_TYPE), state, ignoreNoOuter);
+    }
+
     @Override
     public boolean isStatic() {
         return getParentContext().isStatic();

compiler/testData/codegen/box/extensionFunctions/kt5467.kt

@@ -0,0 +1,20 @@
+fun String.foo() : String {
+    fun Int.bar() : String {
+        fun Long.baz() : String {
+            val x = this@foo
+            val y = this@bar
+            val z = this@baz
+            return "$x $y $z"
+        }
+        return 0L.baz()
+    }
+    return 42.bar()
+}
+
+fun box() : String {
+    val result = "OK".foo()
+
+    if (result != "OK 42 0") return "fail: $result"
+
+    return "OK"
+}

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

@@ -2813,6 +2813,12 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
             doTest(fileName);
         }
         
+        @TestMetadata("kt5467.kt")
+        public void testKt5467() throws Exception {
+            String fileName = JetTestUtils.navigationMetadata("compiler/testData/codegen/box/extensionFunctions/kt5467.kt");
+            doTest(fileName);
+        }
+        
         @TestMetadata("kt606.kt")
         public void testKt606() throws Exception {
             String fileName = JetTestUtils.navigationMetadata("compiler/testData/codegen/box/extensionFunctions/kt606.kt");