Support nested chain of non-inline lambdas inside inline one

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

@@ -37,10 +37,18 @@ public class ConstructorInvocation {
 
     private Map<String, LambdaInfo> capturedLambdasToInline;
 
-    ConstructorInvocation(String ownerInternalName, Map<Integer, LambdaInfo> lambdasToInline, boolean isSameModule) {
+    private boolean capturedOuterRegenerated;
+
+    ConstructorInvocation(
+            String ownerInternalName,
+            Map<Integer, LambdaInfo> lambdasToInline,
+            boolean isSameModule,
+            boolean capturedOuterRegenerated
+    ) {
         this.ownerInternalName = ownerInternalName;
         this.lambdasToInline = lambdasToInline;
         this.isSameModule = isSameModule;
+        this.capturedOuterRegenerated = capturedOuterRegenerated;
     }
 
     public String getOwnerInternalName() {
@@ -48,7 +56,7 @@ public class ConstructorInvocation {
     }
 
     public boolean shouldRegenerate() {
-        return !lambdasToInline.isEmpty() || !isSameModule;
+        return !lambdasToInline.isEmpty() || !isSameModule || capturedOuterRegenerated;
     }
 
     public Map<Integer, LambdaInfo> getLambdasToInline() {
@@ -86,4 +94,5 @@ public class ConstructorInvocation {
     public void setCapturedLambdasToInline(Map<String, LambdaInfo> capturedLambdasToInline) {
         this.capturedLambdasToInline = capturedLambdasToInline;
     }
+
 }

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

@@ -211,7 +211,7 @@ public class InlineCodegen implements ParentCodegenAware, CallGenerator {
         InliningContext info =
                 new InliningContext(expressionMap, null, null, null, state,
                                  codegen.getInlineNameGenerator().subGenerator(functionDescriptor.getName().asString()),
-                                 codegen.getContext(), call, Collections.<String, String>emptyMap(), false);
+                                 codegen.getContext(), call, Collections.<String, String>emptyMap(), false, false);
 
         MethodInliner inliner = new MethodInliner(node, parameters, info, null, new LambdaFieldRemapper(), isSameModule); //with captured
 

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

@@ -47,6 +47,8 @@ public class InliningContext {
 
     public final boolean isInliningLambda;
 
+    public final boolean classRegeneration;
+
     public InliningContext(
             Map<Integer, LambdaInfo> map,
             List<InvokeCall> accesses,
@@ -57,7 +59,8 @@ public class InliningContext {
             CodegenContext startContext,
             Call call,
             Map<String, String> typeMapping,
-            boolean isInliningLambda
+            boolean isInliningLambda,
+            boolean classRegeneration
     ) {
         expressionMap = map;
         invokeCalls = accesses;
@@ -69,14 +72,17 @@ public class InliningContext {
         this.call = call;
         this.typeMapping = typeMapping;
         this.isInliningLambda = isInliningLambda;
+        this.classRegeneration = classRegeneration;
     }
 
     public InliningContext subInline(NameGenerator generator) {
         return subInline(generator, Collections.<String, String>emptyMap());
     }
 
-    public InliningContext subInlineLambda(NameGenerator generator) {
+    public InliningContext subInlineLambda(LambdaInfo lambdaInfo) {
-        return subInline(generator, Collections.<String, String>emptyMap(), true);
+        Map<String, String> map = new HashMap();
+        map.put(lambdaInfo.getLambdaClassType().getInternalName(), null); //mark lambda inlined
+        return subInline(nameGenerator.subGenerator("lambda"), map, true);
     }
 
     public InliningContext subInline(NameGenerator generator, Map<String, String> additionalTypeMappings) {
@@ -87,7 +93,11 @@ public class InliningContext {
         Map<String, String> newTypeMappings = new HashMap<String, String>(typeMapping);
         newTypeMappings.putAll(additionalTypeMappings);
         return new InliningContext(expressionMap, invokeCalls, constructorInvocation, remapper, state, generator, startContext, call,
-                                newTypeMappings, isInliningLambda);
+                                newTypeMappings, isInliningLambda, classRegeneration);
     }
 
+    public InliningContext classRegeneration() {
+        return new InliningContext(expressionMap, invokeCalls, constructorInvocation, remapper, state, nameGenerator, startContext, call,
+                                   typeMapping, isInliningLambda, true);
+    }
 }

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

@@ -55,7 +55,7 @@ public class LambdaTransformer {
 
     private final MethodNode bridge;
 
-    private final InliningContext info;
+    private final InliningContext inliningContext;
 
     private final Type oldLambdaType;
 
@@ -67,11 +67,11 @@ public class LambdaTransformer {
     private String[] interfaces;
     private final boolean isSameModule;
 
-    public LambdaTransformer(String lambdaInternalName, InliningContext info, boolean isSameModule, Type newLambdaType) {
+    public LambdaTransformer(String lambdaInternalName, InliningContext inliningContext, boolean isSameModule, Type newLambdaType) {
         this.isSameModule = isSameModule;
-        this.state = info.state;
+        this.state = inliningContext.state;
         this.typeMapper = state.getTypeMapper();
-        this.info = info;
+        this.inliningContext = inliningContext;
         this.oldLambdaType = Type.getObjectType(lambdaInternalName);
         this.newLambdaType = newLambdaType;
 
@@ -130,7 +130,7 @@ public class LambdaTransformer {
         MethodVisitor invokeVisitor = newMethod(classBuilder, invoke);
         RegeneratedLambdaFieldRemapper
                 remapper = new RegeneratedLambdaFieldRemapper(oldLambdaType.getInternalName(), newLambdaType.getInternalName(), parameters, invocation.getCapturedLambdasToInline());
-        MethodInliner inliner = new MethodInliner(invoke, parameters, info.subInline(info.nameGenerator.subGenerator("lambda")), oldLambdaType,
+        MethodInliner inliner = new MethodInliner(invoke, parameters, inliningContext.subInline(inliningContext.nameGenerator.subGenerator("lambda")), oldLambdaType,
                                                   remapper, isSameModule);
         InlineResult result = inliner.doInline(invokeVisitor, new VarRemapper.ParamRemapper(parameters, 0), remapper, false);
         invokeVisitor.visitMaxs(-1, -1);
@@ -182,8 +182,8 @@ public class LambdaTransformer {
     }
 
     private ClassBuilder createClassBuilder() {
-        return new RemappingClassBuilder(state.getFactory().forLambdaInlining(newLambdaType, info.call.getCallElement().getContainingFile()),
+        return new RemappingClassBuilder(state.getFactory().forLambdaInlining(newLambdaType, inliningContext.call.getCallElement().getContainingFile()),
-                     new TypeRemapper(info.typeMapping));
+                     new TypeRemapper(inliningContext.typeMapping));
     }
 
     private static MethodVisitor newMethod(ClassBuilder builder, MethodNode original) {

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

@@ -13,6 +13,7 @@ import org.jetbrains.asm4.commons.Method;
 import org.jetbrains.asm4.commons.RemappingMethodAdapter;
 import org.jetbrains.asm4.tree.*;
 import org.jetbrains.asm4.tree.analysis.*;
+import org.jetbrains.jet.codegen.AsmUtil;
 import org.jetbrains.jet.codegen.ClosureCodegen;
 import org.jetbrains.jet.codegen.StackValue;
 import org.jetbrains.jet.codegen.state.JetTypeMapper;
@@ -28,7 +29,7 @@ public class MethodInliner {
 
     private final Parameters parameters;
 
-    private final InliningContext parent;
+    private final InliningContext inliningContext;
 
     @Nullable
     private final Type lambdaType;
@@ -52,7 +53,7 @@ public class MethodInliner {
      *
      * @param node
      * @param parameters
-     * @param parent
+     * @param inliningContext
      * @param lambdaType - in case on lambda 'invoke' inlining
      */
     public MethodInliner(
@@ -65,7 +66,7 @@ public class MethodInliner {
     ) {
         this.node = node;
         this.parameters = parameters;
-        this.parent = parent;
+        this.inliningContext = parent;
         this.lambdaType = lambdaType;
         this.lambdaFieldRemapper = lambdaFieldRemapper;
         this.isSameModule = isSameModule;
@@ -125,15 +126,16 @@ public class MethodInliner {
 
                     if (invocation.shouldRegenerate()) {
                         //TODO: need poping of type but what to do with local funs???
-                        Type newLambdaType = Type.getObjectType(parent.nameGenerator.genLambdaClassName());
+                        Type newLambdaType = Type.getObjectType(inliningContext.nameGenerator.genLambdaClassName());
                         currentTypeMapping.put(invocation.getOwnerInternalName(), newLambdaType.getInternalName());
                         LambdaTransformer transformer = new LambdaTransformer(invocation.getOwnerInternalName(),
-                                                                              parent.subInline(parent.nameGenerator, currentTypeMapping),
+                                                                              inliningContext.subInline(inliningContext.nameGenerator, currentTypeMapping).classRegeneration(),
                                                                               isSameModule, newLambdaType);
 
-                        transformer.doTransform(invocation);
+                        InlineResult transformResult = transformer.doTransform(invocation);
+                        result.addAllClassesToRemove(transformResult);
 
-                        if (parent.isInliningLambda) {
+                        if (inliningContext.isInliningLambda) {
                             //this class is transformed and original not used so we should remove original one after inlining
                             result.addClassToRemove(invocation.getOwnerInternalName());
                         }
@@ -163,8 +165,9 @@ public class MethodInliner {
                     Parameters lambdaParameters = info.addAllParameters(capturedRemapper);
 
                     setInlining(true);
-                    MethodInliner inliner = new MethodInliner(info.getNode(), lambdaParameters, parent.subInlineLambda(
+                    MethodInliner inliner = new MethodInliner(info.getNode(), lambdaParameters,
-                            parent.nameGenerator.subGenerator("lambda")), info.getLambdaClassType(),
+                                                              inliningContext.subInlineLambda(info),
+                                                              info.getLambdaClassType(),
                                                               capturedRemapper, true /*cause all calls in same module as lambda*/
                     );
 
@@ -331,7 +334,7 @@ public class MethodInliner {
                             }
                         }
 
-                        constructorInvocations.add(new ConstructorInvocation(owner, lambdaMapping, isSameModule));
+                        constructorInvocations.add(new ConstructorInvocation(owner, lambdaMapping, isSameModule, inliningContext.classRegeneration));
                     }
                 }
             }
@@ -423,6 +426,34 @@ public class MethodInliner {
                     } else {
                         cur = this.lambdaFieldRemapper.doTransform(node, fieldInsnNode, result);
                     }
+                } else {
+                    Type type1 = Type.getType(fieldInsnNode.desc);
+                    if (!AsmUtil.isPrimitive(type1)) {
+                        String type = type1.getInternalName();
+                        if (inliningContext.typeMapping.containsKey(type)) {
+                            //TODO value could be null
+                            String newTypeOrSkip = inliningContext.typeMapping.get(type);
+                            if (newTypeOrSkip != null) {
+                                fieldInsnNode.owner = newTypeOrSkip;
+                            }
+                            else {
+                                //generate owner of next instruction
+                                AbstractInsnNode previous = fieldInsnNode.getPrevious();
+                                AbstractInsnNode nextInstruction = fieldInsnNode.getNext();
+                                if (!(nextInstruction instanceof FieldInsnNode)) {
+                                    throw new IllegalStateException(
+                                            "Instruction after inlined one should be field access: " + nextInstruction);
+                                }
+                                if (!(previous instanceof FieldInsnNode)) {
+                                    throw new IllegalStateException("Instruction before inlined one should be field access: " + previous);
+                                }
+                                cur = nextInstruction;
+                                node.instructions.remove(cur.getPrevious());
+                                ((FieldInsnNode) cur).owner = Type.getType(((FieldInsnNode) previous).desc).getInternalName();
+                                ((FieldInsnNode) cur).name = LambdaTransformer.getNewFieldName(((FieldInsnNode) cur).name);
+                            }
+                        }
+                    }
                 }
             }
             cur = cur.getNext();