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