Pass owner to navigator when it can be inferred from the target

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

@@ -29,6 +29,7 @@ import com.intellij.codeInsight.completion.JavaLookupElementBuilder
 import com.intellij.codeInsight.lookup.LookupElementBuilder
 import com.intellij.openapi.editor.Editor
 import com.intellij.openapi.project.Project
+import com.intellij.psi.JavaPsiFacade
 import com.intellij.psi.PsiAnnotation
 import com.intellij.psi.PsiArrayAccessExpression
 import com.intellij.psi.PsiClass
@@ -37,6 +38,7 @@ import com.intellij.psi.PsiLiteral
 import com.intellij.psi.PsiMethodReferenceExpression
 import com.intellij.psi.PsiModifier
 import com.intellij.psi.PsiReferenceExpression
+import com.intellij.psi.search.GlobalSearchScope
 import com.intellij.psi.util.PsiUtil
 import com.intellij.util.ArrayUtilRt
 import org.objectweb.asm.Opcodes
@@ -87,7 +89,13 @@ class FieldInjectionPoint : QualifiedInjectionPoint<PsiField>() {
             ?.takeIf { it in Const.VALID_OPCODES } ?: -1
         val args = AtResolver.getArgs(at)
         val arrayAccess = getArrayAccessType(args)
-        return target?.let { MyNavigationVisitor(targetClass, it, opcode, arrayAccess) }
+        val ownerOrTarget = target?.owner?.let { ownerName ->
+            JavaPsiFacade.getInstance(targetClass.project).findClass(
+                ownerName.replace('/', '.'),
+                GlobalSearchScope.allScope(targetClass.project)
+            )
+        } ?: targetClass
+        return target?.let { MyNavigationVisitor(ownerOrTarget, it, opcode, arrayAccess) }
     }
 
     override fun doCreateCollectVisitor(

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

@@ -28,6 +28,7 @@ import com.intellij.psi.PsiAnnotation
 import com.intellij.psi.PsiClass
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiExpression
+import com.intellij.psi.PsiModifierList
 import com.intellij.psi.PsiStatement
 import org.objectweb.asm.tree.ClassNode
 import org.objectweb.asm.tree.MethodNode
@@ -72,6 +73,9 @@ class HeadInjectionPoint : InjectionPoint<PsiElement>() {
     private class MyNavigationVisitor : NavigationVisitor() {
         private var firstStatement = true
 
+        override fun visitAnnotation(annotation: PsiAnnotation) {
+        }
+
         override fun visitStatement(statement: PsiStatement) {
             if (firstStatement) {
                 firstStatement = false

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

@@ -25,12 +25,14 @@ import com.demonwav.mcdev.util.MemberReference
 import com.intellij.openapi.editor.Editor
 import com.intellij.openapi.project.Project
 import com.intellij.psi.CommonClassNames
+import com.intellij.psi.JavaPsiFacade
 import com.intellij.psi.PsiAnnotation
 import com.intellij.psi.PsiClass
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiLiteral
 import com.intellij.psi.PsiMethod
 import com.intellij.psi.PsiMethodCallExpression
+import com.intellij.psi.search.GlobalSearchScope
 import com.intellij.util.ArrayUtilRt
 import org.objectweb.asm.Opcodes
 import org.objectweb.asm.tree.ClassNode
@@ -91,7 +93,13 @@ class InvokeAssignInjectionPoint : AbstractMethodInjectionPoint() {
         target: MixinSelector?,
         targetClass: PsiClass,
     ): NavigationVisitor? {
-        return target?.let { MyNavigationVisitor(targetClass, it) }
+        val ownerOrTarget = target?.owner?.let { ownerName ->
+            JavaPsiFacade.getInstance(targetClass.project).findClass(
+                ownerName.replace('/', '.'),
+                GlobalSearchScope.allScope(targetClass.project)
+            )
+        } ?: targetClass
+        return target?.let { MyNavigationVisitor(ownerOrTarget, it) }
     }
 
     override fun doCreateCollectVisitor(

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

@@ -35,6 +35,7 @@ import com.intellij.psi.PsiLiteral
 import com.intellij.psi.PsiMethod
 import com.intellij.psi.PsiMethodCallExpression
 import com.intellij.psi.PsiNewExpression
+import com.intellij.psi.search.GlobalSearchScope
 import org.objectweb.asm.tree.ClassNode
 import org.objectweb.asm.tree.MethodInsnNode
 import org.objectweb.asm.tree.MethodNode
@@ -54,7 +55,14 @@ class InvokeInjectionPoint : AbstractMethodInjectionPoint() {
         target: MixinSelector?,
         targetClass: PsiClass,
     ): NavigationVisitor? {
-        return target?.let { MyNavigationVisitor(targetClass, it) }
+        if (target == null) return null
+        val ownerOrTarget = target.owner?.let { ownerName ->
+            JavaPsiFacade.getInstance(targetClass.project).findClass(
+                ownerName.replace('/', '.'),
+                GlobalSearchScope.allScope(targetClass.project)
+            )
+        } ?: targetClass
+        return MyNavigationVisitor(ownerOrTarget, target)
     }
 
     override fun doCreateCollectVisitor(

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

@@ -451,7 +451,8 @@ fun ClassNode.findSourceClass(project: Project, scope: GlobalSearchScope, canDec
             }
         }
         if (canDecompile) {
-            ((stubFile as? PsiCompiledFile)?.decompiledPsiFile as? PsiJavaFile)?.classes?.firstOrNull()
+            val javaFile = (stubFile as? PsiCompiledFile)?.decompiledPsiFile as? PsiJavaFile ?: stubFile as? PsiJavaFile
+            javaFile?.classes?.firstOrNull()
         } else {
             stubClass
         }