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