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