JVM: generalize getMethodNode

to allow matching by arbitrary predicates.

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

@@ -11,6 +11,7 @@ import org.jetbrains.kotlin.types.KotlinType
 import org.jetbrains.org.objectweb.asm.ClassReader
 import org.jetbrains.org.objectweb.asm.Type
 import org.jetbrains.org.objectweb.asm.Label
+import org.jetbrains.org.objectweb.asm.Opcodes
 import org.jetbrains.org.objectweb.asm.commons.Method
 import org.jetbrains.org.objectweb.asm.tree.FieldInsnNode
 
@@ -98,10 +99,10 @@ abstract class DefaultLambda(
         isPropertyReference = superName in PROPERTY_REFERENCE_SUPER_CLASSES
         isFunctionReference = superName == FUNCTION_REFERENCE.internalName || superName == FUNCTION_REFERENCE_IMPL.internalName
 
-        val constructorDescriptor = Type.getMethodDescriptor(Type.VOID_TYPE, *capturedArgs)
+        val constructorMethod = Method("<init>", Type.VOID_TYPE, capturedArgs)
-        val constructor = getMethodNode(classBytes, "<init>", constructorDescriptor, lambdaClassType)?.node
+        val constructor = getMethodNode(classBytes, lambdaClassType, constructorMethod)?.node
         assert(constructor != null || capturedArgs.isEmpty()) {
-            "Can't find non-default constructor <init>$constructorDescriptor for default lambda $lambdaClassType"
+            "can't find constructor '$constructorMethod' for default lambda '${lambdaClassType.internalName}'"
         }
         // This only works for primitives but not inline classes, since information about the Kotlin type of the bound
         // receiver is not present anywhere. This is why with JVM_IR the constructor argument of bound references
@@ -127,15 +128,18 @@ abstract class DefaultLambda(
 
     // Returns whether the loaded invoke is erased, i.e. the name equals the fallback and all types are `Object`.
     protected fun loadInvoke(sourceCompiler: SourceCompilerForInline, erasedName: String, actualMethod: Method): Boolean {
-        val classBytes = loadClass(sourceCompiler)
+        node = getMethodNodeImprecise(loadClass(sourceCompiler), lambdaClassType, actualMethod, erasedName)
-        // TODO: `signatureAmbiguity = true` ignores the argument types from `invokeMethod` and only looks at the count.
+            ?: error("Can't find method '$actualMethod' in '${lambdaClassType.internalName}'")
-        node = getMethodNode(classBytes, actualMethod.name, actualMethod.descriptor, lambdaClassType, signatureAmbiguity = true)
-            ?: getMethodNode(classBytes, erasedName, actualMethod.descriptor, lambdaClassType, signatureAmbiguity = true)
-                    ?: error("Can't find method '$actualMethod' in '${lambdaClassType.internalName}'")
         return invokeMethod.run { name == erasedName && returnType == OBJECT_TYPE && argumentTypes.all { it == OBJECT_TYPE } }
     }
 
     private companion object {
+        fun getMethodNodeImprecise(classBytes: ByteArray, classType: Type, method: Method, erasedName: String) =
+            getMethodNode(classBytes, classType) { it, access ->
+                (it.name == method.name || it.name == erasedName) &&
+                        it.argumentTypes.size == method.argumentTypes.size && access.and(Opcodes.ACC_SYNTHETIC) == 0
+            }
+
         val PROPERTY_REFERENCE_SUPER_CLASSES =
             listOf(
                 PROPERTY_REFERENCE0, PROPERTY_REFERENCE1, PROPERTY_REFERENCE2,

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

@@ -78,12 +78,12 @@ fun loadCompiledInlineFunction(
     state: GenerationState
 ): SMAPAndMethodNode {
     val containerType = AsmUtil.asmTypeByClassId(containerId)
-    val bytes = state.inlineCache.classBytes.getOrPut(containerType.internalName) {
-        findVirtualFile(state, containerId)?.contentsToByteArray()
-            ?: throw IllegalStateException("Couldn't find declaration file for $containerId")
-    }
     val resultInCache = state.inlineCache.methodNodeById.getOrPut(MethodId(containerType.descriptor, asmMethod)) {
-        getMethodNode(containerType, bytes, asmMethod.name, asmMethod.descriptor, isSuspend, isMangled)
+        val bytes = state.inlineCache.classBytes.getOrPut(containerType.internalName) {
+            findVirtualFile(state, containerId)?.contentsToByteArray()
+                ?: throw IllegalStateException("Couldn't find declaration file for $containerId")
+        }
+        getMethodNode(containerType, bytes, asmMethod, isSuspend, isMangled)
     }
     return SMAPAndMethodNode(cloneMethodNode(resultInCache.node), resultInCache.classSMAP)
 }
@@ -91,25 +91,24 @@ fun loadCompiledInlineFunction(
 private fun getMethodNode(
     owner: Type,
     bytes: ByteArray,
-    name: String,
+    method: Method,
-    descriptor: String,
     isSuspend: Boolean,
     isMangled: Boolean
 ): SMAPAndMethodNode {
-    getMethodNode(owner, bytes, name, descriptor, isSuspend)?.let { return it }
+    getMethodNode(owner, bytes, method, isSuspend)?.let { return it }
     if (isMangled) {
         // Compatibility with old inline class ABI versions.
-        val dashIndex = name.indexOf('-')
+        val dashIndex = method.name.indexOf('-')
-        val nameWithoutManglingSuffix = if (dashIndex > 0) name.substring(0, dashIndex) else name
+        val nameWithoutManglingSuffix = if (dashIndex > 0) method.name.substring(0, dashIndex) else method.name
-        if (nameWithoutManglingSuffix != name) {
+        if (nameWithoutManglingSuffix != method.name) {
-            getMethodNode(owner, bytes, nameWithoutManglingSuffix, descriptor, isSuspend)?.let { return it }
+            getMethodNode(owner, bytes, Method(nameWithoutManglingSuffix, method.descriptor), isSuspend)?.let { return it }
         }
-        getMethodNode(owner, bytes, "$nameWithoutManglingSuffix-impl", descriptor, isSuspend)?.let { return it }
+        getMethodNode(owner, bytes, Method("$nameWithoutManglingSuffix-impl", method.descriptor), isSuspend)?.let { return it }
     }
-    throw IllegalStateException("couldn't find inline method $owner.$name$descriptor")
+    throw IllegalStateException("couldn't find inline method $owner.$method")
 }
 
 // If an `inline suspend fun` has a state machine, it should have a `$$forInline` version without one.
-private fun getMethodNode(owner: Type, bytes: ByteArray, name: String, descriptor: String, isSuspend: Boolean) =
+private fun getMethodNode(owner: Type, bytes: ByteArray, method: Method, isSuspend: Boolean) =
-    (if (isSuspend) getMethodNode(bytes, name + FOR_INLINE_SUFFIX, descriptor, owner) else null)
+    (if (isSuspend) getMethodNode(bytes, owner, Method(method.name + FOR_INLINE_SUFFIX, method.descriptor)) else null)
-        ?: getMethodNode(bytes, name, descriptor, owner)
+        ?: getMethodNode(bytes, owner, method)

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

@@ -43,6 +43,7 @@ import org.jetbrains.kotlin.util.OperatorNameConventions
 import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 import org.jetbrains.org.objectweb.asm.*
 import org.jetbrains.org.objectweb.asm.commons.InstructionAdapter
+import org.jetbrains.org.objectweb.asm.commons.Method
 import org.jetbrains.org.objectweb.asm.tree.*
 import org.jetbrains.org.objectweb.asm.util.Printer
 import org.jetbrains.org.objectweb.asm.util.Textifier
@@ -84,40 +85,18 @@ private const val INLINE_MARKER_AFTER_INLINE_SUSPEND_ID = 7
 private const val INLINE_MARKER_BEFORE_UNBOX_INLINE_CLASS = 8
 private const val INLINE_MARKER_AFTER_UNBOX_INLINE_CLASS = 9
 
-internal fun getMethodNode(
+internal inline fun getMethodNode(classData: ByteArray, classType: Type, crossinline match: (Method, Int) -> Boolean): SMAPAndMethodNode? {
-    classData: ByteArray,
-    methodName: String,
-    methodDescriptor: String,
-    classType: Type,
-    signatureAmbiguity: Boolean = false
-): SMAPAndMethodNode? {
-    val cr = ClassReader(classData)
     var node: MethodNode? = null
-    val debugInfo = arrayOfNulls<String>(2)
+    var sourceFile: String? = null
-
+    var sourceMap: String? = null
-    cr.accept(object : ClassVisitor(Opcodes.API_VERSION) {
+    ClassReader(classData).accept(object : ClassVisitor(Opcodes.API_VERSION) {
         override fun visitSource(source: String?, debug: String?) {
-            super.visitSource(source, debug)
+            sourceFile = source
-            debugInfo[0] = source
+            sourceMap = debug
-            debugInfo[1] = debug
         }
 
-        override fun visitMethod(
+        override fun visitMethod(access: Int, name: String, desc: String, signature: String?, exceptions: Array<String>?): MethodVisitor? {
-            access: Int,
+            if (!match(Method(name, desc), access)) return null
-            name: String,
-            desc: String,
-            signature: String?,
-            exceptions: Array<String>?
-        ): MethodVisitor? {
-            if (methodName != name || (signatureAmbiguity && access.and(Opcodes.ACC_SYNTHETIC) != 0)) return null
-
-            if (methodDescriptor != desc) {
-                val sameNumberOfParameters = Type.getArgumentTypes(methodDescriptor).size == Type.getArgumentTypes(desc).size
-                if (!signatureAmbiguity || !sameNumberOfParameters) {
-                    return null
-                }
-            }
-
             node?.let { existing ->
                 throw AssertionError("Can't find proper '$name' method for inline: ambiguity between '${existing.name + existing.desc}' and '${name + desc}'")
             }
@@ -126,15 +105,15 @@ internal fun getMethodNode(
         }
     }, ClassReader.SKIP_FRAMES or if (GENERATE_SMAP) 0 else ClassReader.SKIP_DEBUG)
 
-    if (node == null) {
+    return node?.let{
-        return null
+        val (first, last) = listOfNotNull(it).lineNumberRange()
+        SMAPAndMethodNode(it, SMAPParser.parseOrCreateDefault(sourceMap, sourceFile, classType.internalName, first, last))
     }
-
-    val (first, last) = listOfNotNull(node).lineNumberRange()
-    val smap = SMAPParser.parseOrCreateDefault(debugInfo[1], debugInfo[0], classType.internalName, first, last)
-    return SMAPAndMethodNode(node!!, smap)
 }
 
+internal fun getMethodNode(classData: ByteArray, classType: Type, method: Method): SMAPAndMethodNode? =
+    getMethodNode(classData, classType) { it, _ -> it == method }
+
 internal fun Collection<MethodNode>.lineNumberRange(): Pair<Int, Int> {
     var minLine = Int.MAX_VALUE
     var maxLine = Int.MIN_VALUE