From 91d44de7eceeb275ca9a6e1ee2099e3788d245f3 Mon Sep 17 00:00:00 2001
From: joe <burtonjae@hotmail.co.uk>
Date: Sun, 21 Jul 2024 22:46:25 +0100
Subject: [PATCH] Fix `@ModifyVariable` method signature checking with `STORE`.
 Closes #2163

---
 .../platform/mixin/handlers/ModifyVariableHandler.kt      | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/src/main/kotlin/platform/mixin/handlers/ModifyVariableHandler.kt b/src/main/kotlin/platform/mixin/handlers/ModifyVariableHandler.kt
index 42398c23d..ffc8799c0 100644
--- a/src/main/kotlin/platform/mixin/handlers/ModifyVariableHandler.kt
+++ b/src/main/kotlin/platform/mixin/handlers/ModifyVariableHandler.kt
@@ -23,6 +23,7 @@ package com.demonwav.mcdev.platform.mixin.handlers
 import com.demonwav.mcdev.platform.mixin.handlers.injectionPoint.AbstractLoadInjectionPoint
 import com.demonwav.mcdev.platform.mixin.handlers.injectionPoint.CollectVisitor
 import com.demonwav.mcdev.platform.mixin.handlers.injectionPoint.InjectionPoint
+import com.demonwav.mcdev.platform.mixin.handlers.injectionPoint.StoreInjectionPoint
 import com.demonwav.mcdev.platform.mixin.inspection.injector.MethodSignature
 import com.demonwav.mcdev.platform.mixin.inspection.injector.ParameterGroup
 import com.demonwav.mcdev.platform.mixin.util.LocalInfo
@@ -47,7 +48,9 @@ class ModifyVariableHandler : InjectorAnnotationHandler() {
 
         val at = annotation.findAttributeValue("at") as? PsiAnnotation
         val atCode = at?.findAttributeValue("value")?.constantStringValue
-        val isLoadStore = atCode != null && InjectionPoint.byAtCode(atCode) is AbstractLoadInjectionPoint
+        val injectionPoint = atCode?.let { InjectionPoint.byAtCode(atCode) }
+        val isLoadStore = injectionPoint is AbstractLoadInjectionPoint
+        val isStore = injectionPoint is StoreInjectionPoint
         val mode = if (isLoadStore) CollectVisitor.Mode.COMPLETION else CollectVisitor.Mode.MATCH_ALL
         val targets = resolveInstructions(annotation, targetClass, targetMethod, mode)
 
@@ -63,7 +66,8 @@ class ModifyVariableHandler : InjectorAnnotationHandler() {
 
         val possibleTypes = mutableSetOf<String>()
         for (insn in targets) {
-            val locals = info.getLocals(module, targetClass, targetMethod, insn.insn) ?: continue
+            val actualInsn = if (isStore) insn.insn.next ?: insn.insn else insn.insn
+            val locals = info.getLocals(module, targetClass, targetMethod, actualInsn) ?: continue
             val matchedLocals = info.matchLocals(locals, CollectVisitor.Mode.COMPLETION, matchType = false)
             for (local in matchedLocals) {
                 possibleTypes += local.desc!!