Support plugin-defined instrinsics that can be non-typeOf operation types.

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

@@ -66,7 +66,12 @@ class ReifiedTypeInliner<KT : KotlinTypeMarker>(
         fun reportSuspendTypeUnsupported()
         fun reportNonReifiedTypeParameterWithRecursiveBoundUnsupported(typeParameterName: Name)
 
-        fun rewritePluginDefinedOperationMarker(v: InstructionAdapter, stubConstNull: AbstractInsnNode, instructions: InsnList, type: KT): Boolean =
+        fun rewritePluginDefinedOperationMarker(
+            v: InstructionAdapter,
+            reifiedInsn: AbstractInsnNode,
+            instructions: InsnList,
+            type: KT
+        ): Boolean =
             false
     }
 
@@ -170,16 +175,19 @@ class ReifiedTypeInliner<KT : KotlinTypeMarker>(
 
             val kotlinType = intrinsicsSupport.toKotlinType(type)
 
-            if (when (operationKind) {
+            var processed = if (isPluginNext(insn)) processPlugin(insn, instructions, type) else false
-                    OperationKind.NEW_ARRAY -> processNewArray(insn, asmType)
+
-                    OperationKind.AS -> processAs(insn, instructions, kotlinType, asmType, safe = false)
+            if (!processed) processed = when (operationKind) {
-                    OperationKind.SAFE_AS -> processAs(insn, instructions, kotlinType, asmType, safe = true)
+                OperationKind.NEW_ARRAY -> processNewArray(insn, asmType)
-                    OperationKind.IS -> processIs(insn, instructions, kotlinType, asmType)
+                OperationKind.AS -> processAs(insn, instructions, kotlinType, asmType, safe = false)
-                    OperationKind.JAVA_CLASS -> processJavaClass(insn, asmType)
+                OperationKind.SAFE_AS -> processAs(insn, instructions, kotlinType, asmType, safe = true)
-                    OperationKind.ENUM_REIFIED -> processSpecialEnumFunction(insn, instructions, asmType)
+                OperationKind.IS -> processIs(insn, instructions, kotlinType, asmType)
-                    OperationKind.TYPE_OF -> processTypeOfOrPlugin(insn, instructions, type)
+                OperationKind.JAVA_CLASS -> processJavaClass(insn, asmType)
-                }
+                OperationKind.ENUM_REIFIED -> processSpecialEnumFunction(insn, instructions, asmType)
-            ) {
+                OperationKind.TYPE_OF -> processTypeOfOrPlugin(insn, instructions, type)
+            }
+
+            if (processed) {
                 instructions.remove(insn.previous.previous!!) // PUSH operation ID
                 instructions.remove(insn.previous!!) // PUSH type parameter
                 instructions.remove(insn) // INVOKESTATIC marker method
@@ -267,15 +275,7 @@ class ReifiedTypeInliner<KT : KotlinTypeMarker>(
         type: KT
     ) = rewriteNextTypeInsn(insn, Opcodes.ACONST_NULL) { stubConstNull: AbstractInsnNode ->
         val newMethodNode = newMethodNodeWithCorrectStackSize {
-            if (!isPluginNext(stubConstNull) || !intrinsicsSupport.rewritePluginDefinedOperationMarker(
+            typeSystem.generateTypeOf(it, type, intrinsicsSupport)
-                    it,
-                    stubConstNull,
-                    instructions,
-                    type,
-                )
-            ) {
-                typeSystem.generateTypeOf(it, type, intrinsicsSupport)
-            }
         }
 
         instructions.insert(insn, newMethodNode.instructions)
@@ -285,13 +285,36 @@ class ReifiedTypeInliner<KT : KotlinTypeMarker>(
         return true
     }
 
+    private fun processPlugin(insn: MethodInsnNode, instructions: InsnList, type: KT): Boolean {
+        val reifiedInsn = insn.next ?: return false
+        val newMethodNode = newMethodNodeWithCorrectStackSize {
+            if (!intrinsicsSupport.rewritePluginDefinedOperationMarker(
+                    it,
+                    reifiedInsn,
+                    instructions,
+                    type,
+                )
+            ) return false
+        }
+
+        instructions.insert(insn, newMethodNode.instructions)
+
+        maxStackSize = max(maxStackSize, newMethodNode.maxStack)
+        return true
+    }
+
+    /** insn: INVOKESTATIC reifiedOperationMarker
+     *  insn.next: operation to be reified
+     *  insn.next.next: ldc(pluginMarker)
+     *  insn.next.next.next: INVOKESTATIC voidMagicApiCall
+     */
     private fun isPluginNext(insn: AbstractInsnNode): Boolean {
-        val magicInsn = insn.next?.next ?: return false
+        val magicInsn = insn.next?.next?.next ?: return false
         return magicInsn is MethodInsnNode && magicInsn.opcode == Opcodes.INVOKESTATIC
                 && magicInsn.owner == pluginIntrinsicsMarkerOwner
                 && magicInsn.name == pluginIntrinsicsMarkerMethod
                 && magicInsn.desc == pluginIntrinsicsMarkerSignature
-                && insn.next is LdcInsnNode
+                && magicInsn.previous is LdcInsnNode
     }
 
     private inline fun rewriteNextTypeInsn(

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

@@ -229,7 +229,7 @@ internal fun isAnonymousClass(internalName: String) =
 fun wrapWithMaxLocalCalc(methodNode: MethodNode) =
     MaxStackFrameSizeAndLocalsCalculator(Opcodes.API_VERSION, methodNode.access, methodNode.desc, methodNode)
 
-fun newMethodNodeWithCorrectStackSize(block: (InstructionAdapter) -> Unit): MethodNode {
+inline fun newMethodNodeWithCorrectStackSize(block: (InstructionAdapter) -> Unit): MethodNode {
     val newMethodNode = MethodNode(Opcodes.API_VERSION, "fake", "()V", null, null)
     val mv = wrapWithMaxLocalCalc(newMethodNode)
     block(InstructionAdapter(mv))

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

@@ -127,7 +127,7 @@ class IrInlineIntrinsicsSupport(
             .report(JvmBackendErrors.TYPEOF_NON_REIFIED_TYPE_PARAMETER_WITH_RECURSIVE_BOUND, typeParameterName.asString())
     }
 
-    override fun rewritePluginDefinedOperationMarker(v: InstructionAdapter, stubConstNull: AbstractInsnNode, instructions: InsnList, type: IrType): Boolean {
+    override fun rewritePluginDefinedOperationMarker(v: InstructionAdapter, reifiedInsn: AbstractInsnNode, instructions: InsnList, type: IrType): Boolean {
-        return pluginExtensions.any { it.rewritePluginDefinedOperationMarker(v, stubConstNull, instructions, type) }
+        return pluginExtensions.any { it.rewritePluginDefinedOperationMarker(v, reifiedInsn, instructions, type) }
     }
 }

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

@@ -14,17 +14,45 @@ import org.jetbrains.org.objectweb.asm.tree.AbstractInsnNode
 import org.jetbrains.org.objectweb.asm.tree.InsnList
 
 interface JvmIrIntrinsicExtension : IrIntrinsicExtension {
+    /**
+     * Returns [IntrinsicMethod] that should emit specific bytecode instead of a regular bytecode for calling [symbol],
+     * or `null` if [symbol]'s call should not be replaced.
+     */
     fun getIntrinsic(symbol: IrFunctionSymbol): IntrinsicMethod?
 
     /**
-     * Should return `true` if marker was processed.
+     * Allows to process plugin-defined reified operation marker.
-     * If this method returns `false`, a regular `TYPE_OF` intrinsic would be inserted.
+     * Reified operation marker is an INVOKESTATIC call to IntrinsicsSupport.reifiedOperationMarker(operationType, typeVariableName) followed
+     * by actual operation to be reified.
+     * For marker to be determined as plugin-defined, another special call should be inserted directly afterwards the operation.
+     * This call is MagicApiIntrinsics.voidMagicApiCall(object). Object should be a string loaded by LDC instruction. Contents of the string is plugin-defined
+     * and recommended way to pass data to a plugin.
      *
-     * This is plugin's responsibility to remove any calls to MagicApiIntrinsics.voidMagicApiCall and its arguments.
+     * This function must return `true` if marker was processed.
+     * If this method returns `false`, other plugins would be queried, and if others also return `false`, a regular intrinsic determined by operationType would be inserted.
+     *
+     * If marker was processed, this is plugin's responsibility to remove any calls to MagicApiIntrinsics.voidMagicApiCall and its arguments.
+     *
+     * Example of plugin-defined reified operation marker:
+     *
+     * ```
+     * iconst(6) // operationType=6
+     * aconst(T) // typeParamName=T
+     * invokestatic(kotlin/jvm/internal/Intrinsics.reifiedOperationMarker)
+     * aconst(null) // This is operation to be reified. In case of operationType=6, this is typeOf<T>().
+     *              // A KType instance would normally be generated on stack instead of null.
+     * aconst("pluginDataString") // arbitrary constant string
+     * invokestatic(kotlin/jvm/internal/MagicApiIntrinsics.voidMagicApiCall(Ljava/lang/Object;)V) // plugin marker call
+     * ```
+     *
+     * Such approach with two markers was chosen mainly for compatibility reasons:
+     * Call to voidMagicApiCall should be inserted directly after operationType-specific instruction (aconst(null) in the example with typeOf).
+     * If we form bytecode this way, old compilers would be able to correctly inline normal reified operation here, even if they do not know anything about plugin reifications.
+     * They won't remove invokestatic(voidMagicApiCall), but it is not a problem, since kotlin-stdlib has this function, and it is a no-op.
      */
     fun rewritePluginDefinedOperationMarker(
         v: InstructionAdapter,
-        stubConstNull: AbstractInsnNode,
+        reifiedInsn: AbstractInsnNode,
         instructions: InsnList,
         type: IrType
     ): Boolean

plugins/kotlinx-serialization/kotlinx-serialization.backend/src/org/jetbrains/kotlinx/serialization/compiler/backend/ir/SerializationJvmIrIntrinsicSupport.kt

@@ -181,28 +181,29 @@ class SerializationJvmIrIntrinsicSupport(val jvmBackendContext: JvmBackendContex
      * 4: swap // if withModule
      * 5: invokestatic(kotlinx.serialization.serializer(module?, kType)
      *
-     * We need to remove instructions from 1 to 5
+     * We need to remove instructions from 0 to 5
-     * Instructions 0, -1 -2 and -3 would be removed by inliner.
+     * Instructions -1, -2 and -3 would be removed by inliner.
      */
     override fun rewritePluginDefinedOperationMarker(
         v: InstructionAdapter,
-        stubConstNull: AbstractInsnNode,
+        reifiedInsn: AbstractInsnNode,
         instructions: InsnList,
         type: IrType
     ): Boolean {
-        val operationTypeStr = (stubConstNull.next as LdcInsnNode).cst as String
+        val operationTypeStr = (reifiedInsn.next as LdcInsnNode).cst as String
         if (!operationTypeStr.startsWith(magicMarkerStringPrefix)) return false
         val operationType = if (operationTypeStr.endsWith("withModule")) {
-            val aload = stubConstNull.next.next.next as VarInsnNode
+            val aload = reifiedInsn.next.next.next as VarInsnNode
             val storedVar = aload.`var`
             instructions.remove(aload.next)
             instructions.remove(aload)
             IntrinsicType.WithModule(storedVar)
         } else IntrinsicType.Simple
         // Remove other instructions
-        instructions.remove(stubConstNull.next.next.next)
+        instructions.remove(reifiedInsn.next.next.next)
-        instructions.remove(stubConstNull.next.next)
+        instructions.remove(reifiedInsn.next.next)
-        instructions.remove(stubConstNull.next)
+        instructions.remove(reifiedInsn.next)
+        instructions.remove(reifiedInsn)
         // generate serializer
         generateSerializerForType(type, v, operationType)
         return true