User: joe
Date: 20 Feb 25 16:46
Revision: 92368e6d84b8a5cb3828293e6507e436c7e4a1a9
Summary:
Don't recurse into lambdas etc with new source navigation


TeamCity URL: https://ci.mcdev.io/viewModification.html?tab=vcsModificationFiles&modId=9836&personal=false

Index: src/main/kotlin/platform/mixin/handlers/injectionPoint/AtResolver.kt
===================================================================
--- src/main/kotlin/platform/mixin/handlers/injectionPoint/AtResolver.kt	(revision c061f2894b0ab81fb7b67667517e4c367ef08dd9)
+++ src/main/kotlin/platform/mixin/handlers/injectionPoint/AtResolver.kt	(revision 92368e6d84b8a5cb3828293e6507e436c7e4a1a9)
@@ -222,27 +222,30 @@
         val bytecodeResults = resolveInstructions()
 
         // Then attempt to find the corresponding source elements using the navigation visitor
-        val targetElements = targetMethod.findBodyElements(
+        val mainTargetElement = targetMethod.findSourceElement(
             getTargetClass(target),
             at.project,
             GlobalSearchScope.allScope(at.project),
-        ).ifEmpty {
-            return listOfNotNull(
-                targetMethod.findSourceElement(
+            canDecompile = true,
+        )
+        val targetElements = targetMethod.findBodyElements(
-                    getTargetClass(target),
-                    at.project,
-                    GlobalSearchScope.allScope(at.project),
+            getTargetClass(target),
+            at.project,
+            GlobalSearchScope.allScope(at.project),
-                    canDecompile = true,
-                )
+        )
-            )
+        if (mainTargetElement == null && targetElements.isEmpty()) {
+            return emptyList()
         }
 
-        val targetPsiClass = targetElements.first().parentOfType<PsiClass>() ?: return emptyList()
+        val targetPsiClass = (mainTargetElement ?: 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)
+        navigationVisitor.visitStart(mainTargetElement ?: targetElements.first())
         targetElements.forEach { it.accept(navigationVisitor) }
+        navigationVisitor.visitEnd(mainTargetElement ?: targetElements.last())
 
         return bytecodeResults.mapNotNull { bytecodeResult ->
             val matcher = bytecodeResult.sourceLocationInfo.createMatcher<PsiElement>(targetPsiFile)
Index: src/main/kotlin/platform/mixin/handlers/injectionPoint/InjectionPoint.kt
===================================================================
--- src/main/kotlin/platform/mixin/handlers/injectionPoint/InjectionPoint.kt	(revision c061f2894b0ab81fb7b67667517e4c367ef08dd9)
+++ src/main/kotlin/platform/mixin/handlers/injectionPoint/InjectionPoint.kt	(revision 92368e6d84b8a5cb3828293e6507e436c7e4a1a9)
@@ -302,7 +302,6 @@
 
 abstract class NavigationVisitor : JavaRecursiveElementVisitor() {
     val result = mutableListOf<PsiElement>()
-    private var hasVisitedAnything = false
 
     protected fun addResult(element: PsiElement) {
         result += element
@@ -317,57 +316,22 @@
     open fun visitEnd(executableElement: PsiElement) {
     }
 
-    override fun visitElement(element: PsiElement) {
-        hasVisitedAnything = true
-        super.visitElement(element)
-    }
-
     override fun visitMethod(method: PsiMethod) {
-        if (!hasVisitedAnything) {
-            visitStart(method)
-            super.visitMethod(method)
-            visitEnd(method)
+        // do not recurse into methods
-        }
+    }
-    }
 
     override fun visitAnonymousClass(aClass: PsiAnonymousClass) {
         // do not recurse into anonymous classes
-        if (!hasVisitedAnything) {
-            visitStart(aClass)
-            super.visitAnonymousClass(aClass)
-            visitEnd(aClass)
-        }
+    }
-    }
 
     override fun visitClass(aClass: PsiClass) {
         // do not recurse into inner classes
-        if (!hasVisitedAnything) {
-            visitStart(aClass)
-            super.visitClass(aClass)
-            visitEnd(aClass)
-        }
+    }
-    }
 
-    override fun visitMethodReferenceExpression(expression: PsiMethodReferenceExpression) {
-        val hadVisitedAnything = hasVisitedAnything
-        if (!hadVisitedAnything) {
-            visitStart(expression)
-        }
-        super.visitMethodReferenceExpression(expression)
-        if (!hadVisitedAnything) {
-            visitEnd(expression)
-        }
-    }
-
     override fun visitLambdaExpression(expression: PsiLambdaExpression) {
         // do not recurse into lambda expressions
-        if (!hasVisitedAnything) {
-            visitStart(expression)
-            super.visitLambdaExpression(expression)
-            visitEnd(expression)
-        }
-    }
+    }
+}
-}
 
 abstract class CollectVisitor<T : PsiElement>(protected val mode: Mode) {
     fun visit(methodNode: MethodNode) {