User: joe
Date: 20 Feb 25 04:12
Revision: 713e13b28cf2c36369b4ae084518af1092d99d0b
Summary:
Improve accuracy of source matching using line numbers


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

Index: src/main/kotlin/platform/mixin/handlers/injectionPoint/AtResolver.kt
===================================================================
--- src/main/kotlin/platform/mixin/handlers/injectionPoint/AtResolver.kt	(revision e7c276d995bdbf0ba977af6dfe5a9454f834e28e)
+++ src/main/kotlin/platform/mixin/handlers/injectionPoint/AtResolver.kt	(revision 713e13b28cf2c36369b4ae084518af1092d99d0b)
@@ -228,13 +228,16 @@
             canDecompile = true,
         ) ?: return emptyList()
         val targetPsiClass = targetElement.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)
 
         return bytecodeResults.mapNotNull { bytecodeResult ->
-            navigationVisitor.result.getOrNull(bytecodeResult.index)
+            val matcher = bytecodeResult.sourceLocationInfo.createMatcher<PsiElement>(targetPsiFile)
+            navigationVisitor.result.forEach(matcher::accept)
+            matcher.result
         }
     }
 
Index: src/main/kotlin/platform/mixin/handlers/injectionPoint/InjectionPoint.kt
===================================================================
--- src/main/kotlin/platform/mixin/handlers/injectionPoint/InjectionPoint.kt	(revision e7c276d995bdbf0ba977af6dfe5a9454f834e28e)
+++ src/main/kotlin/platform/mixin/handlers/injectionPoint/InjectionPoint.kt	(revision 713e13b28cf2c36369b4ae084518af1092d99d0b)
@@ -22,6 +22,7 @@
 
 import com.demonwav.mcdev.platform.mixin.reference.MixinSelector
 import com.demonwav.mcdev.platform.mixin.reference.toMixinString
+import com.demonwav.mcdev.platform.mixin.util.SourceCodeLocationInfo
 import com.demonwav.mcdev.platform.mixin.util.fakeResolve
 import com.demonwav.mcdev.platform.mixin.util.findOrConstructSourceMethod
 import com.demonwav.mcdev.util.constantStringValue
@@ -60,6 +61,7 @@
 import org.objectweb.asm.tree.AbstractInsnNode
 import org.objectweb.asm.tree.ClassNode
 import org.objectweb.asm.tree.InsnList
+import org.objectweb.asm.tree.LineNumberNode
 import org.objectweb.asm.tree.MethodInsnNode
 import org.objectweb.asm.tree.MethodNode
 
@@ -385,6 +387,7 @@
 
     private lateinit var method: MethodNode
     private var nextIndex = 0
+    private val nextIndexByLine = mutableMapOf<Int, Int>()
     val result = mutableListOf<Result<T>>()
     private val resultFilters = mutableListOf<Pair<String, CollectResultFilter<T>>>()
     var filterToBlame: String? = null
@@ -416,14 +419,18 @@
             }
         }
 
+        val index = nextIndex++
+        val lineNumber = getLineNumber(insn)
+        val indexInLineNumber = lineNumber?.let { nextIndexByLine.merge(it, 1, Int::plus)!! - 1 } ?: index
         val result = Result(
-            nextIndex++,
+            SourceCodeLocationInfo(index, lineNumber, indexInLineNumber),
             insn,
             shiftedInsn ?: return,
             element,
             qualifier,
             if (insn === shiftedInsn) decorations else emptyMap()
         )
+
         var isFiltered = false
         for ((name, filter) in resultFilters) {
             if (!filter(result, method)) {
@@ -442,6 +449,18 @@
         }
     }
 
+    private fun getLineNumber(insn: AbstractInsnNode): Int? {
+        var i: AbstractInsnNode? = insn
+        while (i != null) {
+            if (i is LineNumberNode) {
+                return i.line
+            }
+            i = i.previous
+        }
+
+        return null
+    }
+
     @Suppress("MemberVisibilityCanBePrivate")
     protected fun stopWalking() {
         throw StopWalkingException()
@@ -454,13 +473,15 @@
     }
 
     data class Result<T : PsiElement>(
-        val index: Int,
+        val sourceLocationInfo: SourceCodeLocationInfo,
         val originalInsn: AbstractInsnNode,
         val insn: AbstractInsnNode,
         val target: T,
         val qualifier: String? = null,
         val decorations: Map<String, Any?>
-    )
+    ) {
+        val index: Int get() = sourceLocationInfo.index
+    }
 
     enum class Mode { MATCH_ALL, MATCH_FIRST, COMPLETION }
 }
Index: src/main/kotlin/platform/mixin/util/AsmUtil.kt
===================================================================
--- src/main/kotlin/platform/mixin/util/AsmUtil.kt	(revision e7c276d995bdbf0ba977af6dfe5a9454f834e28e)
+++ src/main/kotlin/platform/mixin/util/AsmUtil.kt	(revision 713e13b28cf2c36369b4ae084518af1092d99d0b)
@@ -30,7 +30,6 @@
 import com.demonwav.mcdev.util.findModule
 import com.demonwav.mcdev.util.findQualifiedClass
 import com.demonwav.mcdev.util.fullQualifiedName
-import com.demonwav.mcdev.util.hasSyntheticMethod
 import com.demonwav.mcdev.util.isErasureEquivalentTo
 import com.demonwav.mcdev.util.lockedCached
 import com.demonwav.mcdev.util.loggerForTopLevel
@@ -96,6 +95,7 @@
 import org.objectweb.asm.tree.InsnList
 import org.objectweb.asm.tree.InsnNode
 import org.objectweb.asm.tree.InvokeDynamicInsnNode
+import org.objectweb.asm.tree.LineNumberNode
 import org.objectweb.asm.tree.MethodInsnNode
 import org.objectweb.asm.tree.MethodNode
 import org.objectweb.asm.tree.VarInsnNode
@@ -660,13 +660,18 @@
     return null
 }
 
-private fun findContainingMethod(clazz: ClassNode, lambdaMethod: MethodNode): Pair<MethodNode, Int>? {
+private fun findContainingMethod(clazz: ClassNode, lambdaMethod: MethodNode): Pair<MethodNode, SourceCodeLocationInfo>? {
     if (!lambdaMethod.hasAccess(Opcodes.ACC_SYNTHETIC)) {
         return null
     }
     clazz.methods?.forEach { method ->
         var lambdaCount = 0
+        var lineNumber: Int? = null
+        val lambdaCountPerLine = mutableMapOf<Int, Int>()
         method.instructions?.iterator()?.forEach nextInsn@{ insn ->
+            if (insn is LineNumberNode) {
+                lineNumber = insn.line
+            }
             if (insn !is InvokeDynamicInsnNode) return@nextInsn
             if (insn.bsm.owner != "java/lang/invoke/LambdaMetafactory") return@nextInsn
             val invokedMethod = when (insn.bsm.name) {
@@ -691,9 +696,13 @@
             }
 
             lambdaCount++
+            val lambdaCountThisLine =
+                lineNumber?.let { lambdaCountPerLine.merge(it, 1, Int::plus) } ?: lambdaCount
 
             if (invokedMethod.name == lambdaMethod.name && invokedMethod.desc == lambdaMethod.desc) {
-                return@findContainingMethod method to (lambdaCount - 1)
+                val locationInfo =
+                    SourceCodeLocationInfo(lambdaCount - 1, lineNumber, lambdaCountThisLine - 1)
+                return@findContainingMethod method to locationInfo
             }
         }
     }
@@ -704,12 +713,13 @@
 private fun findAssociatedLambda(psiClass: PsiClass, clazz: ClassNode, lambdaMethod: MethodNode): PsiElement? {
     return RecursionManager.doPreventingRecursion(lambdaMethod, false) {
         val pair = findContainingMethod(clazz, lambdaMethod) ?: return@doPreventingRecursion null
-        val (containingMethod, index) = pair
+        val (containingMethod, locationInfo) = pair
         val parent = findAssociatedLambda(psiClass, clazz, containingMethod)
             ?: psiClass.findMethods(containingMethod.memberReference).firstOrNull()
             ?: return@doPreventingRecursion null
-        var i = 0
-        var result: PsiElement? = null
+
+        val psiFile = psiClass.containingFile ?: return@doPreventingRecursion null
+        val matcher = locationInfo.createMatcher<PsiElement>(psiFile)
         parent.accept(
             object : JavaRecursiveElementWalkingVisitor() {
                 override fun visitAnonymousClass(aClass: PsiAnonymousClass) {
@@ -721,8 +731,7 @@
                 }
 
                 override fun visitLambdaExpression(expression: PsiLambdaExpression) {
-                    if (i++ == index) {
-                        result = expression
+                    if (matcher.accept(expression)) {
                         stopWalking()
                     }
                     // skip walking inside the lambda
@@ -732,16 +741,14 @@
                     // walk inside the reference first, visits the qualifier first (it's first in the bytecode)
                     super.visitMethodReferenceExpression(expression)
 
-                    if (expression.hasSyntheticMethod) {
-                        if (i++ == index) {
-                            result = expression
+                    if (matcher.accept(expression)) {
-                            stopWalking()
-                        }
-                    }
+                        stopWalking()
+                    }
+                }
-                }
             },
         )
-        result
+
+        matcher.result
     }
 }
 
Index: src/main/kotlin/platform/mixin/util/SourceCodeLocationInfo.kt
===================================================================
--- src/main/kotlin/platform/mixin/util/SourceCodeLocationInfo.kt	(revision 713e13b28cf2c36369b4ae084518af1092d99d0b)
+++ src/main/kotlin/platform/mixin/util/SourceCodeLocationInfo.kt	(revision 713e13b28cf2c36369b4ae084518af1092d99d0b)
@@ -0,0 +1,85 @@
+/*
+ * Minecraft Development for IntelliJ
+ *
+ * https://mcdev.io/
+ *
+ * Copyright (C) 2025 minecraft-dev
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Lesser General Public License as published
+ * by the Free Software Foundation, version 3.0 only.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public License
+ * along with this program.  If not, see <https://www.gnu.org/licenses/>.
+ */
+
+package com.demonwav.mcdev.platform.mixin.util
+
+import com.demonwav.mcdev.util.findDocument
+import com.demonwav.mcdev.util.lineNumber
+import com.demonwav.mcdev.util.remapLineNumber
+import com.intellij.psi.PsiElement
+import com.intellij.psi.PsiFile
+
+/**
+ * Info returned from the bytecode to help locate an element from the source code.
+ *
+ * For example, if searching for a lambda in the source code, this contains the index of the lambda in the method (i.e.
+ * how many lambdas were before this one), the starting line number of the lambda (or `null` if there were no line
+ * numbers), and the index of the lambda within this line number (i.e. how many lambdas there were before this one in
+ * the same line).
+ *
+ * The line number stored is unmapped, and may need remapping via [remapLineNumber].
+ * [createMatcher] does this internally.
+ */
+class SourceCodeLocationInfo(val index: Int, val lineNumber: Int?, val indexInLineNumber: Int) {
+    interface Matcher<T: PsiElement> {
+        fun accept(t: T): Boolean
+
+        val result: T?
+    }
+
+    fun <T: PsiElement> createMatcher(psiFile: PsiFile): Matcher<T> {
+        val lineNumber = this.lineNumber?.let(psiFile::remapLineNumber)
+        val document = psiFile.findDocument()
+
+        return object : Matcher<T> {
+            private var count = 0
+            private var currentLine: Int? = null
+            private var countThisLine = 0
+            private var myResult: T? = null
+
+            override fun accept(t: T): Boolean {
+                val line = document?.let(t::lineNumber)
+                if (line != null) {
+                    if (line != currentLine) {
+                        countThisLine = 0
+                        currentLine = line
+                    }
+
+                    countThisLine++
+                    if (line == lineNumber && countThisLine == indexInLineNumber + 1) {
+                        myResult = t
+                        return true
+                    }
+                }
+
+                if (count++ == index) {
+                    myResult = t
+                    if (lineNumber == null) {
+                        return true
+                    }
+                }
+
+                return false
+            }
+
+            override val result get() = myResult
+        }
+    }
+}
Index: src/main/kotlin/util/psi-utils.kt
===================================================================
--- src/main/kotlin/util/psi-utils.kt	(revision e7c276d995bdbf0ba977af6dfe5a9454f834e28e)
+++ src/main/kotlin/util/psi-utils.kt	(revision 713e13b28cf2c36369b4ae084518af1092d99d0b)
@@ -24,7 +24,10 @@
 import com.demonwav.mcdev.platform.mcp.McpModule
 import com.demonwav.mcdev.platform.mcp.McpModuleType
 import com.intellij.codeInsight.lookup.LookupElementBuilder
+import com.intellij.debugger.impl.DebuggerUtilsEx
+import com.intellij.ide.highlighter.JavaClassFileType
 import com.intellij.lang.injection.InjectedLanguageManager
+import com.intellij.openapi.editor.Document
 import com.intellij.openapi.module.Module
 import com.intellij.openapi.module.ModuleManager
 import com.intellij.openapi.module.ModuleUtilCore
@@ -41,6 +44,7 @@
 import com.intellij.psi.PsiAnnotation
 import com.intellij.psi.PsiClass
 import com.intellij.psi.PsiDirectory
+import com.intellij.psi.PsiDocumentManager
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiElementFactory
 import com.intellij.psi.PsiElementResolveResult
@@ -222,6 +226,41 @@
     return modifier in ACCESS_MODIFIERS
 }
 
+fun PsiElement.findDocument(containingFile: PsiFile = this.containingFile): Document? {
+    return containingFile.viewProvider.document ?: PsiDocumentManager.getInstance(project).getDocument(containingFile)
+}
+
+/**
+ * Remaps line numbers if the file is decompiled. Line numbers are 1-indexed
+ */
+fun PsiFile.remapLineNumber(lineNumber: Int): Int {
+    val originalFile = this.originalFile
+    if (originalFile.virtualFile?.fileType != JavaClassFileType.INSTANCE) {
+        // not decompiled
+        return lineNumber
+    }
+
+    val mappedLineNumber = DebuggerUtilsEx.bytecodeToSourceLine(originalFile, lineNumber - 1)
+    if (mappedLineNumber < 0) {
+        return lineNumber
+    }
+
+    return mappedLineNumber + 1
+}
+
+/**
+ * Returns the line number of the start of this `PsiElement`'s text range, with line numbers starting at 1
+ */
+fun PsiElement.lineNumber(): Int? = findDocument()?.let(this::lineNumber)
+
+fun PsiElement.lineNumber(document: Document): Int? {
+    val index = this.textRange.startOffset
+    if (index > document.textLength) {
+        return null
+    }
+    return document.getLineNumber(index) + 1
+}
+
 infix fun PsiElement.equivalentTo(other: PsiElement?): Boolean {
     return manager.areElementsEquivalent(this, other)
 }