Class file factory remembers sources for each class

compiler/backend/src/org/jetbrains/jet/codegen/ClassFileFactory.java

@@ -17,8 +17,11 @@
 package org.jetbrains.jet.codegen;
 
 import com.google.common.collect.Lists;
+import com.intellij.openapi.vfs.VfsUtilCore;
 import com.intellij.openapi.vfs.VirtualFile;
 import com.intellij.psi.PsiFile;
+import com.intellij.util.Function;
+import com.intellij.util.containers.ContainerUtil;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.asm4.Type;
 import org.jetbrains.jet.codegen.state.GenerationState;
@@ -40,7 +43,7 @@ public final class ClassFileFactory extends GenerationStateAware {
     @NotNull private ClassBuilderFactory builderFactory;
 
     private final Map<FqName, NamespaceCodegen> ns2codegen = new HashMap<FqName, NamespaceCodegen>();
-    private final Map<String, ClassBuilder> generators = new LinkedHashMap<String, ClassBuilder>();
+    private final Map<String, ClassBuilderAndSourceFileList> generators = new LinkedHashMap<String, ClassBuilderAndSourceFileList>();
     private boolean isDone = false;
 
     public ClassFileFactory(@NotNull GenerationState state) {
@@ -63,7 +66,7 @@ public final class ClassFileFactory extends GenerationStateAware {
         String outputFilePath = asmType.getInternalName() + ".class";
         state.getProgress().reportOutput(toIoFilesIgnoringNonPhysical(sourceFiles), new File(outputFilePath));
         ClassBuilder answer = builderFactory.newClassBuilder();
-        generators.put(outputFilePath, answer);
+        generators.put(outputFilePath, new ClassBuilderAndSourceFileList(answer, sourceFiles));
         return answer;
     }
 
@@ -78,12 +81,12 @@ public final class ClassFileFactory extends GenerationStateAware {
 
     public String asText(String file) {
         done();
-        return builderFactory.asText(generators.get(file));
+        return builderFactory.asText(generators.get(file).classBuilder);
     }
 
     public byte[] asBytes(String file) {
         done();
-        return builderFactory.asBytes(generators.get(file));
+        return builderFactory.asBytes(generators.get(file).classBuilder);
     }
 
     public List<String> files() {
@@ -91,6 +94,26 @@ public final class ClassFileFactory extends GenerationStateAware {
         return new ArrayList<String>(generators.keySet());
     }
 
+    public List<File> getSourceFiles(String relativeClassFilePath) {
+        ClassBuilderAndSourceFileList pair = generators.get(relativeClassFilePath);
+        if (pair == null) {
+            throw new IllegalStateException("No record for binary file " + relativeClassFilePath);
+        }
+
+        return ContainerUtil.mapNotNull(
+                pair.sourceFiles,
+                new Function<PsiFile, File>() {
+                    @Override
+                    public File fun(PsiFile file) {
+                        VirtualFile virtualFile = file.getVirtualFile();
+                        if (virtualFile == null) return null;
+
+                        return VfsUtilCore.virtualToIoFile(virtualFile);
+                    }
+                }
+        );
+    }
+
     public String createText() {
         StringBuilder answer = new StringBuilder();
 
@@ -155,4 +178,14 @@ public final class ClassFileFactory extends GenerationStateAware {
         return result;
     }
 
+    private static class ClassBuilderAndSourceFileList {
+        private final ClassBuilder classBuilder;
+        private final Collection<? extends PsiFile> sourceFiles;
+
+        private ClassBuilderAndSourceFileList(ClassBuilder classBuilder, Collection<? extends PsiFile> sourceFiles) {
+            this.classBuilder = classBuilder;
+            this.sourceFiles = sourceFiles;
+        }
+    }
+
 }