Fix and refactor ExpressionCodegen#pushClosureOnStack

- change Closure parameter to ClassDescriptor, load closure in the beginning
- invert and rename boolean parameter, also use it only to prevent pushing
  dispatch receiver (not the extension receiver) onto the stack because the
  only non-trivial usage of it is preceded by manual handling of the dispatch
  (not extension) receiver. This fixes innerOfLocalCaptureExtensionReceiver.kt

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

@@ -55,6 +55,7 @@ import static org.jetbrains.org.objectweb.asm.Opcodes.*;
 public class ClosureCodegen extends ParentCodegenAware {
     private final PsiElement fun;
     private final FunctionDescriptor funDescriptor;
+    private final ClassDescriptor classDescriptor;
     private final SamType samType;
     private final JetType superClassType;
     private final List<JetType> superInterfaceTypes;
@@ -87,7 +88,7 @@ public class ClosureCodegen extends ParentCodegenAware {
         this.syntheticClassKind = syntheticClassKind;
         this.strategy = strategy;
 
-        ClassDescriptor classDescriptor = anonymousClassForFunction(bindingContext, funDescriptor);
+        this.classDescriptor = anonymousClassForFunction(bindingContext, funDescriptor);
 
         if (samType == null) {
             this.superInterfaceTypes = new ArrayList<JetType>();
@@ -196,7 +197,7 @@ public class ClosureCodegen extends ParentCodegenAware {
                                    null);
 
 
-        AsmUtil.writeOuterClassAndEnclosingMethod(anonymousClassForFunction(bindingContext, funDescriptor), funDescriptor, typeMapper, cv);
+        AsmUtil.writeOuterClassAndEnclosingMethod(classDescriptor, funDescriptor, typeMapper, cv);
         cv.done();
     }
 
@@ -209,7 +210,7 @@ public class ClosureCodegen extends ParentCodegenAware {
             v.anew(asmType);
             v.dup();
 
-            codegen.pushClosureOnStack(closure, false, codegen.defaultCallGenerator);
+            codegen.pushClosureOnStack(classDescriptor, true, codegen.defaultCallGenerator);
             v.invokespecial(asmType.getInternalName(), "<init>", constructor.getDescriptor(), false);
         }
         return StackValue.onStack(asmType);

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

@@ -29,7 +29,6 @@ import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.jet.codegen.binding.CalculatedClosure;
 import org.jetbrains.jet.codegen.binding.CodegenBinding;
-import org.jetbrains.jet.codegen.binding.MutableClosure;
 import org.jetbrains.jet.codegen.context.*;
 import org.jetbrains.jet.codegen.inline.InlineCodegen;
 import org.jetbrains.jet.codegen.inline.InlineCodegenUtil;
@@ -185,7 +184,7 @@ public class ExpressionCodegen extends JetVisitor<StackValue, StackValue> implem
     }
 
     @NotNull
-    public CalculatedClosure generateObjectLiteral(@NotNull JetObjectLiteralExpression literal) {
+    public ClassDescriptor generateObjectLiteral(@NotNull JetObjectLiteralExpression literal) {
         JetObjectDeclaration objectDeclaration = literal.getObjectDeclaration();
 
         ClassDescriptor classDescriptor = bindingContext.get(CLASS, objectDeclaration);
@@ -198,13 +197,10 @@ public class ExpressionCodegen extends JetVisitor<StackValue, StackValue> implem
                 literal.getContainingFile()
         );
 
-
         ClassContext objectContext = context.intoAnonymousClass(classDescriptor, this, OwnerKind.IMPLEMENTATION);
-
         new ImplementationBodyCodegen(objectDeclaration, objectContext, classBuilder, state, getParentCodegen()).generate();
 
-        //noinspection ConstantConditions
+        return classDescriptor;
-        return bindingContext.get(CLOSURE, classDescriptor);
     }
 
     @NotNull
@@ -1342,20 +1338,16 @@ public class ExpressionCodegen extends JetVisitor<StackValue, StackValue> implem
 
     @Override
     public StackValue visitObjectLiteralExpression(@NotNull JetObjectLiteralExpression expression, StackValue receiver) {
-        CalculatedClosure closure = generateObjectLiteral(expression);
+        ClassDescriptor classDescriptor = generateObjectLiteral(expression);
-
+        Type type = typeMapper.mapType(classDescriptor);
-        ConstructorDescriptor constructorDescriptor = bindingContext.get(CONSTRUCTOR, expression.getObjectDeclaration());
-        assert constructorDescriptor != null : "Unresolved constructor: " + expression.getText();
-        JvmMethodSignature constructor = typeMapper.mapSignature(constructorDescriptor);
-
-        Type type = typeMapper.mapType(constructorDescriptor.getContainingDeclaration());
 
         v.anew(type);
         v.dup();
 
-        pushClosureOnStack(closure, false, defaultCallGenerator);
+        pushClosureOnStack(classDescriptor, true, defaultCallGenerator);
 
-        ResolvedCall<ConstructorDescriptor> superCall = closure.getSuperCall();
+        //noinspection ConstantConditions
+        ResolvedCall<ConstructorDescriptor> superCall = bindingContext.get(CLOSURE, classDescriptor).getSuperCall();
         if (superCall != null) {
             // For an anonymous object, we should also generate all non-default arguments that it captures for its super call
             ConstructorDescriptor superConstructor = superCall.getResultingDescriptor();
@@ -1363,7 +1355,7 @@ public class ExpressionCodegen extends JetVisitor<StackValue, StackValue> implem
             int params = superValueParameters.size();
             List<Type> superMappedTypes = typeMapper.mapToCallableMethod(superConstructor).getValueParameterTypes();
             assert superMappedTypes.size() >= params : String.format("Incorrect number of mapped parameters vs arguments: %d < %d for %s",
-                                                                     superMappedTypes.size(), params, constructorDescriptor);
+                                                                     superMappedTypes.size(), params, classDescriptor);
 
             List<ResolvedValueArgument> valueArguments = new ArrayList<ResolvedValueArgument>(params);
             List<ValueParameterDescriptor> valueParameters = new ArrayList<ValueParameterDescriptor>(params);
@@ -1381,33 +1373,33 @@ public class ExpressionCodegen extends JetVisitor<StackValue, StackValue> implem
             argumentGenerator.generate(valueArguments);
         }
 
+        ConstructorDescriptor constructorDescriptor = bindingContext.get(CONSTRUCTOR, expression.getObjectDeclaration());
+        assert constructorDescriptor != null : "Unresolved constructor: " + expression.getText();
+        JvmMethodSignature constructor = typeMapper.mapSignature(constructorDescriptor);
         v.invokespecial(type.getInternalName(), "<init>", constructor.getAsmMethod().getDescriptor(), false);
         return StackValue.onStack(type);
     }
 
-    public void pushClosureOnStack(
+    public void pushClosureOnStack(@NotNull ClassDescriptor classDescriptor, boolean putThis, @NotNull CallGenerator callGenerator) {
-            @Nullable CalculatedClosure closure,
+        CalculatedClosure closure = bindingContext.get(CLOSURE, classDescriptor);
-            boolean ignoreThisAndReceiver,
-            @NotNull CallGenerator callGenerator
-    ) {
         if (closure == null) return;
 
         int paramIndex = 0;
-        if (!ignoreThisAndReceiver) {
+
+        if (putThis) {
             ClassDescriptor captureThis = closure.getCaptureThis();
             if (captureThis != null) {
                 StackValue thisOrOuter = generateThisOrOuter(captureThis, false);
-
                 assert !isPrimitive(thisOrOuter.type) : "This or outer should be non primitive: " + thisOrOuter.type;
                 callGenerator.putCapturedValueOnStack(thisOrOuter, thisOrOuter.type, paramIndex++);
             }
+        }
 
-            JetType captureReceiver = closure.getCaptureReceiverType();
+        JetType captureReceiver = closure.getCaptureReceiverType();
-            if (captureReceiver != null) {
+        if (captureReceiver != null) {
-                Type asmType = typeMapper.mapType(captureReceiver);
+            Type asmType = typeMapper.mapType(captureReceiver);
-                StackValue.Local capturedReceiver = StackValue.local(AsmUtil.getReceiverIndex(context, context.getContextDescriptor()), asmType);
+            StackValue.Local capturedReceiver = StackValue.local(AsmUtil.getReceiverIndex(context, context.getContextDescriptor()), asmType);
-                callGenerator.putCapturedValueOnStack(capturedReceiver, capturedReceiver.type, paramIndex++);
+            callGenerator.putCapturedValueOnStack(capturedReceiver, capturedReceiver.type, paramIndex++);
-            }
         }
 
         for (Map.Entry<DeclarationDescriptor, EnclosedValueDescriptor> entry : closure.getCaptureVariables().entrySet()) {
@@ -1421,8 +1413,7 @@ public class ExpressionCodegen extends JetVisitor<StackValue, StackValue> implem
 
         ResolvedCall<ConstructorDescriptor> superCall = closure.getSuperCall();
         if (superCall != null) {
-            MutableClosure superClosure = bindingContext.get(CLOSURE, superCall.getResultingDescriptor().getContainingDeclaration());
+            pushClosureOnStack(superCall.getResultingDescriptor().getContainingDeclaration(), putThis, callGenerator);
-            pushClosureOnStack(superClosure, ignoreThisAndReceiver, callGenerator);
         }
     }
 
@@ -3339,11 +3330,9 @@ public class ExpressionCodegen extends JetVisitor<StackValue, StackValue> implem
             generateReceiverValue(resolvedCall.getDispatchReceiver(), receiverType);
         }
 
-        MutableClosure closure = bindingContext.get(CLOSURE, constructor.getContainingDeclaration());
-
         // Resolved call to local class constructor doesn't have dispatchReceiver, so we need to generate closure on stack
         // See StackValue.receiver for more info
-        pushClosureOnStack(closure, dispatchReceiver != null, defaultCallGenerator);
+        pushClosureOnStack(constructor.getContainingDeclaration(), dispatchReceiver == null, defaultCallGenerator);
 
         ConstructorDescriptor originalOfSamAdapter = (ConstructorDescriptor) SamCodegenUtil.getOriginalIfSamAdapter(constructor);
         CallableMethod method = typeMapper.mapToCallableMethod(originalOfSamAdapter == null ? constructor : originalOfSamAdapter);

compiler/backend/src/org/jetbrains/jet/codegen/inline/InlineCodegen.java

@@ -31,7 +31,6 @@ import org.jetbrains.jet.descriptors.serialization.descriptors.DeserializedSimpl
 import org.jetbrains.jet.lang.descriptors.*;
 import org.jetbrains.jet.lang.descriptors.impl.AnonymousFunctionDescriptor;
 import org.jetbrains.jet.lang.psi.*;
-import org.jetbrains.jet.lang.resolve.BindingContext;
 import org.jetbrains.jet.lang.resolve.DescriptorToSourceUtils;
 import org.jetbrains.jet.lang.resolve.DescriptorUtils;
 import org.jetbrains.jet.lang.resolve.calls.model.ResolvedCall;
@@ -54,13 +53,13 @@ import java.util.List;
 import java.util.ListIterator;
 import java.util.Map;
 
-import static org.jetbrains.jet.codegen.AsmUtil.*;
+import static org.jetbrains.jet.codegen.AsmUtil.getMethodAsmFlags;
+import static org.jetbrains.jet.codegen.AsmUtil.isPrimitive;
 import static org.jetbrains.jet.codegen.inline.InlineCodegenUtil.addInlineMarker;
 
 public class InlineCodegen implements CallGenerator {
     private final GenerationState state;
     private final JetTypeMapper typeMapper;
-    private final BindingContext bindingContext;
 
     private final SimpleFunctionDescriptor functionDescriptor;
     private final JvmMethodSignature jvmSignature;
@@ -90,7 +89,6 @@ public class InlineCodegen implements CallGenerator {
         this.codegen = codegen;
         this.callElement = callElement;
         this.functionDescriptor = functionDescriptor.getOriginal();
-        bindingContext = codegen.getBindingContext();
         initialFrameSize = codegen.getFrameMap().getCurrentSize();
 
         context = (MethodContext) getContext(functionDescriptor, state);
@@ -418,7 +416,7 @@ public class InlineCodegen implements CallGenerator {
     private void putClosureParametersOnStack() {
         for (LambdaInfo next : expressionMap.values()) {
             activeLambda = next;
-            codegen.pushClosureOnStack(next.closure, false, this);
+            codegen.pushClosureOnStack(next.getClassDescriptor(), true, this);
         }
         activeLambda = null;
     }

compiler/backend/src/org/jetbrains/jet/codegen/inline/LambdaInfo.java

@@ -18,9 +18,6 @@ package org.jetbrains.jet.codegen.inline;
 
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
-import org.jetbrains.org.objectweb.asm.Type;
-import org.jetbrains.org.objectweb.asm.tree.FieldInsnNode;
-import org.jetbrains.org.objectweb.asm.tree.MethodNode;
 import org.jetbrains.jet.codegen.AsmUtil;
 import org.jetbrains.jet.codegen.binding.CalculatedClosure;
 import org.jetbrains.jet.codegen.context.EnclosedValueDescriptor;
@@ -32,12 +29,15 @@ import org.jetbrains.jet.lang.psi.JetFunctionLiteral;
 import org.jetbrains.jet.lang.psi.JetFunctionLiteralExpression;
 import org.jetbrains.jet.lang.resolve.BindingContext;
 import org.jetbrains.jet.lang.resolve.java.AsmTypeConstants;
+import org.jetbrains.org.objectweb.asm.Type;
+import org.jetbrains.org.objectweb.asm.tree.FieldInsnNode;
+import org.jetbrains.org.objectweb.asm.tree.MethodNode;
 
-import java.util.*;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.List;
 
-import static org.jetbrains.jet.codegen.binding.CodegenBinding.CLOSURE;
+import static org.jetbrains.jet.codegen.binding.CodegenBinding.*;
-import static org.jetbrains.jet.codegen.binding.CodegenBinding.anonymousClassForFunction;
-import static org.jetbrains.jet.codegen.binding.CodegenBinding.asmTypeForAnonymousClass;
 
 public class LambdaInfo implements CapturedParamOwner, LabelOwner {
 
@@ -49,7 +49,7 @@ public class LambdaInfo implements CapturedParamOwner, LabelOwner {
     @Nullable
     public final String labelName;
 
-    public final CalculatedClosure closure;
+    private final CalculatedClosure closure;
 
     private MethodNode node;
 

compiler/testData/codegen/box/localClasses/innerOfLocalCaptureExtensionReceiver.kt

@@ -0,0 +1,15 @@
+fun String.bar(): String {
+    open class Local {
+        fun result() = this@bar
+    }
+
+    class Outer {
+        inner class Inner : Local() {
+            fun outer() = this@Outer
+        }
+    }
+
+    return Outer().Inner().result()
+}
+
+fun box() = "OK".bar()

compiler/testData/codegen/box/localClasses/localClassCaptureExtensionReceiver.kt

@@ -0,0 +1,14 @@
+class Outer {
+    fun String.id(): String {
+        class Local(unused: Long) {
+            fun result() = this@id
+            fun outer() = this@Outer
+        }
+
+        return Local(42L).result()
+    }
+
+    fun result(): String = "OK".id()
+}
+
+fun box() = Outer().result()

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

@@ -4191,6 +4191,12 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
             doTest(fileName);
         }
         
+        @TestMetadata("innerOfLocalCaptureExtensionReceiver.kt")
+        public void testInnerOfLocalCaptureExtensionReceiver() throws Exception {
+            String fileName = JetTestUtils.navigationMetadata("compiler/testData/codegen/box/localClasses/innerOfLocalCaptureExtensionReceiver.kt");
+            doTest(fileName);
+        }
+        
         @TestMetadata("kt2700.kt")
         public void testKt2700() throws Exception {
             String fileName = JetTestUtils.navigationMetadata("compiler/testData/codegen/box/localClasses/kt2700.kt");
@@ -4227,6 +4233,12 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
             doTest(fileName);
         }
         
+        @TestMetadata("localClassCaptureExtensionReceiver.kt")
+        public void testLocalClassCaptureExtensionReceiver() throws Exception {
+            String fileName = JetTestUtils.navigationMetadata("compiler/testData/codegen/box/localClasses/localClassCaptureExtensionReceiver.kt");
+            doTest(fileName);
+        }
+        
         @TestMetadata("localClassInInitializer.kt")
         public void testLocalClassInInitializer() throws Exception {
             String fileName = JetTestUtils.navigationMetadata("compiler/testData/codegen/box/localClasses/localClassInInitializer.kt");