Simplify ClassFileFactory state

Before this change it contained maps to complex object,
after it retains only necessary information:
- sources files for package parts
- PackageParts objects grouped by package fq name

Both of them needed for writing module mappings

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

@@ -20,9 +20,12 @@ import com.google.common.collect.Lists;
 import com.intellij.openapi.vfs.VirtualFile;
 import com.intellij.psi.PsiFile;
 import com.intellij.util.Function;
+import com.intellij.util.SmartList;
 import com.intellij.util.containers.ContainerUtil;
 import com.intellij.util.io.DataOutputStream;
 import kotlin.collections.CollectionsKt;
+import kotlin.collections.MapsKt;
+import kotlin.jvm.functions.Function0;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.annotations.TestOnly;
@@ -49,12 +52,13 @@ import static org.jetbrains.kotlin.codegen.JvmCodegenUtil.getMappingFileName;
 public class ClassFileFactory implements OutputFileCollection {
     private final GenerationState state;
     private final ClassBuilderFactory builderFactory;
-    private final Map<FqName, PackageCodegen> package2codegen = new HashMap<FqName, PackageCodegen>();
-    private final Map<FqName, MultifileClassCodegen> multifileClass2codegen = new HashMap<FqName, MultifileClassCodegen>();
     private final Map<String, OutAndSourceFileList> generators = new LinkedHashMap<String, OutAndSourceFileList>();
 
     private boolean isDone = false;
 
+    private final Set<File> packagePartSourceFiles = new HashSet<File>();
+    private final Map<String, PackageParts> partsGroupedByPackage = new LinkedHashMap<String, PackageParts>();
+
     public ClassFileFactory(@NotNull GenerationState state, @NotNull ClassBuilderFactory builderFactory) {
         this.state = state;
         this.builderFactory = builderFactory;
@@ -84,37 +88,23 @@ public class ClassFileFactory implements OutputFileCollection {
     void done() {
         if (!isDone) {
             isDone = true;
-            Collection<PackageCodegen> packageCodegens = package2codegen.values();
+            writeModuleMappings();
-            Collection<MultifileClassCodegen> multifileClassCodegens = multifileClass2codegen.values();
-            writeModuleMappings(packageCodegens, multifileClassCodegens);
         }
     }
 
-    private void writeModuleMappings(
+    private void writeModuleMappings() {
-            @NotNull Collection<PackageCodegen> packageCodegens,
-            @NotNull Collection<MultifileClassCodegen> multifileClassCodegens
-    ) {
         final JvmPackageTable.PackageTable.Builder builder = JvmPackageTable.PackageTable.newBuilder();
         String outputFilePath = getMappingFileName(state.getModuleName());
 
-        List<PackageParts> parts = collectGeneratedPackageParts(packageCodegens, multifileClassCodegens);
+        List<PackageParts> parts = new SmartList<PackageParts>(partsGroupedByPackage.values());
-
-        Set<File> sourceFiles = new HashSet<File>();
-        // TODO extract common logic
-        for (PackageCodegen codegen : packageCodegens) {
-            sourceFiles.addAll(toIoFilesIgnoringNonPhysical(PackagePartClassUtils.getFilesWithCallables(codegen.getFiles())));
-        }
-        for (MultifileClassCodegen codegen : multifileClassCodegens) {
-            sourceFiles.addAll(toIoFilesIgnoringNonPhysical(PackagePartClassUtils.getFilesWithCallables(codegen.getFiles())));
-        }
 
         for (PackageParts part : ClassFileUtilsKt.addCompiledPartsAndSort(parts, state)) {
             PackageParts.Companion.serialize(part, builder);
         }
 
         if (builder.getPackagePartsCount() != 0) {
-            state.getProgress().reportOutput(sourceFiles, new File(outputFilePath));
+            state.getProgress().reportOutput(packagePartSourceFiles, new File(outputFilePath));
-            generators.put(outputFilePath, new OutAndSourceFileList(CollectionsKt.toList(sourceFiles)) {
+            generators.put(outputFilePath, new OutAndSourceFileList(CollectionsKt.toList(packagePartSourceFiles)) {
                 @Override
                 public byte[] asBytes(ClassBuilderFactory factory) {
                     try {
@@ -150,34 +140,6 @@ public class ClassFileFactory implements OutputFileCollection {
         }
     }
 
-    private static List<PackageParts> collectGeneratedPackageParts(
-            @NotNull Collection<PackageCodegen> packageCodegens,
-            @NotNull Collection<MultifileClassCodegen> multifileClassCodegens
-    ) {
-        Map<String, PackageParts> mergedPartsByPackageName = new LinkedHashMap<String, PackageParts>();
-
-        for (PackageCodegen packageCodegen : packageCodegens) {
-            PackageParts generatedParts = packageCodegen.getPackageParts();
-            PackageParts premergedParts = new PackageParts(generatedParts.getPackageFqName());
-            mergedPartsByPackageName.put(generatedParts.getPackageFqName(), premergedParts);
-            premergedParts.getParts().addAll(generatedParts.getParts());
-        }
-
-        for (MultifileClassCodegen multifileClassCodegen : multifileClassCodegens) {
-            PackageParts multifileClassParts = multifileClassCodegen.getPackageParts();
-            PackageParts premergedParts = mergedPartsByPackageName.get(multifileClassParts.getPackageFqName());
-            if (premergedParts == null) {
-                premergedParts = new PackageParts(multifileClassParts.getPackageFqName());
-                mergedPartsByPackageName.put(multifileClassParts.getPackageFqName(), premergedParts);
-            }
-            premergedParts.getParts().addAll(multifileClassParts.getParts());
-        }
-
-        List<PackageParts> result = new ArrayList<PackageParts>();
-        result.addAll(mergedPartsByPackageName.values());
-        return result;
-    }
-
     @NotNull
     @Override
     public List<OutputFile> asList() {
@@ -222,24 +184,34 @@ public class ClassFileFactory implements OutputFileCollection {
     @NotNull
     public PackageCodegen forPackage(@NotNull FqName fqName, @NotNull Collection<KtFile> files) {
         assert !isDone : "Already done!";
-        PackageCodegen codegen = package2codegen.get(fqName);
+        registerPackagePartSourceFiles(files);
-        if (codegen == null) {
+        return new PackageCodegen(state, files, fqName, buildNewPackagePartRegistry(fqName));
-            codegen = new PackageCodegen(state, files, fqName);
-            package2codegen.put(fqName, codegen);
-        }
-
-        return codegen;
     }
 
     @NotNull
     public MultifileClassCodegen forMultifileClass(@NotNull FqName facadeFqName, @NotNull Collection<KtFile> files) {
         assert !isDone : "Already done!";
-        MultifileClassCodegen codegen = multifileClass2codegen.get(facadeFqName);
+        registerPackagePartSourceFiles(files);
-        if (codegen == null) {
+        return new MultifileClassCodegen(state, files, facadeFqName, buildNewPackagePartRegistry(facadeFqName.parent()));
-            codegen = new MultifileClassCodegen(state, files, facadeFqName);
+    }
-            multifileClass2codegen.put(facadeFqName, codegen);
+
-        }
+    private PackagePartRegistry buildNewPackagePartRegistry(@NotNull FqName packageFqName) {
-        return codegen;
+        final String packageFqNameAsString = packageFqName.asString();
+        return new PackagePartRegistry() {
+            @Override
+            public void addPart(@NotNull String partShortName) {
+                MapsKt.getOrPut(partsGroupedByPackage, packageFqNameAsString, new Function0<PackageParts>() {
+                    @Override
+                    public PackageParts invoke() {
+                        return new PackageParts(packageFqNameAsString);
+                    }
+                }).getParts().add(partShortName);
+            }
+        };
+    }
+
+    public void registerPackagePartSourceFiles(Collection<KtFile> files) {
+        packagePartSourceFiles.addAll(toIoFilesIgnoringNonPhysical(PackagePartClassUtils.getFilesWithCallables(files)));
     }
 
     @NotNull

compiler/backend/src/org/jetbrains/kotlin/codegen/MultifileClassCodegen.kt

@@ -58,7 +58,8 @@ import java.util.*
 class MultifileClassCodegen(
         private val state: GenerationState,
         val files: Collection<KtFile>,
-        private val facadeFqName: FqName
+        private val facadeFqName: FqName,
+        private val packagePartRegistry: PackagePartRegistry
 ) {
     private val facadeClassType = AsmUtil.asmTypeByFqNameWithoutInnerClasses(facadeFqName)
 
@@ -75,8 +76,6 @@ class MultifileClassCodegen(
     private fun getDeserializedCallables(compiledPackageFragment: PackageFragmentDescriptor) =
             compiledPackageFragment.getMemberScope().getContributedDescriptors(DescriptorKindFilter.CALLABLES, MemberScope.ALL_NAME_FILTER).filterIsInstance<DeserializedCallableMemberDescriptor>()
 
-    val packageParts = PackageParts(facadeFqName.parent().asString())
-
     private val classBuilder = ClassBuilderOnDemand {
         val originFile = files.firstOrNull()
         val actualPackageFragment = packageFragment ?:
@@ -202,7 +201,7 @@ class MultifileClassCodegen(
         partFqNames.add(partClassInfo.fileClassFqName)
 
         val name = partType.internalName
-        packageParts.parts.add(name.substring(name.lastIndexOf('/') + 1))
+        packagePartRegistry.addPart(name.substring(name.lastIndexOf('/') + 1))
 
         val builder = state.factory.newVisitor(MultifileClassPart(file, packageFragment, facadeFqName), partType, file)
 

compiler/backend/src/org/jetbrains/kotlin/codegen/PackageCodegen.java

@@ -41,17 +41,18 @@ public class PackageCodegen {
     private final GenerationState state;
     private final Collection<KtFile> files;
     private final PackageFragmentDescriptor packageFragment;
-    private final PackageParts packageParts;
+    private final PackagePartRegistry packagePartRegistry;
 
     public PackageCodegen(
             @NotNull GenerationState state,
             @NotNull Collection<KtFile> files,
-            @NotNull FqName packageFqName
+            @NotNull FqName packageFqName,
+            @NotNull PackagePartRegistry registry
     ) {
         this.state = state;
         this.files = files;
         this.packageFragment = getOnlyPackageFragment(packageFqName);
-        packageParts = new PackageParts(packageFqName.asString());
+        packagePartRegistry = registry;
     }
 
     public void generate(@NotNull CompilationErrorHandler errorHandler) {
@@ -110,7 +111,7 @@ public class PackageCodegen {
         if (!generatePackagePart || !state.getGenerateDeclaredClassFilter().shouldGeneratePackagePart(file)) return null;
 
         String name = fileClassType.getInternalName();
-        packageParts.getParts().add(name.substring(name.lastIndexOf('/') + 1));
+        packagePartRegistry.addPart(name.substring(name.lastIndexOf('/') + 1));
 
         ClassBuilder builder = state.getFactory().newVisitor(JvmDeclarationOriginKt.PackagePart(file, packageFragment), fileClassType, file);
 
@@ -147,10 +148,6 @@ public class PackageCodegen {
         MemberCodegen.genClassOrObject(packagePartContext, classOrObject, state, null);
     }
 
-    public PackageParts getPackageParts() {
-        return packageParts;
-    }
-
     public Collection<KtFile> getFiles() {
         return files;
     }

compiler/backend/src/org/jetbrains/kotlin/codegen/PackagePartRegistry.kt

@@ -0,0 +1,21 @@
+/*
+ * Copyright 2010-2016 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.jetbrains.kotlin.codegen
+
+interface PackagePartRegistry {
+    fun addPart(partShortName: String)
+}