Mixin @At navigation fixes for transformed sources

src/main/kotlin/platform/mixin/handlers/injectionPoint/AtResolver.kt

@@ -44,6 +44,7 @@ import com.demonwav.mcdev.util.countIsAtLeast
 import com.demonwav.mcdev.util.descriptor
 import com.demonwav.mcdev.util.equivalentTo
 import com.demonwav.mcdev.util.findMethods
+import com.demonwav.mcdev.util.findContainingClass
 import com.demonwav.mcdev.util.internalName
 import com.intellij.codeInsight.lookup.LookupElementBuilder
 import com.intellij.psi.JavaPsiFacade
@@ -93,6 +94,7 @@ class AtResolver(
     private val at: PsiAnnotation,
     private val targetClass: ClassNode,
     private val targetMethod: MethodNode,
+    private val sourceElement: PsiElement? = null,
 ) {
     companion object {
         fun getInjectionPoint(at: PsiAnnotation): InjectionPoint<*>? {
@@ -241,8 +243,8 @@ class AtResolver(
         val project = at.project
 
         // 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 +254,16 @@ class AtResolver(
 
         // 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),

src/main/kotlin/platform/mixin/handlers/injectionPoint/FieldInjectionPoint.kt

@@ -26,6 +26,7 @@ import com.demonwav.mcdev.platform.mixin.util.fakeResolve
 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 @@ class FieldInjectionPoint : QualifiedInjectionPoint<PsiField>() {
         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

src/main/kotlin/platform/mixin/reference/target/TargetReference.kt

@@ -29,6 +29,7 @@ import com.demonwav.mcdev.platform.mixin.util.ClassAndMethodNode
 import com.demonwav.mcdev.platform.mixin.util.MethodTargetMember
 import com.demonwav.mcdev.platform.mixin.util.MixinConstants.Annotations.AT
 import com.demonwav.mcdev.util.findQualifiedClass
+import com.demonwav.mcdev.util.computeStringArray
 import com.demonwav.mcdev.util.ifEmpty
 import com.demonwav.mcdev.util.insideAnnotationAttribute
 import com.demonwav.mcdev.util.reference.PolyReferenceResolver
@@ -85,14 +86,17 @@ object TargetReference : PolyReferenceResolver(), MixinReference {
      * Null is returned when no parent annotation handler could be found, in which case we shouldn't mark this
      * reference as unresolved.
      */
-    private fun getTargets(at: PsiAnnotation, forUnresolved: Boolean): List<ClassAndMethodNode>? {
-        val (handler, annotation) = generateSequence(at.parent) { it.parent }
+    private fun getInjector(at: PsiAnnotation): Pair<InjectorAnnotationHandler, PsiAnnotation>? =
+        generateSequence(at.parent) { it.parent }
             .filterIsInstance<PsiAnnotation>()
             .flatMap { it.owner?.annotations?.asSequence() ?: emptySequence() }
             .mapNotNull { annotation ->
                 (MixinAnnotationHandler.forMixinAnnotation(annotation) as? InjectorAnnotationHandler)
                     ?.let { it to annotation }
-            }.firstOrNull() ?: return null
+            }.firstOrNull()
+
+    private fun getTargets(at: PsiAnnotation, forUnresolved: Boolean): List<ClassAndMethodNode>? {
+        val (handler, annotation) = getInjector(at) ?: return null
         if (forUnresolved && handler.isSoft) {
             return null
         }
@@ -113,7 +117,25 @@ object TargetReference : PolyReferenceResolver(), MixinReference {
     fun resolveNavigationTargets(context: PsiElement): Array<PsiElement>? {
         val at = context.parentOfType<PsiAnnotation>() ?: return null
         val targets = getTargets(at, false) ?: return null
-        return targets.flatMap { AtResolver(at, it.clazz, it.method).resolveNavigationTargets() }.toTypedArray()
+        val injector = getInjector(at)?.second
+        val methodAttr = injector?.findAttributeValue("method")
+        val selectors = methodAttr?.computeStringArray()?.mapNotNull { parseMixinSelector(it, methodAttr) }.orEmpty()
+        return targets.flatMap { target ->
+            val sourceElement = selectors.firstNotNullOfOrNull { selector ->
+                if (selector.matchMethod(target.method, target.clazz)) {
+                    selector.findSourceMethod(
+                        target.clazz,
+                        target.method,
+                        context.project,
+                        context.resolveScope,
+                        canDecompile = true,
+                    )
+                } else {
+                    null
+                }
+            }
+            AtResolver(at, target.clazz, target.method, sourceElement).resolveNavigationTargets()
+        }.toTypedArray()
     }
 
     override fun resolveReference(context: PsiElement): Array<ResolveResult> {

src/main/kotlin/platform/mixin/util/AsmUtil.kt

@@ -1028,6 +1028,9 @@ fun MethodNode.findSourceElement(
         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

src/main/kotlin/platform/mixin/util/SourceCodeLocationInfo.kt

@@ -53,6 +53,7 @@ class SourceCodeLocationInfo(val index: Int, val lineNumber: Int?, val indexInLi
             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 @@ class SourceCodeLocationInfo(val index: Int, val lineNumber: Int?, val indexInLi
                 }
 
                 if (count++ == index) {
+                    indexedResult = t
                     if (lineNumber == null) {
                         myResult = t
                         return true
@@ -79,7 +81,7 @@ class SourceCodeLocationInfo(val index: Int, val lineNumber: Int?, val indexInLi
                 return false
             }
 
-            override val result get() = myResult
+            override val result get() = myResult ?: indexedResult
         }
     }
 }