From 151f2eac15cbbd813ac61f3aa786a8b984f7bdda Mon Sep 17 00:00:00 2001
From: Joe <burtonjae@hotmail.co.uk>
Date: Fri, 24 Nov 2023 15:21:23 +0000
Subject: [PATCH] Make INVOKE_ASSIGN target the instruction after the
 invocation

---
 ...int.kt => AbstractInvokeInjectionPoint.kt} | 24 ++++++++++++-------
 src/main/resources/META-INF/plugin.xml        |  4 ++--
 2 files changed, 18 insertions(+), 10 deletions(-)
 rename src/main/kotlin/platform/mixin/handlers/injectionPoint/{MethodInjectionPoint.kt => AbstractInvokeInjectionPoint.kt} (90%)

diff --git a/src/main/kotlin/platform/mixin/handlers/injectionPoint/MethodInjectionPoint.kt b/src/main/kotlin/platform/mixin/handlers/injectionPoint/AbstractInvokeInjectionPoint.kt
similarity index 90%
rename from src/main/kotlin/platform/mixin/handlers/injectionPoint/MethodInjectionPoint.kt
rename to src/main/kotlin/platform/mixin/handlers/injectionPoint/AbstractInvokeInjectionPoint.kt
index e83b1a36a..dd62f10c3 100644
--- a/src/main/kotlin/platform/mixin/handlers/injectionPoint/MethodInjectionPoint.kt
+++ b/src/main/kotlin/platform/mixin/handlers/injectionPoint/AbstractInvokeInjectionPoint.kt
@@ -37,7 +37,7 @@ import org.objectweb.asm.tree.ClassNode
 import org.objectweb.asm.tree.MethodInsnNode
 import org.objectweb.asm.tree.MethodNode
 
-class MethodInjectionPoint : AbstractMethodInjectionPoint() {
+abstract class AbstractInvokeInjectionPoint(private val assign: Boolean) : AbstractMethodInjectionPoint() {
     override fun createNavigationVisitor(
         at: PsiAnnotation,
         target: MixinSelector?,
@@ -53,9 +53,9 @@ class MethodInjectionPoint : AbstractMethodInjectionPoint() {
         mode: CollectVisitor.Mode,
     ): CollectVisitor<PsiMethod>? {
         if (mode == CollectVisitor.Mode.COMPLETION) {
-            return MyCollectVisitor(mode, at.project, MemberReference(""))
+            return MyCollectVisitor(mode, at.project, MemberReference(""), assign)
         }
-        return target?.let { MyCollectVisitor(mode, at.project, it) }
+        return target?.let { MyCollectVisitor(mode, at.project, it, assign) }
     }
 
     private class MyNavigationVisitor(
@@ -141,6 +141,7 @@ class MethodInjectionPoint : AbstractMethodInjectionPoint() {
         mode: Mode,
         private val project: Project,
         private val selector: MixinSelector,
+        private val assign: Boolean,
     ) : CollectVisitor<PsiMethod>(mode) {
         override fun accept(methodNode: MethodNode) {
             val insns = methodNode.instructions ?: return
@@ -150,12 +151,19 @@ class MethodInjectionPoint : AbstractMethodInjectionPoint() {
                 }
 
                 val sourceMethod = nodeMatchesSelector(insn, mode, selector, project) ?: return@forEachRemaining
-                addResult(
-                    insn,
-                    sourceMethod,
-                    qualifier = insn.owner.replace('/', '.'),
-                )
+                val actualInsn = if (assign) insn.next else insn
+                if (actualInsn != null) {
+                    addResult(
+                        actualInsn,
+                        sourceMethod,
+                        qualifier = insn.owner.replace('/', '.'),
+                    )
+                }
             }
         }
     }
 }
+
+class InvokeInjectionPoint : AbstractInvokeInjectionPoint(false)
+
+class InvokeAssignInjectionPoint : AbstractInvokeInjectionPoint(true)
diff --git a/src/main/resources/META-INF/plugin.xml b/src/main/resources/META-INF/plugin.xml
index 6a2e8fef5..c7d130852 100644
--- a/src/main/resources/META-INF/plugin.xml
+++ b/src/main/resources/META-INF/plugin.xml
@@ -142,8 +142,8 @@
         <mixinAnnotationHandler annotation="com.llamalad7.mixinextras.injector.wrapoperation.WrapOperation" implementation="com.demonwav.mcdev.platform.mixin.handlers.mixinextras.WrapOperationHandler" />
         <mixinAnnotationHandler annotation="com.llamalad7.mixinextras.injector.WrapWithCondition" implementation="com.demonwav.mcdev.platform.mixin.handlers.mixinextras.WrapWithConditionHandler" />
 
-        <injectionPoint atCode="INVOKE" implementation="com.demonwav.mcdev.platform.mixin.handlers.injectionPoint.MethodInjectionPoint" />
-        <injectionPoint atCode="INVOKE_ASSIGN" implementation="com.demonwav.mcdev.platform.mixin.handlers.injectionPoint.MethodInjectionPoint" />
+        <injectionPoint atCode="INVOKE" implementation="com.demonwav.mcdev.platform.mixin.handlers.injectionPoint.InvokeInjectionPoint" />
+        <injectionPoint atCode="INVOKE_ASSIGN" implementation="com.demonwav.mcdev.platform.mixin.handlers.injectionPoint.InvokeAssignInjectionPoint" />
         <injectionPoint atCode="INVOKE_STRING" implementation="com.demonwav.mcdev.platform.mixin.handlers.injectionPoint.ConstantStringMethodInjectionPoint" />
         <injectionPoint atCode="FIELD" implementation="com.demonwav.mcdev.platform.mixin.handlers.injectionPoint.FieldInjectionPoint" />
         <injectionPoint atCode="NEW" implementation="com.demonwav.mcdev.platform.mixin.handlers.injectionPoint.NewInsnInjectionPoint" />