JVM_IR: only serialize inline functions themselves with -Xserialize-ir=inline

compiler/ir/serialization.jvm/src/JvmIr.proto

@@ -14,23 +14,11 @@ message XStatementOrExpression {
   }
 }
 
-message AuxTables {
-  // TODO: optimize the representation.
-  repeated common.serialization.proto.IrType type = 1;
-  repeated common.serialization.proto.IdSignature signature = 2;
-  repeated string string = 3;
-  repeated XStatementOrExpression body = 4;
-  repeated string debug_info = 5;
-}
-
-message JvmIrFile {
+message ClassOrFile {
   repeated common.serialization.proto.IrDeclaration declaration = 1;
-  repeated common.serialization.proto.IrConstructorCall annotation = 2;
-  repeated int32 facade_fq_name = 3;
-  required AuxTables aux_tables = 4;
-}
-
-message JvmIrClass {
-  required common.serialization.proto.IrClass ir_class = 1;
-  required AuxTables aux_tables = 2;
+  repeated common.serialization.proto.IrType type = 2;
+  repeated common.serialization.proto.IdSignature signature = 3;
+  repeated string string = 4;
+  repeated XStatementOrExpression body = 5;
+  repeated string debug_info = 6;
 }

compiler/ir/serialization.jvm/src/org/jetbrains/kotlin/backend/jvm/serialization/JvmIrSerializerSession.kt

@@ -11,17 +11,18 @@ import org.jetbrains.kotlin.backend.common.serialization.IrFileSerializer
 import org.jetbrains.kotlin.backend.jvm.serialization.proto.JvmIr
 import org.jetbrains.kotlin.config.JvmSerializeIrMode
 import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
-import org.jetbrains.kotlin.ir.declarations.IrClass
-import org.jetbrains.kotlin.ir.declarations.IrFile
+import org.jetbrains.kotlin.ir.IrElement
+import org.jetbrains.kotlin.ir.declarations.*
 import org.jetbrains.kotlin.ir.symbols.IrSymbol
 import org.jetbrains.kotlin.ir.util.IrMessageLogger
-import org.jetbrains.kotlin.protobuf.ByteString
+import org.jetbrains.kotlin.ir.visitors.IrElementVisitor
+import org.jetbrains.kotlin.synthetic.isVisibleOutside
 
 class JvmIrSerializerSession(
     messageLogger: IrMessageLogger,
     private val declarationTable: DeclarationTable,
     expectDescriptorToSymbol: MutableMap<DeclarationDescriptor, IrSymbol>,
-    mode: JvmSerializeIrMode,
+    private val mode: JvmSerializeIrMode,
     skipExpects: Boolean = false,
 ) : IrFileSerializer(
     messageLogger, declarationTable, expectDescriptorToSymbol, CompatibilityMode.CURRENT,
@@ -34,42 +35,88 @@ class JvmIrSerializerSession(
 
     // Usage protocol: construct an instance, call only one of `serializeIrFile()` and `serializeTopLevelClass()` only once.
 
-    fun serializeJvmIrFile(irFile: IrFile): JvmIr.JvmIrFile {
-        val proto = JvmIr.JvmIrFile.newBuilder()
+    fun serializeJvmIrFile(irFile: IrFile): JvmIr.ClassOrFile? {
+        var anySaved = false
+        val proto = JvmIr.ClassOrFile.newBuilder()
 
         declarationTable.inFile(irFile) {
-            irFile.declarations.filter { it !is IrClass }.forEach { declaration ->
+            irFile.declarations.filter { it !is IrClass }.forEach { topDeclaration ->
+                forEveryDeclarationToSerialize(topDeclaration, mode) { declaration ->
+                    proto.addDeclaration(serializeDeclaration(declaration))
+                    anySaved = true
+                }
+            }
+        }
+        if (!anySaved) return null
+
+        serializeAuxTables(proto)
+
+        return proto.build()
+    }
+
+    fun serializeTopLevelClass(irClass: IrClass): JvmIr.ClassOrFile? {
+        val proto = JvmIr.ClassOrFile.newBuilder()
+        declarationTable.inFile(irClass.parent as IrFile) {
+            forEveryDeclarationToSerialize(irClass, mode) { declaration ->
                 proto.addDeclaration(serializeDeclaration(declaration))
             }
-            proto.addAllAnnotation(serializeAnnotations(irFile.annotations))
         }
-
-        proto.auxTables = serializeAuxTables()
-
+        serializeAuxTables(proto)
         return proto.build()
     }
 
-    fun serializeTopLevelClass(irClass: IrClass): JvmIr.JvmIrClass {
-        val proto = JvmIr.JvmIrClass.newBuilder()
-        declarationTable.inFile(irClass.parent as IrFile) {
-            proto.irClass = serializeIrClass(irClass)
-        }
-        proto.auxTables = serializeAuxTables()
-        return proto.build()
-    }
-
-    private fun serializeAuxTables(): JvmIr.AuxTables {
-        val proto = JvmIr.AuxTables.newBuilder()
+    private fun serializeAuxTables(proto: JvmIr.ClassOrFile.Builder) {
         protoTypeArray.forEach(proto::addType)
         protoIdSignatureArray.forEach(proto::addSignature)
         protoStringArray.forEach(proto::addString)
         protoBodyArray.forEach { proto.addBody(it.toProto()) }
         protoDebugInfoArray.forEach(proto::addDebugInfo)
-        return proto.build()
     }
 
     fun XStatementOrExpression.toProto(): JvmIr.XStatementOrExpression = when (this) {
         is XStatementOrExpression.XStatement -> JvmIr.XStatementOrExpression.newBuilder().setStatement(toProtoStatement()).build()
         is XStatementOrExpression.XExpression -> JvmIr.XStatementOrExpression.newBuilder().setExpression(toProtoExpression()).build()
     }
-}
+}
+
+private fun forEveryDeclarationToSerialize(topDeclaration: IrDeclaration, mode: JvmSerializeIrMode, action: (IrDeclaration) -> Unit) {
+    when (mode) {
+        JvmSerializeIrMode.NONE -> error("should not even be called with serialization mode NONE")
+        JvmSerializeIrMode.ALL -> action(topDeclaration)
+        JvmSerializeIrMode.INLINE ->
+            topDeclaration.accept(ForVisibleInlineFunctionsVisitor, action)
+    }
+}
+
+private object ForVisibleInlineFunctionsVisitor : IrElementVisitor<Unit, (IrDeclaration) -> Unit> {
+    override fun visitElement(element: IrElement, data: (IrDeclaration) -> Unit) {
+        error("Visitor only for nonlocal declarations")
+    }
+
+    override fun visitDeclaration(declaration: IrDeclarationBase, data: (IrDeclaration) -> Unit) {
+        return
+    }
+
+    override fun visitClass(declaration: IrClass, data: (IrDeclaration) -> Unit) {
+        if (!declaration.visibility.isVisibleOutside()) return
+        for (child in declaration.declarations) {
+            child.accept(this, data)
+        }
+    }
+
+    override fun visitProperty(declaration: IrProperty, data: (IrDeclaration) -> Unit) {
+        if (!declaration.visibility.isVisibleOutside()) return
+        declaration.getter?.accept(this, data)
+        declaration.setter?.accept(this, data)
+    }
+
+    override fun visitSimpleFunction(declaration: IrSimpleFunction, data: (IrDeclaration) -> Unit) {
+        val action = data
+        if (declaration.visibility.isVisibleOutside() &&
+            declaration.isInline &&
+            !declaration.isFakeOverride
+        ) {
+            action(declaration)
+        }
+    }
+}

compiler/ir/serialization.jvm/src/org/jetbrains/kotlin/backend/jvm/serialization/deserializeLazyDeclarations.kt

@@ -20,6 +20,7 @@ import org.jetbrains.kotlin.ir.backend.jvm.serialization.JvmIrMangler
 import org.jetbrains.kotlin.ir.declarations.IrClass
 import org.jetbrains.kotlin.ir.declarations.IrDeclaration
 import org.jetbrains.kotlin.ir.declarations.impl.IrFileImpl
+import org.jetbrains.kotlin.ir.declarations.lazy.IrLazyDeclarationBase
 import org.jetbrains.kotlin.ir.declarations.lazy.LazyIrFactory
 import org.jetbrains.kotlin.ir.symbols.IrPropertySymbol
 import org.jetbrains.kotlin.ir.symbols.IrSimpleFunctionSymbol
@@ -35,22 +36,22 @@ import org.jetbrains.kotlin.backend.common.serialization.proto.IrExpression as P
 import org.jetbrains.kotlin.backend.common.serialization.proto.IrStatement as ProtoStatement
 import org.jetbrains.kotlin.backend.common.serialization.proto.IrType as ProtoType
 
-fun deserializeClassFromByteArray(
+fun deserializeFromByteArray(
     byteArray: ByteArray,
     stubGenerator: DeclarationStubGenerator,
-    irClass: IrClass,
+    toplevelParent: IrClass,
     typeSystemContext: IrTypeSystemContext,
     allowErrorNodes: Boolean,
 ) {
     val irBuiltIns = stubGenerator.irBuiltIns
     val symbolTable = stubGenerator.symbolTable
-    val irProto = JvmIr.JvmIrClass.parseFrom(byteArray.codedInputStream)
+    val irProto = JvmIr.ClassOrFile.parseFrom(byteArray.codedInputStream)
     val irLibraryFile = IrLibraryFileFromAnnotation(
-        irProto.auxTables.typeList,
-        irProto.auxTables.signatureList,
-        irProto.auxTables.stringList,
-        irProto.auxTables.bodyList,
-        irProto.auxTables.debugInfoList
+        irProto.typeList,
+        irProto.signatureList,
+        irProto.stringList,
+        irProto.bodyList,
+        irProto.debugInfoList
     )
     val descriptorFinder =
         DescriptorByIdSignatureFinder(
@@ -60,65 +61,7 @@ fun deserializeClassFromByteArray(
         )
 
     // Only needed for local signature computation.
-    val dummyIrFile = IrFileImpl(NaiveSourceBasedFileEntryImpl("<unknown>"), IrFileSymbolImpl(), irClass.packageFqName!!)
-
-    val symbolDeserializer = IrSymbolDeserializer(
-        symbolTable,
-        irLibraryFile,
-        fileSymbol = dummyIrFile.symbol,
-        /* TODO */ actuals = emptyList(),
-        enqueueLocalTopLevelDeclaration = {}, // just link to it in symbolTable
-        handleExpectActualMapping = { _, _ -> TODO() },
-        deserializePublicSymbol = { idSignature, symbolKind ->
-            referencePublicSymbol(symbolTable, descriptorFinder, idSignature, symbolKind)
-        }
-    )
-
-    val lazyIrFactory = LazyIrFactory(irBuiltIns.irFactory)
-
-    val deserializer = IrDeclarationDeserializer(
-        irBuiltIns, symbolTable, lazyIrFactory, irLibraryFile, irClass.parent,
-        allowErrorNodes = allowErrorNodes,
-        deserializeInlineFunctions = true,
-        deserializeBodies = true,
-        symbolDeserializer,
-        DefaultFakeOverrideClassFilter,
-        makeSimpleFakeOverrideBuilder(symbolTable, typeSystemContext, symbolDeserializer),
-        compatibilityMode = CompatibilityMode.CURRENT,
-    )
-
-    deserializer.deserializeIrClass(irProto.irClass)
-
-    ExternalDependenciesGenerator(stubGenerator.symbolTable, listOf(stubGenerator)).generateUnboundSymbolsAsDependencies()
-    buildFakeOverridesForLocalClasses(stubGenerator.symbolTable, typeSystemContext, symbolDeserializer, irClass)
-}
-
-fun deserializeIrFileFromByteArray(
-    byteArray: ByteArray,
-    stubGenerator: DeclarationStubGenerator,
-    facadeClass: IrClass,
-    typeSystemContext: IrTypeSystemContext,
-    allowErrorNodes: Boolean,
-) {
-    val irBuiltIns = stubGenerator.irBuiltIns
-    val symbolTable = stubGenerator.symbolTable
-    val irProto = JvmIr.JvmIrFile.parseFrom(byteArray.codedInputStream)
-    val irLibraryFile = IrLibraryFileFromAnnotation(
-        irProto.auxTables.typeList,
-        irProto.auxTables.signatureList,
-        irProto.auxTables.stringList,
-        irProto.auxTables.bodyList,
-        irProto.auxTables.debugInfoList
-    )
-    val descriptorFinder =
-        DescriptorByIdSignatureFinder(
-            stubGenerator.moduleDescriptor,
-            JvmDescriptorMangler(null),
-            DescriptorByIdSignatureFinder.LookupMode.MODULE_WITH_DEPENDENCIES
-        )
-
-    // Only needed for local signature computation.
-    val dummyIrFile = IrFileImpl(NaiveSourceBasedFileEntryImpl("<unknown>"), IrFileSymbolImpl(), facadeClass.packageFqName!!)
+    val dummyIrFile = IrFileImpl(NaiveSourceBasedFileEntryImpl("<unknown>"), IrFileSymbolImpl(), toplevelParent.packageFqName!!)
 
     val symbolDeserializer = IrSymbolDeserializer(
         symbolTable,
@@ -136,8 +79,10 @@ fun deserializeIrFileFromByteArray(
 
     val fakeOverrideBuilder = makeSimpleFakeOverrideBuilder(symbolTable, typeSystemContext, symbolDeserializer)
 
+    // We have to supply topLevelParent here, but this results in wrong values for parent fields in deeply embedded declarations.
+    // Patching will be needed.
     val deserializer = IrDeclarationDeserializer(
-        irBuiltIns, symbolTable, lazyIrFactory, irLibraryFile, facadeClass,
+        irBuiltIns, symbolTable, lazyIrFactory, irLibraryFile, toplevelParent,
         allowErrorNodes = allowErrorNodes,
         deserializeInlineFunctions = true,
         deserializeBodies = true,
@@ -147,11 +92,17 @@ fun deserializeIrFileFromByteArray(
         compatibilityMode = CompatibilityMode.CURRENT,
     )
     for (declarationProto in irProto.declarationList) {
-        deserializer.deserializeDeclaration(declarationProto)
+        val declaration = deserializer.deserializeDeclaration(declarationProto)
+        // Either declaration is lazy, supplied by LazyIrFactory,
+        // or, if it is newly created, there must have been no references to it from the current module.
+        // These newly created declarations will never be used, so it does not matter if we patch their parent or not.
+        if (declaration is IrLazyDeclarationBase) {
+            declaration.parent = declaration.lazyParent()
+        }
     }
 
     ExternalDependenciesGenerator(stubGenerator.symbolTable, listOf(stubGenerator)).generateUnboundSymbolsAsDependencies()
-    buildFakeOverridesForLocalClasses(stubGenerator.symbolTable, typeSystemContext, symbolDeserializer, facadeClass)
+    buildFakeOverridesForLocalClasses(stubGenerator.symbolTable, typeSystemContext, symbolDeserializer, toplevelParent)
 }
 
 private class IrLibraryFileFromAnnotation(
@@ -245,7 +196,8 @@ private fun buildFakeOverridesForLocalClasses(
     toplevel: IrClass
 ) {
     val builder = makeSimpleFakeOverrideBuilder(symbolTable, typeSystemContext, symbolDeserializer)
-    toplevel.acceptChildrenVoid(object : IrElementVisitorVoid {
+    toplevel.acceptChildrenVoid(
+        object : IrElementVisitorVoid {
         override fun visitElement(element: IrElement) {
             element.acceptChildrenVoid(this)
         }
@@ -268,4 +220,4 @@ class PrePopulatedDeclarationTable(
         symbol2Sig[declaration.symbol]?.let { return it }
         return super.tryComputeBackendSpecificSignature(declaration)
     }
-}
+}