diff options
Diffstat (limited to 'src/main/kotlin/com/replaymod')
3 files changed, 276 insertions, 10 deletions
diff --git a/src/main/kotlin/com/replaymod/gradle/remap/AutoImports.kt b/src/main/kotlin/com/replaymod/gradle/remap/AutoImports.kt new file mode 100644 index 0000000..eb57486 --- /dev/null +++ b/src/main/kotlin/com/replaymod/gradle/remap/AutoImports.kt @@ -0,0 +1,141 @@ +package com.replaymod.gradle.remap + +import org.jetbrains.kotlin.cli.jvm.compiler.KotlinCoreEnvironment +import org.jetbrains.kotlin.com.intellij.psi.* +import org.jetbrains.kotlin.psi.psiUtil.endOffset +import org.jetbrains.kotlin.psi.psiUtil.startOffset + +internal class AutoImports(private val environment: KotlinCoreEnvironment) { + + private val shortClassNames = ShortNameIndex(environment) + + fun apply(originalFile: PsiFile, mappedFile: String, processedFile: String): String = + apply(originalFile, originalFile.text.lines(), mappedFile.lines(), processedFile.lines()) + + private fun apply( + originalFile: PsiFile, + originalLines: List<String>, + mappedLines: List<String>, + processedLines: List<String>, + ): String { + if (originalLines.size != mappedLines.size || originalLines.size != processedLines.size) { + return mappedLines.joinToString("\n") + } + + val inputLines = processedLines.mapIndexed { index, processedLine -> + if (originalLines[index] == processedLine) { + mappedLines[index] + } else { + processedLine + } + } + val inputText = inputLines.joinToString("\n") + + val psiFileFactory = PsiFileFactory.getInstance(environment.project) + val psiFile = + psiFileFactory.createFileFromText(originalFile.language, inputText) as? PsiJavaFile ?: return inputText + val pkg = psiFile.packageStatement?.packageReference?.resolve() as? PsiPackage + + val references = findOutgoingReferences(psiFile) + + val imports = psiFile.importList?.importStatements ?: emptyArray() + val onDemandImports = imports.filter { it.isOnDemand }.mapNotNull { it.qualifiedName }.map { "$it." }.toSet() + val existingImports = imports.filter { !it.isOnDemand }.mapNotNull { it.qualifiedName }.toSet() + val unusedImports = existingImports.filter { it.substringAfterLast(".") !in references }.toSet() + + val implicitReferenceSources = listOfNotNull( + psiFile.classes.flatMap { it.allInnerClasses.asIterable() }, + pkg?.classes?.asIterable(), + ) + val implicitReferences = implicitReferenceSources.flatten().mapNotNull { it.name }.toSet() + val importedReferences = existingImports.map { it.substringAfterLast(".") }.toSet() + val missingReferences = references.asSequence() - importedReferences - implicitReferences + val newImports = missingReferences.mapNotNull { shortClassNames[it].singleOrNull()?.qualifiedName } + .filter { ref -> onDemandImports.none { ref.startsWith(it) } } + .filter { !it.startsWith("java.lang.") } + + val finalImports = existingImports.toSet() - unusedImports.toSet() + newImports + onDemandImports.map { "$it*" } + + val textBuilder = StringBuilder(inputText) + + imports.map { it.textRange }.sortedByDescending { it.startOffset }.forEach { importRange -> + textBuilder.replace(importRange.startOffset, importRange.endOffset, "") + + val start = importRange.startOffset + val whiteSpaceRange = start - 1..start + if (whiteSpaceRange.first in textBuilder.indices && whiteSpaceRange.last in textBuilder.indices) { + val whiteSpaceReplacement = when (textBuilder.substring(whiteSpaceRange)) { + "\n\n" -> "\n" + "\n " -> "\n" + " \n" -> "\n" + " " -> " " + else -> null + } + if (whiteSpaceReplacement != null) { + textBuilder.replace(whiteSpaceRange.first, whiteSpaceRange.last + 1, whiteSpaceReplacement) + } + } + } + + val startOfImports = psiFile.importList?.takeIf { it.textLength > 0 }?.startOffset + val endOfPackage = psiFile.packageStatement?.endOffset ?: 0 + + val removedLineCount = inputLines.size - textBuilder.lineSequence().count() + textBuilder.insert(startOfImports ?: endOfPackage, "\n".repeat(removedLineCount)) + + var index = startOfImports ?: endOfPackage + + if (startOfImports == null) { + repeat(2) { + if (textBuilder[index + 1] == '\n' && textBuilder[index + 2] == '\n') { + index++ + } + } + } + + val javaImports = finalImports.filter { it.startsWith("java.") || it.startsWith("javax.") }.toSet() + val otherImports = finalImports - javaImports + val importGroups = listOf(otherImports, javaImports).filter { it.isNotEmpty() } + + for ((importGroupIndex, importGroup) in importGroups.withIndex()) { + val hasMoreGroups = importGroupIndex + 1 in importGroups.indices + + for (import in importGroup.sorted()) { + val hasPrecedingStatement = index > 0 && textBuilder[index - 1] != '\n' + val canAdvanceToNextLine = textBuilder[index + 1] == '\n' && textBuilder[index + 2] == '\n' + + val str = (if (hasPrecedingStatement) " " else "") + "import $import;" + textBuilder.insert(index, str) + index += str.length + if (canAdvanceToNextLine) 1 else 0 + } + + if (hasMoreGroups && textBuilder[index + 1] == '\n' && textBuilder[index + 2] == '\n') { + index++ + } + } + + return textBuilder.toString() + } + + private fun findOutgoingReferences(file: PsiJavaFile): Set<String> { + val references = mutableSetOf<String>() + + fun recordReference(reference: PsiJavaCodeReferenceElement) { + if (reference.isQualified) return + val name = reference.referenceName ?: return + if (!name.first().isUpperCase()) return + val resolved = reference.resolve() + if (resolved is PsiTypeParameter) return + if (resolved is PsiVariable) return + references.add(name) + } + + file.accept(object : JavaRecursiveElementVisitor() { + override fun visitReferenceElement(reference: PsiJavaCodeReferenceElement) { + recordReference(reference) + super.visitReferenceElement(reference) + } + }) + return references + } +}
\ No newline at end of file diff --git a/src/main/kotlin/com/replaymod/gradle/remap/ShortNameIndex.kt b/src/main/kotlin/com/replaymod/gradle/remap/ShortNameIndex.kt new file mode 100644 index 0000000..c78256a --- /dev/null +++ b/src/main/kotlin/com/replaymod/gradle/remap/ShortNameIndex.kt @@ -0,0 +1,100 @@ +package com.replaymod.gradle.remap + +import org.jetbrains.kotlin.cli.jvm.compiler.KotlinCoreEnvironment +import org.jetbrains.kotlin.cli.jvm.config.javaSourceRoots +import org.jetbrains.kotlin.cli.jvm.config.jvmClasspathRoots +import org.jetbrains.kotlin.com.intellij.lang.jvm.JvmModifier +import org.jetbrains.kotlin.com.intellij.openapi.vfs.StandardFileSystems +import org.jetbrains.kotlin.com.intellij.openapi.vfs.VirtualFile +import org.jetbrains.kotlin.com.intellij.openapi.vfs.VirtualFileManager +import org.jetbrains.kotlin.com.intellij.psi.PsiClass +import org.jetbrains.kotlin.com.intellij.psi.PsiJavaFile +import org.jetbrains.kotlin.com.intellij.psi.PsiManager +import org.jetbrains.kotlin.config.JVMConfigurationKeys + +class ShortNameIndex(private val environment: KotlinCoreEnvironment) { + private val psiManager = PsiManager.getInstance(environment.project) + + private val entries: Map<String, ShortNameEntry> = mutableMapOf<String, ShortNameEntry>().apply { + val localFileSystem = VirtualFileManager.getInstance().getFileSystem(StandardFileSystems.FILE_PROTOCOL) + val jarFileSystem = VirtualFileManager.getInstance().getFileSystem(StandardFileSystems.JAR_PROTOCOL) + val jrtFileSystem = VirtualFileManager.getInstance().getFileSystem(StandardFileSystems.JRT_PROTOCOL) + + val classpathRoots = environment.configuration.jvmClasspathRoots.mapNotNull { file -> + if (file.isFile) { + jarFileSystem.findFileByPath("${file.absolutePath}!/") + } else { + localFileSystem.findFileByPath(file.absolutePath) + } + } + + val jdkHome = environment.configuration[JVMConfigurationKeys.JDK_HOME] + val allModuleRoots = jrtFileSystem.findFileByPath("$jdkHome!/modules")?.children ?: emptyArray() + val javaModuleRoots = allModuleRoots.filter { it.name.startsWith("java.") } + + val sourcesRoots = environment.configuration.javaSourceRoots.mapNotNull { localFileSystem.findFileByPath(it) } + + fun index(file: VirtualFile, pkgPrefix: String) { + if (file.isDirectory) { + val pkg = "$pkgPrefix${file.name}." + file.children.forEach { index(it, pkg) } + } else if (file.extension == "class") { + val fileName = file.nameWithoutExtension + val shortName = if ('$' in fileName) { + val innerName = fileName.substringAfterLast('$') + if (!innerName.first().isJavaIdentifierStart()) { + return + } + innerName + } else { + fileName + } + getOrPut(shortName, ::ShortNameEntry).files.add(file) + } else if (file.extension == "java") { + val psi = psiManager.findFile(file) as? PsiJavaFile ?: return + psi.classes.flatMap { listOf(it) + it.allInnerClasses }.forEach { psiClass -> + getOrPut(psiClass.name ?: return@forEach, ::ShortNameEntry).files.add(file) + } + } + } + + (classpathRoots + javaModuleRoots + sourcesRoots).forEach { root -> + root.children.forEach { index(it, "") } + } + } + + operator fun get(shortName: String): Set<PsiClass> { + val entry = entries[shortName] ?: return emptySet() + return entry.resolve(psiManager, shortName) + } + + private class ShortNameEntry { + var files = mutableListOf<VirtualFile>() + private var classes: Set<PsiClass>? = null + + fun resolve(psiManager: PsiManager, shortName: String): Set<PsiClass> { + return classes ?: resolveClasses(psiManager, shortName) + } + + private fun resolveClasses(psiManager: PsiManager, shortName: String): Set<PsiClass> { + val result = files.flatMap { file -> + if (file.extension == "java" && file.nameWithoutExtension != shortName) { + val psi = psiManager.findFile(file) as? PsiJavaFile ?: return@flatMap emptyList() + psi.classes.flatMap { sequenceOf(it) + it.allInnerClasses.asIterable() } + .filter { it.qualifiedName?.endsWith(shortName) == true } + } else if ('$' in file.name) { + val className = file.nameWithoutExtension.replace('$', '.') + val outerName = className.substringBefore(".") + val outerFile = file.parent.findChild("$outerName.class") ?: return@flatMap emptyList() + val outerPsi = psiManager.findFile(outerFile) as? PsiJavaFile ?: return@flatMap emptyList() + outerPsi.classes.flatMap { it.allInnerClasses.asIterable() } + .filter { it.qualifiedName?.endsWith(className) == true } + } else { + (psiManager.findFile(file) as? PsiJavaFile)?.classes?.asIterable() ?: emptyList() + } + }.filter { it.hasModifier(JvmModifier.PUBLIC) }.toSet() + classes = result + return result + } + } +}
\ No newline at end of file diff --git a/src/main/kotlin/com/replaymod/gradle/remap/Transformer.kt b/src/main/kotlin/com/replaymod/gradle/remap/Transformer.kt index 67b065e..7fbec6f 100644 --- a/src/main/kotlin/com/replaymod/gradle/remap/Transformer.kt +++ b/src/main/kotlin/com/replaymod/gradle/remap/Transformer.kt @@ -17,7 +17,6 @@ import org.jetbrains.kotlin.com.intellij.mock.MockProject import org.jetbrains.kotlin.com.intellij.openapi.Disposable import org.jetbrains.kotlin.com.intellij.openapi.extensions.ExtensionPoint import org.jetbrains.kotlin.com.intellij.openapi.extensions.Extensions -import org.jetbrains.kotlin.com.intellij.openapi.project.Project import org.jetbrains.kotlin.com.intellij.openapi.util.Disposer import org.jetbrains.kotlin.com.intellij.openapi.vfs.StandardFileSystems import org.jetbrains.kotlin.com.intellij.openapi.vfs.VirtualFileManager @@ -34,6 +33,7 @@ import java.io.InputStreamReader import java.lang.Exception import java.nio.charset.StandardCharsets import java.nio.file.Files +import java.nio.file.Path import java.nio.file.StandardOpenOption import java.util.* import kotlin.system.exitProcess @@ -42,20 +42,27 @@ class Transformer(private val map: MappingSet) { var classpath: Array<String>? = null var remappedClasspath: Array<String>? = null var patternAnnotation: String? = null + var manageImports = false @Throws(IOException::class) fun remap(sources: Map<String, String>): Map<String, Pair<String, List<Pair<Int, String>>>> = remap(sources, emptyMap()) @Throws(IOException::class) - fun remap(sources: Map<String, String>, processedSource: Map<String, String>): Map<String, Pair<String, List<Pair<Int, String>>>> { + fun remap(sources: Map<String, String>, processedSources: Map<String, String>): Map<String, Pair<String, List<Pair<Int, String>>>> { val tmpDir = Files.createTempDirectory("remap") + val processedTmpDir = Files.createTempDirectory("remap-processed") val disposable = Disposer.newDisposable() try { for ((unitName, source) in sources) { val path = tmpDir.resolve(unitName) Files.createDirectories(path.parent) Files.write(path, source.toByteArray(StandardCharsets.UTF_8), StandardOpenOption.CREATE) + + val processedSource = processedSources[unitName] ?: source + val processedPath = processedTmpDir.resolve(unitName) + Files.createDirectories(processedPath.parent) + Files.write(processedPath, processedSource.toByteArray(), StandardOpenOption.CREATE) } val config = CompilerConfiguration() @@ -93,7 +100,9 @@ class Transformer(private val map: MappingSet) { analyze1620(environment, ktFiles) } - val remappedProject = remappedClasspath?.let { setupRemappedProject(disposable, it) } + val remappedEnv = remappedClasspath?.let { + setupRemappedProject(disposable, it, processedTmpDir) + } val patterns = patternAnnotation?.let { annotationFQN -> val patterns = PsiPatterns(annotationFQN) @@ -103,7 +112,7 @@ class Transformer(private val map: MappingSet) { try { val patternFile = vfs.findFileByIoFile(tmpDir.resolve(unitName).toFile())!! val patternPsiFile = psiManager.findFile(patternFile)!! - patterns.read(patternPsiFile, processedSource[unitName]!!) + patterns.read(patternPsiFile, processedSources[unitName]!!) } catch (e: Exception) { throw RuntimeException("Failed to read patterns from file \"$unitName\".", e) } @@ -111,29 +120,45 @@ class Transformer(private val map: MappingSet) { patterns } + val autoImports = if (manageImports && remappedEnv != null) { + AutoImports(remappedEnv) + } else { + null + } + val results = HashMap<String, Pair<String, List<Pair<Int, String>>>>() for (name in sources.keys) { val file = vfs.findFileByIoFile(tmpDir.resolve(name).toFile())!! val psiFile = psiManager.findFile(file)!! - val mapped = try { - PsiMapper(map, remappedProject, psiFile, analysis.bindingContext, patterns).remapFile() + var (text, errors) = try { + PsiMapper(map, remappedEnv?.project, psiFile, analysis.bindingContext, patterns).remapFile() } catch (e: Exception) { throw RuntimeException("Failed to map file \"$name\".", e) } - results[name] = mapped + + if (autoImports != null && "/* remap: no-manage-imports */" !in text) { + val processedText = processedSources[name] ?: text + text = autoImports.apply(psiFile, text, processedText) + } + + results[name] = text to errors } return results } finally { - Files.walk(tmpDir).map<File> { it.toFile() }.sorted(Comparator.reverseOrder()).forEach { it.delete() } + Files.walk(tmpDir).sorted(Comparator.reverseOrder()).forEach { Files.delete(it) } + Files.walk(processedTmpDir).sorted(Comparator.reverseOrder()).forEach { Files.delete(it) } Disposer.dispose(disposable) } } - private fun setupRemappedProject(disposable: Disposable, classpath: Array<String>): Project { + private fun setupRemappedProject(disposable: Disposable, classpath: Array<String>, sourceRoot: Path): KotlinCoreEnvironment { val config = CompilerConfiguration() config.put(CommonConfigurationKeys.MODULE_NAME, "main") config.addAll(CLIConfigurationKeys.CONTENT_ROOTS, classpath.map { JvmClasspathRoot(File(it)) }) + if (manageImports) { + config.add(CLIConfigurationKeys.CONTENT_ROOTS, JavaSourceRoot(sourceRoot.toFile(), "")) + } config.put(CLIConfigurationKeys.MESSAGE_COLLECTOR_KEY, PrintingMessageCollector(System.err, MessageRenderer.GRADLE_STYLE, true)) val environment = KotlinCoreEnvironment.createForProduction( @@ -146,7 +171,7 @@ class Transformer(private val map: MappingSet) { } catch (e: NoSuchMethodError) { analyze1620(environment, emptyList()) } - return environment.project + return environment } companion object { |