Fix source navigation to field initializers and other elements outside constructors and static blocks

Earthcomputer · Earthcomputer · commit c061f2894b0a · 2025-02-20T16:35:56.000Z
diff --git a/src/main/kotlin/platform/mixin/handlers/injectionPoint/AtResolver.kt b/src/main/kotlin/platform/mixin/handlers/injectionPoint/AtResolver.kt
@@ -26,6 +26,7 @@ import com.demonwav.mcdev.platform.mixin.reference.parseMixinSelector
 import com.demonwav.mcdev.platform.mixin.reference.target.TargetReference
 import com.demonwav.mcdev.platform.mixin.util.MixinConstants.Annotations.SLICE
 import com.demonwav.mcdev.platform.mixin.util.MixinConstants.Classes.SHIFT
+import com.demonwav.mcdev.platform.mixin.util.findBodyElements
 import com.demonwav.mcdev.platform.mixin.util.findSourceElement
 import com.demonwav.mcdev.util.computeStringArray
 import com.demonwav.mcdev.util.constantStringValue
@@ -221,18 +222,27 @@ class AtResolver(
         val bytecodeResults = resolveInstructions()
 
         // Then attempt to find the corresponding source elements using the navigation visitor
-        val targetElement = targetMethod.findSourceElement(
+        val targetElements = targetMethod.findBodyElements(
             getTargetClass(target),
             at.project,
             GlobalSearchScope.allScope(at.project),
-            canDecompile = true,
-        ) ?: return emptyList()
-        val targetPsiClass = targetElement.parentOfType<PsiClass>() ?: return emptyList()
+        ).ifEmpty {
+            return listOfNotNull(
+                targetMethod.findSourceElement(
+                    getTargetClass(target),
+                    at.project,
+                    GlobalSearchScope.allScope(at.project),
+                    canDecompile = true,
+                )
+            )
+        }
+
+        val targetPsiClass = targetElements.first().parentOfType<PsiClass>() ?: return emptyList()
         val targetPsiFile = targetPsiClass.containingFile ?: return emptyList()
 
         val navigationVisitor = injectionPoint.createNavigationVisitor(at, target, targetPsiClass) ?: return emptyList()
         navigationVisitor.configureBytecodeTarget(targetClass, targetMethod)
-        targetElement.accept(navigationVisitor)
+        targetElements.forEach { it.accept(navigationVisitor) }
 
         return bytecodeResults.mapNotNull { bytecodeResult ->
             val matcher = bytecodeResult.sourceLocationInfo.createMatcher<PsiElement>(targetPsiFile)
diff --git a/src/main/kotlin/platform/mixin/handlers/injectionPoint/CtorHeadInjectionPoint.kt b/src/main/kotlin/platform/mixin/handlers/injectionPoint/CtorHeadInjectionPoint.kt
@@ -23,7 +23,7 @@ package com.demonwav.mcdev.platform.mixin.handlers.injectionPoint
 import com.demonwav.mcdev.platform.mixin.inspection.injector.CtorHeadNoUnsafeInspection
 import com.demonwav.mcdev.platform.mixin.reference.MixinSelector
 import com.demonwav.mcdev.platform.mixin.util.findOrConstructSourceMethod
-import com.demonwav.mcdev.platform.mixin.util.findSuperConstructorCall
+import com.demonwav.mcdev.platform.mixin.util.findDelegateConstructorCall
 import com.demonwav.mcdev.platform.mixin.util.isConstructor
 import com.demonwav.mcdev.platform.mixin.util.isFabricMixin
 import com.demonwav.mcdev.util.createLiteralExpression
@@ -132,13 +132,13 @@ class CtorHeadInjectionPoint : InjectionPoint<PsiElement>() {
                 return
             }
 
-            val superCtorCall = methodNode.findSuperConstructorCall() ?: run {
+            val delegateCtorCall = methodNode.findDelegateConstructorCall() ?: run {
                 super.accept(methodNode)
                 return
             }
 
             if (enforce == EnforceMode.POST_DELEGATE) {
-                val insn = superCtorCall.next ?: return
+                val insn = delegateCtorCall.next ?: return
                 addResult(insn, methodNode.findOrConstructSourceMethod(clazz, project))
                 return
             }
@@ -149,11 +149,11 @@ class CtorHeadInjectionPoint : InjectionPoint<PsiElement>() {
             // doesn't want to change the implementation in case of breaking mixins that rely on this
             // behavior, so it is now effectively intended, so it's what we'll use here.
             val lastFieldStore = generateSequence(insns.last) { it.previous }
-                .takeWhile { it !== superCtorCall }
+                .takeWhile { it !== delegateCtorCall }
                 .firstOrNull { insn ->
                     insn.opcode == Opcodes.PUTFIELD &&
                         (insn as FieldInsnNode).owner == clazz.name
-                } ?: superCtorCall
+                } ?: delegateCtorCall
 
             val lastFieldStoreNext = lastFieldStore.next ?: return
             addResult(lastFieldStoreNext, methodNode.findOrConstructSourceMethod(clazz, project))
diff --git a/src/main/kotlin/platform/mixin/inspection/injector/CancellableBeforeSuperCallInspection.kt b/src/main/kotlin/platform/mixin/inspection/injector/CancellableBeforeSuperCallInspection.kt
@@ -26,7 +26,7 @@ import com.demonwav.mcdev.platform.mixin.inspection.MixinCancellableInspection
 import com.demonwav.mcdev.platform.mixin.inspection.MixinInspection
 import com.demonwav.mcdev.platform.mixin.util.MethodTargetMember
 import com.demonwav.mcdev.platform.mixin.util.MixinConstants
-import com.demonwav.mcdev.platform.mixin.util.findSuperConstructorCall
+import com.demonwav.mcdev.platform.mixin.util.findDelegateConstructorCall
 import com.demonwav.mcdev.platform.mixin.util.isConstructor
 import com.demonwav.mcdev.util.constantValue
 import com.intellij.codeInspection.ProblemsHolder
@@ -72,10 +72,10 @@ class CancellableBeforeSuperCallInspection : MixinInspection() {
                     continue
                 }
                 val methodInsns = target.classAndMethod.method.instructions ?: continue
-                val superCtorCall = target.classAndMethod.method.findSuperConstructorCall() ?: continue
+                val delegateCtorCall = target.classAndMethod.method.findDelegateConstructorCall() ?: continue
                 val instructions =
                     handler.resolveInstructions(annotation, target.classAndMethod.clazz, target.classAndMethod.method)
-                if (instructions.any { methodInsns.indexOf(it.insn) <= methodInsns.indexOf(superCtorCall) }) {
+                if (instructions.any { methodInsns.indexOf(it.insn) <= methodInsns.indexOf(delegateCtorCall) }) {
                     return true
                 }
             }
diff --git a/src/main/kotlin/platform/mixin/inspection/injector/InjectIntoConstructorInspection.kt b/src/main/kotlin/platform/mixin/inspection/injector/InjectIntoConstructorInspection.kt
@@ -27,7 +27,7 @@ import com.demonwav.mcdev.platform.mixin.inspection.MixinInspection
 import com.demonwav.mcdev.platform.mixin.inspection.fix.AnnotationAttributeFix
 import com.demonwav.mcdev.platform.mixin.util.MethodTargetMember
 import com.demonwav.mcdev.platform.mixin.util.MixinConstants.Annotations.INJECT
-import com.demonwav.mcdev.platform.mixin.util.findSuperConstructorCall
+import com.demonwav.mcdev.platform.mixin.util.findDelegateConstructorCall
 import com.demonwav.mcdev.platform.mixin.util.isConstructor
 import com.demonwav.mcdev.platform.mixin.util.isFabricMixin
 import com.demonwav.mcdev.util.constantValue
@@ -100,12 +100,12 @@ class InjectIntoConstructorInspection : MixinInspection() {
                             return
                         }
 
-                        val superCtorCall = targetMethod.findSuperConstructorCall()
-                        if (superCtorCall != null &&
+                        val delegateCtorCall = targetMethod.findDelegateConstructorCall()
+                        if (delegateCtorCall != null &&
                             instructions.any {
                                 val insnIndex = targetMethod.instructions.indexOf(it.insn)
-                                val superCtorIndex = targetMethod.instructions.indexOf(superCtorCall)
-                                insnIndex <= superCtorIndex
+                                val delegateCtorIndex = targetMethod.instructions.indexOf(delegateCtorCall)
+                                insnIndex <= delegateCtorIndex
                             }
                         ) {
                             holder.registerProblem(
diff --git a/src/main/kotlin/platform/mixin/inspection/injector/InvalidInjectorMethodSignatureInspection.kt b/src/main/kotlin/platform/mixin/inspection/injector/InvalidInjectorMethodSignatureInspection.kt
@@ -26,7 +26,7 @@ import com.demonwav.mcdev.platform.mixin.inspection.MixinInspection
 import com.demonwav.mcdev.platform.mixin.reference.MethodReference
 import com.demonwav.mcdev.platform.mixin.util.MixinConstants
 import com.demonwav.mcdev.platform.mixin.util.MixinConstants.Annotations.COERCE
-import com.demonwav.mcdev.platform.mixin.util.findSuperConstructorCall
+import com.demonwav.mcdev.platform.mixin.util.findDelegateConstructorCall
 import com.demonwav.mcdev.platform.mixin.util.hasAccess
 import com.demonwav.mcdev.platform.mixin.util.isAssignable
 import com.demonwav.mcdev.platform.mixin.util.isConstructor
@@ -108,15 +108,15 @@ class InvalidInjectorMethodSignatureInspection : MixinInspection() {
                         if (!shouldBeStatic && targetMethod.method.isConstructor) {
                             // before the superclass constructor call, everything must be static
                             val methodInsns = targetMethod.method.instructions
-                            val superCtorCall = targetMethod.method.findSuperConstructorCall()
-                            if (methodInsns != null && superCtorCall != null) {
+                            val delegateCtorCall = targetMethod.method.findDelegateConstructorCall()
+                            if (methodInsns != null && delegateCtorCall != null) {
                                 val insns = handler.resolveInstructions(
                                     annotation,
                                     targetMethod.clazz,
                                     targetMethod.method,
                                 )
                                 shouldBeStatic = insns.any {
-                                    methodInsns.indexOf(it.insn) <= methodInsns.indexOf(superCtorCall)
+                                    methodInsns.indexOf(it.insn) <= methodInsns.indexOf(delegateCtorCall)
                                 }
                             }
                         }
diff --git a/src/main/kotlin/platform/mixin/util/AsmUtil.kt b/src/main/kotlin/platform/mixin/util/AsmUtil.kt
@@ -56,12 +56,15 @@ import com.intellij.psi.PsiCompiledFile
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiElementFactory
 import com.intellij.psi.PsiEllipsisType
+import com.intellij.psi.PsiEnumConstant
 import com.intellij.psi.PsiField
 import com.intellij.psi.PsiFileFactory
 import com.intellij.psi.PsiJavaFile
+import com.intellij.psi.PsiKeyword
 import com.intellij.psi.PsiLambdaExpression
 import com.intellij.psi.PsiManager
 import com.intellij.psi.PsiMethod
+import com.intellij.psi.PsiMethodCallExpression
 import com.intellij.psi.PsiMethodReferenceExpression
 import com.intellij.psi.PsiModifier
 import com.intellij.psi.PsiModifierList
@@ -636,18 +639,18 @@ val MethodNode.isClinit
     get() = this.name == "<clinit>"
 
 /**
- * Finds the super() call in this method node, assuming it is a constructor
+ * Finds the `this()` or `super()` call in this method node, assuming it is a constructor
  */
-fun MethodNode.findSuperConstructorCall(): AbstractInsnNode? {
+fun MethodNode.findDelegateConstructorCall(): MethodInsnNode? {
     val insns = instructions ?: return null
     var superCtorCall = insns.first
     var newCount = 0
     while (superCtorCall != null) {
         if (superCtorCall.opcode == Opcodes.NEW) {
             newCount++
         } else if (superCtorCall.opcode == Opcodes.INVOKESPECIAL) {
-            val methodCall = superCtorCall as MethodInsnNode
-            if (methodCall.name == "<init>") {
+            superCtorCall as MethodInsnNode
+            if (superCtorCall.name == "<init>") {
                 if (newCount == 0) {
                     return superCtorCall
                 } else {
@@ -976,6 +979,67 @@ fun MethodNode.findSourceElement(
     return findAssociatedLambda(psiClass, clazz, this)
 }
 
+fun MethodNode.findBodyElements(clazz: ClassNode, project: Project, scope: GlobalSearchScope): List<PsiElement> {
+    if (isClinit) {
+        val psiClass = clazz.findSourceClass(project, scope, canDecompile = true) ?: return emptyList()
+        val result = mutableListOf<PsiElement>()
+        for (element in psiClass.children) {
+            when (element) {
+                is PsiEnumConstant -> element.argumentList?.expressions?.let { result += it }
+                is PsiField -> {
+                    if (element.hasModifierProperty(PsiModifier.STATIC)) {
+                        element.initializer?.let { result += it }
+                    }
+                }
+                is PsiClassInitializer -> {
+                    if (element.hasModifierProperty(PsiModifier.STATIC)) {
+                        result += element.body
+                    }
+                }
+            }
+        }
+        return result
+    }
+
+    val sourceMethod = findSourceElement(clazz, project, scope, canDecompile = true) ?: return emptyList()
+
+    if (isConstructor && findDelegateConstructorCall()?.owner != clazz.name && sourceMethod is PsiMethod) {
+        val result = mutableListOf<PsiElement>()
+        val body = sourceMethod.body
+        if (body != null) {
+            val children = body.children
+            val superCtorIndex = children.indexOfFirst {
+                it is PsiMethodCallExpression && it.methodExpression.text == PsiKeyword.SUPER
+            }
+            result += children.take(superCtorIndex + 1)
+            sourceMethod.containingClass?.children?.forEach { element ->
+                when (element) {
+                    is PsiField -> {
+                        if (!element.hasModifierProperty(PsiModifier.STATIC)) {
+                            element.initializer?.let { result += it }
+                        }
+                    }
+                    is PsiClassInitializer -> {
+                        if (!element.hasModifierProperty(PsiModifier.STATIC)) {
+                            result += element.body
+                        }
+                    }
+                }
+            }
+            result += children.drop(superCtorIndex + 1)
+            return result
+        }
+    }
+
+    val body = when (sourceMethod) {
+        is PsiMethod -> sourceMethod.body
+        is PsiLambdaExpression -> sourceMethod.body
+        else -> null
+    }
+
+    return listOfNotNull(body)
+}
+
 /**
  * Constructs a fake method node which could have been reached via this method instruction
  */
