User: joe
Date: 07 May 25 14:58
Revision: af7e57e5cacfb12afce5059ff777281162865205
Summary:
Implement infrastructure to desugar Java code before passing to the navigation visitor


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

Index: src/main/kotlin/platform/mixin/handlers/desugar/DesugarUtil.kt
===================================================================
--- src/main/kotlin/platform/mixin/handlers/desugar/DesugarUtil.kt	(revision af7e57e5cacfb12afce5059ff777281162865205)
+++ src/main/kotlin/platform/mixin/handlers/desugar/DesugarUtil.kt	(revision af7e57e5cacfb12afce5059ff777281162865205)
@@ -0,0 +1,90 @@
+/*
+ * 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.handlers.desugar
+
+import com.intellij.openapi.project.Project
+import com.intellij.openapi.util.Key
+import com.intellij.psi.JavaRecursiveElementWalkingVisitor
+import com.intellij.psi.PsiClass
+import com.intellij.psi.PsiElement
+import com.intellij.psi.PsiFile
+import com.intellij.psi.util.PsiTreeUtil
+
+object DesugarUtil {
+    private val ORIGINAL_ELEMENT_KEY = Key.create<PsiElement>("mcdev.desugar.originalElement")
+
+    private val DESUGARERS = arrayOf<Desugarer>(
+        FieldAssignmentDesugarer,
+    )
+
+    fun getOriginalElement(desugared: PsiElement): PsiElement? {
+        return desugared.getCopyableUserData(ORIGINAL_ELEMENT_KEY)
+    }
+
+    fun setOriginalElement(desugared: PsiElement, original: PsiElement) {
+        desugared.putCopyableUserData(ORIGINAL_ELEMENT_KEY, original)
+    }
+
+    fun desugar(project: Project, clazz: PsiClass): PsiClass {
+        val desugaredFile = clazz.containingFile.copy() as PsiFile
+        val desugaredClass = PsiTreeUtil.findSameElementInCopy(clazz, desugaredFile)
+        setOriginalRecursive(desugaredClass, clazz)
+        for (desugarer in DESUGARERS) {
+            desugarer.desugar(project, desugaredClass)
+        }
+        return desugaredClass
+    }
+
+    private fun setOriginalRecursive(desugared: PsiElement, original: PsiElement) {
+        val desugaredElements = mutableListOf<PsiElement>()
+        desugared.accept(object : JavaRecursiveElementWalkingVisitor() {
+            override fun visitElement(element: PsiElement) {
+                super.visitElement(element)
+                desugaredElements.add(element)
+            }
+        })
+
+        val originalElements = mutableListOf<PsiElement>()
+        original.accept(object : JavaRecursiveElementWalkingVisitor() {
+            override fun visitElement(element: PsiElement) {
+                super.visitElement(element)
+                originalElements.add(element)
+            }
+        })
+
+        for ((originalElement, desugaredElement) in originalElements.zip(desugaredElements)) {
+            setOriginalElement(desugaredElement, originalElement)
+        }
+    }
+
+    fun getOriginalToDesugaredMap(desugared: PsiElement): Map<PsiElement, List<PsiElement>> {
+        val result = mutableMapOf<PsiElement, MutableList<PsiElement>>()
+        desugared.accept(object : JavaRecursiveElementWalkingVisitor() {
+            override fun visitElement(element: PsiElement) {
+                super.visitElement(element)
+                getOriginalElement(element)?.let { original ->
+                    result.getOrPut(original) { mutableListOf() } += desugared
+                }
+            }
+        })
+        return result
+    }
+}
Index: src/main/kotlin/platform/mixin/handlers/desugar/Desugarer.kt
===================================================================
--- src/main/kotlin/platform/mixin/handlers/desugar/Desugarer.kt	(revision af7e57e5cacfb12afce5059ff777281162865205)
+++ src/main/kotlin/platform/mixin/handlers/desugar/Desugarer.kt	(revision af7e57e5cacfb12afce5059ff777281162865205)
@@ -0,0 +1,28 @@
+/*
+ * 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.handlers.desugar
+
+import com.intellij.openapi.project.Project
+import com.intellij.psi.PsiClass
+
+interface Desugarer {
+    fun desugar(project: Project, clazz: PsiClass)
+}
Index: src/main/kotlin/platform/mixin/handlers/desugar/FieldAssignmentDesugarer.kt
===================================================================
--- src/main/kotlin/platform/mixin/handlers/desugar/FieldAssignmentDesugarer.kt	(revision af7e57e5cacfb12afce5059ff777281162865205)
+++ src/main/kotlin/platform/mixin/handlers/desugar/FieldAssignmentDesugarer.kt	(revision af7e57e5cacfb12afce5059ff777281162865205)
@@ -0,0 +1,141 @@
+/*
+ * 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.handlers.desugar
+
+import com.demonwav.mcdev.util.constantValue
+import com.intellij.openapi.project.Project
+import com.intellij.psi.JavaPsiFacade
+import com.intellij.psi.PsiAssignmentExpression
+import com.intellij.psi.PsiClass
+import com.intellij.psi.PsiClassInitializer
+import com.intellij.psi.PsiExpressionStatement
+import com.intellij.psi.PsiField
+import com.intellij.psi.PsiMethod
+import com.intellij.psi.PsiModifier
+import com.intellij.psi.PsiStatement
+import com.intellij.psi.PsiType
+import com.intellij.util.JavaPsiConstructorUtil
+
+object FieldAssignmentDesugarer : Desugarer {
+    override fun desugar(project: Project, clazz: PsiClass) {
+        val staticStatementsToInsertPre = mutableListOf<PsiStatement>()
+        val staticStatementsToInsertPost = mutableListOf<PsiStatement>()
+        val nonStaticStatementsToInsert = mutableListOf<PsiStatement>()
+        var seenStaticInitializer = false
+
+        for (child in clazz.children) {
+            when (child) {
+                is PsiField -> {
+                    val initializer = child.initializer ?: continue
+
+                    if (child.hasModifierProperty(PsiModifier.STATIC)) {
+                        // check if the field is a ConstantValue with no initializer in the bytecode
+                        val constantValue = initializer.constantValue
+                        if (constantValue != null && constantValue !is PsiType) {
+                            continue
+                        }
+
+                        val fieldInitializer = JavaPsiFacade.getElementFactory(project)
+                            .createStatementFromText("${child.name} = null;", child) as PsiExpressionStatement
+                        (fieldInitializer.expression as PsiAssignmentExpression).rExpression!!.replace(initializer)
+                        DesugarUtil.setOriginalElement(fieldInitializer, child)
+
+                        if (seenStaticInitializer) {
+                            staticStatementsToInsertPost += fieldInitializer
+                        } else {
+                            staticStatementsToInsertPre += fieldInitializer
+                        }
+                    } else {
+                        val fieldInitializer = JavaPsiFacade.getElementFactory(project)
+                            .createStatementFromText("this.${child.name} = null;", child) as PsiExpressionStatement
+                        (fieldInitializer.expression as PsiAssignmentExpression).rExpression!!.replace(initializer)
+                        DesugarUtil.setOriginalElement(fieldInitializer, child)
+
+                        nonStaticStatementsToInsert += fieldInitializer
+                    }
+
+                    initializer.delete()
+                }
+                is PsiClassInitializer -> {
+                    if (child.hasModifierProperty(PsiModifier.STATIC)) {
+                        seenStaticInitializer = true
+                    } else {
+                        nonStaticStatementsToInsert += child.body.statements
+                        child.delete()
+                    }
+                }
+            }
+        }
+
+        if (staticStatementsToInsertPre.isNotEmpty() || staticStatementsToInsertPost.isNotEmpty()) {
+            val staticBlock = findStaticBlock(project, clazz)
+            for (statement in staticStatementsToInsertPre) {
+                staticBlock.body.addAfter(statement, staticBlock.body.lBrace)
+            }
+            for (statement in staticStatementsToInsertPost) {
+                staticBlock.body.addBefore(statement, staticBlock.body.rBrace)
+            }
+        }
+
+        if (nonStaticStatementsToInsert.isNotEmpty()) {
+            for (constructor in findConstructorsCallingSuper(project, clazz)) {
+                val body = constructor.body ?: continue
+                val delegateCtorCall = JavaPsiConstructorUtil.findThisOrSuperCallInConstructor(constructor)
+                for (statement in nonStaticStatementsToInsert.asReversed()) {
+                    body.addAfter(statement, delegateCtorCall?.parent ?: body.lBrace)
+                }
+            }
+        }
+    }
+
+    private fun findStaticBlock(project: Project, clazz: PsiClass): PsiClassInitializer {
+        for (initializer in clazz.initializers) {
+            if (initializer.hasModifierProperty(PsiModifier.STATIC)) {
+                return initializer
+            }
+        }
+
+        val initializer = JavaPsiFacade.getElementFactory(project)
+            .createClass("class _Dummy_ { static {} }")
+            .initializers
+            .first()
+        DesugarUtil.setOriginalElement(initializer, clazz)
+        return clazz.add(initializer) as PsiClassInitializer
+    }
+
+    private fun findConstructorsCallingSuper(project: Project, clazz: PsiClass): List<PsiMethod> {
+        val className = clazz.name ?: return emptyList()
+
+        val constructors = clazz.constructors.filter {
+            !JavaPsiConstructorUtil.isChainedConstructorCall(
+                JavaPsiConstructorUtil.findThisOrSuperCallInConstructor(it)
+            )
+        }
+
+        if (constructors.isNotEmpty()) {
+            return constructors
+        }
+
+        val constructor = JavaPsiFacade.getElementFactory(project).createConstructor(className)
+        DesugarUtil.setOriginalElement(constructor, clazz)
+        return listOf(clazz.add(constructor) as PsiMethod)
+    }
+}
Index: src/main/kotlin/platform/mixin/handlers/injectionPoint/AtResolver.kt
===================================================================
--- src/main/kotlin/platform/mixin/handlers/injectionPoint/AtResolver.kt	(revision d1e8b460e193d919889bf7ae16bee9bff5778097)
+++ src/main/kotlin/platform/mixin/handlers/injectionPoint/AtResolver.kt	(revision af7e57e5cacfb12afce5059ff777281162865205)
@@ -20,18 +20,22 @@
 
 package com.demonwav.mcdev.platform.mixin.handlers.injectionPoint
 
+import com.demonwav.mcdev.platform.mixin.handlers.desugar.DesugarUtil
 import com.demonwav.mcdev.platform.mixin.reference.MixinSelector
 import com.demonwav.mcdev.platform.mixin.reference.isMiscDynamicSelector
 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.findSourceClass
 import com.demonwav.mcdev.platform.mixin.util.findSourceElement
+import com.demonwav.mcdev.platform.mixin.util.isClinit
+import com.demonwav.mcdev.platform.mixin.util.memberReference
 import com.demonwav.mcdev.util.computeStringArray
 import com.demonwav.mcdev.util.constantStringValue
 import com.demonwav.mcdev.util.constantValue
 import com.demonwav.mcdev.util.equivalentTo
+import com.demonwav.mcdev.util.findMethods
 import com.demonwav.mcdev.util.fullQualifiedName
 import com.intellij.codeInsight.lookup.LookupElementBuilder
 import com.intellij.psi.JavaPsiFacade
@@ -43,7 +47,9 @@
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiEnumConstant
 import com.intellij.psi.PsiExpression
+import com.intellij.psi.PsiModifier
 import com.intellij.psi.PsiModifierList
+import com.intellij.psi.PsiParameterListOwner
 import com.intellij.psi.PsiQualifiedReference
 import com.intellij.psi.PsiReference
 import com.intellij.psi.PsiReferenceExpression
@@ -221,35 +227,56 @@
         val target = targetAttr?.let { parseMixinSelector(it) }
         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),
-        )
-        if (mainTargetElement == null && targetElements.isEmpty()) {
-            return emptyList()
-        }
+        val project = at.project
 
-        val targetPsiClass = (mainTargetElement ?: targetElements.first()).parentOfType<PsiClass>()
+        // Resolve the target source class
+        val targetPsiClass = getTargetClass(target)
+            .findSourceClass(project, GlobalSearchScope.allScope(project), canDecompile = true)
             ?: return emptyList()
         val targetPsiFile = targetPsiClass.containingFile ?: return emptyList()
 
+        // Desugar the target class
+        val desugaredTargetClass = DesugarUtil.desugar(project, targetPsiClass)
+
+        // Find the element in the desugared class, first by directly searching and then by searching in the original
+        // and reverse mapping it into the desugared class.
+        val desugaredTargetElement = when {
+            targetMethod.isClinit -> desugaredTargetClass.initializers.firstOrNull {
+                it.hasModifierProperty(PsiModifier.STATIC)
+            }
+            else -> desugaredTargetClass.findMethods(targetMethod.memberReference).firstOrNull()
+        } ?: run {
+            val originalTargetElement = targetMethod.findSourceElement(
+                getTargetClass(target),
+                project,
+                GlobalSearchScope.allScope(project),
+                canDecompile = true,
+            ) ?: return emptyList()
+            DesugarUtil.getOriginalToDesugaredMap(desugaredTargetClass)[originalTargetElement]
+                ?.firstOrNull { it is PsiParameterListOwner }
+                ?: return listOf(originalTargetElement)
+        }
+
+        // Find the source element in the desugared class
         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())
+        navigationVisitor.visitStart(desugaredTargetElement)
+        if (desugaredTargetElement is PsiParameterListOwner) {
+            desugaredTargetElement.acceptChildren(navigationVisitor)
+        } else {
+            desugaredTargetElement.accept(navigationVisitor)
+        }
+        navigationVisitor.visitEnd(desugaredTargetElement)
 
+        // Map the desugared results back into the original source class
+        val sourceResults = navigationVisitor.result.mapNotNull { desugaredResult ->
+            desugaredResult.parents(true).firstNotNullOfOrNull(DesugarUtil::getOriginalElement)
+        }
+
+        // Match the bytecode results to the source results
         return bytecodeResults.mapNotNull { bytecodeResult ->
             val matcher = bytecodeResult.sourceLocationInfo.createMatcher<PsiElement>(targetPsiFile)
-            navigationVisitor.result.forEach(matcher::accept)
+            sourceResults.forEach(matcher::accept)
             matcher.result
         }
     }
Index: src/main/kotlin/platform/mixin/util/AsmUtil.kt
===================================================================
--- src/main/kotlin/platform/mixin/util/AsmUtil.kt	(revision d1e8b460e193d919889bf7ae16bee9bff5778097)
+++ src/main/kotlin/platform/mixin/util/AsmUtil.kt	(revision af7e57e5cacfb12afce5059ff777281162865205)
@@ -71,6 +71,7 @@
 import com.intellij.psi.PsiModifierList
 import com.intellij.psi.PsiParameter
 import com.intellij.psi.PsiParameterList
+import com.intellij.psi.PsiParameterListOwner
 import com.intellij.psi.PsiType
 import com.intellij.psi.PsiTypes
 import com.intellij.psi.impl.compiled.ClsElementImpl
@@ -1036,8 +1037,7 @@
     }
 
     val body = when (sourceMethod) {
-        is PsiMethod -> sourceMethod.body
-        is PsiLambdaExpression -> sourceMethod.body
+        is PsiParameterListOwner -> sourceMethod.body
         else -> null
     }