JVM: extract some helpers for SMAP inlining

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

@@ -118,11 +118,8 @@ public abstract class AbstractClassBuilder implements ClassBuilder {
     public void done(boolean generateSmapCopyToAnnotation) {
         getVisitor().visitSource(sourceName, debugInfo);
         if (generateSmapCopyToAnnotation && debugInfo != null) {
-            AnnotationVisitor v =
+            AnnotationVisitor v = getVisitor().visitAnnotation(JvmAnnotationNames.SOURCE_DEBUG_EXTENSION_DESC, false);
-                    getVisitor().visitAnnotation(JvmAnnotationNames.SOURCE_DEBUG_EXTENSION_DESC, false).visitArray("value");
+            CodegenUtilKt.visitWithSplitting(v, "value", debugInfo);
-            for (String part : CodegenUtilKt.splitStringConstant(debugInfo)) {
-                v.visit(null, part);
-            }
             v.visitEnd();
         }
         getVisitor().visitEnd();

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

@@ -19,6 +19,7 @@ import org.jetbrains.kotlin.codegen.binding.CodegenBinding;
 import org.jetbrains.kotlin.codegen.context.*;
 import org.jetbrains.kotlin.codegen.coroutines.CoroutineCodegenUtilKt;
 import org.jetbrains.kotlin.codegen.coroutines.SuspendFunctionGenerationStrategy;
+import org.jetbrains.kotlin.codegen.inline.InlineCodegenUtilsKt;
 import org.jetbrains.kotlin.codegen.state.GenerationState;
 import org.jetbrains.kotlin.codegen.state.KotlinTypeMapper;
 import org.jetbrains.kotlin.codegen.state.TypeMapperUtilsKt;
@@ -1305,7 +1306,7 @@ public class FunctionCodegen {
         if (((ClassDescriptor) container).getModality() == Modality.FINAL) return;
 
         Label end = new Label();
-        int handleIndex = (Type.getArgumentsAndReturnSizes(defaultMethod.getDescriptor()) >> 2) - 2; /*-1 for this, and -1 for handle*/
+        int handleIndex = InlineCodegenUtilsKt.argumentsSize(defaultMethod.getDescriptor(), true) - 1; // last argument
         iv.load(handleIndex, OBJECT_TYPE);
         iv.ifnull(end);
         AsmUtil.genThrow(

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

@@ -52,6 +52,7 @@ import org.jetbrains.kotlin.types.checker.KotlinTypeChecker
 import org.jetbrains.kotlin.types.model.KotlinTypeMarker
 import org.jetbrains.kotlin.types.model.TypeParameterMarker
 import org.jetbrains.kotlin.utils.DFS
+import org.jetbrains.org.objectweb.asm.AnnotationVisitor
 import org.jetbrains.org.objectweb.asm.Label
 import org.jetbrains.org.objectweb.asm.Opcodes.*
 import org.jetbrains.org.objectweb.asm.Type
@@ -738,6 +739,14 @@ fun splitStringConstant(value: String): List<String> {
     }
 }
 
+fun AnnotationVisitor.visitWithSplitting(name: String?, value: String) {
+    val av = visitArray(name)
+    for (part in splitStringConstant(value)) {
+        av.visit(null, part)
+    }
+    av.visitEnd()
+}
+
 fun String.encodedUTF8Size(): Int {
     var result = 0
     for (char in this) {

compiler/backend/src/org/jetbrains/kotlin/codegen/inline/InlineCodegenForDefaultBody.kt

@@ -39,20 +39,15 @@ class InlineCodegenForDefaultBody(
     override fun genCallInner(callableMethod: Callable, resolvedCall: ResolvedCall<*>?, callDefault: Boolean, codegen: ExpressionCodegen) {
         assert(!callDefault) { "inlining default stub into another default stub" }
         val (node, smap) = sourceCompilerForInline.compileInlineFunction(jvmSignature)
-        val childSourceMapper = SourceMapCopier(sourceMapper, smap)
 
-        val argsSize =
+        val argsSize = argumentsSize(jvmSignature.asmMethod.descriptor, callableMethod.isStaticCall())
-            (Type.getArgumentsAndReturnSizes(jvmSignature.asmMethod.descriptor) ushr 2) - if (callableMethod.isStaticCall()) 1 else 0
+        val mv = object : MethodBodyVisitor(codegen.visitor) {
-        // `$default` is only for Kotlin use so it has no `$$forInline` version - this *is* what the inliner will use.
-        node.preprocessSuspendMarkers(forInline = true, keepFakeContinuation = false)
-        node.accept(object : MethodBodyVisitor(codegen.visitor) {
-            // The LVT was not generated at all, so move the start of parameters to the start of the method.
             override fun visitLocalVariable(name: String, desc: String, signature: String?, start: Label, end: Label, index: Int) =
                 super.visitLocalVariable(name, desc, signature, if (index < argsSize) methodStartLabel else start, end, index)
-
+        }
-            override fun visitLineNumber(line: Int, start: Label) =
+        // `$default` is only for Kotlin use so it has no `$$forInline` version - this *is* what the inliner will use.
-                super.visitLineNumber(childSourceMapper.mapLineNumber(line), start)
+        node.preprocessSuspendMarkers(forInline = true, keepFakeContinuation = false)
-        })
+        node.accept(SourceMapCopyingMethodVisitor(sourceMapper, smap, mv))
     }
 
     override fun genValueAndPut(

compiler/backend/src/org/jetbrains/kotlin/codegen/inline/InlineScopesGenerator.kt

@@ -340,13 +340,10 @@ class InlineScopesGenerator {
     }
 }
 
-fun updateCallSiteLineNumber(name: String, lineNumberMapping: Map<Int, Int>): String =
-    updateCallSiteLineNumber(name) { lineNumberMapping[it] ?: it }
-
 fun updateCallSiteLineNumber(name: String, newLineNumber: Int): String =
     updateCallSiteLineNumber(name) { newLineNumber }
 
-private fun updateCallSiteLineNumber(name: String, calculate: (Int) -> Int): String {
+fun updateCallSiteLineNumber(name: String, calculate: (Int) -> Int): String {
     val (scopeNumber, callSiteLineNumber, surroundingScopeNumber) = name.getInlineScopeInfo() ?: return name
     if (callSiteLineNumber == null) {
         return name

compiler/backend/src/org/jetbrains/kotlin/codegen/inline/MaxLocalsCalculator.java

@@ -20,7 +20,6 @@ import org.jetbrains.annotations.NotNull;
 import org.jetbrains.org.objectweb.asm.Label;
 import org.jetbrains.org.objectweb.asm.MethodVisitor;
 import org.jetbrains.org.objectweb.asm.Opcodes;
-import org.jetbrains.org.objectweb.asm.Type;
 
 public class MaxLocalsCalculator extends MethodVisitor {
 
@@ -28,14 +27,7 @@ public class MaxLocalsCalculator extends MethodVisitor {
 
     public MaxLocalsCalculator(int api, int access, String descriptor, MethodVisitor mv) {
         super(api, mv);
-
+        maxLocals = InlineCodegenUtilsKt.argumentsSize(descriptor, (access & Opcodes.ACC_STATIC) != 0);
-        // updates maxLocals
-        int size = Type.getArgumentsAndReturnSizes(descriptor) >> 2;
-        if ((access & Opcodes.ACC_STATIC) != 0) {
-            --size;
-        }
-
-        maxLocals = size;
     }
 
     @Override

compiler/backend/src/org/jetbrains/kotlin/codegen/inline/MethodInliner.kt

@@ -156,7 +156,6 @@ class MethodInliner(
         val fakeContinuationName = CoroutineTransformer.findFakeContinuationConstructorClassName(node)
         val markerShift = calcMarkerShift(parameters, node)
         var currentLineNumber = if (isInlineOnlyMethod) sourceMapper.callSite!!.line else -1
-        val lineNumberMapping = mutableMapOf<Int, Int>()
         val lambdaInliner = object : InlineAdapter(remappingMethodAdapter, parameters.argsSizeOnStack, sourceMapper) {
             private var transformationInfo: TransformationInfo? = null
             private var currentLabel: Label? = null
@@ -170,11 +169,6 @@ class MethodInliner(
                 if (!isInlineOnlyMethod) {
                     currentLineNumber = line
                 }
-
-                if (GENERATE_SMAP) {
-                    lineNumberMapping[line] = sourceMapper.mapLineNumber(line)
-                }
-
                 super.visitLineNumber(line, start)
             }
 
@@ -428,14 +422,14 @@ class MethodInliner(
 
         surroundInvokesWithSuspendMarkersIfNeeded(resultNode)
 
-        if (inliningContext.inlineScopesGenerator != null) {
+        if (inliningContext.inlineScopesGenerator != null && GENERATE_SMAP) {
-            updateCallSiteLineNumbers(resultNode, node, lineNumberMapping)
+            updateCallSiteLineNumbers(resultNode, node)
         }
 
         return resultNode
     }
 
-    private fun updateCallSiteLineNumbers(resultNode: MethodNode, inlinedNode: MethodNode, lineNumberMapping: Map<Int, Int>) {
+    private fun updateCallSiteLineNumbers(resultNode: MethodNode, inlinedNode: MethodNode) {
         val inlinedNodeLocalVariables = inlinedNode.localVariables ?: return
         val resultNodeLocalVariables = resultNode.localVariables ?: return
         if (inlinedNodeLocalVariables.isEmpty() || resultNodeLocalVariables.isEmpty()) {
@@ -462,7 +456,7 @@ class MethodInliner(
         for (variable in resultNodeLocalVariables) {
             val name = variable.name
             if (isFakeLocalVariableForInline(name) && name in markerVariableNamesFromInlinedNode) {
-                variable.name = updateCallSiteLineNumber(name, lineNumberMapping)
+                variable.name = updateCallSiteLineNumber(name) { sourceMapper.mapLineNumber(it) }
             }
         }
     }

compiler/backend/src/org/jetbrains/kotlin/codegen/inline/SMAP.kt

@@ -8,6 +8,9 @@ package org.jetbrains.kotlin.codegen.inline
 import it.unimi.dsi.fastutil.ints.Int2IntOpenHashMap
 import org.jetbrains.kotlin.codegen.SourceInfo
 import org.jetbrains.kotlin.codegen.optimization.common.asSequence
+import org.jetbrains.org.objectweb.asm.Label
+import org.jetbrains.org.objectweb.asm.MethodVisitor
+import org.jetbrains.org.objectweb.asm.Opcodes
 import org.jetbrains.org.objectweb.asm.tree.LineNumberNode
 import org.jetbrains.org.objectweb.asm.tree.MethodNode
 import java.util.*
@@ -78,6 +81,19 @@ class SourceMapCopier(val parent: SourceMapper, private val smap: SMAP, val call
     }
 }
 
+class SourceMapCopyingMethodVisitor(private val smapCopier: SourceMapCopier, mv: MethodVisitor) : MethodVisitor(Opcodes.API_VERSION, mv) {
+    constructor(target: SourceMapper, source: SMAP, mv: MethodVisitor) : this(SourceMapCopier(target, source), mv)
+
+    override fun visitLineNumber(line: Int, start: Label) =
+        super.visitLineNumber(smapCopier.mapLineNumber(line), start)
+
+    override fun visitLocalVariable(name: String, descriptor: String, signature: String?, start: Label, end: Label, index: Int) =
+        if (isFakeLocalVariableForInline(name))
+            super.visitLocalVariable(updateCallSiteLineNumber(name, smapCopier::mapLineNumber), descriptor, signature, start, end, index)
+        else
+            super.visitLocalVariable(name, descriptor, signature, start, end, index)
+}
+
 data class SourcePosition(val line: Int, val file: String, val path: String)
 
 class SourceMapper(val sourceInfo: SourceInfo?) {

compiler/backend/src/org/jetbrains/kotlin/codegen/inline/inlineCodegenUtils.kt

@@ -113,6 +113,9 @@ internal inline fun getMethodNode(classData: ByteArray, classType: Type, crossin
 internal fun getMethodNode(classData: ByteArray, classType: Type, method: Method): SMAPAndMethodNode? =
     getMethodNode(classData, classType) { it == method }
 
+fun argumentsSize(descriptor: String, isStatic: Boolean): Int =
+    (Type.getArgumentsAndReturnSizes(descriptor) shr 2) - (if (isStatic) 1 else 0)
+
 internal fun findVirtualFile(state: GenerationState, classId: ClassId): VirtualFile? {
     return VirtualFileFinder.getInstance(state.project, state.module).findVirtualFileWithHeader(classId)
 }

compiler/ir/backend.jvm/codegen/src/org/jetbrains/kotlin/backend/jvm/codegen/ClassCodegen.kt

@@ -405,25 +405,7 @@ class ClassCodegen private constructor(
             method.origin == JvmLoweredDeclarationOrigin.FOR_INLINE_STATE_MACHINE_TEMPLATE_CAPTURES_CROSSINLINE
         )
         val mv = with(node) { visitor.newMethod(method.descriptorOrigin, access, name, desc, signature, exceptions.toTypedArray()) }
-        val smapCopier = SourceMapCopier(classSMAP, smap)
+        val smapCopyingVisitor = SourceMapCopyingMethodVisitor(classSMAP, smap, mv)
-
-        val smapCopyingVisitor = object : MethodVisitor(Opcodes.API_VERSION, mv) {
-            private val lineNumberMapping = mutableMapOf<Int, Int>()
-
-            override fun visitLineNumber(line: Int, start: Label) {
-                val newLine = smapCopier.mapLineNumber(line)
-                lineNumberMapping[line] = newLine
-                super.visitLineNumber(newLine, start)
-            }
-
-            override fun visitLocalVariable(name: String, descriptor: String, signature: String?, start: Label, end: Label, index: Int) {
-                if (state.configuration.getBoolean(JVMConfigurationKeys.USE_INLINE_SCOPES_NUMBERS) && isFakeLocalVariableForInline(name)) {
-                    val newName = updateCallSiteLineNumber(name, lineNumberMapping)
-                    return super.visitLocalVariable(newName, descriptor, signature, start, end, index)
-                }
-                super.visitLocalVariable(name, descriptor, signature, start, end, index)
-            }
-        }
 
         if (method.hasContinuation()) {
             // Generate a state machine within this method. The continuation class for it should be generated

compiler/ir/backend.jvm/codegen/src/org/jetbrains/kotlin/backend/jvm/codegen/IrInlineDefaultCodegen.kt

@@ -7,7 +7,8 @@ package org.jetbrains.kotlin.backend.jvm.codegen
 
 import org.jetbrains.kotlin.backend.jvm.mapping.IrCallableMethod
 import org.jetbrains.kotlin.codegen.inline.MethodBodyVisitor
-import org.jetbrains.kotlin.codegen.inline.SourceMapCopier
+import org.jetbrains.kotlin.codegen.inline.SourceMapCopyingMethodVisitor
+import org.jetbrains.kotlin.codegen.inline.argumentsSize
 import org.jetbrains.kotlin.ir.declarations.IrValueParameter
 import org.jetbrains.kotlin.ir.expressions.IrExpression
 import org.jetbrains.kotlin.ir.expressions.IrFunctionAccessExpression
@@ -42,20 +43,14 @@ object IrInlineDefaultCodegen : IrInlineCallGenerator {
         isInsideIfCondition: Boolean
     ) {
         val function = expression.symbol.owner
-        val nodeAndSmap = codegen.classCodegen.generateMethodNode(function)
+        val (node, smap) = codegen.classCodegen.generateMethodNode(function)
-        val childSourceMapper = SourceMapCopier(codegen.smap, nodeAndSmap.classSMAP)
+        val argsSize = argumentsSize(callableMethod.asmMethod.descriptor, function.isStatic)
-
+        val mv = object : MethodBodyVisitor(codegen.visitor) {
-        val argsSize =
-            (Type.getArgumentsAndReturnSizes(callableMethod.asmMethod.descriptor) ushr 2) - if (function.isStatic) 1 else 0
-        nodeAndSmap.node.accept(object : MethodBodyVisitor(codegen.visitor) {
             override fun visitLocalVariable(name: String, desc: String, signature: String?, start: Label, end: Label, index: Int) {
-                // We only copy LVT entries for local variables, since we already generated entries for the method parameters,
+                // We only copy LVT entries for local variables, since we already generated entries for the method parameters.
                 if (index >= argsSize) super.visitLocalVariable(name, desc, signature, start, end, index)
             }
-
+        }
-            override fun visitLineNumber(line: Int, start: Label?) {
+        node.accept(SourceMapCopyingMethodVisitor(codegen.smap, smap, mv))
-                super.visitLineNumber(childSourceMapper.mapLineNumber(line), start)
-            }
-        })
     }
 }