JVM: generalize getMethodNode

to allow matching by arbitrary predicates.
2021-06-27 20:49:44 +02:00
parent f986591ba9
commit d0b9c4ae6b
3 changed files with 43 additions and 61 deletions
@@ -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 constructor = getMethodNode(classBytes, "<init>", constructorDescriptor, lambdaClassType)?.node
+        val constructorMethod = Method("<init>", Type.VOID_TYPE, capturedArgs)
+        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)
-        // TODO: `signatureAmbiguity = true` ignores the argument types from `invokeMethod` and only looks at the count.
-        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}'")
+        node = getMethodNodeImprecise(loadClass(sourceCompiler), lambdaClassType, actualMethod, erasedName)
+            ?: 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,
@@ -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,
-    descriptor: String,
+    method: Method,
    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 nameWithoutManglingSuffix = if (dashIndex > 0) name.substring(0, dashIndex) else name
-        if (nameWithoutManglingSuffix != name) {
-            getMethodNode(owner, bytes, nameWithoutManglingSuffix, descriptor, isSuspend)?.let { return it }
+        val dashIndex = method.name.indexOf('-')
+        val nameWithoutManglingSuffix = if (dashIndex > 0) method.name.substring(0, dashIndex) else method.name
+        if (nameWithoutManglingSuffix != method.name) {
+            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) =
-    (if (isSuspend) getMethodNode(bytes, name + FOR_INLINE_SUFFIX, descriptor, owner) else null)
-        ?: getMethodNode(bytes, name, descriptor, owner)
+private fun getMethodNode(owner: Type, bytes: ByteArray, method: Method, isSuspend: Boolean) =
+    (if (isSuspend) getMethodNode(bytes, owner, Method(method.name + FOR_INLINE_SUFFIX, method.descriptor)) else null)
+        ?: getMethodNode(bytes, owner, method)
@@ -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(
-    classData: ByteArray,
-    methodName: String,
-    methodDescriptor: String,
-    classType: Type,
-    signatureAmbiguity: Boolean = false
-): SMAPAndMethodNode? {
-    val cr = ClassReader(classData)
+internal inline fun getMethodNode(classData: ByteArray, classType: Type, crossinline match: (Method, Int) -> Boolean): SMAPAndMethodNode? {
    var node: MethodNode? = null
-    val debugInfo = arrayOfNulls<String>(2)
-
-    cr.accept(object : ClassVisitor(Opcodes.API_VERSION) {
+    var sourceFile: String? = null
+    var sourceMap: String? = null
+    ClassReader(classData).accept(object : ClassVisitor(Opcodes.API_VERSION) {
        override fun visitSource(source: String?, debug: String?) {
-            super.visitSource(source, debug)
-            debugInfo[0] = source
-            debugInfo[1] = debug
+            sourceFile = source
+            sourceMap = debug
        }

-        override fun visitMethod(
-            access: Int,
-            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
-                }
-            }
-
+        override fun visitMethod(access: Int, name: String, desc: String, signature: String?, exceptions: Array<String>?): MethodVisitor? {
+            if (!match(Method(name, desc), access)) 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 null
+    return node?.let{
+        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