Lambda captured field remapper

compiler/backend/src/org/jetbrains/jet/codegen/asm/CapturedParamInfo.java

@@ -26,6 +26,8 @@ public class CapturedParamInfo extends ParameterInfo {
 
     private int shift = 0;
 
+    private LambdaInfo recapturedFrom;
+
     public static final CapturedParamInfo STUB = new CapturedParamInfo("", AsmTypeConstants.OBJECT_TYPE, true, -1, -1);
 
     public CapturedParamInfo(@NotNull String fieldName, @NotNull Type type, boolean skipped, int remapIndex, int index) {
@@ -45,4 +47,12 @@ public class CapturedParamInfo extends ParameterInfo {
     public void setShift(int shift) {
         this.shift = shift;
     }
+
+    public LambdaInfo getRecapturedFrom() {
+        return recapturedFrom;
+    }
+
+    public void setRecapturedFrom(LambdaInfo recapturedFrom) {
+        this.recapturedFrom = recapturedFrom;
+    }
 }

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

@@ -33,6 +33,8 @@ public class ConstructorInvocation {
 
     private List<CapturedParamInfo> recaptured;
 
+    private Map<String, LambdaInfo> recapturedLambdas;
+
     ConstructorInvocation(String ownerInternalName, Map<Integer, InlinableAccess> access) {
         this.ownerInternalName = ownerInternalName;
         this.access = access;
@@ -74,4 +76,12 @@ public class ConstructorInvocation {
     public void setRecaptured(List<CapturedParamInfo> recaptured) {
         this.recaptured = recaptured;
     }
+
+    public Map<String, LambdaInfo> getRecapturedLambdas() {
+        return recapturedLambdas;
+    }
+
+    public void setRecapturedLambdas(Map<String, LambdaInfo> recapturedLambdas) {
+        this.recapturedLambdas = recapturedLambdas;
+    }
 }

compiler/backend/src/org/jetbrains/jet/codegen/asm/InlineFieldRemapper.java

@@ -16,30 +16,47 @@
 
 package org.jetbrains.jet.codegen.asm;
 
+import org.jetbrains.annotations.Nullable;
 import org.jetbrains.asm4.Opcodes;
 import org.jetbrains.asm4.Type;
 import org.jetbrains.asm4.tree.AbstractInsnNode;
 import org.jetbrains.asm4.tree.FieldInsnNode;
 import org.jetbrains.asm4.tree.MethodNode;
-import org.jetbrains.asm4.tree.VarInsnNode;
+
+import java.util.Collection;
+import java.util.List;
+import java.util.Map;
 
 import static org.jetbrains.jet.codegen.asm.MethodInliner.getPreviousNoLabelNoLine;
 
 public class InlineFieldRemapper extends LambdaFieldRemapper {
 
-    private String oldOwnerType;
+    private final String oldOwnerType;
 
-    private String newOwnerType;
+    private final String newOwnerType;
 
-    public InlineFieldRemapper(String oldOwnerType, String newOwnerType) {
+    private final Parameters parameters;
+
+    private final Map<String, LambdaInfo> recapturedLambdas;
+
+    public InlineFieldRemapper(String oldOwnerType, String newOwnerType, Parameters parameters, Map<String, LambdaInfo> recapturedLambdas) {
         this.oldOwnerType = oldOwnerType;
         this.newOwnerType = newOwnerType;
+        this.parameters = parameters;
+        this.recapturedLambdas = recapturedLambdas;
     }
 
     @Override
     public AbstractInsnNode doTransform(
             MethodNode node, FieldInsnNode fieldInsnNode, CapturedParamInfo capturedField
     ) {
+        boolean isRecaptured = isRecapruredLambdaType(fieldInsnNode.owner);
+
+        if (!isRecaptured && capturedField.getLambda() != null) {
+            //strict inlining
+            return super.doTransform(node, fieldInsnNode, capturedField);
+        }
+
         AbstractInsnNode prev = getPreviousNoLabelNoLine(fieldInsnNode);
 
         assert prev.getType() == AbstractInsnNode.VAR_INSN;
@@ -56,8 +73,44 @@ public class InlineFieldRemapper extends LambdaFieldRemapper {
         node.instructions.remove(loadThis);
 
         fieldInsnNode.owner = newOwnerType;
-        fieldInsnNode.name = LambdaTransformer.getNewFieldName(fieldInsnNode.name);
+        fieldInsnNode.name = isRecaptured || capturedField.getRecapturedFrom() != null ? LambdaTransformer.getNewFieldName(capturedField.getFieldName()) : capturedField.getFieldName();
 
         return fieldInsnNode;
     }
+
+    @Override
+    public List<CapturedParamInfo> markRecaptured(List<CapturedParamInfo> originalCaptured, LambdaInfo lambda) {
+        List<CapturedParamInfo> captured = parameters.getCaptured();
+        for (CapturedParamInfo originalField : originalCaptured) {
+            for (CapturedParamInfo capturedParamInfo : captured) {
+                if (capturedParamInfo.getRecapturedFrom() == lambda) {
+                    if (capturedParamInfo.getFieldName().equals(LambdaTransformer.getNewFieldName(originalField.getFieldName()))) {
+                        originalField.setRecapturedFrom(lambda);//just mark recaptured
+                    }
+                }
+            }
+        }
+        return originalCaptured;
+    }
+
+    @Override
+    public boolean canProcess(String owner, String currentLambdaType) {
+        return super.canProcess(owner, currentLambdaType) || isRecapruredLambdaType(owner);
+    }
+
+    private boolean isRecapruredLambdaType(String owner) {
+        return recapturedLambdas.containsKey(owner);
+    }
+
+
+    @Nullable
+    @Override
+    public CapturedParamInfo findField(FieldInsnNode fieldInsnNode, Collection<CapturedParamInfo> captured) {
+        if (!isRecapruredLambdaType(fieldInsnNode.owner)) {
+            return super.findField(fieldInsnNode, captured);
+        } else {
+            LambdaInfo info = recapturedLambdas.get(fieldInsnNode.owner);
+            return super.findField(fieldInsnNode, info.getCapturedVars());
+        }
+    }
 }

compiler/backend/src/org/jetbrains/jet/codegen/asm/LambdaFieldRemapper.java

@@ -16,12 +16,16 @@
 
 package org.jetbrains.jet.codegen.asm;
 
+import org.jetbrains.annotations.Nullable;
 import org.jetbrains.asm4.Opcodes;
 import org.jetbrains.asm4.tree.AbstractInsnNode;
 import org.jetbrains.asm4.tree.FieldInsnNode;
 import org.jetbrains.asm4.tree.MethodNode;
 import org.jetbrains.asm4.tree.VarInsnNode;
 
+import java.util.Collection;
+import java.util.List;
+
 import static org.jetbrains.jet.codegen.asm.MethodInliner.getPreviousNoLabelNoLine;
 
 public class LambdaFieldRemapper {
@@ -43,4 +47,25 @@ public class LambdaFieldRemapper {
         return insn;
     }
 
+    public List<CapturedParamInfo> markRecaptured(List<CapturedParamInfo> originalCaptured, LambdaInfo lambda) {
+        return originalCaptured;
+    }
+
+
+    public boolean canProcess(String owner, String currentLambdaType) {
+        return owner.equals(currentLambdaType);
+    }
+
+    @Nullable
+    public CapturedParamInfo findField(FieldInsnNode fieldInsnNode, Collection<CapturedParamInfo> captured) {
+        String name = fieldInsnNode.name;
+        CapturedParamInfo result = null;
+        for (CapturedParamInfo valueDescriptor : captured) {
+            if (valueDescriptor.getFieldName().equals(name)) {
+                result = valueDescriptor;
+                break;
+            }
+        }
+        return result;
+    }
 }

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

@@ -122,9 +122,10 @@ public class LambdaTransformer {
         Parameters parameters = getLambdaParameters(builder, invocation);
 
         MethodVisitor invokeVisitor = newMethod(classBuilder, invoke);
+        InlineFieldRemapper remapper = new InlineFieldRemapper(oldLambdaType.getInternalName(), newLambdaType.getInternalName(), parameters, invocation.getRecapturedLambdas());
         MethodInliner inliner = new MethodInliner(invoke, parameters, info.subInline(info.nameGenerator.subGenerator("lambda")), oldLambdaType,
-                                                  new LambdaFieldRemapper());
+                                                  remapper);
-        inliner.doTransformAndMerge(invokeVisitor, new VarRemapper.ParamRemapper(parameters, null), new InlineFieldRemapper(oldLambdaType.getInternalName(), newLambdaType.getInternalName()), false);
+        inliner.doTransformAndMerge(invokeVisitor, new VarRemapper.ParamRemapper(parameters, null), remapper, false);
         invokeVisitor.visitMaxs(-1, -1);
 
         generateConstructorAndFields(classBuilder, builder, invocation);
@@ -212,16 +213,20 @@ public class LambdaTransformer {
             cur = cur.getNext();
         }
 
+        Map<String, LambdaInfo> recapturedLambdas = new HashMap<String, LambdaInfo>(); //captured var of inlined parameter
         List<CapturedParamInfo> recaptured = new ArrayList<CapturedParamInfo>();
         for (LambdaInfo info : additionalCaptured) {
             List<CapturedParamInfo> vars = info.getCapturedVars();
             for (CapturedParamInfo var : vars) {
                 CapturedParamInfo recapturedParamInfo = builder.addCapturedParam(getNewFieldName(var.getFieldName()), var.getType(), true, var);
+                recapturedParamInfo.setRecapturedFrom(info);
                 recaptured.add(var);
             }
+            recapturedLambdas.put(info.getLambdaClassType().getInternalName(), info);
         }
 
         invocation.setRecaptured(recaptured);
+        invocation.setRecapturedLambdas(recapturedLambdas);
     }
 
     @Nullable

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

@@ -139,7 +139,7 @@ public class MethodInliner {
 
                     List<ParameterInfo> lambdaParameters = inlinableAccess.getParameters();
 
-                    Parameters params = new Parameters(lambdaParameters, Parameters.transformList(info.getCapturedVars(), lambdaParameters.size()));
+                    Parameters params = new Parameters(lambdaParameters, Parameters.transformList(capturedRemapper.markRecaptured(info.getCapturedVars(), info), lambdaParameters.size()));
 
                     MethodInliner inliner = new MethodInliner(info.getNode(), params, parent.subInline(parent.nameGenerator.subGenerator("lambda")), info.getLambdaClassType(),
                                                               capturedRemapper);
@@ -383,18 +383,13 @@ public class MethodInliner {
                 FieldInsnNode fieldInsnNode = (FieldInsnNode) cur;
                 //TODO check closure
                 String owner = fieldInsnNode.owner;
-                if (lambdaClassType.getInternalName().equals(fieldInsnNode.owner)) {
+                if (lambdaFieldRemapper.canProcess(fieldInsnNode.owner, lambdaClassType.getInternalName())) {
                     String name = fieldInsnNode.name;
                     String desc = fieldInsnNode.desc;
 
                     Collection<CapturedParamInfo> vars = paramsToSearch.getCaptured();
-                    CapturedParamInfo result = null;
+                    CapturedParamInfo result = lambdaFieldRemapper.findField(fieldInsnNode, paramsToSearch.getCaptured());
-                    for (CapturedParamInfo valueDescriptor : vars) {
+
-                        if (valueDescriptor.getFieldName().equals(name)) {
-                            result = valueDescriptor;
-                            break;
-                        }
-                    }
                     if (result == null) {
                         throw new UnsupportedOperationException("Coudn't find field " +
                                                                 owner +

compiler/backend/src/org/jetbrains/jet/codegen/asm/Parameters.java

@@ -68,6 +68,7 @@ public class Parameters implements Iterable<ParameterInfo> {
                                           capturedParamInfo.getIndex(), result.size() + realSize);
 
             newInfo.setLambda(capturedParamInfo.getLambda());
+            newInfo.setRecapturedFrom(capturedParamInfo.getRecapturedFrom());
 
             result.add(newInfo);
 
@@ -84,6 +85,7 @@ public class Parameters implements Iterable<ParameterInfo> {
             CapturedParamInfo newInfo = new CapturedParamInfo(capturedParamInfo.getFieldName(), capturedParamInfo.getType(), capturedParamInfo.isSkipped,
                                                               capturedParamInfo.getRemapIndex(), result.size() + realSize);
             newInfo.setLambda(capturedParamInfo.getLambda());
+            newInfo.setRecapturedFrom(capturedParamInfo.getRecapturedFrom());
             result.add(newInfo);
             if (capturedParamInfo.getType().getSize() == 2) {
                 result.add(CapturedParamInfo.STUB);