Replace SourceInterpreter with a specific one in coroutines inlining

compiler/backend/src/org/jetbrains/kotlin/codegen/inline/coroutines/CoroutineTransformer.kt

@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.codegen.inline.coroutines
 
 import com.intellij.util.ArrayUtil
+import jdk.internal.org.objectweb.asm.Type
 import org.jetbrains.kotlin.codegen.ClassBuilder
 import org.jetbrains.kotlin.codegen.coroutines.*
 import org.jetbrains.kotlin.codegen.inline.*
@@ -14,12 +15,14 @@ import org.jetbrains.kotlin.codegen.optimization.transformer.MethodTransformer
 import org.jetbrains.kotlin.config.isReleaseCoroutines
 import org.jetbrains.kotlin.descriptors.ClassDescriptor
 import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
+import org.jetbrains.kotlin.resolve.jvm.AsmTypes
 import org.jetbrains.kotlin.resolve.jvm.diagnostics.JvmDeclarationOrigin
-import org.jetbrains.org.objectweb.asm.*
+import org.jetbrains.org.objectweb.asm.MethodVisitor
+import org.jetbrains.org.objectweb.asm.Opcodes
 import org.jetbrains.org.objectweb.asm.tree.*
+import org.jetbrains.org.objectweb.asm.tree.analysis.BasicInterpreter
+import org.jetbrains.org.objectweb.asm.tree.analysis.BasicValue
 import org.jetbrains.org.objectweb.asm.tree.analysis.Frame
-import org.jetbrains.org.objectweb.asm.tree.analysis.SourceInterpreter
-import org.jetbrains.org.objectweb.asm.tree.analysis.SourceValue
 
 const val FOR_INLINE_SUFFIX = "\$\$forInline"
 
@@ -187,8 +190,7 @@ fun markNoinlineLambdaIfSuspend(mv: MethodVisitor, info: FunctionalArgument?) {
     }
 }
 
-private fun Frame<SourceValue>.getSource(offset: Int): AbstractInsnNode? =
+private fun Frame<BasicValue>.getSource(offset: Int): AbstractInsnNode? = (getStack(stackSize - offset - 1) as? PossibleLambdaLoad)?.insn
-    getStack(stackSize - offset - 1)?.insns?.singleOrNull()
 
 fun surroundInvokesWithSuspendMarkersIfNeeded(node: MethodNode) {
     val markers = node.instructions.asSequence().filter {
@@ -196,14 +198,14 @@ fun surroundInvokesWithSuspendMarkersIfNeeded(node: MethodNode) {
     }.toList()
     if (markers.isEmpty()) return
 
-    val sourceFrames = MethodTransformer.analyze("fake", node, SourceInterpreter())
+    val sourceFrames = MethodTransformer.analyze("fake", node, CapturedLambdaInterpreter())
     val loads = markers.map { marker ->
         val arity = (marker.next as MethodInsnNode).owner.removePrefix(NUMBERED_FUNCTION_PREFIX).toInt()
         var receiver = sourceFrames[node.instructions.indexOf(marker) + 1].getSource(arity)
         // Navigate the ALOAD+GETFIELD+... chain to the first instruction. We need to insert a stack
         // spilling marker before it starts.
         while (receiver?.opcode == Opcodes.GETFIELD) {
-            receiver = sourceFrames[node.instructions.indexOf(receiver)].getSource(0)
+            receiver = receiver.previous
         }
         receiver
     }
@@ -244,3 +246,39 @@ fun FieldInsnNode.isSuspendLambdaCapturedByOuterObjectOrLambda(inliningContext:
     }
     return isCapturedSuspendLambda(container, name, inliningContext.state.bindingContext)
 }
+
+// Interpreter, that keeps track of captured functional arguments
+private class PossibleLambdaLoad(val insn: AbstractInsnNode) : BasicValue(AsmTypes.OBJECT_TYPE)
+
+private class CapturedLambdaInterpreter : BasicInterpreter(Opcodes.API_VERSION) {
+    override fun newOperation(insn: AbstractInsnNode): BasicValue? {
+        if (insn.opcode == Opcodes.GETSTATIC) {
+            insn.fieldLoad()?.let { return it }
+        }
+
+        return super.newOperation(insn)
+    }
+
+    private fun AbstractInsnNode.fieldLoad(): PossibleLambdaLoad? {
+        if (this !is FieldInsnNode) return null
+        if (desc.startsWith('L') && Type.getType(desc).internalName.isNumberedFunctionInternalName()) {
+            if ((opcode == Opcodes.GETSTATIC && name.startsWith(CAPTURED_FIELD_FOLD_PREFIX + CAPTURED_FIELD_PREFIX)) ||
+                (opcode == Opcodes.GETFIELD && isCapturedFieldName(name))
+            ) return PossibleLambdaLoad(this)
+        }
+        return null
+    }
+
+    override fun copyOperation(insn: AbstractInsnNode, value: BasicValue?): BasicValue? =
+        if (insn.opcode == Opcodes.ALOAD) PossibleLambdaLoad(insn) else super.copyOperation(insn, value)
+
+    override fun unaryOperation(insn: AbstractInsnNode, value: BasicValue?): BasicValue? {
+        if (insn.opcode == Opcodes.GETFIELD) {
+            insn.fieldLoad()?.let { return it }
+        }
+        return super.unaryOperation(insn, value)
+    }
+
+    override fun merge(v: BasicValue?, w: BasicValue?): BasicValue? =
+        if (v is PossibleLambdaLoad && w is PossibleLambdaLoad && v.insn == w.insn) v else super.merge(v, w)
+}

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

@@ -205,11 +205,12 @@ private fun getInlineName(
 }
 
 internal fun isInvokeOnLambda(owner: String, name: String): Boolean {
-    return OperatorNameConventions.INVOKE.asString() == name &&
+    return OperatorNameConventions.INVOKE.asString() == name && owner.isNumberedFunctionInternalName()
-            owner.startsWith(NUMBERED_FUNCTION_PREFIX) &&
-            owner.substring(NUMBERED_FUNCTION_PREFIX.length).isInteger()
 }
 
+internal fun String.isNumberedFunctionInternalName(): Boolean =
+    startsWith(NUMBERED_FUNCTION_PREFIX) && substring(NUMBERED_FUNCTION_PREFIX.length).isInteger()
+
 internal fun isAnonymousConstructorCall(internalName: String, methodName: String): Boolean =
     isConstructor(methodName) && isAnonymousClass(internalName)