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

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

@@ -25,6 +25,7 @@ import com.demonwav.mcdev.platform.mixin.handlers.injectionPoint.CollectVisitor
 import com.demonwav.mcdev.platform.mixin.handlers.injectionPoint.InsnResolutionInfo
 import com.demonwav.mcdev.platform.mixin.inspection.injector.MethodSignature
 import com.demonwav.mcdev.platform.mixin.reference.DescSelectorParser
+import com.demonwav.mcdev.platform.mixin.reference.DynamicMixinSelector
 import com.demonwav.mcdev.platform.mixin.reference.isMiscDynamicSelector
 import com.demonwav.mcdev.platform.mixin.reference.parseMixinSelector
 import com.demonwav.mcdev.platform.mixin.util.ClassAndMethodNode
@@ -54,19 +55,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 = DynamicMixinSelector.apply(selector, 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 +106,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,8 @@
 
 package com.demonwav.mcdev.platform.mixin.handlers.injectionPoint
 
+import com.demonwav.mcdev.platform.mixin.reference.DynamicMixinSelector
+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 +154,7 @@ class AtResolver(
  val collectVisitor = injectionPoint.createCollectVisitor(
  at,
  target,
- targetClass,
+ getTargetClass(target),
  CollectVisitor.Mode.MATCH_FIRST,
  )
  if (collectVisitor == null) {
@@ -181,7 +183,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 +203,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 +225,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 DynamicMixinSelector.apply(selector, 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 {
+ DynamicMixinSelector.apply(targetMethodInfo, 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 = DynamicMixinSelector.apply(selector, 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
@@ -387,6 +389,25 @@ private class MixinRegexSelector(
 
 // Dynamic selectors
 
+/**
+ * A dynamic selector can change which class to query for members if desired.
+ */
+interface DynamicMixinSelector : MixinSelector {
+ fun redirectOwner(owner: ClassNode): ClassNode {
+ return owner
+ }
+
+ companion object {
+ fun apply(selector: MixinSelector?, owner: ClassNode): ClassNode {
+ return if (selector is DynamicMixinSelector) {
+ selector.redirectOwner(owner)
+ } else {
+ owner
+ }
+ }
+ }
+}
+
 /**
  * Checks if the string uses a dynamic selector that exists in the project but has no special handling
  * in mcdev, used to suppress invalid selector errors.
@@ -417,13 +438,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 +458,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 {
@@ -446,13 +504,13 @@ abstract class DynamicSelectorParser(id: String, vararg aliases: String) : Mixin
  return parseDynamic(matchResult.groups[4]?.value ?: "", context)
  }
 
- abstract fun parseDynamic(args: String, context: PsiElement): MixinSelector?
+ abstract fun parseDynamic(args: String, context: PsiElement): DynamicMixinSelector?
 }
 
 // @Desc
 
 class DescSelectorParser : DynamicSelectorParser("Desc", "mixin:Desc") {
- override fun parseDynamic(args: String, context: PsiElement): MixinSelector? {
+ override fun parseDynamic(args: String, context: PsiElement): DynamicMixinSelector? {
  val descAnnotation = findDescAnnotation(args.lowercase(Locale.ENGLISH), context) ?: return null
  return descSelectorFromAnnotation(descAnnotation)
  }
@@ -593,7 +651,7 @@ data class DescSelector(
  val owners: Set<String>,
  val name: String,
  override val methodDescriptor: String,
-) : MixinSelector {
+) : DynamicMixinSelector {
  override fun matchField(owner: String, name: String, desc: String): Boolean {
  return this.owners.contains(owner) && this.name == name && this.fieldDescriptor.substringBefore("(") == desc
  }

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"