001    /*
002     * Copyright 2010-2013 JetBrains s.r.o.
003     *
004     * Licensed under the Apache License, Version 2.0 (the "License");
005     * you may not use this file except in compliance with the License.
006     * You may obtain a copy of the License at
007     *
008     * http://www.apache.org/licenses/LICENSE-2.0
009     *
010     * Unless required by applicable law or agreed to in writing, software
011     * distributed under the License is distributed on an "AS IS" BASIS,
012     * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
013     * See the License for the specific language governing permissions and
014     * limitations under the License.
015     */
016    
017    package org.jetbrains.jet.codegen.inline;
018    
019    import com.intellij.openapi.util.Pair;
020    import com.intellij.openapi.vfs.VirtualFile;
021    import org.jetbrains.annotations.NotNull;
022    import org.jetbrains.annotations.Nullable;
023    import org.jetbrains.asm4.*;
024    import org.jetbrains.asm4.commons.Method;
025    import org.jetbrains.asm4.tree.AbstractInsnNode;
026    import org.jetbrains.asm4.tree.FieldInsnNode;
027    import org.jetbrains.asm4.tree.MethodNode;
028    import org.jetbrains.asm4.tree.VarInsnNode;
029    import org.jetbrains.jet.OutputFile;
030    import org.jetbrains.jet.codegen.*;
031    import org.jetbrains.jet.codegen.state.GenerationState;
032    import org.jetbrains.jet.codegen.state.JetTypeMapper;
033    
034    import java.io.IOException;
035    import java.util.*;
036    
037    import static org.jetbrains.asm4.Opcodes.ASM4;
038    import static org.jetbrains.asm4.Opcodes.V1_6;
039    
040    public class LambdaTransformer {
041    
042        protected final GenerationState state;
043    
044        protected final JetTypeMapper typeMapper;
045    
046        private final MethodNode constructor;
047    
048        private final MethodNode invoke;
049    
050        private final MethodNode bridge;
051    
052        private final InliningContext inliningContext;
053    
054        private final Type oldLambdaType;
055    
056        private final Type newLambdaType;
057    
058        private int classAccess;
059        private String signature;
060        private String superName;
061        private String[] interfaces;
062        private final boolean isSameModule;
063    
064        private Map<String, List<String>> fieldNames = new HashMap<String, List<String>>();
065    
066        public LambdaTransformer(String lambdaInternalName, InliningContext inliningContext, boolean isSameModule, Type newLambdaType) {
067            this.isSameModule = isSameModule;
068            this.state = inliningContext.state;
069            this.typeMapper = state.getTypeMapper();
070            this.inliningContext = inliningContext;
071            this.oldLambdaType = Type.getObjectType(lambdaInternalName);
072            this.newLambdaType = newLambdaType;
073    
074            //try to find just compiled classes then in dependencies
075            ClassReader reader;
076            try {
077                OutputFile outputFile = state.getFactory().get(lambdaInternalName + ".class");
078                if (outputFile != null) {
079                    reader = new ClassReader(outputFile.asByteArray());
080                } else {
081                    VirtualFile file = InlineCodegenUtil.findVirtualFile(state.getProject(), lambdaInternalName);
082                    if (file == null) {
083                        throw new RuntimeException("Couldn't find virtual file for " + lambdaInternalName);
084                    }
085                    reader = new ClassReader(file.getInputStream());
086                }
087            }
088            catch (IOException e) {
089                throw new RuntimeException(e);
090            }
091    
092            //TODO rewrite to one step
093            constructor = getMethodNode(reader, true, false);
094            invoke = getMethodNode(reader, false, false);
095            bridge = getMethodNode(reader, false, true);
096        }
097    
098        private void buildInvokeParams(ParametersBuilder builder) {
099            builder.addThis(oldLambdaType, false);
100    
101            Type[] types = Type.getArgumentTypes(invoke.desc);
102            for (Type type : types) {
103                builder.addNextParameter(type, false, null);
104            }
105        }
106    
107        public InlineResult doTransform(ConstructorInvocation invocation, FieldRemapper parentRemapper) {
108            ClassBuilder classBuilder = createClassBuilder();
109    
110            //TODO: public visibility for inline function
111            classBuilder.defineClass(null,
112                                     V1_6,
113                                     classAccess,
114                                     newLambdaType.getInternalName(),
115                                     signature,
116                                     superName,
117                                     interfaces
118            );
119    
120            // TODO: load synthetic class kind from the transformed class and write the same kind to the copy of that class here
121            // See AsmUtil.writeKotlinSyntheticClassAnnotation
122    
123            ParametersBuilder builder = ParametersBuilder.newBuilder();
124            Parameters parameters = getLambdaParameters(builder, invocation);
125    
126            MethodVisitor invokeVisitor = newMethod(classBuilder, invoke);
127    
128            RegeneratedLambdaFieldRemapper remapper =
129                    new RegeneratedLambdaFieldRemapper(oldLambdaType.getInternalName(), newLambdaType.getInternalName(),
130                                                       parameters, invocation.getCapturedLambdasToInline(),
131                                                       parentRemapper);
132    
133            MethodInliner inliner = new MethodInliner(invoke, parameters, inliningContext.subInline(inliningContext.nameGenerator.subGenerator("lambda")),
134                                                      remapper, isSameModule, "Transformer for " + invocation.getOwnerInternalName());
135            InlineResult result = inliner.doInline(invokeVisitor, new LocalVarRemapper(parameters, 0), false);
136            invokeVisitor.visitMaxs(-1, -1);
137    
138            generateConstructorAndFields(classBuilder, builder, invocation);
139    
140            if (bridge != null) {
141                MethodVisitor invokeBridge = newMethod(classBuilder, bridge);
142                bridge.accept(new MethodVisitor(ASM4, invokeBridge) {
143                    @Override
144                    public void visitMethodInsn(int opcode, String owner, String name, String desc) {
145                        if (owner.equals(oldLambdaType.getInternalName())) {
146                            super.visitMethodInsn(opcode, newLambdaType.getInternalName(), name, desc);
147                        } else {
148                            super.visitMethodInsn(opcode, owner, name, desc);
149                        }
150                    }
151                });
152            }
153    
154            classBuilder.done();
155    
156            invocation.setNewLambdaType(newLambdaType);
157            return result;
158        }
159    
160        private void generateConstructorAndFields(@NotNull ClassBuilder classBuilder, @NotNull ParametersBuilder builder, @NotNull ConstructorInvocation invocation) {
161            List<CapturedParamInfo> infos = builder.buildCaptured();
162            List<Pair<String, Type>> newConstructorSignature = new ArrayList<Pair<String, Type>>();
163            for (CapturedParamInfo capturedParamInfo : infos) {
164                if (capturedParamInfo.getLambda() == null) { //not inlined
165                    newConstructorSignature.add(new Pair<String, Type>(capturedParamInfo.getNewFieldName(), capturedParamInfo.getType()));
166                }
167            }
168    
169            List<FieldInfo> fields = AsmUtil.transformCapturedParams(newConstructorSignature, newLambdaType);
170    
171            AsmUtil.genClosureFields(newConstructorSignature, classBuilder);
172    
173            //TODO for inline method make public class
174            Method newConstructor = ClosureCodegen.generateConstructor(classBuilder, fields, null, Type.getObjectType(superName), state, AsmUtil.NO_FLAG_PACKAGE_PRIVATE);
175            invocation.setNewConstructorDescriptor(newConstructor.getDescriptor());
176        }
177    
178        private Parameters getLambdaParameters(ParametersBuilder builder, ConstructorInvocation invocation) {
179            buildInvokeParams(builder);
180            extractParametersMapping(constructor, builder, invocation);
181            return builder.buildParameters();
182        }
183    
184        private ClassBuilder createClassBuilder() {
185            return new RemappingClassBuilder(state.getFactory().forLambdaInlining(newLambdaType, inliningContext.call.getCallElement().getContainingFile()),
186                         new TypeRemapper(inliningContext.typeMapping));
187        }
188    
189        private static MethodVisitor newMethod(ClassBuilder builder, MethodNode original) {
190            return builder.newMethod(
191                    null,
192                    original.access,
193                    original.name,
194                    original.desc,
195                    original.signature,
196                    null //TODO: change signature to list
197            );
198        }
199    
200        private void extractParametersMapping(MethodNode constructor, ParametersBuilder builder, final ConstructorInvocation invocation) {
201            Map<Integer, LambdaInfo> indexToLambda = invocation.getLambdasToInline();
202    
203            AbstractInsnNode cur = constructor.instructions.getFirst();
204            cur = cur.getNext(); //skip super call
205            List<LambdaInfo> capturedLambdas = new ArrayList<LambdaInfo>(); //captured var of inlined parameter
206            CapturedParamOwner owner = new CapturedParamOwner() {
207                @Override
208                public Type getType() {
209                    return Type.getObjectType(invocation.getOwnerInternalName());
210                }
211            };
212    
213            while (cur != null) {
214                if (cur.getType() == AbstractInsnNode.FIELD_INSN) {
215                    FieldInsnNode fieldNode = (FieldInsnNode) cur;
216                    CapturedParamInfo info = builder.addCapturedParam(fieldNode.name, Type.getType(fieldNode.desc), false, null, owner);
217    
218                    assert fieldNode.getPrevious() instanceof VarInsnNode : "Previous instruction should be VarInsnNode but was " + fieldNode.getPrevious();
219                    VarInsnNode previous = (VarInsnNode) fieldNode.getPrevious();
220                    int varIndex = previous.var;
221                    LambdaInfo lambdaInfo = indexToLambda.get(varIndex);
222                    if (lambdaInfo != null) {
223                        info.setLambda(lambdaInfo);
224                        capturedLambdas.add(lambdaInfo);
225                    }
226    
227                    addUniqueField(info.getOriginalFieldName());
228                }
229                cur = cur.getNext();
230            }
231    
232            //For all inlined lambdas add their captured parameters
233            //TODO: some of such parameters could be skipped - we should perform additional analysis
234            Map<String, LambdaInfo> capturedLambdasToInline = new HashMap<String, LambdaInfo>(); //captured var of inlined parameter
235            List<CapturedParamInfo> allRecapturedParameters = new ArrayList<CapturedParamInfo>();
236            for (LambdaInfo info : capturedLambdas) {
237                for (CapturedParamInfo var : info.getCapturedVars()) {
238                    CapturedParamInfo recapturedParamInfo = builder.addCapturedParam(var, getNewFieldName(var.getOriginalFieldName()));
239                    StackValue composed = StackValue.composed(StackValue.local(0, oldLambdaType),
240                                                              StackValue.field(var.getType(),
241                                                                               oldLambdaType, /*TODO owner type*/
242                                                                               recapturedParamInfo.getNewFieldName(), false)
243                    );
244                    recapturedParamInfo.setRemapValue(composed);
245                    allRecapturedParameters.add(var);
246                }
247                capturedLambdasToInline.put(info.getLambdaClassType().getInternalName(), info);
248            }
249    
250            invocation.setAllRecapturedParameters(allRecapturedParameters);
251            invocation.setCapturedLambdasToInline(capturedLambdasToInline);
252        }
253    
254        @Nullable
255        public MethodNode getMethodNode(@NotNull ClassReader reader, final boolean findConstructor, final boolean findBridge) {
256            final MethodNode[] methodNode = new MethodNode[1];
257            reader.accept(new ClassVisitor(InlineCodegenUtil.API) {
258    
259                @Override
260                public void visit(int version, int access, String name, String signature, String superName, String[] interfaces) {
261                    super.visit(version, access, name, signature, superName, interfaces);
262                    LambdaTransformer.this.classAccess = access;
263                    LambdaTransformer.this.signature = signature;
264                    LambdaTransformer.this.superName = superName;
265                    LambdaTransformer.this.interfaces = interfaces;
266                }
267    
268                @Override
269                public MethodVisitor visitMethod(int access, String name, String desc, String signature, String[] exceptions) {
270                    boolean isConstructorMethod = "<init>".equals(name);
271                    boolean isBridge = (access & Opcodes.ACC_BRIDGE) != 0;
272                    if (findConstructor && isConstructorMethod || (!findConstructor && !isConstructorMethod && (isBridge == findBridge))) {
273                        assert methodNode[0] == null : "Wrong lambda/sam structure: " + methodNode[0].name + " conflicts with " + name;
274                        return methodNode[0] = new MethodNode(access, name, desc, signature, exceptions);
275                    }
276                    return null;
277                }
278            }, ClassReader.SKIP_DEBUG | ClassReader.SKIP_FRAMES);
279    
280            if (methodNode[0] == null && !findBridge) {
281                throw new RuntimeException("Couldn't find operation method of lambda/sam class " + oldLambdaType.getInternalName() + ": findConstructor = " + findConstructor);
282            }
283    
284            return methodNode[0];
285        }
286    
287        @NotNull
288        public String getNewFieldName(@NotNull String oldName) {
289            if (oldName.equals("this$0")) {
290                //"this$0" couldn't clash and we should keep this name invariant for further transformations
291                return oldName;
292            }
293            return addUniqueField(oldName + "$inlined");
294        }
295    
296        @NotNull
297        private String addUniqueField(@NotNull String name) {
298            List<String> existNames = fieldNames.get(name);
299            if (existNames == null) {
300                existNames = new LinkedList<String>();
301                fieldNames.put(name, existNames);
302            }
303            String suffix = existNames.isEmpty() ? "" : "$" + existNames.size();
304            String newName = name + suffix;
305            existNames.add(newName);
306            return newName;
307        }
308    }