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