Refactor .kotlin_module-writing code, simplify PackagePartRegistry

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

@@ -49,7 +49,7 @@ import org.jetbrains.kotlin.serialization.StringTableImpl;
 import org.jetbrains.org.objectweb.asm.Type;
 
 import java.io.File;
-import java.io.UnsupportedEncodingException;
+import java.nio.charset.StandardCharsets;
 import java.util.*;
 
 import static org.jetbrains.kotlin.codegen.JvmCodegenUtil.getMappingFileName;
@@ -62,7 +62,7 @@ public class ClassFileFactory implements OutputFileCollection {
     private boolean isDone = false;
 
     private final Set<File> sourceFiles = new HashSet<>();
-    private final Map<String, PackageParts> partsGroupedByPackage = new LinkedHashMap<>();
+    private final PackagePartRegistry packagePartRegistry = new PackagePartRegistry();
 
     public ClassFileFactory(@NotNull GenerationState state, @NotNull ClassBuilderFactory builderFactory) {
         this.state = state;
@@ -73,6 +73,11 @@ public class ClassFileFactory implements OutputFileCollection {
         return state;
     }
 
+    @NotNull
+    public PackagePartRegistry getPackagePartRegistry() {
+        return packagePartRegistry;
+    }
+
     @NotNull
     public ClassBuilder newVisitor(
             @NotNull JvmDeclarationOrigin origin,
@@ -110,7 +115,7 @@ public class ClassFileFactory implements OutputFileCollection {
         JvmModuleProtoBuf.Module.Builder builder = JvmModuleProtoBuf.Module.newBuilder();
         String outputFilePath = getMappingFileName(state.getModuleName());
 
-        for (PackageParts part : ClassFileUtilsKt.addCompiledPartsAndSort(partsGroupedByPackage.values(), state)) {
+        for (PackageParts part : ClassFileUtilsKt.addCompiledPartsAndSort(packagePartRegistry.getParts().values(), state)) {
             part.addTo(builder);
         }
 
@@ -133,12 +138,7 @@ public class ClassFileFactory implements OutputFileCollection {
 
             @Override
             public String asText(ClassBuilderFactory factory) {
-                try {
+                return new String(asBytes(factory), StandardCharsets.UTF_8);
-                    return new String(asBytes(factory), "UTF-8");
-                }
-                catch (UnsupportedEncodingException e) {
-                    throw new RuntimeException(e);
-                }
             }
         });
     }
@@ -239,22 +239,14 @@ public class ClassFileFactory implements OutputFileCollection {
     public PackageCodegen forPackage(@NotNull FqName fqName, @NotNull Collection<KtFile> files) {
         assert !isDone : "Already done!";
         registerSourceFiles(files);
-        return state.getCodegenFactory().createPackageCodegen(state, files, fqName, buildNewPackagePartRegistry(fqName));
+        return state.getCodegenFactory().createPackageCodegen(state, files, fqName);
     }
 
     @NotNull
     public MultifileClassCodegen forMultifileClass(@NotNull FqName facadeFqName, @NotNull Collection<KtFile> files) {
         assert !isDone : "Already done!";
         registerSourceFiles(files);
-        return state.getCodegenFactory().createMultifileClassCodegen(state, files, facadeFqName, buildNewPackagePartRegistry(facadeFqName.parent()));
+        return state.getCodegenFactory().createMultifileClassCodegen(state, files, facadeFqName);
-    }
-
-    private PackagePartRegistry buildNewPackagePartRegistry(@NotNull FqName packageFqName) {
-        String packageFqNameAsString = packageFqName.asString();
-        return (partInternalName, facadeInternalName) -> {
-            PackageParts packageParts = partsGroupedByPackage.computeIfAbsent(packageFqNameAsString, PackageParts::new);
-            packageParts.addPart(partInternalName, facadeInternalName);
-        };
     }
 
     private void registerSourceFiles(Collection<KtFile> files) {

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

@@ -26,8 +26,10 @@ import org.jetbrains.kotlin.psi.KtFile
 
 interface CodegenFactory {
     fun generateModule(state: GenerationState, files: Collection<KtFile?>, errorHandler: CompilationErrorHandler)
-    fun createPackageCodegen(state: GenerationState, files: Collection<KtFile>, fqName: FqName, registry: PackagePartRegistry): PackageCodegen
+
-    fun createMultifileClassCodegen(state: GenerationState, files: Collection<KtFile>, fqName: FqName, registry: PackagePartRegistry): MultifileClassCodegen
+    fun createPackageCodegen(state: GenerationState, files: Collection<KtFile>, fqName: FqName): PackageCodegen
+
+    fun createMultifileClassCodegen(state: GenerationState, files: Collection<KtFile>, fqName: FqName): MultifileClassCodegen
 
     companion object {
         fun doCheckCancelled(state: GenerationState) {
@@ -69,11 +71,11 @@ object DefaultCodegenFactory : CodegenFactory {
         }
     }
 
-    override fun createPackageCodegen(state: GenerationState, files: Collection<KtFile>, fqName: FqName, registry: PackagePartRegistry) =
+    override fun createPackageCodegen(state: GenerationState, files: Collection<KtFile>, fqName: FqName) =
-            PackageCodegenImpl(state, files, fqName, registry)
+            PackageCodegenImpl(state, files, fqName)
 
-    override fun createMultifileClassCodegen(state: GenerationState, files: Collection<KtFile>, fqName: FqName, registry: PackagePartRegistry) =
+    override fun createMultifileClassCodegen(state: GenerationState, files: Collection<KtFile>, fqName: FqName) =
-            MultifileClassCodegenImpl(state, files, fqName, registry)
+            MultifileClassCodegenImpl(state, files, fqName)
 
     private fun generateMultifileClass(
             state: GenerationState,
@@ -96,5 +98,4 @@ object DefaultCodegenFactory : CodegenFactory {
         val codegen = state.factory.forPackage(packageFqName, jetFiles)
         codegen.generate(errorHandler)
     }
-
+}
-}

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

@@ -58,10 +58,9 @@ interface MultifileClassCodegen {
 }
 
 class MultifileClassCodegenImpl(
-        private val state: GenerationState,
+    private val state: GenerationState,
-        private val files: Collection<KtFile>,
+    private val files: Collection<KtFile>,
-        private val facadeFqName: FqName,
+    private val facadeFqName: FqName
-        private val packagePartRegistry: PackagePartRegistry
 ) : MultifileClassCodegen {
     private val facadeClassType = AsmUtil.asmTypeByFqNameWithoutInnerClasses(facadeFqName)
 
@@ -214,7 +213,7 @@ class MultifileClassCodegenImpl(
             !file.hasDeclarationsForPartClass(state.bindingContext)
         ) return
 
-        packagePartRegistry.addPart(partType.internalName, facadeClassType.internalName)
+        state.factory.packagePartRegistry.addPart(packageFragment.fqName, partType.internalName, facadeClassType.internalName)
 
         val builder = state.factory.newVisitor(MultifileClassPart(file, packageFragment), partType, file)
 

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

@@ -44,18 +44,15 @@ public class PackageCodegenImpl implements PackageCodegen {
     private final GenerationState state;
     private final Collection<KtFile> files;
     private final PackageFragmentDescriptor packageFragment;
-    private final PackagePartRegistry packagePartRegistry;
 
     public PackageCodegenImpl(
             @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);
-        packagePartRegistry = registry;
     }
 
     @Override
@@ -120,7 +117,7 @@ public class PackageCodegenImpl implements PackageCodegen {
 
         if (!generatePackagePart || !state.getGenerateDeclaredClassFilter().shouldGeneratePackagePart(file)) return;
 
-        packagePartRegistry.addPart(fileClassType.getInternalName(), null);
+        state.getFactory().getPackagePartRegistry().addPart(packageFragment.getFqName(), fileClassType.getInternalName(), null);
 
         ClassBuilder builder = state.getFactory().newVisitor(JvmDeclarationOriginKt.PackagePart(file, packageFragment), fileClassType, file);
 

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

@@ -16,6 +16,13 @@
 
 package org.jetbrains.kotlin.codegen
 
-interface PackagePartRegistry {
+import org.jetbrains.kotlin.metadata.jvm.deserialization.PackageParts
-    fun addPart(partInternalName: String, facadeInternalName: String?)
+import org.jetbrains.kotlin.name.FqName
+
+internal class PackagePartRegistry {
+    val parts = mutableMapOf<FqName, PackageParts>()
+
+    fun addPart(packageFqName: FqName, partInternalName: String, facadeInternalName: String?) {
+        parts.computeIfAbsent(packageFqName) { PackageParts(it.asString()) }.addPart(partInternalName, facadeInternalName)
+    }
 }

compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/JvmIrCodegenFactory.kt

@@ -36,8 +36,8 @@ object JvmIrCodegenFactory : CodegenFactory {
         JvmBackendFacade.doGenerateFilesInternal(state, errorHandler, irModuleFragment, psi2irContext)
     }
 
-    override fun createPackageCodegen(state: GenerationState, files: Collection<KtFile>, fqName: FqName, registry: PackagePartRegistry): PackageCodegen {
+    override fun createPackageCodegen(state: GenerationState, files: Collection<KtFile>, fqName: FqName): PackageCodegen {
-        val impl = PackageCodegenImpl(state, files, fqName, registry)
+        val impl = PackageCodegenImpl(state, files, fqName)
 
         return object : PackageCodegen {
             override fun generate(errorHandler: CompilationErrorHandler) {
@@ -54,7 +54,7 @@ object JvmIrCodegenFactory : CodegenFactory {
         }
     }
 
-    override fun createMultifileClassCodegen(state: GenerationState, files: Collection<KtFile>, fqName: FqName, registry: PackagePartRegistry): MultifileClassCodegen {
+    override fun createMultifileClassCodegen(state: GenerationState, files: Collection<KtFile>, fqName: FqName): MultifileClassCodegen {
         TODO()
     }
-}
+}