User: joe
Date: 11 May 25 14:53
Revision: 98d2b2f7e6d2e5261202c6b33c58229125147855
Summary:
Add tests for desugaring, fix a couple of bugs


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

Index: src/main/kotlin/platform/mixin/handlers/desugar/DesugarUtil.kt
===================================================================
--- src/main/kotlin/platform/mixin/handlers/desugar/DesugarUtil.kt	(revision 4d4f9682c28b89c9e514f7e9ec2bd2eac4bb0272)
+++ src/main/kotlin/platform/mixin/handlers/desugar/DesugarUtil.kt	(revision 98d2b2f7e6d2e5261202c6b33c58229125147855)
@@ -29,6 +29,7 @@
 import com.intellij.psi.PsiJavaFile
 import com.intellij.psi.util.PsiTreeUtil
 import com.intellij.psi.util.parents
+import org.jetbrains.annotations.VisibleForTesting
 
 object DesugarUtil {
     private val ORIGINAL_ELEMENT_KEY = Key.create<PsiElement>("mcdev.desugar.originalElement")
@@ -46,19 +47,34 @@
         desugared.putCopyableUserData(ORIGINAL_ELEMENT_KEY, original)
     }
 
+    fun getOriginalToDesugaredMap(desugared: PsiElement): Map<PsiElement, List<PsiElement>> {
+        val desugaredFile = desugared.containingFile ?: return emptyMap()
+        return desugaredFile.cached {
+            val result = mutableMapOf<PsiElement, MutableList<PsiElement>>()
+            PsiTreeUtil.processElements(desugaredFile) { desugaredElement ->
+                desugaredElement.getCopyableUserData(ORIGINAL_ELEMENT_KEY)?.let { original ->
+                    result.getOrPut(original) { mutableListOf() } += desugaredElement
+                }
+                true
+            }
+            result
+        }
+    }
+
     fun desugar(project: Project, clazz: PsiClass): PsiClass? {
-        return clazz.cached {
-            val desugaredFile = clazz.containingFile.copy() as? PsiJavaFile ?: return@cached null
-            var desugaredClass = PsiTreeUtil.findSameElementInCopy(clazz, desugaredFile)
-            setOriginalRecursive(desugaredClass, clazz)
+        val file = clazz.containingFile as? PsiJavaFile ?: return null
+        return file.cached {
+            val desugaredFile = file.copy() as PsiJavaFile
+            setOriginalRecursive(desugaredFile, file)
             for (desugarer in DESUGARERS) {
-                desugaredClass = desugarer.desugar(project, desugaredFile, desugaredClass)
+                desugarer.desugar(project, desugaredFile)
             }
-            desugaredClass
+            getOriginalToDesugaredMap(desugaredFile)[clazz]?.filterIsInstance<PsiClass>()?.firstOrNull()
         }
     }
 
-    private fun setOriginalRecursive(desugared: PsiElement, original: PsiElement) {
+    @VisibleForTesting
+    fun setOriginalRecursive(desugared: PsiElement, original: PsiElement) {
         val desugaredElements = mutableListOf<PsiElement>()
         desugared.accept(object : JavaRecursiveElementWalkingVisitor() {
             override fun visitElement(element: PsiElement) {
@@ -79,17 +95,4 @@
             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 4d4f9682c28b89c9e514f7e9ec2bd2eac4bb0272)
+++ src/main/kotlin/platform/mixin/handlers/desugar/Desugarer.kt	(revision 98d2b2f7e6d2e5261202c6b33c58229125147855)
@@ -26,7 +26,7 @@
 import com.intellij.psi.util.childrenOfType
 
 abstract class Desugarer {
-    abstract fun desugar(project: Project, file: PsiJavaFile, clazz: PsiClass): PsiClass
+    abstract fun desugar(project: Project, file: PsiJavaFile)
 
     companion object {
         @JvmStatic
Index: src/main/kotlin/platform/mixin/handlers/desugar/FieldAssignmentDesugarer.kt
===================================================================
--- src/main/kotlin/platform/mixin/handlers/desugar/FieldAssignmentDesugarer.kt	(revision 4d4f9682c28b89c9e514f7e9ec2bd2eac4bb0272)
+++ src/main/kotlin/platform/mixin/handlers/desugar/FieldAssignmentDesugarer.kt	(revision 98d2b2f7e6d2e5261202c6b33c58229125147855)
@@ -36,7 +36,7 @@
 import com.intellij.util.JavaPsiConstructorUtil
 
 object FieldAssignmentDesugarer : Desugarer() {
-    override fun desugar(project: Project, file: PsiJavaFile, clazz: PsiClass): PsiClass {
+    override fun desugar(project: Project, file: PsiJavaFile) {
         val staticStatementsToInsertPre = mutableListOf<PsiStatement>()
         val staticStatementsToInsertPost = mutableListOf<PsiStatement>()
         val nonStaticStatementsToInsert = mutableListOf<PsiStatement>()
@@ -108,8 +108,6 @@
                 }
             }
         }
-
-        return clazz
     }
 
     private fun findStaticBlock(project: Project, clazz: PsiClass): PsiClassInitializer {
Index: src/main/kotlin/platform/mixin/handlers/desugar/RemoveVarArgsDesugarer.kt
===================================================================
--- src/main/kotlin/platform/mixin/handlers/desugar/RemoveVarArgsDesugarer.kt	(revision 4d4f9682c28b89c9e514f7e9ec2bd2eac4bb0272)
+++ src/main/kotlin/platform/mixin/handlers/desugar/RemoveVarArgsDesugarer.kt	(revision 98d2b2f7e6d2e5261202c6b33c58229125147855)
@@ -26,7 +26,6 @@
 import com.intellij.psi.JavaPsiFacade
 import com.intellij.psi.PsiCall
 import com.intellij.psi.PsiCapturedWildcardType
-import com.intellij.psi.PsiClass
 import com.intellij.psi.PsiEllipsisType
 import com.intellij.psi.PsiJavaFile
 import com.intellij.psi.PsiMethod
@@ -39,7 +38,7 @@
 import com.siyeh.ig.psiutils.MethodCallUtils
 
 object RemoveVarArgsDesugarer : Desugarer() {
-    override fun desugar(project: Project, file: PsiJavaFile, clazz: PsiClass): PsiClass {
+    override fun desugar(project: Project, file: PsiJavaFile) {
         val varArgsStarts = mutableListOf<Pair<PsiCall, Int>>()
         PsiTreeUtil.processElements(file) { element ->
             if (element is PsiCall && MethodCallUtils.isVarArgCall(element)) {
@@ -75,8 +74,10 @@
                 newExpr to newExpr
             } else {
                 val erasure = TypeConversionUtil.erasure(type)
-                val castExpr = elementFactory.createExpressionFromText("(${type.canonicalText}) new ${erasure.canonicalText}[] {}", call)
-                    as PsiTypeCastExpression
+                val castExpr = elementFactory.createExpressionFromText(
+                    "(${type.createArrayType().canonicalText}) new ${erasure.canonicalText}[] {}",
+                    call
+                ) as PsiTypeCastExpression
                 castExpr to castExpr.operand as PsiNewExpression
             }
 
@@ -105,7 +106,5 @@
             val newTypeElement = ellipsisType.replace(elementFactory.createTypeElement(newType))
             DesugarUtil.setOriginalElement(newTypeElement, DesugarUtil.getOriginalElement(ellipsisType))
         }
-
-        return clazz
     }
 }
Index: src/test/kotlin/platform/mixin/desugar/AbstractDesugarTest.kt
===================================================================
--- src/test/kotlin/platform/mixin/desugar/AbstractDesugarTest.kt	(revision 98d2b2f7e6d2e5261202c6b33c58229125147855)
+++ src/test/kotlin/platform/mixin/desugar/AbstractDesugarTest.kt	(revision 98d2b2f7e6d2e5261202c6b33c58229125147855)
@@ -0,0 +1,99 @@
+/*
+ * 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.desugar
+
+import com.demonwav.mcdev.framework.BaseMinecraftTest
+import com.demonwav.mcdev.platform.mixin.handlers.desugar.DesugarUtil
+import com.demonwav.mcdev.platform.mixin.handlers.desugar.Desugarer
+import com.intellij.openapi.command.WriteCommandAction
+import com.intellij.psi.PsiClass
+import com.intellij.psi.PsiJavaFile
+import com.intellij.psi.codeStyle.CodeStyleManager
+import com.intellij.psi.codeStyle.JavaCodeStyleManager
+import com.intellij.psi.util.PsiTreeUtil
+import com.intellij.psi.util.childrenOfType
+import com.intellij.testFramework.IndexingTestUtil
+import org.intellij.lang.annotations.Language
+import org.junit.jupiter.api.Assertions.assertEquals
+import org.junit.jupiter.api.Assertions.assertInstanceOf
+import org.junit.jupiter.api.Assertions.assertTrue
+
+abstract class AbstractDesugarTest : BaseMinecraftTest() {
+    abstract val desugarer: Desugarer
+
+    protected fun doTest(@Language("JAVA") before: String, @Language("JAVA") after: String) {
+        WriteCommandAction.runWriteCommandAction(project) {
+            val codeStyleManager = CodeStyleManager.getInstance(project)
+            val javaCodeStyleManager = JavaCodeStyleManager.getInstance(project)
+
+            val expectedFile = fixture.addClass(after).containingFile
+            assertEquals(
+                expectedFile,
+                codeStyleManager.reformat(expectedFile),
+                "Reformatting changed the file!",
+            )
+            val expectedText = expectedFile.text
+            expectedFile.delete()
+
+            val testFile = assertInstanceOf(
+                PsiJavaFile::class.java,
+                fixture.configureByText("Test.java", before)
+            )
+            assertEquals(
+                testFile,
+                codeStyleManager.reformat(testFile),
+                "Reformatting changed the file!",
+            )
+
+            IndexingTestUtil.waitUntilIndexesAreReady(project)
+
+            val desugaredFile = testFile.copy() as PsiJavaFile
+            DesugarUtil.setOriginalRecursive(desugaredFile, testFile)
+            desugarer.desugar(project, desugaredFile)
+            assertEquals(
+                expectedText,
+                codeStyleManager.reformat(javaCodeStyleManager.shortenClassReferences(desugaredFile.copy())).text
+            )
+
+            PsiTreeUtil.processElements(desugaredFile) { desugaredElement ->
+                val originalElement = DesugarUtil.getOriginalElement(desugaredElement)
+                if (originalElement != null) {
+                    assertTrue(
+                        PsiTreeUtil.isAncestor(testFile, originalElement, false)
+                    ) {
+                        "The original element of $desugaredElement is not from the original file"
+                    }
+                }
+                true
+            }
+
+            val originalClasses = testFile.childrenOfType<PsiClass>()
+            val desugaredClassesSet = mutableSetOf<PsiClass>()
+            val originalToDesugaredMap = DesugarUtil.getOriginalToDesugaredMap(desugaredFile)
+            for (clazz in originalClasses) {
+                val desugaredClasses = originalToDesugaredMap[clazz]?.filterIsInstance<PsiClass>() ?: emptyList()
+                assertEquals(1, desugaredClasses.size) { "Unexpected number of desugared classes for ${clazz.name}" }
+                desugaredClassesSet += desugaredClasses.first()
+            }
+            assertEquals(originalClasses.size, desugaredClassesSet.size, "Unexpected number of desugared classes")
+        }
+    }
+}
Index: src/test/kotlin/platform/mixin/desugar/RemoveVarArgsDesugarTest.kt
===================================================================
--- src/test/kotlin/platform/mixin/desugar/RemoveVarArgsDesugarTest.kt	(revision 98d2b2f7e6d2e5261202c6b33c58229125147855)
+++ src/test/kotlin/platform/mixin/desugar/RemoveVarArgsDesugarTest.kt	(revision 98d2b2f7e6d2e5261202c6b33c58229125147855)
@@ -0,0 +1,136 @@
+/*
+ * 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.desugar
+
+import com.demonwav.mcdev.platform.mixin.handlers.desugar.RemoveVarArgsDesugarer
+import org.junit.jupiter.api.Test
+
+class RemoveVarArgsDesugarTest : AbstractDesugarTest() {
+    override val desugarer = RemoveVarArgsDesugarer
+
+    @Test
+    fun testSimple() {
+        doTest(
+            """
+                class Test {
+                    static void foo(String... args) {
+                    }
+                    
+                    static void test() {
+                        foo("a", "b", "c")
+                    }
+                }
+            """.trimIndent(),
+            """
+                class Test {
+                    static void foo(String[] args) {
+                    }
+                    
+                    static void test() {
+                        foo(new String[] {"a", "b", "c"})
+                    }
+                }
+            """.trimIndent()
+        )
+    }
+
+    @Test
+    fun testNoArgs() {
+        doTest(
+            """
+                class Test {
+                    static void foo(String... args) {
+                    }
+                    
+                    static void test() {
+                        foo();
+                    }
+                }
+            """.trimIndent(),
+            """
+                class Test {
+                    static void foo(String[] args) {                
+                    }
+                    
+                    static void test() {
+                        foo(new String[] {});
+                    }
+                }
+            """.trimIndent()
+        )
+    }
+
+    @Test
+    fun testGenericMethod() {
+        doTest(
+            """
+                class Test {
+                    @SafeVarargs
+                    static <T> void foo(T... args) {
+                    }
+                    
+                    static void test() {
+                        foo("a", "b", "c");
+                    }
+                }
+            """.trimIndent(),
+            """
+                class Test {
+                    @SafeVarargs
+                    static <T> void foo(T[] args) {
+                    }
+                    
+                    static void test() {
+                        foo(new String[] {"a", "b", "c"});
+                    }
+                }
+            """.trimIndent()
+        )
+    }
+
+    @Test
+    fun testGenericClass() {
+        doTest(
+            """
+                class Test<T> {
+                    @SafeVarargs
+                    static void foo(T... args) {
+                    }
+                    
+                    static void test() {
+                        foo("a", "b", "c");
+                    }
+                }
+            """.trimIndent(),
+            """
+                class Test<T> {
+                    @SafeVarargs
+                    static void foo(T[] args) {
+                    }
+                    
+                    static void test() {
+                        foo((T[]) new Object[] {"a", "b", "c"});
+                    }
+                }
+            """.trimIndent()
+        )
+    }
+}