Support JvmPackageName annotation in binary format

The main changes are in jvm_package_table.proto and ModuleMapping.kt.
With JvmPackageName, package parts can now have a JVM package name that
differs from their Kotlin name. So, in addition to the old package parts
which were stored as short names + short name of multifile facade (we
can't change this because of compatibility with old compilers), we now
store separately those package parts, which have a different JVM package
name. The format is optimized to avoid storing any package name more
than once as a string.

Another notable change is in KotlinCliJavaFileManagerImpl, where we now
load .kotlin_module files when determining whether or not a package
exists. Before this change, no PsiPackage (and thus, no JavaPackage and
eventually, no LazyJavaPackageFragment) was created unless there was at
least one file in the corresponding directory. Now we also create
packages if they are "mapped" to other JVM packages, i.e. if all package
parts in them have been annotated with JvmPackageName.

Most of the other changes are refactorings to allow internal names of
package parts/multifile classes where previously there were only short
names.

core/descriptor.loader.java/src/jvm_package_table.proto

@@ -25,20 +25,37 @@ message PackageTable {
 
   // Names of .kotlin_metadata files for each package
   repeated PackageParts metadata_parts = 2;
+
+  // Values of @JvmPackageName annotation used in this module; can be referenced in PackageParts#class_with_jvm_package_name_package_id.
+  // The names here are dot-separated, e.g. "org.foo.bar"
+  repeated string jvm_package_name = 3;
 }
 
 message PackageParts {
     required string package_fq_name = 1;
 
     // Short names of files, without extension, present in this package. Only single file facades and multi-file _parts_ are listed here
-    // (multi-file facades are not present in this list, they are defined below)
-    repeated string class_name = 2;
+    // (multi-file facades are not present in this list, they are defined below). Only files whose JVM package name is equal to the
+    // Kotlin package name (i.e. it has not been changed with @JvmPackageName) are listed here.
+    repeated string short_class_name = 2;
 
-    // For each class name listed above, index of the name of the corresponding multi-file facade class in multifile_facade_name + 1,
+    // For each name in short_class_name, index of the name of the corresponding multi-file facade class in multifile_facade_short_name + 1,
     // or 0 if the class is not a multi-file part. If there's no value in this list at some index, the value is assumed to be 0.
     // (e.g. if there are no multi-file classes in the module, this list is not going to exist at all)
-    repeated int32 multifile_facade_id = 3 [packed=true];
+    repeated int32 multifile_facade_short_name_id = 3 [packed = true];
 
-    // Short names of multi-file facades, used in multifile_facade_id to store the part -> facade mapping.
-    repeated string multifile_facade_name = 4;
+    // Short names of multi-file facades, used in multifile_facade_short_name_id to store the part -> facade mapping.
+    repeated string multifile_facade_short_name = 4;
+
+    // Short names of files (single file facades), whose JVM package differs from the Kotlin package because of @JvmPackageName.
+    // The JVM package name of each file is stored at the same index in class_with_jvm_package_name_package_id.
+    repeated string class_with_jvm_package_name_short_name = 5;
+
+    // For each name in class_with_jvm_package_name_short_name, the index (into PackageTable#jvm_package_name) of the JVM package name.
+    // This list should have at least one element, otherwise classes with JVM package names are going to be ignored completely.
+    //
+    // If there's no value in this list at some index other than 0, the value is assumed to be the same as the value of the last element
+    // of this list. The intended use case for this optimization is to have just a list of a single element in the most frequent case
+    // when a bunch of files from the same Kotlin package have the same JVM package name.
+    repeated int32 class_with_jvm_package_name_package_id = 6 [packed = true];
 }

core/descriptor.loader.java/src/org/jetbrains/kotlin/load/java/descriptors/util.kt

@@ -22,9 +22,9 @@ import org.jetbrains.kotlin.descriptors.ValueParameterDescriptor
 import org.jetbrains.kotlin.descriptors.impl.ValueParameterDescriptorImpl
 import org.jetbrains.kotlin.load.java.lazy.descriptors.LazyJavaStaticClassScope
 import org.jetbrains.kotlin.load.kotlin.JvmPackagePartSource
-import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.resolve.descriptorUtil.getSuperClassNotAny
 import org.jetbrains.kotlin.resolve.descriptorUtil.module
+import org.jetbrains.kotlin.resolve.jvm.JvmClassName
 import org.jetbrains.kotlin.serialization.deserialization.descriptors.DeserializedMemberDescriptor
 import org.jetbrains.kotlin.types.KotlinType
 
@@ -66,9 +66,8 @@ fun ClassDescriptor.getParentJavaStaticClassScope(): LazyJavaStaticClassScope? {
     return staticScope
 }
 
-fun DeserializedMemberDescriptor.getImplClassNameForDeserialized(): Name? {
-    return (containerSource as? JvmPackagePartSource)?.simpleName
-}
+fun DeserializedMemberDescriptor.getImplClassNameForDeserialized(): JvmClassName? =
+        (containerSource as? JvmPackagePartSource)?.className
 
 fun DeserializedMemberDescriptor.isFromJvmPackagePart(): Boolean =
         containerSource is JvmPackagePartSource

core/descriptor.loader.java/src/org/jetbrains/kotlin/load/java/lazy/descriptors/LazyJavaPackageFragment.kt

@@ -29,7 +29,7 @@ import org.jetbrains.kotlin.load.kotlin.KotlinJvmBinaryPackageSourceElement
 import org.jetbrains.kotlin.load.kotlin.header.KotlinClassHeader
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
-import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.resolve.jvm.JvmClassName
 import org.jetbrains.kotlin.storage.getValue
 
 class LazyJavaPackageFragment(
@@ -40,7 +40,7 @@ class LazyJavaPackageFragment(
 
     internal val binaryClasses by c.storageManager.createLazyValue {
         c.components.packageMapper.findPackageParts(fqName.asString()).mapNotNull { partName ->
-            val classId = ClassId(fqName, Name.identifier(partName))
+            val classId = ClassId.topLevel(JvmClassName.byInternalName(partName).fqNameForTopLevelClassMaybeWithDollars)
             c.components.kotlinClassFinder.findKotlinClass(classId)?.let { partName to it }
         }.toMap()
     }
@@ -63,13 +63,13 @@ class LazyJavaPackageFragment(
     internal fun findClassifierByJavaClass(jClass: JavaClass): ClassDescriptor? = scope.javaScope.findClassifierByJavaClass(jClass)
 
     private val partToFacade by c.storageManager.createLazyValue {
-        val result = hashMapOf<String, String>()
-        kotlinClasses@for ((partName, kotlinClass) in binaryClasses) {
+        val result = hashMapOf<JvmClassName, JvmClassName>()
+        kotlinClasses@for ((partInternalName, kotlinClass) in binaryClasses) {
+            val partName = JvmClassName.byInternalName(partInternalName)
             val header = kotlinClass.classHeader
             when (header.kind) {
                 KotlinClassHeader.Kind.MULTIFILE_CLASS_PART -> {
-                    val facadeName = header.multifileClassName ?: continue@kotlinClasses
-                    result[partName] = facadeName.substringAfterLast('/')
+                    result[partName] = JvmClassName.byInternalName(header.multifileClassName ?: continue@kotlinClasses)
                 }
                 KotlinClassHeader.Kind.FILE_FACADE -> {
                     result[partName] = partName
@@ -80,7 +80,7 @@ class LazyJavaPackageFragment(
         result
     }
 
-    fun getFacadeSimpleNameForPartSimpleName(partName: String): String? = partToFacade[partName]
+    fun getFacadeNameForPartName(partName: JvmClassName): JvmClassName? = partToFacade[partName]
 
     override fun getMemberScope() = scope
 

core/descriptor.loader.java/src/org/jetbrains/kotlin/load/kotlin/JvmMetadataVersion.kt

@@ -29,7 +29,7 @@ class JvmMetadataVersion(vararg numbers: Int) : BinaryVersion(*numbers) {
 
     companion object {
         @JvmField
-        val INSTANCE = JvmMetadataVersion(1, 1, 7)
+        val INSTANCE = JvmMetadataVersion(1, 1, 8)
 
         @JvmField
         val INVALID_VERSION = JvmMetadataVersion()

core/descriptor.loader.java/src/org/jetbrains/kotlin/load/kotlin/KotlinJvmBinaryPackageSourceElement.kt

@@ -35,6 +35,6 @@ class KotlinJvmBinaryPackageSourceElement(
 
     fun getContainingBinaryClass(descriptor: DeserializedMemberDescriptor): KotlinJvmBinaryClass? {
         val name = descriptor.getImplClassNameForDeserialized() ?: return null
-        return packageFragment.binaryClasses[name.asString()]
+        return packageFragment.binaryClasses[name.internalName]
     }
 }

core/descriptor.loader.java/src/org/jetbrains/kotlin/load/kotlin/ModuleMapping.kt

@@ -16,6 +16,9 @@
 
 package org.jetbrains.kotlin.load.kotlin
 
+import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.resolve.jvm.JvmClassName
 import org.jetbrains.kotlin.serialization.deserialization.DeserializationConfiguration
 import org.jetbrains.kotlin.serialization.jvm.JvmPackageTable
 import java.io.ByteArrayInputStream
@@ -64,16 +67,28 @@ class ModuleMapping private constructor(val packageFqName2Parts: Map<String, Pac
                 val result = linkedMapOf<String, PackageParts>()
 
                 for (proto in table.packagePartsList) {
-                    val packageParts = result.getOrPut(proto.packageFqName) { PackageParts(proto.packageFqName) }
-                    for ((index, partShortName) in proto.classNameList.withIndex()) {
-                        val multifileFacadeId = proto.multifileFacadeIdList.getOrNull(index)?.minus(1)
-                        packageParts.addPart(partShortName, multifileFacadeId?.let(proto.multifileFacadeNameList::getOrNull))
+                    val packageFqName = proto.packageFqName
+                    val packageParts = result.getOrPut(packageFqName) { PackageParts(packageFqName) }
+
+                    for ((index, partShortName) in proto.shortClassNameList.withIndex()) {
+                        val multifileFacadeId = proto.multifileFacadeShortNameIdList.getOrNull(index)?.minus(1)
+                        val facadeShortName = multifileFacadeId?.let(proto.multifileFacadeShortNameList::getOrNull)
+                        val facadeInternalName = facadeShortName?.let { internalNameOf(packageFqName, it) }
+                        packageParts.addPart(internalNameOf(packageFqName, partShortName), facadeInternalName)
+                    }
+
+                    for ((index, partShortName) in proto.classWithJvmPackageNameShortNameList.withIndex()) {
+                        val packageId = proto.classWithJvmPackageNamePackageIdList.getOrNull(index)
+                                        ?: proto.classWithJvmPackageNamePackageIdList.lastOrNull()
+                                        ?: continue
+                        val jvmPackageName = table.jvmPackageNameList.getOrNull(packageId) ?: continue
+                        packageParts.addPart(internalNameOf(jvmPackageName, partShortName), null)
                     }
                 }
 
                 for (proto in table.metadataPartsList) {
                     val packageParts = result.getOrPut(proto.packageFqName) { PackageParts(proto.packageFqName) }
-                    proto.classNameList.forEach(packageParts::addMetadataPart)
+                    proto.shortClassNameList.forEach(packageParts::addMetadataPart)
                 }
 
                 return ModuleMapping(result, debugName)
@@ -87,21 +102,23 @@ class ModuleMapping private constructor(val packageFqName2Parts: Map<String, Pac
     }
 }
 
-class PackageParts(val packageFqName: String) {
-    // Short name of package part -> short name of the corresponding multifile facade (or null, if it's not a multifile part)
-    private val packageParts = linkedMapOf<String, String?>()
+private fun internalNameOf(packageFqName: String, className: String): String =
+        JvmClassName.byFqNameWithoutInnerClasses(FqName(packageFqName).child(Name.identifier(className))).internalName
 
-    // See JvmPackageTable.PackageTable.package_parts
+class PackageParts(val packageFqName: String) {
+    // JVM internal name of package part -> JVM internal name of the corresponding multifile facade (or null, if it's not a multifile part)
+    private val packageParts = linkedMapOf<String, String?>()
     val parts: Set<String> get() = packageParts.keys
-    // See JvmPackageTable.PackageTable.metadata_parts
+
+    // Short names of .kotlin_metadata package parts
     val metadataParts: Set<String> = linkedSetOf()
 
-    fun addPart(partShortName: String, facadeShortName: String?) {
-        packageParts[partShortName] = facadeShortName
+    fun addPart(partInternalName: String, facadeInternalName: String?) {
+        packageParts[partInternalName] = facadeInternalName
     }
 
-    fun removePart(shortName: String) {
-        packageParts.remove(shortName)
+    fun removePart(internalName: String) {
+        packageParts.remove(internalName)
     }
 
     fun addMetadataPart(shortName: String) {
@@ -113,36 +130,76 @@ class PackageParts(val packageFqName: String) {
             builder.addPackageParts(JvmPackageTable.PackageParts.newBuilder().apply {
                 packageFqName = this@PackageParts.packageFqName
 
-                val facadeNameToId = mutableMapOf<String, Int>()
-                for ((facadeName, partNames) in parts.groupBy { getMultifileFacadeName(it) }.toSortedMap(nullsLast())) {
-                    for (partName in partNames.sorted()) {
-                        addClassName(partName)
-                        if (facadeName != null) {
-                            addMultifileFacadeId(1 + facadeNameToId.getOrPut(facadeName) { facadeNameToId.size })
-                        }
-                    }
+                val packageInternalName = packageFqName.replace('.', '/')
+                val (partsWithinPackage, partsOutsidePackage) = parts.partition { partInternalName ->
+                    partInternalName.packageName == packageInternalName
                 }
 
-                for ((facadeId, facadeName) in facadeNameToId.values.zip(facadeNameToId.keys).sortedBy(Pair<Int, String>::first)) {
-                    assert(facadeId == multifileFacadeNameCount) { "Multifile facades are loaded incorrectly: $facadeNameToId" }
-                    addMultifileFacadeName(facadeName)
-                }
+                writePartsWithinPackage(partsWithinPackage)
+
+                writePartsOutsidePackage(partsOutsidePackage, builder)
             })
         }
 
         if (metadataParts.isNotEmpty()) {
             builder.addMetadataParts(JvmPackageTable.PackageParts.newBuilder().apply {
                 packageFqName = this@PackageParts.packageFqName
-                addAllClassName(metadataParts.sorted())
+                addAllShortClassName(metadataParts.sorted())
             })
         }
     }
 
-    fun getMultifileFacadeName(partShortName: String): String? = packageParts[partShortName]
+    private fun JvmPackageTable.PackageParts.Builder.writePartsWithinPackage(parts: List<String>) {
+        val facadeNameToId = mutableMapOf<String, Int>()
+        for ((facadeInternalName, partInternalNames) in parts.groupBy { getMultifileFacadeName(it) }.toSortedMap(nullsLast())) {
+            for (partInternalName in partInternalNames.sorted()) {
+                addShortClassName(partInternalName.className)
+                if (facadeInternalName != null) {
+                    addMultifileFacadeShortNameId(1 + facadeNameToId.getOrPut(facadeInternalName.className) { facadeNameToId.size })
+                }
+            }
+        }
+
+        for ((facadeId, facadeName) in facadeNameToId.values.zip(facadeNameToId.keys).sortedBy(Pair<Int, String>::first)) {
+            assert(facadeId == multifileFacadeShortNameCount) { "Multifile facades are loaded incorrectly: $facadeNameToId" }
+            addMultifileFacadeShortName(facadeName)
+        }
+    }
+
+    // Writes information about package parts which have a different JVM package from the Kotlin package (with the help of @JvmPackageName)
+    private fun JvmPackageTable.PackageParts.Builder.writePartsOutsidePackage(
+            parts: List<String>,
+            packageTableBuilder: JvmPackageTable.PackageTable.Builder
+    ) {
+        val packageIds = mutableListOf<Int>()
+        for ((packageInternalName, partsInPackage) in parts.groupBy { it.packageName }.toSortedMap()) {
+            val packageFqName = packageInternalName.replace('/', '.')
+            if (packageFqName !in packageTableBuilder.jvmPackageNameList) {
+                packageTableBuilder.addJvmPackageName(packageFqName)
+            }
+            val packageId = packageTableBuilder.jvmPackageNameList.indexOf(packageFqName)
+            for (part in partsInPackage.map { it.className }.sorted()) {
+                addClassWithJvmPackageNameShortName(part)
+                packageIds.add(packageId)
+            }
+        }
+
+        // See PackageParts#class_with_jvm_package_name_package_id in jvm_package_table.proto for description of this optimization
+        while (packageIds.size > 1 && packageIds[packageIds.size - 1] == packageIds[packageIds.size - 2]) {
+            packageIds.removeAt(packageIds.size - 1)
+        }
+
+        addAllClassWithJvmPackageNamePackageId(packageIds)
+    }
+
+    private val String.packageName: String get() = substringBeforeLast('/', "")
+    private val String.className: String get() = substringAfterLast('/')
+
+    fun getMultifileFacadeName(partInternalName: String): String? = packageParts[partInternalName]
 
     operator fun plusAssign(other: PackageParts) {
-        for ((partShortName, facadeShortName) in other.packageParts) {
-            addPart(partShortName, facadeShortName)
+        for ((partInternalName, facadeInternalName) in other.packageParts) {
+            addPart(partInternalName, facadeInternalName)
         }
         other.metadataParts.forEach(this::addMetadataPart)
     }

core/descriptor.loader.java/src/org/jetbrains/kotlin/resolve/jvm/JvmClassName.java

@@ -89,7 +89,12 @@ public class JvmClassName {
     }
 
     /**
-     * WARNING: internal name cannot be converted to FQ name for a class which contains dollars in the name
+     * WARNING: internal name cannot be reliably converted to FQ name.
+     *
+     * This method treats all dollar characters ('$') in the internal name as inner class separators.
+     * So it _will work incorrectly_ for classes where dollar characters are a part of the identifier.
+     *
+     * E.g. JvmClassName("org/foo/bar/Baz$quux").getFqNameForClassNameWithoutDollars() -> FqName("org.foo.bar.Baz.quux")
      */
     @NotNull
     public FqName getFqNameForClassNameWithoutDollars() {
@@ -99,6 +104,19 @@ public class JvmClassName {
         return fqName;
     }
 
+    /**
+     * WARNING: internal name cannot be reliably converted to FQ name.
+     *
+     * This method treats all dollar characters ('$') in the internal name as a part of the identifier.
+     * So it _will work incorrectly_ for inner classes.
+     *
+     * E.g. JvmClassName("org/foo/bar/Baz$quux").getFqNameForTopLevelClassMaybeWithDollars() -> FqName("org.foo.bar.Baz$quux")
+     */
+    @NotNull
+    public FqName getFqNameForTopLevelClassMaybeWithDollars() {
+        return new FqName(internalName.replace('/', '.'));
+    }
+
     @NotNull
     public FqName getPackageFqName() {
         int lastSlash = internalName.lastIndexOf("/");