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.AsmUtil;
031    import org.jetbrains.jet.codegen.ClassBuilder;
032    import org.jetbrains.jet.codegen.ClosureCodegen;
033    import org.jetbrains.jet.codegen.FieldInfo;
034    import org.jetbrains.jet.codegen.state.GenerationState;
035    import org.jetbrains.jet.codegen.state.JetTypeMapper;
036    
037    import java.io.IOException;
038    import java.util.ArrayList;
039    import java.util.HashMap;
040    import java.util.List;
041    import java.util.Map;
042    
043    import static org.jetbrains.asm4.Opcodes.ASM4;
044    import static org.jetbrains.asm4.Opcodes.V1_6;
045    
046    public class LambdaTransformer {
047    
048        protected final GenerationState state;
049    
050        protected final JetTypeMapper typeMapper;
051    
052        private final MethodNode constructor;
053    
054        private final MethodNode invoke;
055    
056        private final MethodNode bridge;
057    
058        private final InliningContext info;
059    
060        private final Type oldLambdaType;
061    
062        private final Type newLambdaType;
063    
064        private int classAccess;
065        private String signature;
066        private String superName;
067        private String[] interfaces;
068        private final boolean isSameModule;
069    
070        public LambdaTransformer(String lambdaInternalName, InliningContext info, boolean isSameModule, Type newLambdaType) {
071            this.isSameModule = isSameModule;
072            this.state = info.state;
073            this.typeMapper = state.getTypeMapper();
074            this.info = info;
075            this.oldLambdaType = Type.getObjectType(lambdaInternalName);
076            this.newLambdaType = newLambdaType;
077    
078            //try to find just compiled classes then in dependencies
079            ClassReader reader;
080            try {
081                OutputFile outputFile = state.getFactory().get(lambdaInternalName + ".class");
082                if (outputFile != null) {
083                    reader = new ClassReader(outputFile.asByteArray());
084                } else {
085                    VirtualFile file = InlineCodegenUtil.findVirtualFile(state.getProject(), lambdaInternalName);
086                    if (file == null) {
087                        throw new RuntimeException("Couldn't find virtual file for " + lambdaInternalName);
088                    }
089                    reader = new ClassReader(file.getInputStream());
090                }
091            }
092            catch (IOException e) {
093                throw new RuntimeException(e);
094            }
095    
096            //TODO rewrite to one step
097            constructor = getMethodNode(reader, true, false);
098            invoke = getMethodNode(reader, false, false);
099            bridge = getMethodNode(reader, false, true);
100        }
101    
102        private void buildInvokeParams(ParametersBuilder builder) {
103            builder.addThis(oldLambdaType, false);
104    
105            Type[] types = Type.getArgumentTypes(invoke.desc);
106            for (Type type : types) {
107                builder.addNextParameter(type, false, null);
108            }
109        }
110    
111        public void doTransform(ConstructorInvocation invocation) {
112            ClassBuilder classBuilder = createClassBuilder();
113    
114            //TODO: public visibility for inline function
115            classBuilder.defineClass(null,
116                                     V1_6,
117                                     classAccess,
118                                     newLambdaType.getInternalName(),
119                                     signature,
120                                     superName,
121                                     interfaces
122            );
123            ParametersBuilder builder = ParametersBuilder.newBuilder();
124            Parameters parameters = getLambdaParameters(builder, invocation);
125    
126            MethodVisitor invokeVisitor = newMethod(classBuilder, invoke);
127            RegeneratedLambdaFieldRemapper
128                    remapper = new RegeneratedLambdaFieldRemapper(oldLambdaType.getInternalName(), newLambdaType.getInternalName(), parameters, invocation.getCapturedLambdasToInline());
129            MethodInliner inliner = new MethodInliner(invoke, parameters, info.subInline(info.nameGenerator.subGenerator("lambda")), oldLambdaType,
130                                                      remapper, isSameModule);
131            inliner.doInline(invokeVisitor, new VarRemapper.ParamRemapper(parameters, 0), remapper, false);
132            invokeVisitor.visitMaxs(-1, -1);
133    
134            generateConstructorAndFields(classBuilder, builder, invocation);
135    
136            if (bridge != null) {
137                MethodVisitor invokeBridge = newMethod(classBuilder, bridge);
138                bridge.accept(new MethodVisitor(ASM4, invokeBridge) {
139                    @Override
140                    public void visitMethodInsn(int opcode, String owner, String name, String desc) {
141                        if (owner.equals(oldLambdaType.getInternalName())) {
142                            super.visitMethodInsn(opcode, newLambdaType.getInternalName(), name, desc);
143                        } else {
144                            super.visitMethodInsn(opcode, owner, name, desc);
145                        }
146                    }
147                });
148            }
149    
150            classBuilder.done();
151    
152            invocation.setNewLambdaType(newLambdaType);
153        }
154    
155        private void generateConstructorAndFields(@NotNull ClassBuilder classBuilder, @NotNull ParametersBuilder builder, @NotNull ConstructorInvocation invocation) {
156            List<CapturedParamInfo> infos = builder.buildCaptured();
157            List<Pair<String, Type>> newConstructorSignature = new ArrayList<Pair<String, Type>>();
158            for (CapturedParamInfo capturedParamInfo : infos) {
159                if (capturedParamInfo.getLambda() == null) { //not inlined
160                    newConstructorSignature.add(new Pair<String, Type>(capturedParamInfo.getFieldName(), capturedParamInfo.getType()));
161                }
162            }
163    
164            List<FieldInfo> fields = AsmUtil.transformCapturedParams(newConstructorSignature, newLambdaType);
165    
166            AsmUtil.genClosureFields(newConstructorSignature, classBuilder);
167    
168            //TODO for inline method make public class
169            Method newConstructor = ClosureCodegen.generateConstructor(classBuilder, fields, null, Type.getObjectType(superName), state, AsmUtil.NO_FLAG_PACKAGE_PRIVATE);
170            invocation.setNewConstructorDescriptor(newConstructor.getDescriptor());
171        }
172    
173        private Parameters getLambdaParameters(ParametersBuilder builder, ConstructorInvocation invocation) {
174            buildInvokeParams(builder);
175            extractParametersMapping(constructor, builder, invocation);
176            return builder.buildParameters();
177        }
178    
179        private ClassBuilder createClassBuilder() {
180            return new RemappingClassBuilder(state.getFactory().forLambdaInlining(newLambdaType, info.call.getCallElement().getContainingFile()),
181                         new TypeRemapper(info.typeMapping));
182        }
183    
184        private static MethodVisitor newMethod(ClassBuilder builder, MethodNode original) {
185            return builder.newMethod(
186                    null,
187                    original.access,
188                    original.name,
189                    original.desc,
190                    original.signature,
191                    null //TODO: change signature to list
192            );
193        }
194    
195        private static void extractParametersMapping(MethodNode constructor, ParametersBuilder builder, ConstructorInvocation invocation) {
196            Map<Integer, LambdaInfo> indexToLambda = invocation.getLambdasToInline();
197    
198            AbstractInsnNode cur = constructor.instructions.getFirst();
199            cur = cur.getNext(); //skip super call
200            List<LambdaInfo> capturedLambdas = new ArrayList<LambdaInfo>(); //captured var of inlined parameter
201            while (cur != null) {
202                if (cur.getType() == AbstractInsnNode.FIELD_INSN) {
203                    FieldInsnNode fieldNode = (FieldInsnNode) cur;
204                    CapturedParamInfo info = builder.addCapturedParam(fieldNode.name, Type.getType(fieldNode.desc), false, null);
205    
206                    assert fieldNode.getPrevious() instanceof VarInsnNode : "Previous instruction should be VarInsnNode but was " + fieldNode.getPrevious();
207                    VarInsnNode previous = (VarInsnNode) fieldNode.getPrevious();
208                    int varIndex = previous.var;
209                    LambdaInfo lambdaInfo = indexToLambda.get(varIndex);
210                    if (lambdaInfo != null) {
211                        info.setLambda(lambdaInfo);
212                        capturedLambdas.add(lambdaInfo);
213                    }
214                }
215                cur = cur.getNext();
216            }
217    
218            //For all inlined lambdas add their captured parameters
219            //TODO: some of such parameters could be skipped - we should perform additional analysis
220            Map<String, LambdaInfo> capturedLambdasToInline = new HashMap<String, LambdaInfo>(); //captured var of inlined parameter
221            List<CapturedParamInfo> allRecapturedParameters = new ArrayList<CapturedParamInfo>();
222            for (LambdaInfo info : capturedLambdas) {
223                for (CapturedParamInfo var : info.getCapturedVars()) {
224                    CapturedParamInfo recapturedParamInfo = builder.addCapturedParam(getNewFieldName(var.getFieldName()), var.getType(), true, var);
225                    recapturedParamInfo.setRecapturedFrom(info);
226                    allRecapturedParameters.add(var);
227                }
228                capturedLambdasToInline.put(info.getLambdaClassType().getInternalName(), info);
229            }
230    
231            invocation.setAllRecapturedParameters(allRecapturedParameters);
232            invocation.setCapturedLambdasToInline(capturedLambdasToInline);
233        }
234    
235        @Nullable
236        public MethodNode getMethodNode(@NotNull ClassReader reader, final boolean findConstructor, final boolean findBridge) {
237            final MethodNode[] methodNode = new MethodNode[1];
238            reader.accept(new ClassVisitor(InlineCodegenUtil.API) {
239    
240                @Override
241                public void visit(int version, int access, String name, String signature, String superName, String[] interfaces) {
242                    super.visit(version, access, name, signature, superName, interfaces);
243                    LambdaTransformer.this.classAccess = access;
244                    LambdaTransformer.this.signature = signature;
245                    LambdaTransformer.this.superName = superName;
246                    LambdaTransformer.this.interfaces = interfaces;
247                }
248    
249                @Override
250                public MethodVisitor visitMethod(int access, String name, String desc, String signature, String[] exceptions) {
251                    boolean isConstructorMethod = "<init>".equals(name);
252                    boolean isBridge = (access & Opcodes.ACC_BRIDGE) != 0;
253                    if (findConstructor && isConstructorMethod || (!findConstructor && !isConstructorMethod && (isBridge == findBridge))) {
254                        assert methodNode[0] == null : "Wrong lambda/sam structure: " + methodNode[0].name + " conflicts with " + name;
255                        return methodNode[0] = new MethodNode(access, name, desc, signature, exceptions);
256                    }
257                    return null;
258                }
259            }, ClassReader.SKIP_DEBUG | ClassReader.SKIP_FRAMES);
260    
261            if (methodNode[0] == null && !findBridge) {
262                throw new RuntimeException("Couldn't find operation method of lambda/sam class " + oldLambdaType.getInternalName() + ": findConstructor = " + findConstructor);
263            }
264    
265            return methodNode[0];
266        }
267    
268        public static String getNewFieldName(String oldName) {
269            return oldName + "$inlined";
270        }
271    }