Don't recurse into lambdas etc with new source navigation

Author: Earthcomputer

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

@@ -222,27 +222,30 @@ class AtResolver(
         val bytecodeResults = resolveInstructions()
 
         // Then attempt to find the corresponding source elements using the navigation visitor
+        val mainTargetElement = targetMethod.findSourceElement(
+            getTargetClass(target),
+            at.project,
+            GlobalSearchScope.allScope(at.project),
+            canDecompile = true,
+        )
         val targetElements = targetMethod.findBodyElements(
             getTargetClass(target),
             at.project,
             GlobalSearchScope.allScope(at.project),
-        ).ifEmpty {
-            return listOfNotNull(
-                targetMethod.findSourceElement(
-                    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)

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

@@ -302,7 +302,6 @@ abstract class AbstractMethodInjectionPoint : QualifiedInjectionPoint<PsiMethod>
 
 abstract class NavigationVisitor : JavaRecursiveElementVisitor() {
     val result = mutableListOf<PsiElement>()
-    private var hasVisitedAnything = false
 
     protected fun addResult(element: PsiElement) {
         result += element
@@ -317,55 +316,20 @@ abstract class NavigationVisitor : JavaRecursiveElementVisitor() {
     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)
-        }
     }
 }