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 {
282                OutputUtilsPackage.writeAll(outputFiles, outputDir == null ? new File(".") : outputDir, messageCollector);
283            }
284        }
285    
286        @NotNull
287        public static List<JetFile> getJetFiles(
288                @NotNull final Project project,
289                @NotNull List<String> sourceRoots,
290                @NotNull Function1<String, Unit> reportError
291        ) {
292            final VirtualFileSystem localFileSystem = VirtualFileManager.getInstance().getFileSystem(StandardFileSystems.FILE_PROTOCOL);
293    
294            final List<JetFile> result = Lists.newArrayList();
295    
296            for (String sourceRootPath : sourceRoots) {
297                if (sourceRootPath == null) {
298                    continue;
299                }
300    
301                VirtualFile vFile = localFileSystem.findFileByPath(sourceRootPath);
302                if (vFile == null) {
303                    reportError.invoke("Source file or directory not found: " + sourceRootPath);
304                    continue;
305                }
306                if (!vFile.isDirectory() && vFile.getFileType() != JetFileType.INSTANCE) {
307                    reportError.invoke("Source entry is not a Kotlin file: " + sourceRootPath);
308                    continue;
309                }
310    
311                IoPackage.recurse(new File(sourceRootPath), new Function1<File, Unit>() {
312                    @Override
313                    public Unit invoke(File file) {
314                        if (file.isFile()) {
315                            VirtualFile fileByPath = localFileSystem.findFileByPath(file.getAbsolutePath());
316                            if (fileByPath != null) {
317                                PsiFile psiFile = PsiManager.getInstance(project).findFile(fileByPath);
318                                if (psiFile instanceof JetFile) {
319                                    result.add((JetFile) psiFile);
320                                }
321                            }
322                        }
323                        return Unit.VALUE;
324                    }
325                });
326            }
327    
328            return result;
329        }
330    
331        public static class ModuleScriptData {
332            @Nullable
333            private final String incrementalCacheDir;
334            @NotNull
335            private final List<Module> modules;
336    
337            @NotNull
338            public List<Module> getModules() {
339                return modules;
340            }
341    
342            @Nullable
343            public String getIncrementalCacheDir() {
344                return incrementalCacheDir;
345            }
346    
347            private ModuleScriptData(@NotNull List<Module> modules, @Nullable String incrementalCacheDir) {
348                this.incrementalCacheDir = incrementalCacheDir;
349                this.modules = modules;
350            }
351        }
352    
353        private static class DescriptionToModuleAdapter implements Module {
354            private final ModuleDescription description;
355    
356            public DescriptionToModuleAdapter(ModuleDescription description) {
357                this.description = description;
358            }
359    
360            @NotNull
361            @Override
362            public String getModuleName() {
363                return description.getModuleName();
364            }
365    
366            @NotNull
367            @Override
368            public String getOutputDirectory() {
369                return description.getOutputDir();
370            }
371    
372            @NotNull
373            @Override
374            public List<String> getSourceFiles() {
375                return description.getSourceFiles();
376            }
377    
378            @NotNull
379            @Override
380            public List<String> getClasspathRoots() {
381                return description.getClasspathRoots();
382            }
383    
384            @NotNull
385            @Override
386            public List<String> getAnnotationsRoots() {
387                return description.getAnnotationsRoots();
388            }
389        }
390    }