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