Add support for adding missing and removing unused imports

This adds a post-process step which automatically adds unambiguous imports,
removes unused imports and sorts the import list (formatting matches standard
IntelliJ settings).
This will preserve line count across versions at all cost.

Java only for now because it's a lot more tricky with Kotlin and we don't yet
use Kotlin ourselves (and won't be preprocessing it in the future either).

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 {