Allow dynamic selectors to choose the class to perform member lookups…

… in (#2293)

* * fix dynamic selector missing namespace when the selector declares it at registration
* allow dynamic selectors to change the target class that is used when looking up members

* move custom owner over to mixin selector

src/main/kotlin/platform/mixin/handlers/InjectorAnnotationHandler.kt

@@ -54,19 +54,25 @@ import org.objectweb.asm.tree.MethodNode
 
 abstract class InjectorAnnotationHandler : MixinAnnotationHandler {
     override fun resolveTarget(annotation: PsiAnnotation, targetClass: ClassNode): List<MixinTargetMember> {
-        val targetClassMethods = targetClass.methods ?: return emptyList()
-
         val methodAttr = annotation.findAttributeValue("method")
         val method = methodAttr?.computeStringArray() ?: emptyList()
         val desc = annotation.findAttributeValue("desc")?.findAnnotations() ?: emptyList()
         val selectors = method.mapNotNull { parseMixinSelector(it, methodAttr!!) } +
             desc.mapNotNull { DescSelectorParser.descSelectorFromAnnotation(it) }
 
-        return targetClassMethods.mapNotNull { targetMethod ->
-            if (selectors.any { it.matchMethod(targetMethod, targetClass) }) {
-                MethodTargetMember(targetClass, targetMethod)
-            } else {
-                null
+        val targetClassMethods = selectors.associateWith { selector ->
+            val actualTarget = selector.getCustomOwner(targetClass)
+            (actualTarget to actualTarget.methods)
+        }
+
+        return targetClassMethods.mapNotNull { (selector, pair) ->
+            val (clazz, methods) = pair
+            methods.firstNotNullOfOrNull { method ->
+                if (selector.matchMethod(method, clazz)) {
+                    MethodTargetMember(clazz, method)
+                } else {
+                    null
+                }
             }
         }
     }
@@ -99,7 +105,7 @@ abstract class InjectorAnnotationHandler : MixinAnnotationHandler {
     override fun resolveForNavigation(annotation: PsiAnnotation, targetClass: ClassNode): List<PsiElement> {
         return resolveTarget(annotation, targetClass).flatMap { targetMember ->
             val targetMethod = targetMember as? MethodTargetMember ?: return@flatMap emptyList()
-            resolveForNavigation(annotation, targetClass, targetMethod.classAndMethod.method)
+            resolveForNavigation(annotation, targetMethod.classAndMethod.clazz, targetMethod.classAndMethod.method)
         }
     }
 

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

@@ -20,6 +20,7 @@
 
 package com.demonwav.mcdev.platform.mixin.handlers.injectionPoint
 
+import com.demonwav.mcdev.platform.mixin.reference.MixinSelector
 import com.demonwav.mcdev.platform.mixin.reference.isMiscDynamicSelector
 import com.demonwav.mcdev.platform.mixin.reference.parseMixinSelector
 import com.demonwav.mcdev.platform.mixin.reference.target.TargetReference
@@ -152,7 +153,7 @@ class AtResolver(
         val collectVisitor = injectionPoint.createCollectVisitor(
             at,
             target,
-            targetClass,
+            getTargetClass(target),
             CollectVisitor.Mode.MATCH_FIRST,
         )
         if (collectVisitor == null) {
@@ -181,7 +182,7 @@ class AtResolver(
         val targetAttr = at.findAttributeValue("target")
         val target = targetAttr?.let { parseMixinSelector(it) }
 
-        val collectVisitor = injectionPoint.createCollectVisitor(at, target, targetClass, mode)
+        val collectVisitor = injectionPoint.createCollectVisitor(at, target, getTargetClass(target), mode)
             ?: return InsnResolutionInfo.Failure()
         collectVisitor.visit(targetMethod)
         val result = collectVisitor.result
@@ -201,7 +202,7 @@ class AtResolver(
 
         // Then attempt to find the corresponding source elements using the navigation visitor
         val targetElement = targetMethod.findSourceElement(
-            targetClass,
+            getTargetClass(target),
             at.project,
             GlobalSearchScope.allScope(at.project),
             canDecompile = true,
@@ -223,16 +224,23 @@ class AtResolver(
 
         // Collect all possible targets
         fun <T : PsiElement> doCollectVariants(injectionPoint: InjectionPoint<T>): List<Any> {
-            val visitor = injectionPoint.createCollectVisitor(at, target, targetClass, CollectVisitor.Mode.COMPLETION)
+            val visitor = injectionPoint.createCollectVisitor(
+                at, target, getTargetClass(target),
+                CollectVisitor.Mode.COMPLETION
+            )
                 ?: return emptyList()
             visitor.visit(targetMethod)
             return visitor.result
                 .mapNotNull { result ->
-                    injectionPoint.createLookup(targetClass, result)?.let { completionHandler(it) }
+                    injectionPoint.createLookup(getTargetClass(target), result)?.let { completionHandler(it) }
                 }
         }
         return doCollectVariants(injectionPoint)
     }
+
+    private fun getTargetClass(selector: MixinSelector?): ClassNode {
+        return selector?.getCustomOwner(targetClass) ?: targetClass
+    }
 }
 
 sealed class InsnResolutionInfo {

src/main/kotlin/platform/mixin/reference/AbstractMethodReference.kt

@@ -83,7 +83,9 @@ abstract class AbstractMethodReference : PolyReferenceResolver(), MixinReference
         val stringValue = context.constantStringValue ?: return false
         val targetMethodInfo = parseSelector(stringValue, context) ?: return false
         val targets = getTargets(context) ?: return false
-        return !targets.asSequence().flatMap { it.findMethods(targetMethodInfo) }.any()
+        return !targets.asSequence().flatMap {
+            targetMethodInfo.getCustomOwner(it).findMethods(targetMethodInfo)
+        }.any()
     }
 
     fun getReferenceIfAmbiguous(context: PsiElement): MemberReference? {
@@ -125,7 +127,10 @@ abstract class AbstractMethodReference : PolyReferenceResolver(), MixinReference
         selector: MixinSelector,
     ): Sequence<ClassAndMethodNode> {
         return targets.asSequence()
-            .flatMap { target -> target.findMethods(selector).map { ClassAndMethodNode(target, it) } }
+            .flatMap { target ->
+                val actualTarget = selector.getCustomOwner(target)
+                actualTarget.findMethods(selector).map { ClassAndMethodNode(actualTarget, it) }
+            }
     }
 
     fun resolveIfUnique(context: PsiElement): ClassAndMethodNode? {

src/main/kotlin/platform/mixin/reference/MixinSelectors.kt

@@ -52,6 +52,7 @@ import com.intellij.openapi.util.text.StringUtil
 import com.intellij.psi.CommonClassNames
 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.PsiField
@@ -61,6 +62,7 @@ import com.intellij.psi.PsiNameValuePair
 import com.intellij.psi.PsiTypes
 import com.intellij.psi.search.GlobalSearchScope
 import com.intellij.psi.search.searches.AnnotatedMembersSearch
+import com.intellij.psi.search.searches.MethodReferencesSearch
 import com.intellij.psi.util.InheritanceUtil
 import com.intellij.psi.util.PsiModificationTracker
 import com.intellij.psi.util.PsiTreeUtil
@@ -123,6 +125,10 @@ interface MixinSelector {
         return matchMethod(qualifier.name, method.name, method.desc)
     }
 
+    fun getCustomOwner(owner: ClassNode): ClassNode {
+        return owner
+    }
+
     /**
      * Implement this to return false for early-out optimizations, so you don't need to resolve the member in the
      * navigation visitor
@@ -417,13 +423,18 @@ private fun getAllDynamicSelectors(project: Project): Set<String> {
             val annotation = member.findAnnotation(MixinConstants.Classes.SELECTOR_ID) ?: return@flatMap emptySequence()
             val value = annotation.findAttributeValue("value")?.constantStringValue
                 ?: return@flatMap emptySequence()
-            val namespace = annotation.findAttributeValue("namespace")?.constantStringValue
+            var namespace = annotation.findAttributeValue("namespace")?.constantStringValue
             if (namespace.isNullOrEmpty()) {
                 val builtinPrefix = "org.spongepowered.asm.mixin.injection.selectors."
                 if (member.qualifiedName?.startsWith(builtinPrefix) == true) {
                     sequenceOf(value, "mixin:$value")
                 } else {
-                    sequenceOf(value)
+                    namespace = findNamespace(project, member)
+                    if (namespace != null) {
+                        sequenceOf("$namespace:$value")
+                    } else {
+                        sequenceOf(value)
+                    }
                 }
             } else {
                 sequenceOf("$namespace:$value")
@@ -432,6 +443,38 @@ private fun getAllDynamicSelectors(project: Project): Set<String> {
     }
 }
 
+/**
+ * Dynamic selectors don't have to declare their namespace in the annotation,
+ * so instead we look for the registration call and extract the namespace from there.
+ */
+private fun findNamespace(
+    project: Project,
+    member: PsiClass
+): String? {
+    val targetSelector = JavaPsiFacade.getInstance(project)
+        .findClass(MixinConstants.Classes.TARGET_SELECTOR, GlobalSearchScope.allScope(project))
+    val registerMethod = targetSelector?.findMethodsByName("register", false)?.firstOrNull() ?: return null
+
+    val query = MethodReferencesSearch.search(registerMethod)
+    val usages = query.findAll()
+    for (usage in usages) {
+        val element = usage.element
+        val callExpression = PsiTreeUtil.getParentOfType(element, PsiCallExpression::class.java) ?: continue
+        val args = callExpression.argumentList ?: continue
+        if (args.expressions.size != 2) continue
+
+        // is the registered selector the one we're checking?
+        val selectorName = args.expressions[0].text.removeSuffix(".class")
+        if (selectorName != member.name) continue
+
+        val namespaceArg = args.expressions[1].text.removeSurrounding("\"")
+        if (namespaceArg.isEmpty()) continue
+
+        return namespaceArg
+    }
+    return null
+}
+
 private val DYNAMIC_SELECTOR_PATTERN = "(?i)^@([a-z]+(:[a-z]+)?)(\\((.*)\\))?$".toRegex()
 
 abstract class DynamicSelectorParser(id: String, vararg aliases: String) : MixinSelectorParser {

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

@@ -38,6 +38,7 @@ object MixinConstants {
         const val CONSTANT_CONDITION = "org.spongepowered.asm.mixin.injection.Constant.Condition"
         const val INJECTION_POINT = "org.spongepowered.asm.mixin.injection.InjectionPoint"
         const val SELECTOR = "org.spongepowered.asm.mixin.injection.InjectionPoint.Selector"
+        const val TARGET_SELECTOR = "org.spongepowered.asm.mixin.injection.selectors.TargetSelector"
         const val MIXIN_AGENT = "org.spongepowered.tools.agent.MixinAgent"
         const val MIXIN_CONFIG = "org.spongepowered.asm.mixin.transformer.MixinConfig"
         const val MIXIN_PLUGIN = "org.spongepowered.asm.mixin.extensibility.IMixinConfigPlugin"