From fb0e54ad548e63a8e847a074cbf8998683d6b595 Mon Sep 17 00:00:00 2001 From: Jonas Herzig Date: Sun, 15 Sep 2019 18:25:45 +0200 Subject: Add support for remapping Kotlin files --- .../kotlin/com/replaymod/gradle/remap/PsiMapper.kt | 116 +++++++++++++++++++-- .../com/replaymod/gradle/remap/Transformer.kt | 23 ++-- 2 files changed, 122 insertions(+), 17 deletions(-) (limited to 'src/main/kotlin/com/replaymod/gradle/remap') 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() 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() { + 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(CLIConfigurationKeys.CONTENT_ROOTS, JavaSourceRoot(tmpDir.toFile(), "")) config.add(CLIConfigurationKeys.CONTENT_ROOTS, KotlinSourceRoot(tmpDir.toAbsolutePath().toString(), false)) + config.addAll(CLIConfigurationKeys.CONTENT_ROOTS, classpath!!.map { JvmClasspathRoot(File(it)) }) config.put(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() - 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() 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 -- cgit