Add support for remapping Kotlin files

2 files changed, 122 insertions, 17 deletions

diff --git a/src/main/kotlin/com/replaymod/gradle/remap/PsiMapper.kt b/src/main/kotlin/com/replaymod/gradle/remap/PsiMapper.kt
index ca11e96..44efec6 100644
--- a/src/main/kotlin/com/replaymod/gradle/remap/PsiMapper.kt
+++ b/src/main/kotlin/com/replaymod/gradle/remap/PsiMapper.kt
@@ -4,9 +4,20 @@ import com.replaymod.gradle.remap.PsiUtils.getSignature
 import org.cadixdev.bombe.type.signature.MethodSignature
 import org.cadixdev.lorenz.MappingSet
 import org.cadixdev.lorenz.model.ClassMapping
+import org.jetbrains.kotlin.asJava.getRepresentativeLightMethod
+import org.jetbrains.kotlin.com.intellij.lang.ASTNode
 import org.jetbrains.kotlin.com.intellij.openapi.util.TextRange
 import org.jetbrains.kotlin.com.intellij.openapi.util.text.StringUtil
 import org.jetbrains.kotlin.com.intellij.psi.*
+import org.jetbrains.kotlin.descriptors.CallableMemberDescriptor
+import org.jetbrains.kotlin.descriptors.FunctionDescriptor
+import org.jetbrains.kotlin.js.resolve.diagnostics.findPsi
+import org.jetbrains.kotlin.lexer.KtTokens
+import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.resolve.BindingContext
+import org.jetbrains.kotlin.synthetic.SyntheticJavaPropertyDescriptor
+import org.jetbrains.kotlin.synthetic.SyntheticJavaPropertyDescriptor.Companion.propertyNameByGetMethodName
 import java.util.*
 
 internal class PsiMapper(private val map: MappingSet, private val file: PsiFile) {
@@ -25,11 +36,14 @@ internal class PsiMapper(private val map: MappingSet, private val file: PsiFile)
     }
 
     private fun replaceIdentifier(parent: PsiElement, with: String) {
-        for (child in parent.children) {
-            if (child is PsiIdentifier) {
+        var child = parent.firstChild
+        while (child != null) {
+            if (child is PsiIdentifier
+                    || child is ASTNode && child.elementType == KtTokens.IDENTIFIER) {
                 replace(child, with)
                 return
             }
+            child = child.nextSibling
         }
     }
 
@@ -74,9 +88,50 @@ internal class PsiMapper(private val map: MappingSet, private val file: PsiFile)
     }
 
     private fun map(expr: PsiElement, method: PsiMethod) {
-        if (method.isConstructor) return
+        if (method.isConstructor) {
+            if (expr is KtSimpleNameExpression) {
+                map(expr, method.containingClass ?: return)
+            }
+            return
+        }
+
+        val mapped = findMapping(method)
+        if (mapped != null && mapped != method.name) {
+            replaceIdentifier(expr, mapped)
+        }
+    }
+
+    private fun map(expr: PsiElement, method: KtNamedFunction) {
+        val psiMethod = method.getRepresentativeLightMethod()
+        val mapped = findMapping(psiMethod ?: return)
+        if (mapped != null && mapped != method.name) {
+            replaceIdentifier(expr, mapped)
+        }
+    }
+
+    private fun map(expr: PsiElement, property: SyntheticJavaPropertyDescriptor) {
+        val getter = property.getMethod.findPsi() as? PsiMethod ?: return
+        val mappedGetter = findMapping(getter) ?: return
+        if (mappedGetter != getter.name) {
+            val mapped = propertyNameByGetMethodName(Name.identifier(mappedGetter))!!.identifier
+            replaceIdentifier(expr, mapped)
+        }
+    }
+
+    // See caller for why this exists
+    private fun map(expr: PsiElement, method: FunctionDescriptor) {
+        for (overriddenDescriptor in method.overriddenDescriptors) {
+            val overriddenPsi = overriddenDescriptor.findPsi()
+            if (overriddenPsi != null) {
+                map(expr, overriddenPsi) // found a psi element, continue as usually
+            } else {
+                map(expr, overriddenDescriptor) // recursion
+            }
+        }
+    }
 
-        var declaringClass: PsiClass? = method.containingClass ?: return
+    private fun findMapping(method: PsiMethod): String? {
+        var declaringClass: PsiClass? = method.containingClass ?: return null
         val parentQueue = ArrayDeque<PsiClass>()
         parentQueue.offer(declaringClass)
         var mapping: ClassMapping<*, *>? = null
@@ -89,16 +144,13 @@ internal class PsiMapper(private val map: MappingSet, private val file: PsiFile)
             if (mapping != null) {
                 val mapped = mapping.findMethodMapping(getSignature(method))?.deobfuscatedName
                 if (mapped != null) {
-                    if (mapped != method.name) {
-                        replaceIdentifier(expr, mapped)
-                    }
-                    return
+                    return mapped
                 }
                 mapping = null
             }
             while (mapping == null) {
                 declaringClass = parentQueue.poll()
-                if (declaringClass == null) return
+                if (declaringClass == null) return null
 
                 val superClass = declaringClass.superClass
                 if (superClass != null) {
@@ -126,6 +178,15 @@ internal class PsiMapper(private val map: MappingSet, private val file: PsiFile)
             replace(expr, mapped)
             return
         }
+        val parent: PsiElement? = expr.parent
+        if ((parent is KtUserType || parent is KtQualifiedExpression) && parent.text == name) {
+            replace(parent, mapped)
+            return
+        }
+        // FIXME this incorrectly filters things like "Packet<?>" and doesn't filter same-name type aliases
+        // if (expr.text != name.substring(name.lastIndexOf('.') + 1)) {
+        //     return // type alias, will be remapped at its definition
+        // }
         replaceIdentifier(expr, mapped.substring(mapped.lastIndexOf('.') + 1))
     }
 
@@ -133,6 +194,7 @@ internal class PsiMapper(private val map: MappingSet, private val file: PsiFile)
         when (resolved) {
             is PsiField -> map(expr, resolved)
             is PsiMethod -> map(expr, resolved)
+            is KtNamedFunction -> map(expr, resolved.getRepresentativeLightMethod())
             is PsiClass, is PsiPackage -> map(expr, resolved as PsiQualifiedNamedElement)
         }
     }
@@ -313,7 +375,7 @@ internal class PsiMapper(private val map: MappingSet, private val file: PsiFile)
         })
     }
 
-    fun remapFile(): String? {
+    fun remapFile(bindingContext: BindingContext): String? {
         file.accept(object : JavaRecursiveElementVisitor() {
             override fun visitClass(psiClass: PsiClass) {
                 val annotation = psiClass.getAnnotation(CLASS_MIXIN) ?: return
@@ -359,6 +421,40 @@ internal class PsiMapper(private val map: MappingSet, private val file: PsiFile)
             }
         })
 
+        if (file is KtFile) {
+            file.accept(object : KtTreeVisitor<Void>() {
+                override fun visitNamedFunction(function: KtNamedFunction, data: Void?): Void? {
+                    if (valid(function)) {
+                        map(function, function)
+                    }
+                    return super.visitNamedFunction(function, data)
+                }
+
+                override fun visitReferenceExpression(expression: KtReferenceExpression, data: Void?): Void? {
+                    if (valid(expression)) {
+                        val target = bindingContext[BindingContext.REFERENCE_TARGET, expression]
+                        if (target is SyntheticJavaPropertyDescriptor) {
+                            map(expression, target)
+                        } else if (target != null && (target as? CallableMemberDescriptor)?.kind != CallableMemberDescriptor.Kind.SYNTHESIZED) {
+                            val targetPsi = target.findPsi()
+                            if (targetPsi != null) {
+                                map(expression, targetPsi)
+                            } else if (target is FunctionDescriptor) {
+                                // Appears to be the case if we're referencing an overwritten function in a previously
+                                // compiled kotlin file
+                                // E.g. A.f overwrites B.f overwrites C.f
+                                //      C is a Minecraft class, B is a previously compiled (and already remapped) kotlin
+                                //      class and we're currently in A.f trying to call `super.f()`.
+                                //      `target` is a DeserializedSimpleFunctionDescriptor with no linked PSI element.
+                                map(expression, target)
+                            }
+                        }
+                    }
+                    return super.visitReferenceExpression(expression, data)
+                }
+            }, null)
+        }
+
         return getResult(file.text)
     }
 
diff --git a/src/main/kotlin/com/replaymod/gradle/remap/Transformer.kt b/src/main/kotlin/com/replaymod/gradle/remap/Transformer.kt
index 53713fd..90296d3 100644
--- a/src/main/kotlin/com/replaymod/gradle/remap/Transformer.kt
+++ b/src/main/kotlin/com/replaymod/gradle/remap/Transformer.kt
@@ -26,10 +26,12 @@ import org.jetbrains.kotlin.com.intellij.psi.PsiManager
 import org.jetbrains.kotlin.com.intellij.psi.search.GlobalSearchScope
 import org.jetbrains.kotlin.config.CommonConfigurationKeys
 import org.jetbrains.kotlin.config.CompilerConfiguration
+import org.jetbrains.kotlin.psi.KtFile
 import java.io.BufferedReader
 import java.io.File
 import java.io.IOException
 import java.io.InputStreamReader
+import java.lang.Exception
 import java.nio.charset.StandardCharsets
 import java.nio.file.Files
 import java.nio.file.StandardOpenOption
@@ -54,6 +56,7 @@ class Transformer(private val map: MappingSet) {
             config.put(CommonConfigurationKeys.MODULE_NAME, "main")
             config.add<ContentRoot>(CLIConfigurationKeys.CONTENT_ROOTS, JavaSourceRoot(tmpDir.toFile(), ""))
             config.add<ContentRoot>(CLIConfigurationKeys.CONTENT_ROOTS, KotlinSourceRoot(tmpDir.toAbsolutePath().toString(), false))
+            config.addAll<ContentRoot>(CLIConfigurationKeys.CONTENT_ROOTS, classpath!!.map { JvmClasspathRoot(File(it)) })
             config.put<MessageCollector>(CLIConfigurationKeys.MESSAGE_COLLECTOR_KEY, PrintingMessageCollector(System.err, MessageRenderer.GRADLE_STYLE, true))
 
             val environment = KotlinCoreEnvironment.createForProduction(
@@ -67,24 +70,30 @@ class Transformer(private val map: MappingSet) {
             }
 
             val project = environment.project as MockProject
+            val psiManager = PsiManager.getInstance(project)
+            val vfs = VirtualFileManager.getInstance().getFileSystem(StandardFileSystems.FILE_PROTOCOL) as CoreLocalFileSystem
+            val virtualFiles = sources.mapValues { vfs.findFileByIoFile(tmpDir.resolve(it.key).toFile())!! }
+            val psiFiles = virtualFiles.mapValues { psiManager.findFile(it.value)!! }
+            val ktFiles = psiFiles.values.filterIsInstance<KtFile>()
 
-            environment.updateClasspath(classpath!!.map { JvmClasspathRoot(File(it)) })
-
-            TopDownAnalyzerFacadeForJVM.analyzeFilesWithJavaIntegration(
+            val analysis = TopDownAnalyzerFacadeForJVM.analyzeFilesWithJavaIntegration(
                     project,
-                    emptyList(),
+                    ktFiles,
                     NoScopeRecordCliBindingTrace(),
                     environment.configuration,
                     { scope: GlobalSearchScope -> environment.createPackagePartProvider(scope) }
             )
 
-            val vfs = VirtualFileManager.getInstance().getFileSystem(StandardFileSystems.FILE_PROTOCOL) as CoreLocalFileSystem
             val results = HashMap<String, String>()
             for (name in sources.keys) {
                 val file = vfs.findFileByIoFile(tmpDir.resolve(name).toFile())!!
-                val psiFile = PsiManager.getInstance(project).findFile(file)!!
+                val psiFile = psiManager.findFile(file)!!
 
-                val mapped = PsiMapper(map, psiFile).remapFile()
+                val mapped = try {
+                    PsiMapper(map, psiFile).remapFile(analysis.bindingContext)
+                } catch (e: Exception) {
+                    throw RuntimeException("Failed to map file \"$name\".", e)
+                }
                 if (mapped == null) {
                     fail = true
                     continue