001    /*
002     * Copyright 2010-2015 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.kotlin.codegen;
018    
019    import com.google.common.collect.Lists;
020    import com.intellij.openapi.vfs.VirtualFile;
021    import com.intellij.psi.PsiFile;
022    import com.intellij.util.Function;
023    import com.intellij.util.containers.ContainerUtil;
024    import com.intellij.util.io.DataOutputStream;
025    import kotlin.CollectionsKt;
026    import org.jetbrains.annotations.NotNull;
027    import org.jetbrains.annotations.Nullable;
028    import org.jetbrains.annotations.TestOnly;
029    import org.jetbrains.kotlin.backend.common.output.OutputFile;
030    import org.jetbrains.kotlin.backend.common.output.OutputFileCollection;
031    import org.jetbrains.kotlin.codegen.state.GenerationState;
032    import org.jetbrains.kotlin.load.java.JvmAbi;
033    import org.jetbrains.kotlin.load.kotlin.PackagePartClassUtils;
034    import org.jetbrains.kotlin.load.kotlin.PackageParts;
035    import org.jetbrains.kotlin.name.FqName;
036    import org.jetbrains.kotlin.psi.KtFile;
037    import org.jetbrains.kotlin.resolve.jvm.diagnostics.JvmDeclarationOrigin;
038    import org.jetbrains.kotlin.serialization.jvm.JvmPackageTable;
039    import org.jetbrains.org.objectweb.asm.Type;
040    
041    import java.io.ByteArrayOutputStream;
042    import java.io.File;
043    import java.io.IOException;
044    import java.io.UnsupportedEncodingException;
045    import java.util.*;
046    
047    import static org.jetbrains.kotlin.codegen.JvmCodegenUtil.getMappingFileName;
048    
049    public class ClassFileFactory implements OutputFileCollection {
050        private final GenerationState state;
051        private final ClassBuilderFactory builderFactory;
052        private final Map<FqName, PackageCodegen> package2codegen = new HashMap<FqName, PackageCodegen>();
053        private final Map<FqName, MultifileClassCodegen> multifileClass2codegen = new HashMap<FqName, MultifileClassCodegen>();
054        private final Map<String, OutAndSourceFileList> generators = new LinkedHashMap<String, OutAndSourceFileList>();
055    
056        private boolean isDone = false;
057    
058        public ClassFileFactory(@NotNull GenerationState state, @NotNull ClassBuilderFactory builderFactory) {
059            this.state = state;
060            this.builderFactory = builderFactory;
061        }
062    
063        @NotNull
064        public ClassBuilder newVisitor(
065                @NotNull JvmDeclarationOrigin origin,
066                @NotNull Type asmType,
067                @NotNull PsiFile sourceFile) {
068            return newVisitor(origin, asmType, Collections.singletonList(sourceFile));
069        }
070    
071        @NotNull
072        public ClassBuilder newVisitor(
073                @NotNull JvmDeclarationOrigin origin,
074                @NotNull Type asmType,
075                @NotNull Collection<? extends PsiFile> sourceFiles) {
076            String outputFilePath = asmType.getInternalName() + ".class";
077            List<File> ioSourceFiles = toIoFilesIgnoringNonPhysical(sourceFiles);
078            state.getProgress().reportOutput(ioSourceFiles, new File(outputFilePath));
079            ClassBuilder answer = builderFactory.newClassBuilder(origin);
080            generators.put(outputFilePath, new ClassBuilderAndSourceFileList(answer, ioSourceFiles));
081            return answer;
082        }
083    
084        void done() {
085            if (!isDone) {
086                isDone = true;
087                Collection<PackageCodegen> packageCodegens = package2codegen.values();
088                Collection<MultifileClassCodegen> multifileClassCodegens = multifileClass2codegen.values();
089                for (MultifileClassCodegen codegen : multifileClassCodegens) {
090                    codegen.done();
091                }
092                writeModuleMappings(packageCodegens, multifileClassCodegens);
093            }
094        }
095    
096        private void writeModuleMappings(
097                @NotNull Collection<PackageCodegen> packageCodegens,
098                @NotNull Collection<MultifileClassCodegen> multifileClassCodegens
099        ) {
100            final JvmPackageTable.PackageTable.Builder builder = JvmPackageTable.PackageTable.newBuilder();
101            String outputFilePath = getMappingFileName(state.getModuleName());
102    
103            List<PackageParts> parts = collectGeneratedPackageParts(packageCodegens, multifileClassCodegens);
104    
105            Set<File> sourceFiles = new HashSet<File>();
106            // TODO extract common logic
107            for (PackageCodegen codegen : packageCodegens) {
108                sourceFiles.addAll(toIoFilesIgnoringNonPhysical(PackagePartClassUtils.getFilesWithCallables(codegen.getFiles())));
109            }
110            for (MultifileClassCodegen codegen : multifileClassCodegens) {
111                sourceFiles.addAll(toIoFilesIgnoringNonPhysical(PackagePartClassUtils.getFilesWithCallables(codegen.getFiles())));
112            }
113    
114            for (PackageParts part : ClassFileUtilsKt.addCompiledPartsAndSort(parts, state)) {
115                PackageParts.Companion.serialize(part, builder);
116            }
117    
118            if (builder.getPackagePartsCount() != 0) {
119                state.getProgress().reportOutput(sourceFiles, new File(outputFilePath));
120                generators.put(outputFilePath, new OutAndSourceFileList(CollectionsKt.toList(sourceFiles)) {
121                    @Override
122                    public byte[] asBytes(ClassBuilderFactory factory) {
123                        try {
124                            ByteArrayOutputStream moduleMapping = new ByteArrayOutputStream(4096);
125                            DataOutputStream dataOutStream = new DataOutputStream(moduleMapping);
126                            int[] version = JvmAbi.VERSION.toArray();
127                            dataOutStream.writeInt(version.length);
128                            for (int number : version) {
129                                dataOutStream.writeInt(number);
130                            }
131                            builder.build().writeTo(dataOutStream);
132                            dataOutStream.flush();
133                            return moduleMapping.toByteArray();
134                        }
135                        catch (UnsupportedEncodingException e) {
136                            throw new RuntimeException(e);
137                        }
138                        catch (IOException e) {
139                            throw new RuntimeException(e);
140                        }
141                    }
142    
143                    @Override
144                    public String asText(ClassBuilderFactory factory) {
145                        try {
146                            return new String(asBytes(factory), "UTF-8");
147                        }
148                        catch (UnsupportedEncodingException e) {
149                            throw new RuntimeException(e);
150                        }
151                    }
152                });
153            }
154        }
155    
156        private static List<PackageParts> collectGeneratedPackageParts(
157                @NotNull Collection<PackageCodegen> packageCodegens,
158                @NotNull Collection<MultifileClassCodegen> multifileClassCodegens
159        ) {
160            Map<String, PackageParts> mergedPartsByPackageName = new LinkedHashMap<String, PackageParts>();
161    
162            for (PackageCodegen packageCodegen : packageCodegens) {
163                PackageParts generatedParts = packageCodegen.getPackageParts();
164                PackageParts premergedParts = new PackageParts(generatedParts.getPackageFqName());
165                mergedPartsByPackageName.put(generatedParts.getPackageFqName(), premergedParts);
166                premergedParts.getParts().addAll(generatedParts.getParts());
167            }
168    
169            for (MultifileClassCodegen multifileClassCodegen : multifileClassCodegens) {
170                PackageParts multifileClassParts = multifileClassCodegen.getPackageParts();
171                PackageParts premergedParts = mergedPartsByPackageName.get(multifileClassParts.getPackageFqName());
172                if (premergedParts == null) {
173                    premergedParts = new PackageParts(multifileClassParts.getPackageFqName());
174                    mergedPartsByPackageName.put(multifileClassParts.getPackageFqName(), premergedParts);
175                }
176                premergedParts.getParts().addAll(multifileClassParts.getParts());
177            }
178    
179            List<PackageParts> result = new ArrayList<PackageParts>();
180            result.addAll(mergedPartsByPackageName.values());
181            return result;
182        }
183    
184        @NotNull
185        @Override
186        public List<OutputFile> asList() {
187            done();
188            return ContainerUtil.map(generators.keySet(), new Function<String, OutputFile>() {
189                @Override
190                public OutputFile fun(String relativeClassFilePath) {
191                    return new OutputClassFile(relativeClassFilePath);
192                }
193            });
194        }
195    
196        @Override
197        @Nullable
198        public OutputFile get(@NotNull String relativePath) {
199            return generators.containsKey(relativePath) ? new OutputClassFile(relativePath) : null;
200        }
201    
202        @NotNull
203        @TestOnly
204        public String createText() {
205            StringBuilder answer = new StringBuilder();
206    
207            for (OutputFile file : asList()) {
208                answer.append("@").append(file.getRelativePath()).append('\n');
209                answer.append(file.asText());
210            }
211    
212            return answer.toString();
213        }
214    
215        @NotNull
216        @TestOnly
217        public Map<String, String> createTextForEachFile() {
218            Map<String, String> answer = new LinkedHashMap<String, String>();
219            for (OutputFile file : asList()) {
220                answer.put(file.getRelativePath(), file.asText());
221            }
222            return answer;
223        }
224    
225        @NotNull
226        public PackageCodegen forPackage(@NotNull FqName fqName, @NotNull Collection<KtFile> files) {
227            assert !isDone : "Already done!";
228            PackageCodegen codegen = package2codegen.get(fqName);
229            if (codegen == null) {
230                codegen = new PackageCodegen(state, files, fqName);
231                package2codegen.put(fqName, codegen);
232            }
233    
234            return codegen;
235        }
236    
237        @NotNull
238        public MultifileClassCodegen forMultifileClass(@NotNull FqName facadeFqName, @NotNull Collection<KtFile> files) {
239            assert !isDone : "Already done!";
240            MultifileClassCodegen codegen = multifileClass2codegen.get(facadeFqName);
241            if (codegen == null) {
242                codegen = new MultifileClassCodegen(state, files, facadeFqName);
243                multifileClass2codegen.put(facadeFqName, codegen);
244            }
245            return codegen;
246        }
247    
248        @NotNull
249        private static List<File> toIoFilesIgnoringNonPhysical(@NotNull Collection<? extends PsiFile> psiFiles) {
250            List<File> result = Lists.newArrayList();
251            for (PsiFile psiFile : psiFiles) {
252                VirtualFile virtualFile = psiFile.getVirtualFile();
253                // We ignore non-physical files here, because this code is needed to tell the make what inputs affect which outputs
254                // a non-physical file cannot be processed by make
255                if (virtualFile != null) {
256                    result.add(new File(virtualFile.getPath()));
257                }
258            }
259            return result;
260        }
261    
262        private class OutputClassFile implements OutputFile {
263            private final String relativeClassFilePath;
264    
265            public OutputClassFile(String relativeClassFilePath) {
266                this.relativeClassFilePath = relativeClassFilePath;
267            }
268    
269            @NotNull
270            @Override
271            public String getRelativePath() {
272                return relativeClassFilePath;
273            }
274    
275            @NotNull
276            @Override
277            public List<File> getSourceFiles() {
278                OutAndSourceFileList pair = generators.get(relativeClassFilePath);
279                if (pair == null) {
280                    throw new IllegalStateException("No record for binary file " + relativeClassFilePath);
281                }
282    
283                return pair.sourceFiles;
284            }
285    
286            @NotNull
287            @Override
288            public byte[] asByteArray() {
289                return generators.get(relativeClassFilePath).asBytes(builderFactory);
290            }
291    
292            @NotNull
293            @Override
294            public String asText() {
295                return generators.get(relativeClassFilePath).asText(builderFactory);
296            }
297    
298            @NotNull
299            @Override
300            public String toString() {
301                return getRelativePath() + " (compiled from " + getSourceFiles() + ")";
302            }
303        }
304    
305        private static final class ClassBuilderAndSourceFileList extends OutAndSourceFileList {
306            private final ClassBuilder classBuilder;
307    
308            private ClassBuilderAndSourceFileList(ClassBuilder classBuilder, List<File> sourceFiles) {
309                super(sourceFiles);
310                this.classBuilder = classBuilder;
311            }
312    
313            @Override
314            public byte[] asBytes(ClassBuilderFactory factory) {
315                return factory.asBytes(classBuilder);
316            }
317    
318            @Override
319            public String asText(ClassBuilderFactory factory) {
320                return factory.asText(classBuilder);
321            }
322        }
323    
324        private static abstract class OutAndSourceFileList {
325    
326            protected final List<File> sourceFiles;
327    
328            private OutAndSourceFileList(List<File> sourceFiles) {
329                this.sourceFiles = sourceFiles;
330            }
331    
332            public abstract byte[] asBytes(ClassBuilderFactory factory);
333    
334            public abstract String asText(ClassBuilderFactory factory);
335        }
336    
337        public void removeClasses(Set<String> classNamesToRemove) {
338            for (String classInternalName : classNamesToRemove) {
339                generators.remove(classInternalName + ".class");
340            }
341        }
342    
343        @TestOnly
344        public List<KtFile> getInputFiles() {
345            return state.getFiles();
346        }
347    }