Expose class passed to AtResolver to enclosing selector

src/main/kotlin/platform/mixin/expression/gui/MEShowFlowAction.kt

@@ -93,8 +93,8 @@ class MEShowFlowAction : AnAction() {
 
         fun resolveMixinMethodString(): Sequence<Resolved> {
             val string = element.parentOfType<PsiLiteralExpression>() ?: return emptySequence()
-            return MethodReference.resolve(string)?.map { (clazz, method) ->
-                Resolved(clazz, method)
+            return MethodReference.resolve(string)?.map { member ->
+                Resolved(member.classAndMethod.clazz, member.classAndMethod.method)
             }.orEmpty()
         }
 

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

@@ -22,6 +22,7 @@ package com.demonwav.mcdev.platform.mixin.handlers
 
 import com.demonwav.mcdev.platform.mixin.inspection.injector.MethodSignature
 import com.demonwav.mcdev.platform.mixin.inspection.injector.ParameterGroup
+import com.demonwav.mcdev.platform.mixin.reference.MixinSelector
 import com.demonwav.mcdev.platform.mixin.util.LocalVariables
 import com.demonwav.mcdev.platform.mixin.util.callbackInfoReturnableType
 import com.demonwav.mcdev.platform.mixin.util.callbackInfoType
@@ -49,6 +50,7 @@ class InjectAnnotationHandler : InjectorAnnotationHandler() {
         annotation: PsiAnnotation,
         targetClass: ClassNode,
         targetMethod: MethodNode,
+        enclosingSelector: MixinSelector?,
     ): List<MethodSignature> {
         val returnType = targetMethod.getGenericReturnType(targetClass, annotation.project)
 
@@ -85,7 +87,7 @@ class InjectAnnotationHandler : InjectorAnnotationHandler() {
         if (localCapture != "NO_CAPTURE") {
             annotation.findModule()?.let { module ->
                 var commonLocalsPrefix: MutableList<LocalVariables.LocalVariable>? = null
-                val resolvedInsns = resolveInstructions(annotation, targetClass, targetMethod).ifEmpty { return@let }
+                val resolvedInsns = resolveInstructions(annotation, targetClass, targetMethod, enclosingSelector).ifEmpty { return@let }
                 for (insn in resolvedInsns) {
                     val locals = LocalVariables.getLocals(module, targetClass, targetMethod, insn.insn)
                         ?.filterNotNull()

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

@@ -26,9 +26,11 @@ 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.MixinSelector
 import com.demonwav.mcdev.platform.mixin.reference.isMiscDynamicSelector
 import com.demonwav.mcdev.platform.mixin.reference.parseMixinSelector
 import com.demonwav.mcdev.platform.mixin.util.ClassAndMethodNode
+import com.demonwav.mcdev.platform.mixin.util.ContextAwareMethodTargetMember
 import com.demonwav.mcdev.platform.mixin.util.MethodTargetMember
 import com.demonwav.mcdev.platform.mixin.util.MixinTargetMember
 import com.demonwav.mcdev.platform.mixin.util.getGenericParameterTypes
@@ -72,7 +74,11 @@ abstract class InjectorAnnotationHandler : MixinAnnotationHandler {
             val (clazz, methods) = pair
             methods.firstNotNullOfOrNull { method ->
                 if (selector.matchMethod(method, clazz)) {
-                    MethodTargetMember(clazz, method)
+                    ContextAwareMethodTargetMember(
+                        clazz,
+                        method,
+                        selector,
+                    )
                 } else {
                     null
                 }
@@ -88,8 +94,13 @@ abstract class InjectorAnnotationHandler : MixinAnnotationHandler {
         }
 
         return resolveTarget(annotation, targetClass).map { targetMember ->
-            val targetMethod = targetMember as? MethodTargetMember ?: return@map InsnResolutionInfo.Failure()
-            isUnresolved(annotation, targetClass, targetMethod.classAndMethod.method) ?: return@isUnresolved null
+            val targetMethod = (targetMember as? MethodTargetMember)?.classAndMethod ?: return@map InsnResolutionInfo.Failure()
+            isUnresolved(
+                annotation,
+                targetClass,
+                targetMethod.method,
+                (targetMember as? ContextAwareMethodTargetMember)?.selector
+            ) ?: return@isUnresolved null
         }.reduceOrNull(InsnResolutionInfo.Failure::combine) ?: InsnResolutionInfo.Failure()
     }
 
@@ -99,27 +110,34 @@ abstract class InjectorAnnotationHandler : MixinAnnotationHandler {
         annotation: PsiAnnotation,
         targetClass: ClassNode,
         targetMethod: MethodNode,
+        enclosingSelector: MixinSelector?
     ): InsnResolutionInfo.Failure? {
         return annotation.findAttributeValue(getAtKey(annotation))?.findAnnotations()
             .ifNullOrEmpty { return InsnResolutionInfo.Failure() }!!
-            .firstNotNullOfOrNull { AtResolver(it, targetClass, targetMethod).isUnresolved() }
+            .firstNotNullOfOrNull { AtResolver(it, targetClass, targetMethod, enclosingSelector).isUnresolved() }
     }
 
     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, targetMethod.classAndMethod.clazz, targetMethod.classAndMethod.method)
+            val targetMethod = (targetMember as? MethodTargetMember)?.classAndMethod ?: return@flatMap emptyList()
+            resolveForNavigation(
+                annotation,
+                targetMethod.clazz,
+                targetMethod.method,
+                (targetMember as? ContextAwareMethodTargetMember)?.selector
+            )
         }
     }
 
     protected open fun resolveForNavigation(
         annotation: PsiAnnotation,
         targetClass: ClassNode,
         targetMethod: MethodNode,
+        enclosingSelector: MixinSelector?
     ): List<PsiElement> {
         return annotation.findAttributeValue(getAtKey(annotation))?.findAnnotations()
             .ifNullOrEmpty { return emptyList() }!!
-            .flatMap { AtResolver(it, targetClass, targetMethod).resolveNavigationTargets() }
+            .flatMap { AtResolver(it, targetClass, targetMethod, enclosingSelector).resolveNavigationTargets() }
     }
 
     fun resolveInstructions(annotation: PsiAnnotation) = annotation.cached(PsiModificationTracker.MODIFICATION_COUNT) {
@@ -131,7 +149,12 @@ abstract class InjectorAnnotationHandler : MixinAnnotationHandler {
         return resolveTarget(annotation, targetClass)
             .flatMap { targetMember ->
                 val targetMethod = (targetMember as? MethodTargetMember)?.classAndMethod ?: return@flatMap emptyList()
-                resolveInstructions(annotation, targetMethod.clazz, targetMethod.method).map { result ->
+                resolveInstructions(
+                    annotation,
+                    targetMethod.clazz,
+                    targetMethod.method,
+                    (targetMember as? ContextAwareMethodTargetMember)?.selector
+                ).map { result ->
                     InsnResult(targetMethod, result)
                 }
             }
@@ -141,6 +164,7 @@ abstract class InjectorAnnotationHandler : MixinAnnotationHandler {
         annotation: PsiAnnotation,
         targetClass: ClassNode,
         targetMethod: MethodNode,
+        enclosingSelector: MixinSelector?,
         mode: CollectVisitor.Mode = CollectVisitor.Mode.RESOLUTION,
     ): List<CollectVisitor.Result<*>> {
         val cache = annotation.cached(PsiModificationTracker.MODIFICATION_COUNT) {
@@ -149,7 +173,7 @@ abstract class InjectorAnnotationHandler : MixinAnnotationHandler {
         return cache.computeIfAbsent(ClassAndMethodNode(targetClass, targetMethod) to mode) {
             annotation.findAttributeValue(getAtKey(annotation))?.findAnnotations()
                 .ifNullOrEmpty { return@computeIfAbsent emptyList() }!!
-                .flatMap { AtResolver(it, targetClass, targetMethod).resolveInstructions(mode) }
+                .flatMap { AtResolver(it, targetClass, targetMethod, enclosingSelector).resolveInstructions(mode) }
         }
     }
 
@@ -162,6 +186,7 @@ abstract class InjectorAnnotationHandler : MixinAnnotationHandler {
         annotation: PsiAnnotation,
         targetClass: ClassNode,
         targetMethod: MethodNode,
+        enclosingSelector: MixinSelector?
     ): List<MethodSignature>?
 
     open fun isInsnAllowed(insn: AbstractInsnNode, decorations: Map<String, Any?>): Boolean {
@@ -229,6 +254,7 @@ object DefaultInjectorAnnotationHandler : InjectorAnnotationHandler() {
         annotation: PsiAnnotation,
         targetClass: ClassNode,
         targetMethod: MethodNode,
+        enclosingSelector: MixinSelector?,
     ) = null
 
     override val isSoft = true

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

@@ -22,6 +22,7 @@ package com.demonwav.mcdev.platform.mixin.handlers
 
 import com.demonwav.mcdev.platform.mixin.inspection.injector.MethodSignature
 import com.demonwav.mcdev.platform.mixin.inspection.injector.ParameterGroup
+import com.demonwav.mcdev.platform.mixin.reference.MixinSelector
 import com.demonwav.mcdev.platform.mixin.util.fakeResolve
 import com.demonwav.mcdev.platform.mixin.util.getParameter
 import com.demonwav.mcdev.platform.mixin.util.toPsiType
@@ -48,12 +49,13 @@ class ModifyArgHandler : InjectorAnnotationHandler() {
         annotation: PsiAnnotation,
         targetClass: ClassNode,
         targetMethod: MethodNode,
+        enclosingSelector: MixinSelector?,
     ): List<MethodSignature>? {
         val index = annotation.findDeclaredAttributeValue("index")?.constantValue as? Int
         val validSingleArgTypes = mutableSetOf<String>()
         var mayHaveValidFullSignature = true
         var validFullSignature: String? = null
-        val insns = resolveInstructions(annotation, targetClass, targetMethod).ifEmpty { return emptyList() }
+        val insns = resolveInstructions(annotation, targetClass, targetMethod, enclosingSelector).ifEmpty { return emptyList() }
         for (insn in insns) {
             if (insn.insn !is MethodInsnNode) return null
 

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

@@ -22,6 +22,7 @@ package com.demonwav.mcdev.platform.mixin.handlers
 
 import com.demonwav.mcdev.platform.mixin.inspection.injector.MethodSignature
 import com.demonwav.mcdev.platform.mixin.inspection.injector.ParameterGroup
+import com.demonwav.mcdev.platform.mixin.reference.MixinSelector
 import com.demonwav.mcdev.platform.mixin.util.MixinConstants.Classes.ARGS
 import com.demonwav.mcdev.util.Parameter
 import com.intellij.psi.JavaPsiFacade
@@ -44,6 +45,7 @@ class ModifyArgsHandler : InjectorAnnotationHandler() {
         annotation: PsiAnnotation,
         targetClass: ClassNode,
         targetMethod: MethodNode,
+        enclosingSelector: MixinSelector?,
     ): List<MethodSignature> {
         val argsType = JavaPsiFacade.getElementFactory(annotation.project)
             .createTypeByFQClassName(ARGS, annotation.resolveScope)

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

@@ -24,6 +24,7 @@ import com.demonwav.mcdev.platform.mixin.handlers.injectionPoint.ConstantInjecti
 import com.demonwav.mcdev.platform.mixin.handlers.injectionPoint.InjectionPoint
 import com.demonwav.mcdev.platform.mixin.inspection.injector.MethodSignature
 import com.demonwav.mcdev.platform.mixin.inspection.injector.ParameterGroup
+import com.demonwav.mcdev.platform.mixin.reference.MixinSelector
 import com.demonwav.mcdev.util.findAnnotations
 import com.intellij.openapi.project.Project
 import com.intellij.psi.JavaPsiFacade
@@ -84,6 +85,7 @@ class ModifyConstantHandler : InjectorAnnotationHandler() {
         annotation: PsiAnnotation,
         targetClass: ClassNode,
         targetMethod: MethodNode,
+        enclosingSelector: MixinSelector?,
     ): List<MethodSignature> {
         val constantInfos = getConstantInfos(annotation)
         if (constantInfos == null) {

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

@@ -25,6 +25,7 @@ import com.demonwav.mcdev.platform.mixin.handlers.injectionPoint.CollectVisitor
 import com.demonwav.mcdev.platform.mixin.handlers.injectionPoint.InjectionPoint
 import com.demonwav.mcdev.platform.mixin.inspection.injector.MethodSignature
 import com.demonwav.mcdev.platform.mixin.inspection.injector.ParameterGroup
+import com.demonwav.mcdev.platform.mixin.reference.MixinSelector
 import com.demonwav.mcdev.platform.mixin.util.LocalInfo
 import com.demonwav.mcdev.platform.mixin.util.toPsiType
 import com.demonwav.mcdev.util.constantStringValue
@@ -42,14 +43,15 @@ class ModifyVariableHandler : InjectorAnnotationHandler() {
         annotation: PsiAnnotation,
         targetClass: ClassNode,
         targetMethod: MethodNode,
+        enclosingSelector: MixinSelector?,
     ): List<MethodSignature>? {
         val module = annotation.findModule() ?: return null
 
         val at = annotation.findAttributeValue("at") as? PsiAnnotation
         val atCode = at?.findAttributeValue("value")?.constantStringValue
         val isLoadStore = atCode != null && InjectionPoint.byAtCode(atCode) is AbstractLoadInjectionPoint
         val mode = if (isLoadStore) CollectVisitor.Mode.COMPLETION else CollectVisitor.Mode.RESOLUTION
-        val targets = resolveInstructions(annotation, targetClass, targetMethod, mode)
+        val targets = resolveInstructions(annotation, targetClass, targetMethod, enclosingSelector, mode)
 
         val targetParamsGroup = ParameterGroup(
             collectTargetMethodParameters(annotation.project, targetClass, targetMethod),

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

@@ -23,6 +23,7 @@ package com.demonwav.mcdev.platform.mixin.handlers
 import com.demonwav.mcdev.platform.mixin.handlers.injectionPoint.NewInsnInjectionPoint
 import com.demonwav.mcdev.platform.mixin.inspection.injector.MethodSignature
 import com.demonwav.mcdev.platform.mixin.inspection.injector.ParameterGroup
+import com.demonwav.mcdev.platform.mixin.reference.MixinSelector
 import com.demonwav.mcdev.platform.mixin.util.AsmDfaUtil
 import com.demonwav.mcdev.platform.mixin.util.FieldTargetMember
 import com.demonwav.mcdev.platform.mixin.util.MethodTargetMember
@@ -88,8 +89,9 @@ class RedirectInjectorHandler : InjectorAnnotationHandler() {
         annotation: PsiAnnotation,
         targetClass: ClassNode,
         targetMethod: MethodNode,
+        enclosingSelector: MixinSelector?,
     ): List<MethodSignature>? {
-        val insns = resolveInstructions(annotation, targetClass, targetMethod).ifEmpty { return emptyList() }
+        val insns = resolveInstructions(annotation, targetClass, targetMethod, enclosingSelector).ifEmpty { return emptyList() }
         return getRedirectType(insns[0].insn)?.expectedMethodSignature(
             annotation,
             targetClass,

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

@@ -91,6 +91,7 @@ class AtResolver(
     private val at: PsiAnnotation,
     private val targetClass: ClassNode,
     private val targetMethod: MethodNode,
+    private val enclosingSelector: MixinSelector?,
 ) {
     companion object {
         fun getInjectionPoint(at: PsiAnnotation): InjectionPoint<*>? {
@@ -186,6 +187,7 @@ class AtResolver(
             at,
             target,
             getTargetClass(target),
+            enclosingSelector,
             CollectVisitor.Mode.RESOLUTION,
         )
         if (collectVisitor == null) {
@@ -211,7 +213,13 @@ class AtResolver(
         val targetAttr = at.findAttributeValue("target")
         val target = targetAttr?.let { parseMixinSelector(it) }
 
-        val collectVisitor = injectionPoint.createCollectVisitor(at, target, getTargetClass(target), mode)
+        val collectVisitor = injectionPoint.createCollectVisitor(
+            at,
+            target,
+            getTargetClass(target),
+            enclosingSelector,
+            mode
+        )
             ?: return InsnResolutionInfo.Failure()
 
         return collectVisitor.visit(targetMethod)
@@ -288,7 +296,10 @@ class AtResolver(
         // Collect all possible targets
         fun <T : PsiElement> doCollectVariants(injectionPoint: InjectionPoint<T>): List<Any> {
             val visitor = injectionPoint.createCollectVisitor(
-                at, target, getTargetClass(target),
+                at,
+                target,
+                getTargetClass(target),
+                enclosingSelector,
                 CollectVisitor.Mode.COMPLETION
             )
                 ?: return emptyList()
@@ -305,7 +316,8 @@ class AtResolver(
     }
 
     private fun getTargetClass(selector: MixinSelector?): ClassNode {
-        return selector?.getCustomOwner(targetClass) ?: targetClass
+        val owner = selector?.getCustomOwner(targetClass) ?: targetClass
+        return enclosingSelector?.transformTargetClass(owner) ?: owner
     }
 }
 

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

@@ -127,12 +127,13 @@ abstract class InjectionPoint<T : PsiElement> {
         at: PsiAnnotation,
         target: MixinSelector?,
         targetClass: ClassNode,
+        enclosingSelector: MixinSelector?,
         mode: CollectVisitor.Mode,
     ): CollectVisitor<T>? {
         return doCreateCollectVisitor(at, target, targetClass, mode)?.also {
             val isInsideSlice = at.parentOfType<PsiAnnotation>()?.hasQualifiedName(SLICE) == true
             val defaultSpecifier = if (isInsideSlice) InjectionPointSpecifier.FIRST else InjectionPointSpecifier.ALL
-            addFilters(at, targetClass, it, defaultSpecifier, mode)
+            addFilters(at, targetClass, it, defaultSpecifier, enclosingSelector, mode)
         }
     }
 
@@ -141,20 +142,22 @@ abstract class InjectionPoint<T : PsiElement> {
         targetClass: ClassNode,
         collectVisitor: CollectVisitor<T>,
         defaultSpecifier: InjectionPointSpecifier,
+        enclosingSelector: MixinSelector?,
         mode: CollectVisitor.Mode,
     ) {
-        addStandardFilters(at, targetClass, collectVisitor, defaultSpecifier, mode)
+        addStandardFilters(at, targetClass, collectVisitor, defaultSpecifier, enclosingSelector, mode)
     }
 
     fun addStandardFilters(
         at: PsiAnnotation,
         targetClass: ClassNode,
         collectVisitor: CollectVisitor<T>,
         defaultSpecifier: InjectionPointSpecifier,
+        enclosingSelector: MixinSelector?,
         mode: CollectVisitor.Mode,
     ) {
         addShiftSupport(at, targetClass, collectVisitor)
-        addSliceFilter(at, targetClass, collectVisitor)
+        addSliceFilter(at, targetClass, enclosingSelector, collectVisitor)
 
         // Make sure the ordinal and specifier filters are last, so that the ordinal only increments once the other
         // filters have passed, and the specifier acts on the result of them.
@@ -170,7 +173,7 @@ abstract class InjectionPoint<T : PsiElement> {
         collectVisitor.shiftBy = AtResolver.getShift(at)
     }
 
-    protected open fun addSliceFilter(at: PsiAnnotation, targetClass: ClassNode, collectVisitor: CollectVisitor<T>) {
+    protected open fun addSliceFilter(at: PsiAnnotation, targetClass: ClassNode, enclosingSelector: MixinSelector?, collectVisitor: CollectVisitor<T>) {
         // resolve slice annotation, take into account slice id if present
         val sliceId = at.findDeclaredAttributeValue("slice")?.constantStringValue
         val parentAnnotation = at.parentOfType<PsiAnnotation>() ?: return
@@ -194,10 +197,10 @@ abstract class InjectionPoint<T : PsiElement> {
 
         fun resolveSliceIndex(
             sliceAt: PsiAnnotation?,
-            method: MethodNode,
+            method: MethodNode
         ): Int? {
             return sliceAt?.let {
-                AtResolver(sliceAt, targetClass, method).resolveInstructions()
+                AtResolver(sliceAt, targetClass, method, enclosingSelector).resolveInstructions()
                     .singleOrNull()
                     ?.let { method.instructions.indexOf(it.insn) }
             }