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