User: bawnorton
Date: 13 Nov 26 20:35
Revision: f7345544478a5e15fa7aefddbfb0aa9d911ad054
Summary:
Mixin `@At` navigation fixes for transformed sources (#2656)

* Allow mixin selectors to resolve source methods for navigation

* Navigate injection targets through selector source methods

* Move selector resolution into at resolver

TeamCity URL: http://ci.mcdev.io:80/viewModification.html?tab=vcsModificationFiles&modId=10759&personal=false

Index: src/main/kotlin/platform/mixin/handlers/injectionPoint/AtResolver.kt
===================================================================
--- src/main/kotlin/platform/mixin/handlers/injectionPoint/AtResolver.kt	(revision 40750da50c6bfa7e43da32e8588ddf4aecfc6da9)
+++ src/main/kotlin/platform/mixin/handlers/injectionPoint/AtResolver.kt	(revision f7345544478a5e15fa7aefddbfb0aa9d911ad054)
@@ -43,6 +43,7 @@
 import com.demonwav.mcdev.util.countIsAtLeast
 import com.demonwav.mcdev.util.descriptor
 import com.demonwav.mcdev.util.equivalentTo
+import com.demonwav.mcdev.util.findContainingClass
 import com.demonwav.mcdev.util.findMethods
 import com.demonwav.mcdev.util.internalName
 import com.intellij.codeInsight.lookup.LookupElementBuilder
@@ -240,9 +241,26 @@
 
         val project = at.project
 
+        val injector = findInjectorAnnotation(at)
+        val methodAttr = injector?.findAttributeValue("method")
+        val selectors = methodAttr?.computeStringArray()?.mapNotNull { parseMixinSelector(it, methodAttr) }.orEmpty()
+        val sourceElement = selectors.firstNotNullOfOrNull { selector ->
+            if (selector.matchMethod(targetMethod, targetClass)) {
+                selector.findSourceMethod(
+                    targetClass,
+                    targetMethod,
+                    project,
+                    at.resolveScope,
+                    canDecompile = true,
+                )
+            } else {
+                null
+            }
+        }
+
         // Resolve the target source class
-        val targetPsiClass = getTargetClass(target)
-            .findSourceClass(project, GlobalSearchScope.allScope(project), canDecompile = true)
+        val targetPsiClass = sourceElement?.findContainingClass()
+            ?: getTargetClass(target).findSourceClass(project, GlobalSearchScope.allScope(project), canDecompile = true)
             ?: return emptyList()
         val targetPsiFile = targetPsiClass.containingFile ?: return emptyList()
 
@@ -252,13 +270,16 @@
 
         // Find the element in the desugared class, first by directly searching and then by searching in the original
         // and reverse mapping it into the desugared class.
-        val desugaredTargetElement = when {
+        val desugaredTargetElement = sourceElement?.let { original ->
+            DesugarUtil.getOriginalToDesugaredMap(desugaredTargetClass)[original]
+                ?.firstOrNull { it is PsiParameterListOwner }
+        } ?: when {
             targetMethod.isClinit -> desugaredTargetClass.initializers.firstOrNull {
                 it.hasModifierProperty(PsiModifier.STATIC)
             }
             else -> desugaredTargetClass.findMethods(targetMethod.memberReference).firstOrNull()
         } ?: run {
-            val originalTargetElement = targetMethod.findSourceElement(
+            val originalTargetElement = sourceElement ?: targetMethod.findSourceElement(
                 getTargetClass(target),
                 project,
                 GlobalSearchScope.allScope(project),
Index: src/main/kotlin/platform/mixin/handlers/injectionPoint/FieldInjectionPoint.kt
===================================================================
--- src/main/kotlin/platform/mixin/handlers/injectionPoint/FieldInjectionPoint.kt	(revision 40750da50c6bfa7e43da32e8588ddf4aecfc6da9)
+++ src/main/kotlin/platform/mixin/handlers/injectionPoint/FieldInjectionPoint.kt	(revision f7345544478a5e15fa7aefddbfb0aa9d911ad054)
@@ -26,6 +26,7 @@
 import com.demonwav.mcdev.platform.mixin.util.findOrConstructSourceField
 import com.demonwav.mcdev.util.Quantifier
 import com.demonwav.mcdev.util.constantValue
+import com.demonwav.mcdev.util.descriptor
 import com.intellij.codeInsight.completion.JavaLookupElementBuilder
 import com.intellij.codeInsight.lookup.LookupElementBuilder
 import com.intellij.openapi.editor.Editor
@@ -126,16 +127,24 @@
         private val opcode: Int,
         private val arrayAccess: ArrayAccessType?,
     ) : NavigationVisitor() {
+        private var bytecodeOwner: String? = null
+
+        override fun configureBytecodeTarget(classNode: ClassNode, methodNode: MethodNode) {
+            bytecodeOwner = classNode.name
+        }
+
         override fun visitReferenceExpression(expression: PsiReferenceExpression) {
             if (expression !is PsiMethodReferenceExpression) {
                 // early out for if the name does not match
                 val name = expression.referenceName
                 if (name == null || selector.canEverMatch(name)) {
                     (expression.resolve() as? PsiField)?.let { resolved ->
-                        var matches = selector.matchField(
-                            resolved,
-                            QualifiedMember.resolveQualifier(expression) ?: targetClass,
-                        )
+                        val qualifier = QualifiedMember.resolveQualifier(expression) ?: targetClass
+                        var matches = if (qualifier.qualifiedName == targetClass.qualifiedName && bytecodeOwner != null) {
+                            resolved.descriptor?.let { selector.matchField(bytecodeOwner!!, resolved.name, it) } == true
+                        } else {
+                            selector.matchField(resolved, qualifier)
+                        }
                         if (matches && opcode != -1) {
                             // check if we match the opcode
                             val isStatic = opcode == Opcodes.GETSTATIC || opcode == Opcodes.PUTSTATIC
Index: src/main/kotlin/platform/mixin/util/AsmUtil.kt
===================================================================
--- src/main/kotlin/platform/mixin/util/AsmUtil.kt	(revision 40750da50c6bfa7e43da32e8588ddf4aecfc6da9)
+++ src/main/kotlin/platform/mixin/util/AsmUtil.kt	(revision f7345544478a5e15fa7aefddbfb0aa9d911ad054)
@@ -1028,6 +1028,9 @@
         return psiClass.childrenOfType<PsiClassInitializer>().firstOrNull { it.hasModifierProperty(PsiModifier.STATIC) }
             ?: psiClass
     }
+    if (hasAccess(Opcodes.ACC_SYNTHETIC) && psiClass !is PsiCompiledElement) {
+        findAssociatedLambda(project, scope, clazz, this)?.let { return it }
+    }
     psiClass.findMethods(memberReference).firstOrNull()?.let { return it }
     if (psiClass is PsiCompiledElement) {
         // don't walk into stub compiled elements to look for lambdas
Index: src/main/kotlin/platform/mixin/util/SourceCodeLocationInfo.kt
===================================================================
--- src/main/kotlin/platform/mixin/util/SourceCodeLocationInfo.kt	(revision 40750da50c6bfa7e43da32e8588ddf4aecfc6da9)
+++ src/main/kotlin/platform/mixin/util/SourceCodeLocationInfo.kt	(revision f7345544478a5e15fa7aefddbfb0aa9d911ad054)
@@ -53,6 +53,7 @@
             private var currentLine: Int? = null
             private var countThisLine = 0
             private var myResult: T? = null
+            private var indexedResult: T? = null
 
             override fun accept(t: T): Boolean {
                 val line = document?.let(t::lineNumber)
@@ -70,6 +71,7 @@
                 }
 
                 if (count++ == index) {
+                    indexedResult = t
                     if (lineNumber == null) {
                         myResult = t
                         return true
@@ -79,7 +81,7 @@
                 return false
             }
 
-            override val result get() = myResult
+            override val result get() = myResult ?: indexedResult
         }
     }
 }