Support new scheme of compilation of OptionalExpectation annotations

Instead of generating these annotation classes as package-private on
JVM, serialize their metadata to the .kotlin_module file, and load it
when compiling dependent multiplatform modules.

The problem with generating them as package-private was that
kotlin-stdlib for JVM would end up declaring symbols from other
platforms, which would include some annotations from package
kotlin.native. But using that package is discouraged by some tools
because it has a Java keyword in its name. In particular, jlink refused
to work with such artifact altogether (KT-21266).

 #KT-38652 Fixed

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

@@ -39,7 +39,6 @@ import org.jetbrains.kotlin.resolve.BindingContext;
 import org.jetbrains.kotlin.resolve.DescriptorUtils;
 import org.jetbrains.kotlin.resolve.InlineClassDescriptorResolver;
 import org.jetbrains.kotlin.resolve.InlineClassesUtilsKt;
-import org.jetbrains.kotlin.resolve.checkers.ExpectedActualDeclarationChecker;
 import org.jetbrains.kotlin.resolve.deprecation.DeprecationResolver;
 import org.jetbrains.kotlin.resolve.inline.InlineUtil;
 import org.jetbrains.kotlin.resolve.jvm.*;
@@ -433,9 +432,6 @@ public class AsmUtil {
         if (descriptor instanceof SyntheticClassDescriptorForLambda) {
             return getVisibilityAccessFlagForAnonymous(descriptor);
         }
-        if (ExpectedActualDeclarationChecker.isOptionalAnnotationClass(descriptor)) {
-            return NO_FLAG_PACKAGE_PRIVATE;
-        }
         if (descriptor.getKind() == ClassKind.ENUM_ENTRY) {
             return NO_FLAG_PACKAGE_PRIVATE;
         }

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

@@ -39,6 +39,7 @@ import org.jetbrains.kotlin.metadata.jvm.JvmModuleProtoBuf;
 import org.jetbrains.kotlin.metadata.jvm.deserialization.ModuleMapping;
 import org.jetbrains.kotlin.metadata.jvm.deserialization.ModuleMappingKt;
 import org.jetbrains.kotlin.metadata.jvm.deserialization.PackageParts;
+import org.jetbrains.kotlin.metadata.serialization.StringTable;
 import org.jetbrains.kotlin.name.ClassId;
 import org.jetbrains.kotlin.name.FqName;
 import org.jetbrains.kotlin.psi.KtFile;
@@ -129,15 +130,18 @@ public class ClassFileFactory implements OutputFileCollection {
         JvmModuleProtoBuf.Module.Builder builder = JvmModuleProtoBuf.Module.newBuilder();
         String outputFilePath = getMappingFileName(state.getModuleName());
 
-        for (PackageParts part : ClassFileUtilsKt.addCompiledPartsAndSort(packagePartRegistry.getParts().values(), state)) {
-            part.addTo(builder);
-        }
+        StringTableImpl stringTable = new StringTableImpl();
+        ClassFileUtilsKt.addDataFromCompiledModule(builder, packagePartRegistry, stringTable, state);
 
         List<String> experimental = state.getLanguageVersionSettings().getFlag(AnalysisFlags.getExperimental());
         if (!experimental.isEmpty()) {
-            writeExperimentalMarkers(state.getModule(), builder, experimental);
+            writeExperimentalMarkers(state.getModule(), builder, experimental, stringTable);
         }
 
+        Pair<ProtoBuf.StringTable, ProtoBuf.QualifiedNameTable> tables = stringTable.buildProto();
+        builder.setStringTable(tables.getFirst());
+        builder.setQualifiedNameTable(tables.getSecond());
+
         JvmModuleProtoBuf.Module moduleProto = builder.build();
 
         generators.put(outputFilePath, new OutAndSourceFileList(CollectionsKt.toList(sourceFiles)) {
@@ -160,9 +164,9 @@ public class ClassFileFactory implements OutputFileCollection {
     private static void writeExperimentalMarkers(
             @NotNull ModuleDescriptor module,
             @NotNull JvmModuleProtoBuf.Module.Builder builder,
-            @NotNull List<String> experimental
+            @NotNull List<String> experimental,
+            @NotNull StringTable stringTable
     ) {
-        StringTableImpl stringTable = new StringTableImpl();
         for (String fqName : experimental) {
             ClassDescriptor descriptor =
                     DescriptorUtilKt.resolveClassByFqName(module, new FqName(fqName), NoLookupLocation.FOR_ALREADY_TRACKED);
@@ -175,9 +179,6 @@ public class ClassFileFactory implements OutputFileCollection {
                 }
             }
         }
-        Pair<ProtoBuf.StringTable, ProtoBuf.QualifiedNameTable> tables = stringTable.buildProto();
-        builder.setStringTable(tables.getFirst());
-        builder.setQualifiedNameTable(tables.getSecond());
     }
 
     @NotNull

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

@@ -90,8 +90,12 @@ public class PackageCodegenImpl implements PackageCodegen {
         for (KtDeclaration declaration : file.getDeclarations()) {
             if (declaration instanceof KtClassOrObject) {
                 ClassDescriptor descriptor = state.getBindingContext().get(BindingContext.CLASS, declaration);
-                if (PsiUtilsKt.hasExpectModifier(declaration) &&
-                    (descriptor == null || !ExpectedActualDeclarationChecker.shouldGenerateExpectClass(descriptor))) {
+                if (PsiUtilsKt.hasExpectModifier(declaration)) {
+                    if (descriptor != null && ExpectedActualDeclarationChecker.shouldGenerateExpectClass(descriptor)) {
+                        assert ExpectedActualDeclarationChecker.isOptionalAnnotationClass(descriptor) :
+                                "Expect class should be generated only if it's an optional annotation: " + descriptor;
+                        state.getFactory().getPackagePartRegistry().getOptionalAnnotations().add(descriptor);
+                    }
                     continue;
                 }
 

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

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

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

@@ -18,10 +18,19 @@ package org.jetbrains.kotlin.codegen
 
 import org.jetbrains.kotlin.backend.common.output.OutputFile
 import org.jetbrains.kotlin.codegen.state.GenerationState
+import org.jetbrains.kotlin.descriptors.findClassAcrossModuleDependencies
 import org.jetbrains.kotlin.load.kotlin.loadModuleMapping
+import org.jetbrains.kotlin.metadata.deserialization.BinaryVersion
+import org.jetbrains.kotlin.metadata.jvm.JvmModuleProtoBuf
+import org.jetbrains.kotlin.metadata.jvm.deserialization.JvmMetadataVersion
 import org.jetbrains.kotlin.metadata.jvm.deserialization.ModuleMapping
 import org.jetbrains.kotlin.metadata.jvm.deserialization.PackageParts
+import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.resolve.jvm.JvmClassName
+import org.jetbrains.kotlin.serialization.DescriptorSerializer
+import org.jetbrains.kotlin.serialization.KotlinSerializerExtensionBase
+import org.jetbrains.kotlin.serialization.StringTableImpl
+import org.jetbrains.kotlin.serialization.deserialization.builtins.BuiltInSerializerProtocol
 
 fun ClassFileFactory.getClassFiles(): Iterable<OutputFile> {
     return asList().filterClassFiles()
@@ -31,27 +40,64 @@ fun List<OutputFile>.filterClassFiles(): List<OutputFile> {
     return filter { it.relativePath.endsWith(".class") }
 }
 
-fun Iterable<PackageParts>.addCompiledPartsAndSort(state: GenerationState): List<PackageParts> =
-        addCompiledParts(state).sortedBy { it.packageFqName }
-
-private fun Iterable<PackageParts>.addCompiledParts(state: GenerationState): List<PackageParts> {
-    val incrementalCache = state.incrementalCacheForThisTarget ?: return this.toList()
-    val moduleMappingData = incrementalCache.getModuleMappingData() ?: return this.toList()
-
-    val mapping = ModuleMapping.loadModuleMapping(moduleMappingData, "<incremental>", state.deserializationConfiguration) { version ->
-        throw IllegalStateException("Version of the generated module cannot be incompatible: $version")
+fun JvmModuleProtoBuf.Module.Builder.addDataFromCompiledModule(
+    registry: PackagePartRegistry, stringTable: StringTableImpl, state: GenerationState
+) {
+    for (part in registry.parts.values.addCompiledPartsAndSort(state)) {
+        part.addTo(this)
     }
 
-    incrementalCache.getObsoletePackageParts().forEach { internalName ->
+    // Take all optional annotation classes from sources, as well as look up all previously compiled optional annotation classes
+    // by FQ name in the current module. The latter is needed because in incremental compilation scenario, the already compiled
+    // classes will not be available via sources.
+    val optionalAnnotationClassDescriptors =
+        registry.optionalAnnotations.toSet() +
+                state.loadCompiledModule()?.moduleData?.run {
+                    optionalAnnotations.mapNotNull { proto ->
+                        state.module.findClassAcrossModuleDependencies(
+                            ClassId.fromString(nameResolver.getQualifiedClassName(proto.fqName))
+                        )
+                    }
+                }.orEmpty()
+
+    val serializer = DescriptorSerializer.createTopLevel(JvmOptionalAnnotationSerializerExtension(stringTable))
+    for (descriptor in optionalAnnotationClassDescriptors) {
+        addOptionalAnnotationClass(serializer.classProto(descriptor))
+    }
+}
+
+class JvmOptionalAnnotationSerializerExtension(
+    override val stringTable: StringTableImpl
+) : KotlinSerializerExtensionBase(BuiltInSerializerProtocol) {
+    override val metadataVersion: BinaryVersion
+        get() = JvmMetadataVersion.INSTANCE
+
+    override fun shouldUseTypeTable(): Boolean = true
+}
+
+private fun Iterable<PackageParts>.addCompiledPartsAndSort(state: GenerationState): List<PackageParts> =
+    addCompiledParts(state).sortedBy { it.packageFqName }
+
+private fun Iterable<PackageParts>.addCompiledParts(state: GenerationState): List<PackageParts> {
+    val mapping = state.loadCompiledModule() ?: return this.toList()
+
+    state.incrementalCacheForThisTarget?.getObsoletePackageParts()?.forEach { internalName ->
         val qualifier = JvmClassName.byInternalName(internalName).packageFqName.asString()
         mapping.findPackageParts(qualifier)?.removePart(internalName)
     }
 
     return (this + mapping.packageFqName2Parts.values)
-            .groupBy { it.packageFqName }
-            .map { (packageFqName, allOldPackageParts) ->
-                PackageParts(packageFqName).apply {
-                    allOldPackageParts.forEach { packageParts -> this += packageParts }
-                }
+        .groupBy { it.packageFqName }
+        .map { (packageFqName, allOldPackageParts) ->
+            PackageParts(packageFqName).apply {
+                allOldPackageParts.forEach { packageParts -> this += packageParts }
             }
+        }
+}
+
+private fun GenerationState.loadCompiledModule(): ModuleMapping? {
+    val moduleMappingData = incrementalCacheForThisTarget?.getModuleMappingData() ?: return null
+    return ModuleMapping.loadModuleMapping(moduleMappingData, "<incremental>", deserializationConfiguration) { version ->
+        throw IllegalStateException("Version of the generated module cannot be incompatible: $version")
+    }
 }