diff --git a/src/main/kotlin/platform/mixin/handlers/InjectorAnnotationHandler.kt b/src/main/kotlin/platform/mixin/handlers/InjectorAnnotationHandler.kt index 5c1079602..39058b76c 100644 --- a/src/main/kotlin/platform/mixin/handlers/InjectorAnnotationHandler.kt +++ b/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 { - 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 { 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) } } diff --git a/src/main/kotlin/platform/mixin/handlers/injectionPoint/AtResolver.kt b/src/main/kotlin/platform/mixin/handlers/injectionPoint/AtResolver.kt index 15b2ab3e7..c1b1df6ac 100644 --- a/src/main/kotlin/platform/mixin/handlers/injectionPoint/AtResolver.kt +++ b/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 doCollectVariants(injectionPoint: InjectionPoint): List { - 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 { diff --git a/src/main/kotlin/platform/mixin/reference/AbstractMethodReference.kt b/src/main/kotlin/platform/mixin/reference/AbstractMethodReference.kt index 67b355f30..89aa5dbd5 100644 --- a/src/main/kotlin/platform/mixin/reference/AbstractMethodReference.kt +++ b/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 { 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? { diff --git a/src/main/kotlin/platform/mixin/reference/MixinSelectors.kt b/src/main/kotlin/platform/mixin/reference/MixinSelectors.kt index 59ab0d8a0..4e8d34f4f 100644 --- a/src/main/kotlin/platform/mixin/reference/MixinSelectors.kt +++ b/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 { 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 { } } +/** + * 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 { diff --git a/src/main/kotlin/platform/mixin/util/MixinConstants.kt b/src/main/kotlin/platform/mixin/util/MixinConstants.kt index b0e6093c6..7d74f6cdb 100644 --- a/src/main/kotlin/platform/mixin/util/MixinConstants.kt +++ b/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"