jvm-abi-gen: remove parts of the SMAP used only in stripped methods

plugins/jvm-abi-gen/src/org/jetbrains/kotlin/jvm/abi/JvmAbiOutputExtension.kt

@@ -13,6 +13,8 @@ import org.jetbrains.kotlin.cli.common.output.writeAllTo
 import org.jetbrains.kotlin.cli.jvm.compiler.CompileEnvironmentUtil
 import org.jetbrains.kotlin.codegen.ClassFileFactory
 import org.jetbrains.kotlin.codegen.extensions.ClassFileFactoryFinalizerExtension
+import org.jetbrains.kotlin.codegen.inline.*
+import org.jetbrains.kotlin.codegen.visitWithSplitting
 import org.jetbrains.kotlin.load.java.JvmAnnotationNames
 import org.jetbrains.org.objectweb.asm.*
 import org.jetbrains.org.objectweb.asm.commons.ClassRemapper
@@ -94,7 +96,11 @@ class JvmAbiOutputExtension(
                         val prune = abiInfo.prune
                         val memberInfo = abiInfo.memberInfo
                         val innerClassesToKeep = mutableSetOf<String>()
-                        val noMethodsKept = memberInfo.none { it.value == AbiMethodInfo.KEEP && it.key is JvmMethodSignature }
+
+                        var sourceFile: String? = null
+                        var sourceMap: SourceMapCopier? = null
+                        var sourceMapAnnotationPresent = false
+
                         val writer = ClassWriter(0)
                         val remapper = ClassRemapper(writer, object : Remapper() {
                             override fun map(internalName: String): String =
@@ -105,6 +111,13 @@ class JvmAbiOutputExtension(
                             private val keptMethods = mutableListOf<MethodNode>()
                             private val innerClassInfos = mutableMapOf<String, InnerClassInfo>()
 
+                            override fun visitSource(source: String?, debug: String?) {
+                                sourceFile = source.takeIf { !removeDebugInfo }
+                                sourceMap = debug.takeIf { !removeDebugInfo && !prune }
+                                    ?.let(SMAPParser::parseOrNull)
+                                    ?.let { SourceMapCopier(SourceMapper(sourceFile, it), it) }
+                            }
+
                             // Strip private fields.
                             override fun visitField(
                                 access: Int,
@@ -134,21 +147,12 @@ class JvmAbiOutputExtension(
                                     keptMethods += it
                                 }
 
-                                if (info == AbiMethodInfo.KEEP || access and (Opcodes.ACC_NATIVE or Opcodes.ACC_ABSTRACT) != 0) {
+                                return when {
-                                    return if (removeDebugInfo) DebugInfoRemovingMethodVisitor(node) else node
+                                    info != AbiMethodInfo.KEEP && access and (Opcodes.ACC_NATIVE or Opcodes.ACC_ABSTRACT) == 0 ->
-                                }
+                                        BodyStrippingMethodVisitor(node)
-
+                                    removeDebugInfo ->
-                                return BodyStrippingMethodVisitor(node)
+                                        DebugInfoRemovingMethodVisitor(node)
-                            }
+                                    else -> node
-
-                            override fun visitSource(source: String?, debug: String?) {
-                                when {
-                                    removeDebugInfo -> super.visitSource(null, null)
-                                    prune || noMethodsKept -> {
-                                        // Strip SourceDebugExtension attribute if there are no inline functions.
-                                        super.visitSource(source, null)
-                                    }
-                                    else -> super.visitSource(source, debug)
                                 }
                             }
 
@@ -160,9 +164,9 @@ class JvmAbiOutputExtension(
                             }
 
                             override fun visitAnnotation(descriptor: String?, visible: Boolean): AnnotationVisitor? {
-                                // Strip @SourceDebugExtension annotation if we're removing debug info.
                                 if (descriptor == JvmAnnotationNames.SOURCE_DEBUG_EXTENSION_DESC) {
-                                    if (removeDebugInfo || noMethodsKept || prune) return null
+                                    sourceMapAnnotationPresent = true
+                                    return null
                                 }
 
                                 val delegate = super.visitAnnotation(descriptor, visible)
@@ -180,19 +184,36 @@ class JvmAbiOutputExtension(
                             }
 
                             override fun visitEnd() {
-                                // Output class members in sorted order so that changes in original ordering don't affect the ABI JAR.
+                                if (!preserveDeclarationOrder) {
+                                    // Output class members in sorted order so that changes in original ordering don't affect the ABI JAR.
+                                    keptFields.sortWith(compareBy(FieldNode::name, FieldNode::desc))
+                                    keptMethods.sortWith(compareBy(MethodNode::name, MethodNode::desc))
+                                }
 
-                                keptFields
+                                for (field in keptFields) {
-                                    .apply { if (!preserveDeclarationOrder) sortWith(compareBy(FieldNode::name, FieldNode::desc)) }
+                                    field.accept(cv)
-                                    .forEach { it.accept(cv) }
+                                }
 
-                                keptMethods
+                                for (method in keptMethods) {
-                                    .apply { if (!preserveDeclarationOrder) sortWith(compareBy(MethodNode::name, MethodNode::desc)) }
+                                    val mv = with(method) { cv.visitMethod(access, name, desc, signature, exceptions?.toTypedArray()) }
-                                    .forEach { it.accept(cv) }
+                                    // Mapping the line numbers should only be done *after* sorting methods, as otherwise the order
+                                    // of inline methods may be visible in their synthetic line numbers.
+                                    method.accept(sourceMap?.let { SourceMapCopyingMethodVisitor(it, mv) } ?: mv)
+                                }
+
+                                val sourceMapText = sourceMap?.parent?.takeIf { !it.isTrivial }
+                                    ?.let { SMAPBuilder.build(it.resultMappings, backwardsCompatibleSyntax = false) }
+                                // This is technically not the right way to use `ClassVisitor` (`visitSource` should be called before
+                                // `visitMethod` and such), but `ClassWriter` doesn't care, and we're a bit constrained here (see above).
+                                cv.visitSource(sourceFile, sourceMapText)
+                                if (sourceMapAnnotationPresent && sourceMapText != null) {
+                                    val av = cv.visitAnnotation(JvmAnnotationNames.SOURCE_DEBUG_EXTENSION_DESC, false)
+                                    av.visitWithSplitting("value", sourceMapText)
+                                    av.visitEnd()
+                                }
 
                                 innerClassesToKeep.addInnerClasses(innerClassInfos, internalName)
                                 innerClassesToKeep.addOuterClasses(innerClassInfos)
-
                                 for (name in innerClassesToKeep.sorted()) {
                                     innerClassInfos[name]?.let { cv.visitInnerClass(it.name, it.outerName, it.innerName, it.access) }
                                 }

plugins/jvm-abi-gen/testData/compare/declarationOrderInlineCall/directives.txt

@@ -1 +0,0 @@
-// IGNORE_BACKEND: JVM_IR

plugins/jvm-abi-gen/testData/compare/declarationOrderInlineCall/sameAbi/test.kt

@@ -4,5 +4,6 @@ fun g() = f()
 
 inline fun f() = 1
 
-// This does not change the line numbers in f, but does change the line
+// The line numbers in f are the same, but the SourceDebug extension section
-// numbers in the SourceDebug extension section which has to be rewritten.
+// contains a different line number for the inline function call in g.
+// g is not inline, though, so the debugging information for its body is removed.