From d8c7ffe77237de8bc7b0123e9770afe436daf398 Mon Sep 17 00:00:00 2001
From: Bawnorton <bawnorton@me.com>
Date: Mon, 9 Nov 2026 14:26:44 +0000
Subject: [PATCH] Allow mixin selectors to resolve source methods for
 navigation

---
 .../reference/AbstractMethodReference.kt      | 42 ++++++++-----------
 .../mixin/reference/MixinSelectors.kt         | 13 +++++-
 2 files changed, 30 insertions(+), 25 deletions(-)

diff --git a/src/main/kotlin/platform/mixin/reference/AbstractMethodReference.kt b/src/main/kotlin/platform/mixin/reference/AbstractMethodReference.kt
index 94f6abb02..ed6f048df 100644
--- a/src/main/kotlin/platform/mixin/reference/AbstractMethodReference.kt
+++ b/src/main/kotlin/platform/mixin/reference/AbstractMethodReference.kt
@@ -27,7 +27,6 @@ import com.demonwav.mcdev.platform.mixin.util.MemberInfo
 import com.demonwav.mcdev.platform.mixin.util.bytecode
 import com.demonwav.mcdev.platform.mixin.util.findMethods
 import com.demonwav.mcdev.platform.mixin.util.findOrConstructSourceMethod
-import com.demonwav.mcdev.platform.mixin.util.findSourceElement
 import com.demonwav.mcdev.platform.mixin.util.findUpstreamMixin
 import com.demonwav.mcdev.platform.mixin.util.mixinTargets
 import com.demonwav.mcdev.util.MemberReference
@@ -116,6 +115,10 @@ abstract class AbstractMethodReference : PolyReferenceResolver(), MixinReference
     }
 
     fun resolve(context: PsiElement): Sequence<ClassAndMethodNode>? {
+        return resolveWithSelectors(context)?.map { it.second }
+    }
+
+    private fun resolveWithSelectors(context: PsiElement): Sequence<Pair<MixinSelector, ClassAndMethodNode>>? {
         val allowStatic = context.parentOfType<PsiMethod>()?.hasModifierProperty(PsiModifier.STATIC) ?: true
         val targets = getTargets(context) ?: return null
         val targetedMethods = when (context) {
@@ -124,30 +127,20 @@ abstract class AbstractMethodReference : PolyReferenceResolver(), MixinReference
         }
         val selectors = targetedMethods.mapNotNull { parseSelector(it, context) }
 
-        return resolve(targets, selectors, allowStatic)
-    }
-
-    private fun resolve(
-        targets: Collection<ClassNode>,
-        selectors: List<MixinSelector>,
-        allowStatic: Boolean,
-    ): Sequence<ClassAndMethodNode> {
-        return targets.asSequence()
-            .flatMap { target ->
-                selectors.asSequence().map { it.getCustomOwner(target) to it }
-            }
-            .groupBy({ it.first }, { it.second })
-            .asSequence()
-            .flatMap { (target, selectors) ->
-                target.findMethods(selectors, allowStatic)
-                    .map { ClassAndMethodNode(target, it) }
+        return targets.asSequence().flatMap { target ->
+            selectors.asSequence().flatMap { selector ->
+                val owner = selector.getCustomOwner(target)
+                owner.findMethods(listOf(selector), allowStatic)
+                    .map { selector to ClassAndMethodNode(owner, it) }
             }
+        }
     }
 
     fun resolveForNavigation(context: PsiElement): Array<PsiElement>? {
-        return resolve(context)?.mapNotNull {
-            it.method.findSourceElement(
-                it.clazz,
+        return resolveWithSelectors(context)?.mapNotNull { (selector, target) ->
+            selector.findSourceMethod(
+                target.clazz,
+                target.method,
                 context.project,
                 scope = context.resolveScope,
                 canDecompile = true,
@@ -156,9 +149,10 @@ abstract class AbstractMethodReference : PolyReferenceResolver(), MixinReference
     }
 
     override fun resolveReference(context: PsiElement): Array<ResolveResult> {
-        return resolve(context)?.mapNotNull {
-            it.method.findSourceElement(
-                it.clazz,
+        return resolveWithSelectors(context)?.mapNotNull { (selector, target) ->
+            selector.findSourceMethod(
+                target.clazz,
+                target.method,
                 context.project,
                 scope = context.resolveScope,
                 canDecompile = false,
diff --git a/src/main/kotlin/platform/mixin/reference/MixinSelectors.kt b/src/main/kotlin/platform/mixin/reference/MixinSelectors.kt
index 7c8d25d3c..7656e6bbc 100644
--- a/src/main/kotlin/platform/mixin/reference/MixinSelectors.kt
+++ b/src/main/kotlin/platform/mixin/reference/MixinSelectors.kt
@@ -24,6 +24,7 @@ import com.demonwav.mcdev.platform.mixin.util.MemberInfo
 import com.demonwav.mcdev.platform.mixin.util.MixinConstants
 import com.demonwav.mcdev.platform.mixin.util.MixinConstants.Annotations.DESC
 import com.demonwav.mcdev.platform.mixin.util.MixinConstants.Annotations.SLICE
+import com.demonwav.mcdev.platform.mixin.util.findSourceElement
 import com.demonwav.mcdev.platform.mixin.util.mixinTargets
 import com.demonwav.mcdev.util.MemberMatcher
 import com.demonwav.mcdev.util.Quantifier
@@ -42,11 +43,11 @@ import com.demonwav.mcdev.util.resolveTypeArray
 import com.intellij.openapi.editor.Editor
 import com.intellij.openapi.extensions.ExtensionPointName
 import com.intellij.openapi.project.Project
+import com.intellij.psi.PsiElement
 import com.intellij.psi.JavaPsiFacade
 import com.intellij.psi.PsiAnnotation
 import com.intellij.psi.PsiCallExpression
 import com.intellij.psi.PsiClass
-import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiLiteral
 import com.intellij.psi.PsiMethod
 import com.intellij.psi.PsiModifierList
@@ -61,6 +62,8 @@ import com.intellij.psi.util.InheritanceUtil
 import com.intellij.psi.util.PsiModificationTracker
 import com.intellij.psi.util.PsiTreeUtil
 import com.intellij.psi.util.parentOfType
+import org.objectweb.asm.tree.ClassNode
+import org.objectweb.asm.tree.MethodNode
 import java.util.Locale
 import java.util.regex.PatternSyntaxException
 import org.objectweb.asm.Type
@@ -98,6 +101,14 @@ interface MixinSelector : MemberMatcher {
     val next: MixinSelector? get() = null
 
     fun withQuantifier(quantifier: Quantifier): MixinSelector
+
+    fun findSourceMethod(
+        owner: ClassNode,
+        method: MethodNode,
+        project: Project,
+        scope: GlobalSearchScope,
+        canDecompile: Boolean,
+    ): PsiElement? = method.findSourceElement(owner, project, scope, canDecompile)
 }
 
 class MixinMemberParser : MixinSelectorParser {