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 com.intellij.util.ArrayUtil;
022    import org.jetbrains.annotations.NotNull;
023    import org.jetbrains.jet.OutputFile;
024    import org.jetbrains.jet.codegen.AsmUtil;
025    import org.jetbrains.jet.codegen.ClassBuilder;
026    import org.jetbrains.jet.codegen.FieldInfo;
027    import org.jetbrains.jet.codegen.StackValue;
028    import org.jetbrains.jet.codegen.state.GenerationState;
029    import org.jetbrains.jet.codegen.state.JetTypeMapper;
030    import org.jetbrains.org.objectweb.asm.*;
031    import org.jetbrains.org.objectweb.asm.commons.InstructionAdapter;
032    import org.jetbrains.org.objectweb.asm.tree.*;
033    
034    import java.io.IOException;
035    import java.util.*;
036    
037    public class AnonymousObjectTransformer {
038    
039        protected final GenerationState state;
040    
041        protected final JetTypeMapper typeMapper;
042    
043        private MethodNode constructor;
044    
045        private final InliningContext inliningContext;
046    
047        private final Type oldObjectType;
048    
049        private final Type newLambdaType;
050    
051        private final ClassReader reader;
052    
053        private final boolean isSameModule;
054    
055        private final Map<String, List<String>> fieldNames = new HashMap<String, List<String>>();
056    
057        public AnonymousObjectTransformer(
058                @NotNull String objectInternalName,
059                @NotNull InliningContext inliningContext,
060                boolean isSameModule,
061                @NotNull Type newLambdaType
062        ) {
063            this.isSameModule = isSameModule;
064            this.state = inliningContext.state;
065            this.typeMapper = state.getTypeMapper();
066            this.inliningContext = inliningContext;
067            this.oldObjectType = Type.getObjectType(objectInternalName);
068            this.newLambdaType = newLambdaType;
069    
070            //try to find just compiled classes then in dependencies
071            try {
072                OutputFile outputFile = state.getFactory().get(objectInternalName + ".class");
073                if (outputFile != null) {
074                    reader = new ClassReader(outputFile.asByteArray());
075                } else {
076                    VirtualFile file = InlineCodegenUtil.findVirtualFile(state.getProject(), objectInternalName);
077                    if (file == null) {
078                        throw new RuntimeException("Couldn't find virtual file for " + objectInternalName);
079                    }
080                    reader = new ClassReader(file.getInputStream());
081                }
082            }
083            catch (IOException e) {
084                throw new RuntimeException(e);
085            }
086        }
087    
088        private void buildInvokeParamsFor(@NotNull ParametersBuilder builder, @NotNull MethodNode node) {
089            builder.addThis(oldObjectType, false);
090    
091            Type[] types = Type.getArgumentTypes(node.desc);
092            for (Type type : types) {
093                builder.addNextParameter(type, false, null);
094            }
095        }
096    
097        @NotNull
098        public InlineResult doTransform(@NotNull ConstructorInvocation invocation, @NotNull FieldRemapper parentRemapper) {
099            ClassBuilder classBuilder = createClassBuilder();
100            final List<MethodNode> methodsToTransform = new ArrayList<MethodNode>();
101            reader.accept(new ClassVisitor(InlineCodegenUtil.API, classBuilder.getVisitor()) {
102    
103                @Override
104                public void visitOuterClass(@NotNull String owner, String name, String desc) {
105                    InliningContext parent = inliningContext.getParent();
106                    assert parent != null : "Context for transformer should have parent one: " + inliningContext;
107    
108                    //we don't write owner info for lamdbas and SAMs just only for objects
109                    if (parent.isRoot() || parent.isInliningLambdaRootContext()) {
110                        //TODO: think about writing method info - there is some problem with new constructor desc calculation
111                        super.visitOuterClass(inliningContext.getParent().getClassNameToInline(), null, null);
112                        return;
113                    }
114    
115                    super.visitOuterClass(owner, name, desc);
116                }
117    
118                @Override
119                public MethodVisitor visitMethod(
120                        int access, @NotNull String name, @NotNull String desc, String signature, String[] exceptions
121                ) {
122                    MethodNode node = new MethodNode(access, name, desc, signature, exceptions);
123                    if (name.equals("<init>")){
124                        if (constructor != null)
125                            throw new RuntimeException("Lambda, SAM or anonymous object should have only one constructor");
126    
127                        constructor = node;
128                    } else {
129                        methodsToTransform.add(node);
130                    }
131                    return node;
132                }
133    
134                @Override
135                public FieldVisitor visitField(
136                        int access, @NotNull String name, @NotNull String desc, String signature, Object value
137                ) {
138                    addUniqueField(name);
139                    if (InlineCodegenUtil.isCapturedFieldName(name)) {
140                        return null;
141                    } else {
142                        return super.visitField(access, name, desc, signature, value);
143                    }
144                }
145            }, ClassReader.SKIP_FRAMES);
146    
147            ParametersBuilder allCapturedParamBuilder = ParametersBuilder.newBuilder();
148            ParametersBuilder constructorParamBuilder = ParametersBuilder.newBuilder();
149            extractParametersMappingAndPatchConstructor(constructor, allCapturedParamBuilder, constructorParamBuilder, invocation);
150    
151            InlineResult result = InlineResult.create();
152            for (MethodNode next : methodsToTransform) {
153                MethodVisitor visitor = newMethod(classBuilder, next);
154                InlineResult funResult = inlineMethod(invocation, parentRemapper, visitor, next, allCapturedParamBuilder);
155                result.addAllClassesToRemove(funResult);
156            }
157    
158            InlineResult constructorResult =
159                    generateConstructorAndFields(classBuilder, allCapturedParamBuilder, constructorParamBuilder, invocation, parentRemapper);
160    
161            result.addAllClassesToRemove(constructorResult);
162    
163            classBuilder.done();
164    
165            invocation.setNewLambdaType(newLambdaType);
166            return result;
167        }
168    
169        @NotNull
170        private InlineResult inlineMethod(
171                @NotNull ConstructorInvocation invocation,
172                @NotNull FieldRemapper parentRemapper,
173                @NotNull MethodVisitor resultVisitor,
174                @NotNull MethodNode sourceNode,
175                @NotNull ParametersBuilder capturedBuilder
176        ) {
177    
178            Parameters parameters = getMethodParametersWithCaptured(capturedBuilder, sourceNode);
179    
180            RegeneratedLambdaFieldRemapper remapper =
181                    new RegeneratedLambdaFieldRemapper(oldObjectType.getInternalName(), newLambdaType.getInternalName(),
182                                                       parameters, invocation.getCapturedLambdasToInline(),
183                                                       parentRemapper);
184    
185            MethodInliner inliner = new MethodInliner(sourceNode, parameters, inliningContext.subInline(inliningContext.nameGenerator.subGenerator("lambda")),
186                                                      remapper, isSameModule, "Transformer for " + invocation.getOwnerInternalName());
187            InlineResult result = inliner.doInline(resultVisitor, new LocalVarRemapper(parameters, 0), false);
188            resultVisitor.visitMaxs(-1, -1);
189            return result;
190        }
191    
192        private InlineResult generateConstructorAndFields(
193                @NotNull ClassBuilder classBuilder,
194                @NotNull ParametersBuilder allCapturedBuilder,
195                @NotNull ParametersBuilder constructorInlineBuilder,
196                @NotNull ConstructorInvocation invocation,
197                @NotNull FieldRemapper parentRemapper
198        ) {
199            List<Type> descTypes = new ArrayList<Type>();
200    
201            Parameters constructorParams = constructorInlineBuilder.buildParameters();
202            int [] capturedIndexes = new int [constructorParams.totalSize()];
203            int index = 0;
204            int size = 0;
205    
206            //complex processing cause it could have super constructor call params
207            for (ParameterInfo info : constructorParams) {
208                if (!info.isSkipped()) { //not inlined
209                    if (info.isCaptured() || info instanceof CapturedParamInfo) {
210                        capturedIndexes[index] = size;
211                        index++;
212                    }
213    
214                    if (size != 0) { //skip this
215                        descTypes.add(info.getType());
216                    }
217                    size += info.getType().getSize();
218                }
219            }
220    
221            List<Pair<String, Type>> capturedFieldsToGenerate = new ArrayList<Pair<String, Type>>();
222            for (CapturedParamInfo capturedParamInfo : allCapturedBuilder.listCaptured()) {
223                if (capturedParamInfo.getLambda() == null) { //not inlined
224                    capturedFieldsToGenerate.add(new Pair<String, Type>(capturedParamInfo.getNewFieldName(), capturedParamInfo.getType()));
225                }
226            }
227    
228            String constructorDescriptor = Type.getMethodDescriptor(Type.VOID_TYPE, descTypes.toArray(new Type[descTypes.size()]));
229    
230            MethodVisitor constructorVisitor = classBuilder.newMethod(null,
231                                                                      AsmUtil.NO_FLAG_PACKAGE_PRIVATE,
232                                                                      "<init>", constructorDescriptor,
233                                                                      null, ArrayUtil.EMPTY_STRING_ARRAY);
234    
235            //initialize captured fields
236            List<FieldInfo> fields = AsmUtil.transformCapturedParams(capturedFieldsToGenerate, newLambdaType);
237            int paramIndex = 0;
238            InstructionAdapter capturedFieldInitializer = new InstructionAdapter(constructorVisitor);
239            for (FieldInfo fieldInfo : fields) {
240                AsmUtil.genAssignInstanceFieldFromParam(fieldInfo, capturedIndexes[paramIndex], capturedFieldInitializer);
241                paramIndex++;
242            }
243    
244            //then transform constructor
245            Parameters constructorParameters = constructorInlineBuilder.buildParameters();
246    
247            RegeneratedLambdaFieldRemapper remapper =
248                    new RegeneratedLambdaFieldRemapper(oldObjectType.getInternalName(), newLambdaType.getInternalName(),
249                                                       constructorParameters, invocation.getCapturedLambdasToInline(),
250                                                       parentRemapper);
251    
252            MethodInliner inliner = new MethodInliner(constructor, constructorParameters, inliningContext.subInline(inliningContext.nameGenerator.subGenerator("lambda")),
253                                                      remapper, isSameModule, "Transformer for constructor of " + invocation.getOwnerInternalName());
254            InlineResult result = inliner.doInline(capturedFieldInitializer, new LocalVarRemapper(constructorParameters, 0), false);
255            constructorVisitor.visitMaxs(-1, -1);
256    
257            AsmUtil.genClosureFields(capturedFieldsToGenerate, classBuilder);
258            //TODO for inline method make public class
259            invocation.setNewConstructorDescriptor(constructorDescriptor);
260            return result;
261        }
262    
263        @NotNull
264        private Parameters getMethodParametersWithCaptured(
265                @NotNull ParametersBuilder capturedBuilder,
266                @NotNull MethodNode sourceNode
267        ) {
268            ParametersBuilder builder = ParametersBuilder.newBuilder();
269            buildInvokeParamsFor(builder, sourceNode);
270            for (CapturedParamInfo param : capturedBuilder.listCaptured()) {
271                builder.addCapturedParamCopy(param);
272            }
273            return builder.buildParameters();
274        }
275    
276        @NotNull
277        private ClassBuilder createClassBuilder() {
278            return new RemappingClassBuilder(state.getFactory().forLambdaInlining(newLambdaType, inliningContext.getRoot().call.getCallElement().getContainingFile()),
279                         new TypeRemapper(inliningContext.typeMapping));
280        }
281    
282        @NotNull
283        private static MethodVisitor newMethod(@NotNull ClassBuilder builder, @NotNull MethodNode original) {
284            return builder.newMethod(
285                    null,
286                    original.access,
287                    original.name,
288                    original.desc,
289                    original.signature,
290                    original.exceptions.toArray(new String [original.exceptions.size()])
291            );
292        }
293    
294        private void extractParametersMappingAndPatchConstructor(
295                @NotNull MethodNode constructor,
296                @NotNull ParametersBuilder capturedParamBuilder,
297                @NotNull ParametersBuilder constructorParamBuilder,
298                @NotNull final ConstructorInvocation invocation
299        ) {
300    
301            CapturedParamOwner owner = new CapturedParamOwner() {
302                @Override
303                public Type getType() {
304                    return Type.getObjectType(invocation.getOwnerInternalName());
305                }
306            };
307    
308            List<LambdaInfo> capturedLambdas = new ArrayList<LambdaInfo>(); //captured var of inlined parameter
309            List<CapturedParamInfo> capturedToInline = new ArrayList<CapturedParamInfo>();
310            Map<Integer, LambdaInfo> indexToLambda = invocation.getLambdasToInline();
311            Set<Integer> capturedParams = new HashSet<Integer>();
312    
313            //load captured parameters and patch instruction list (NB: there is also could be object fields)
314            AbstractInsnNode cur = constructor.instructions.getFirst();
315            while (cur != null) {
316                if (cur instanceof FieldInsnNode) {
317                    FieldInsnNode fieldNode = (FieldInsnNode) cur;
318                    if (fieldNode.getOpcode() == Opcodes.PUTFIELD && InlineCodegenUtil.isCapturedFieldName(fieldNode.name)) {
319    
320                        boolean isPrevVarNode = fieldNode.getPrevious() instanceof VarInsnNode;
321                        boolean isPrevPrevVarNode = isPrevVarNode && fieldNode.getPrevious().getPrevious() instanceof VarInsnNode;
322    
323                        if (isPrevPrevVarNode) {
324                            VarInsnNode node = (VarInsnNode) fieldNode.getPrevious().getPrevious();
325                            if (node.var == 0) {
326                                VarInsnNode previous = (VarInsnNode) fieldNode.getPrevious();
327                                int varIndex = previous.var;
328                                LambdaInfo lambdaInfo = indexToLambda.get(varIndex);
329                                CapturedParamInfo info = capturedParamBuilder.addCapturedParam(owner, fieldNode.name, Type.getType(fieldNode.desc), lambdaInfo != null, null);
330                                if (lambdaInfo != null) {
331                                    info.setLambda(lambdaInfo);
332                                    capturedLambdas.add(lambdaInfo);
333                                    capturedToInline.add(info);
334                                }
335                                capturedParams.add(varIndex);
336    
337                                constructor.instructions.remove(previous.getPrevious());
338                                constructor.instructions.remove(previous);
339                                AbstractInsnNode temp = cur;
340                                cur = cur.getNext();
341                                constructor.instructions.remove(temp);
342                                continue;
343                            }
344                        }
345                    }
346                }
347                cur = cur.getNext();
348            }
349    
350            constructorParamBuilder.addThis(oldObjectType, false);
351            Type [] types = Type.getArgumentTypes(invocation.getDesc());
352            for (Type type : types) {
353                LambdaInfo info = indexToLambda.get(constructorParamBuilder.getNextValueParameterIndex());
354                ParameterInfo parameterInfo = constructorParamBuilder.addNextParameter(type, info != null, null);
355                parameterInfo.setLambda(info);
356                if (capturedParams.contains(parameterInfo.getIndex())) {
357                    parameterInfo.setCaptured(true);
358                } else {
359                    //otherwise it's super constructor parameter
360                }
361            }
362    
363            //For all inlined lambdas add their captured parameters
364            //TODO: some of such parameters could be skipped - we should perform additional analysis
365            Map<String, LambdaInfo> capturedLambdasToInline = new HashMap<String, LambdaInfo>(); //captured var of inlined parameter
366            List<CapturedParamInfo> allRecapturedParameters = new ArrayList<CapturedParamInfo>();
367            for (LambdaInfo info : capturedLambdas) {
368                for (CapturedParamInfo var : info.getCapturedVars()) {
369                    CapturedParamInfo recapturedParamInfo = capturedParamBuilder.addCapturedParam(var,
370                                                                                                  getNewFieldName(var.getOriginalFieldName()));
371                    StackValue composed = StackValue.composed(StackValue.local(0, oldObjectType),
372                                                              StackValue.field(var.getType(),
373                                                                               oldObjectType, /*TODO owner type*/
374                                                                               recapturedParamInfo.getNewFieldName(), false)
375                    );
376                    recapturedParamInfo.setRemapValue(composed);
377                    allRecapturedParameters.add(var);
378    
379                    constructorParamBuilder.addCapturedParam(var, recapturedParamInfo.getNewFieldName()).setRemapValue(composed);
380                }
381                capturedLambdasToInline.put(info.getLambdaClassType().getInternalName(), info);
382            }
383    
384            //HACK: in inlinining into constructor we access inlined lambda with field access not local var but this lambda added no general params not captured one,
385            //so we need to add them to captured params
386            for (CapturedParamInfo info : capturedToInline) {
387                constructorParamBuilder.addCapturedParamCopy(info);
388            }
389    
390            invocation.setAllRecapturedParameters(allRecapturedParameters);
391            invocation.setCapturedLambdasToInline(capturedLambdasToInline);
392        }
393    
394        @NotNull
395        public String getNewFieldName(@NotNull String oldName) {
396            if (InlineCodegenUtil.THIS$0.equals(oldName)) {
397                //"this$0" couldn't clash and we should keep this name invariant for further transformations
398                return oldName;
399            }
400            return addUniqueField(oldName + "$inlined");
401        }
402    
403        @NotNull
404        private String addUniqueField(@NotNull String name) {
405            List<String> existNames = fieldNames.get(name);
406            if (existNames == null) {
407                existNames = new LinkedList<String>();
408                fieldNames.put(name, existNames);
409            }
410            String suffix = existNames.isEmpty() ? "" : "$" + existNames.size();
411            String newName = name + suffix;
412            existNames.add(newName);
413            return newName;
414        }
415    }