Make INVOKE_ASSIGN target the instruction after the invocation

minecraft-dev · Nov 24, 2023 · 151f2ea · 151f2ea
1 parent 4bd9e2b
commit 151f2ea
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 10 deletions.
diff --git a/...rs/injectionPoint/MethodInjectionPoint.kt → ...tionPoint/AbstractInvokeInjectionPoint.kt b/...rs/injectionPoint/MethodInjectionPoint.kt → ...tionPoint/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
@@ -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" />