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.cli.jvm.compiler;
018    
019    import com.google.common.collect.Lists;
020    import com.google.common.collect.Sets;
021    import com.intellij.openapi.project.Project;
022    import com.intellij.openapi.util.io.FileUtil;
023    import com.intellij.openapi.util.io.FileUtilRt;
024    import com.intellij.openapi.vfs.StandardFileSystems;
025    import com.intellij.openapi.vfs.VirtualFile;
026    import com.intellij.openapi.vfs.VirtualFileManager;
027    import com.intellij.openapi.vfs.VirtualFileSystem;
028    import com.intellij.psi.PsiFile;
029    import com.intellij.psi.PsiManager;
030    import kotlin.Unit;
031    import kotlin.io.FilesKt;
032    import kotlin.jvm.functions.Function1;
033    import kotlin.sequences.SequencesKt;
034    import org.jetbrains.annotations.NotNull;
035    import org.jetbrains.annotations.Nullable;
036    import org.jetbrains.kotlin.backend.common.output.OutputFile;
037    import org.jetbrains.kotlin.cli.common.messages.MessageCollector;
038    import org.jetbrains.kotlin.cli.common.modules.ModuleScriptData;
039    import org.jetbrains.kotlin.cli.common.modules.ModuleXmlParser;
040    import org.jetbrains.kotlin.codegen.ClassFileFactory;
041    import org.jetbrains.kotlin.idea.KotlinFileType;
042    import org.jetbrains.kotlin.name.FqName;
043    import org.jetbrains.kotlin.psi.KtFile;
044    import org.jetbrains.kotlin.utils.ExceptionUtilsKt;
045    import org.jetbrains.kotlin.utils.PathUtil;
046    
047    import java.io.*;
048    import java.util.Collection;
049    import java.util.List;
050    import java.util.Set;
051    import java.util.jar.*;
052    
053    import static org.jetbrains.kotlin.cli.common.messages.CompilerMessageLocation.NO_LOCATION;
054    import static org.jetbrains.kotlin.cli.common.messages.CompilerMessageSeverity.ERROR;
055    
056    public class CompileEnvironmentUtil {
057    
058        @NotNull
059        public static ModuleScriptData loadModuleDescriptions(String moduleDefinitionFile, MessageCollector messageCollector) {
060            File file = new File(moduleDefinitionFile);
061            if (!file.exists()) {
062                messageCollector.report(ERROR, "Module definition file does not exist: " + moduleDefinitionFile, NO_LOCATION);
063                return ModuleScriptData.EMPTY;
064            }
065            String extension = FileUtilRt.getExtension(moduleDefinitionFile);
066            if ("xml".equalsIgnoreCase(extension)) {
067                return ModuleXmlParser.parseModuleScript(moduleDefinitionFile, messageCollector);
068            }
069            messageCollector.report(ERROR, "Unknown module definition type: " + moduleDefinitionFile, NO_LOCATION);
070            return ModuleScriptData.EMPTY;
071        }
072    
073        // TODO: includeRuntime should be not a flag but a path to runtime
074        private static void doWriteToJar(ClassFileFactory outputFiles, OutputStream fos, @Nullable FqName mainClass, boolean includeRuntime) {
075            try {
076                Manifest manifest = new Manifest();
077                Attributes mainAttributes = manifest.getMainAttributes();
078                mainAttributes.putValue("Manifest-Version", "1.0");
079                mainAttributes.putValue("Created-By", "JetBrains Kotlin");
080                if (mainClass != null) {
081                    mainAttributes.putValue("Main-Class", mainClass.asString());
082                }
083                JarOutputStream stream = new JarOutputStream(fos, manifest);
084                for (OutputFile outputFile : outputFiles.asList()) {
085                    stream.putNextEntry(new JarEntry(outputFile.getRelativePath()));
086                    stream.write(outputFile.asByteArray());
087                }
088                if (includeRuntime) {
089                    writeRuntimeToJar(stream);
090                }
091                stream.finish();
092            }
093            catch (IOException e) {
094                throw new CompileEnvironmentException("Failed to generate jar file", e);
095            }
096        }
097    
098        public static void writeToJar(File jarPath, boolean jarRuntime, FqName mainClass, ClassFileFactory outputFiles) {
099            FileOutputStream outputStream = null;
100            try {
101                outputStream = new FileOutputStream(jarPath);
102                doWriteToJar(outputFiles, outputStream, mainClass, jarRuntime);
103                outputStream.close();
104            }
105            catch (FileNotFoundException e) {
106                throw new CompileEnvironmentException("Invalid jar path " + jarPath, e);
107            }
108            catch (IOException e) {
109                throw ExceptionUtilsKt.rethrow(e);
110            }
111            finally {
112                ExceptionUtilsKt.closeQuietly(outputStream);
113            }
114        }
115    
116        private static void writeRuntimeToJar(JarOutputStream stream) throws IOException {
117            File runtimePath = PathUtil.getKotlinPathsForCompiler().getRuntimePath();
118            if (!runtimePath.exists()) {
119                throw new CompileEnvironmentException("Couldn't find runtime library");
120            }
121    
122            JarInputStream jis = new JarInputStream(new FileInputStream(runtimePath));
123            try {
124                while (true) {
125                    JarEntry e = jis.getNextJarEntry();
126                    if (e == null) {
127                        break;
128                    }
129                    if (FileUtilRt.extensionEquals(e.getName(), "class")) {
130                        stream.putNextEntry(e);
131                        FileUtil.copy(jis, stream);
132                    }
133                }
134            }
135            finally {
136                jis.close();
137            }
138        }
139    
140        @NotNull
141        public static List<KtFile> getKtFiles(
142                @NotNull final Project project,
143                @NotNull Collection<String> sourceRoots,
144                @NotNull Function1<String, Unit> reportError
145        ) {
146            final VirtualFileSystem localFileSystem = VirtualFileManager.getInstance().getFileSystem(StandardFileSystems.FILE_PROTOCOL);
147    
148            final Set<VirtualFile> processedFiles = Sets.newHashSet();
149            final List<KtFile> result = Lists.newArrayList();
150    
151            for (String sourceRootPath : sourceRoots) {
152                if (sourceRootPath == null) {
153                    continue;
154                }
155    
156                VirtualFile vFile = localFileSystem.findFileByPath(sourceRootPath);
157                if (vFile == null) {
158                    reportError.invoke("Source file or directory not found: " + sourceRootPath);
159                    continue;
160                }
161                if (!vFile.isDirectory() && vFile.getFileType() != KotlinFileType.INSTANCE) {
162                    reportError.invoke("Source entry is not a Kotlin file: " + sourceRootPath);
163                    continue;
164                }
165    
166                SequencesKt.forEach(FilesKt.walkTopDown(new File(sourceRootPath)), new Function1<File, Unit>() {
167                    @Override
168                    public Unit invoke(File file) {
169                        if (file.isFile()) {
170                            VirtualFile virtualFile = localFileSystem.findFileByPath(file.getAbsolutePath());
171                            if (virtualFile != null && !processedFiles.contains(virtualFile)) {
172                                processedFiles.add(virtualFile);
173                                PsiFile psiFile = PsiManager.getInstance(project).findFile(virtualFile);
174                                if (psiFile instanceof KtFile) {
175                                    result.add((KtFile) psiFile);
176                                }
177                            }
178                        }
179                        return Unit.INSTANCE;
180                    }
181                });
182            }
183    
184            return result;
185        }
186    }