001    /*
002     * Copyright 2010-2013 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.jet.cli.jvm.compiler;
018    
019    import com.google.common.collect.Lists;
020    import com.intellij.openapi.Disposable;
021    import com.intellij.openapi.project.Project;
022    import com.intellij.openapi.util.Disposer;
023    import com.intellij.openapi.util.io.FileUtil;
024    import com.intellij.openapi.util.io.FileUtilRt;
025    import com.intellij.openapi.vfs.StandardFileSystems;
026    import com.intellij.openapi.vfs.VirtualFile;
027    import com.intellij.openapi.vfs.VirtualFileManager;
028    import com.intellij.openapi.vfs.VirtualFileSystem;
029    import com.intellij.psi.PsiFile;
030    import com.intellij.psi.PsiManager;
031    import kotlin.Function1;
032    import kotlin.Unit;
033    import kotlin.io.IoPackage;
034    import kotlin.modules.AllModules;
035    import kotlin.modules.Module;
036    import org.jetbrains.annotations.NotNull;
037    import org.jetbrains.annotations.Nullable;
038    import org.jetbrains.jet.OutputFile;
039    import org.jetbrains.jet.cli.common.CLIConfigurationKeys;
040    import org.jetbrains.jet.cli.common.messages.MessageCollector;
041    import org.jetbrains.jet.cli.common.messages.MessageRenderer;
042    import org.jetbrains.jet.cli.common.modules.ModuleScriptData;
043    import org.jetbrains.jet.cli.common.modules.ModuleXmlParser;
044    import org.jetbrains.jet.cli.common.output.outputUtils.OutputUtilsPackage;
045    import org.jetbrains.jet.cli.jvm.JVMConfigurationKeys;
046    import org.jetbrains.jet.codegen.ClassFileFactory;
047    import org.jetbrains.jet.codegen.GeneratedClassLoader;
048    import org.jetbrains.jet.codegen.state.GenerationState;
049    import org.jetbrains.jet.config.CommonConfigurationKeys;
050    import org.jetbrains.jet.config.CompilerConfiguration;
051    import org.jetbrains.jet.lang.psi.JetFile;
052    import org.jetbrains.jet.lang.resolve.java.PackageClassUtils;
053    import org.jetbrains.jet.lang.resolve.name.FqName;
054    import org.jetbrains.jet.plugin.JetFileType;
055    import org.jetbrains.jet.utils.KotlinPaths;
056    import org.jetbrains.jet.utils.PathUtil;
057    import org.jetbrains.jet.utils.UtilsPackage;
058    
059    import java.io.*;
060    import java.lang.reflect.Method;
061    import java.net.MalformedURLException;
062    import java.net.URL;
063    import java.net.URLClassLoader;
064    import java.util.ArrayList;
065    import java.util.List;
066    import java.util.jar.*;
067    
068    import static org.jetbrains.jet.cli.common.messages.CompilerMessageLocation.NO_LOCATION;
069    import static org.jetbrains.jet.cli.common.messages.CompilerMessageSeverity.ERROR;
070    
071    public class CompileEnvironmentUtil {
072    
073        @Nullable
074        private static File getRuntimeJarPath() {
075            File runtimePath = PathUtil.getKotlinPathsForCompiler().getRuntimePath();
076            return runtimePath.exists() ? runtimePath : null;
077        }
078    
079        @NotNull
080        public static ModuleScriptData loadModuleDescriptions(KotlinPaths paths, String moduleDefinitionFile, MessageCollector messageCollector) {
081            File file = new File(moduleDefinitionFile);
082            if (!file.exists()) {
083                messageCollector.report(ERROR, "Module definition file does not exist: " + moduleDefinitionFile, NO_LOCATION);
084                return ModuleScriptData.EMPTY;
085            }
086            String extension = FileUtilRt.getExtension(moduleDefinitionFile);
087            if ("ktm".equalsIgnoreCase(extension)) {
088                return loadModuleScript(paths, moduleDefinitionFile, messageCollector);
089            }
090            if ("xml".equalsIgnoreCase(extension)) {
091                return ModuleXmlParser.parseModuleScript(moduleDefinitionFile, messageCollector);
092            }
093            messageCollector.report(ERROR, "Unknown module definition type: " + moduleDefinitionFile, NO_LOCATION);
094            return ModuleScriptData.EMPTY;
095        }
096    
097        @NotNull
098        private static ModuleScriptData loadModuleScript(KotlinPaths paths, String moduleScriptFile, MessageCollector messageCollector) {
099            CompilerConfiguration configuration = new CompilerConfiguration();
100            File runtimePath = paths.getRuntimePath();
101            if (runtimePath.exists()) {
102                configuration.add(JVMConfigurationKeys.CLASSPATH_KEY, runtimePath);
103            }
104            configuration.addAll(JVMConfigurationKeys.CLASSPATH_KEY, PathUtil.getJdkClassesRoots());
105            File jdkAnnotationsPath = paths.getJdkAnnotationsPath();
106            if (jdkAnnotationsPath.exists()) {
107                configuration.add(JVMConfigurationKeys.ANNOTATIONS_PATH_KEY, jdkAnnotationsPath);
108            }
109            configuration.add(CommonConfigurationKeys.SOURCE_ROOTS_KEY, moduleScriptFile);
110            configuration.put(CLIConfigurationKeys.MESSAGE_COLLECTOR_KEY, messageCollector);
111    
112            List<Module> modules;
113    
114            Disposable disposable = Disposer.newDisposable();
115            try {
116                JetCoreEnvironment scriptEnvironment = JetCoreEnvironment.createForProduction(disposable, configuration);
117                GenerationState generationState = KotlinToJVMBytecodeCompiler.analyzeAndGenerate(scriptEnvironment);
118                if (generationState == null) {
119                    throw new CompileEnvironmentException("Module script " + moduleScriptFile + " analyze failed:\n" +
120                                                          loadModuleScriptText(moduleScriptFile));
121                }
122    
123                modules = runDefineModules(paths, generationState.getFactory());
124            }
125            finally {
126                Disposer.dispose(disposable);
127            }
128    
129            if (modules == null) {
130                throw new CompileEnvironmentException("Module script " + moduleScriptFile + " compilation failed");
131            }
132    
133            if (modules.isEmpty()) {
134                throw new CompileEnvironmentException("No modules where defined by " + moduleScriptFile);
135            }
136            return new ModuleScriptData(modules, null);
137        }
138    
139        private static List<Module> runDefineModules(KotlinPaths paths, ClassFileFactory factory) {
140            File stdlibJar = paths.getRuntimePath();
141            GeneratedClassLoader loader;
142            if (stdlibJar.exists()) {
143                try {
144                    loader = new GeneratedClassLoader(factory, new URLClassLoader(new URL[]{stdlibJar.toURI().toURL()},
145                                                                                  AllModules.class.getClassLoader()));
146                }
147                catch (MalformedURLException e) {
148                    throw new RuntimeException(e);
149                }
150            }
151            else {
152                loader = new GeneratedClassLoader(factory, KotlinToJVMBytecodeCompiler.class.getClassLoader());
153            }
154            try {
155                Class<?> packageClass = loader.loadClass(PackageClassUtils.getPackageClassName(FqName.ROOT));
156                Method method = packageClass.getDeclaredMethod("project");
157    
158                method.setAccessible(true);
159                method.invoke(null);
160    
161                ArrayList<Module> answer = new ArrayList<Module>(AllModules.INSTANCE$.get());
162                AllModules.INSTANCE$.get().clear();
163                return answer;
164            }
165            catch (Exception e) {
166                throw new ModuleExecutionException(e);
167            }
168            finally {
169                loader.dispose();
170            }
171        }
172    
173        // TODO: includeRuntime should be not a flag but a path to runtime
174        private static void doWriteToJar(ClassFileFactory outputFiles, OutputStream fos, @Nullable FqName mainClass, boolean includeRuntime) {
175            try {
176                Manifest manifest = new Manifest();
177                Attributes mainAttributes = manifest.getMainAttributes();
178                mainAttributes.putValue("Manifest-Version", "1.0");
179                mainAttributes.putValue("Created-By", "JetBrains Kotlin");
180                if (mainClass != null) {
181                    mainAttributes.putValue("Main-Class", mainClass.asString());
182                }
183                JarOutputStream stream = new JarOutputStream(fos, manifest);
184                for (OutputFile outputFile : outputFiles.asList()) {
185                    stream.putNextEntry(new JarEntry(outputFile.getRelativePath()));
186                    stream.write(outputFile.asByteArray());
187                }
188                if (includeRuntime) {
189                    writeRuntimeToJar(stream);
190                }
191                stream.finish();
192            }
193            catch (IOException e) {
194                throw new CompileEnvironmentException("Failed to generate jar file", e);
195            }
196        }
197    
198        public static void writeToJar(File jarPath, boolean jarRuntime, FqName mainClass, ClassFileFactory outputFiles) {
199            FileOutputStream outputStream = null;
200            try {
201                outputStream = new FileOutputStream(jarPath);
202                doWriteToJar(outputFiles, outputStream, mainClass, jarRuntime);
203                outputStream.close();
204            }
205            catch (FileNotFoundException e) {
206                throw new CompileEnvironmentException("Invalid jar path " + jarPath, e);
207            }
208            catch (IOException e) {
209                throw UtilsPackage.rethrow(e);
210            }
211            finally {
212                UtilsPackage.closeQuietly(outputStream);
213            }
214        }
215    
216        private static void writeRuntimeToJar(JarOutputStream stream) throws IOException {
217            File runtimeJarPath = getRuntimeJarPath();
218            if (runtimeJarPath != null) {
219                JarInputStream jis = new JarInputStream(new FileInputStream(runtimeJarPath));
220                try {
221                    while (true) {
222                        JarEntry e = jis.getNextJarEntry();
223                        if (e == null) {
224                            break;
225                        }
226                        if (FileUtilRt.extensionEquals(e.getName(), "class")) {
227                            stream.putNextEntry(e);
228                            FileUtil.copy(jis, stream);
229                        }
230                    }
231                }
232                finally {
233                    jis.close();
234                }
235            }
236            else {
237                throw new CompileEnvironmentException("Couldn't find runtime library");
238            }
239        }
240    
241        // Used for debug output only
242        private static String loadModuleScriptText(String moduleScriptFile) {
243            String moduleScriptText;
244            try {
245                moduleScriptText = FileUtil.loadFile(new File(moduleScriptFile));
246            }
247            catch (IOException e) {
248                moduleScriptText = "Can't load module script text:\n" + MessageRenderer.PLAIN.renderException(e);
249            }
250            return moduleScriptText;
251        }
252    
253        static void writeOutputToDirOrJar(
254                @Nullable File jar,
255                @Nullable File outputDir,
256                boolean includeRuntime,
257                @Nullable FqName mainClass,
258                @NotNull ClassFileFactory outputFiles,
259                @NotNull MessageCollector messageCollector
260        ) {
261            if (jar != null) {
262                writeToJar(jar, includeRuntime, mainClass, outputFiles);
263            }
264            else {
265                OutputUtilsPackage.writeAll(outputFiles, outputDir == null ? new File(".") : outputDir, messageCollector);
266            }
267        }
268    
269        @NotNull
270        public static List<JetFile> getJetFiles(
271                @NotNull final Project project,
272                @NotNull List<String> sourceRoots,
273                @NotNull Function1<String, Unit> reportError
274        ) {
275            final VirtualFileSystem localFileSystem = VirtualFileManager.getInstance().getFileSystem(StandardFileSystems.FILE_PROTOCOL);
276    
277            final List<JetFile> result = Lists.newArrayList();
278    
279            for (String sourceRootPath : sourceRoots) {
280                if (sourceRootPath == null) {
281                    continue;
282                }
283    
284                VirtualFile vFile = localFileSystem.findFileByPath(sourceRootPath);
285                if (vFile == null) {
286                    reportError.invoke("Source file or directory not found: " + sourceRootPath);
287                    continue;
288                }
289                if (!vFile.isDirectory() && vFile.getFileType() != JetFileType.INSTANCE) {
290                    reportError.invoke("Source entry is not a Kotlin file: " + sourceRootPath);
291                    continue;
292                }
293    
294                IoPackage.recurse(new File(sourceRootPath), new Function1<File, Unit>() {
295                    @Override
296                    public Unit invoke(File file) {
297                        if (file.isFile()) {
298                            VirtualFile fileByPath = localFileSystem.findFileByPath(file.getAbsolutePath());
299                            if (fileByPath != null) {
300                                PsiFile psiFile = PsiManager.getInstance(project).findFile(fileByPath);
301                                if (psiFile instanceof JetFile) {
302                                    result.add((JetFile) psiFile);
303                                }
304                            }
305                        }
306                        return Unit.INSTANCE$;
307                    }
308                });
309            }
310    
311            return result;
312        }
313    }