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).

9 files changed, 812 insertions, 25 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 {
diff --git a/src/test/kotlin/com/replaymod/gradle/remap/imports/TestJavaAutoImports.kt b/src/test/kotlin/com/replaymod/gradle/remap/imports/TestJavaAutoImports.kt
new file mode 100644
index 0000000..f41767f
--- /dev/null
+++ b/src/test/kotlin/com/replaymod/gradle/remap/imports/TestJavaAutoImports.kt
@@ -0,0 +1,344 @@
+package com.replaymod.gradle.remap.imports
+
+import com.replaymod.gradle.remap.util.TestData
+import io.kotest.matchers.shouldBe
+import org.junit.jupiter.api.Test
+
+class TestJavaAutoImports {
+    @Test
+    fun `should remove unused imports`() {
+        TestData.remap("""
+            package test;
+            
+            import java.util.ArrayList;
+            
+            class Test {
+            }
+        """.trimIndent()) shouldBe """
+            package test;
+            
+            
+            
+            class Test {
+            }
+        """.trimIndent()
+    }
+
+    @Test
+    fun `should add unambiguous missing imports from JDK`() {
+        TestData.remap("""
+            package test;
+            
+            
+            
+            class Test extends ArrayList {
+            }
+        """.trimIndent()) shouldBe """
+            package test;
+            
+            import java.util.ArrayList;
+            
+            class Test extends ArrayList {
+            }
+        """.trimIndent()
+    }
+
+    @Test
+    fun `should add unambiguous missing imports from remapped classpath`() {
+        TestData.remap("""
+            package test;
+            
+            
+            
+            class Test extends B {
+            }
+        """.trimIndent()) shouldBe """
+            package test;
+            
+            import b.pkg.B;
+            
+            class Test extends B {
+            }
+        """.trimIndent()
+    }
+
+    @Test
+    fun `should not add unambiguous missing imports from original classpath`() {
+        TestData.remap("""
+            package test;
+            
+            
+            
+            class Test extends A {
+            }
+        """.trimIndent()) shouldBe """
+            package test;
+            
+            
+            
+            class Test extends A {
+            }
+        """.trimIndent()
+    }
+
+    @Test
+    fun `should add unambiguous missing imports from the same source set`() {
+        val original = """
+            package test;
+            
+            
+            
+            class Test extends Sample {
+            }
+        """.trimIndent()
+        TestData.transformer.remap(mapOf(
+            "test.java" to original,
+            "my/Sample.java" to "package my; public class Sample {}"
+        ))["test.java"]!!.first shouldBe """
+            package test;
+            
+            import my.Sample;
+            
+            class Test extends Sample {
+            }
+        """.trimIndent()
+    }
+
+    @Test
+    fun `should preserve existing ambiguous imports`() {
+        TestData.remap("""
+            package test;
+            
+            import java.awt.List;
+            
+            class Test implements List {
+            }
+        """.trimIndent()) shouldBe """
+            package test;
+            
+            import java.awt.List;
+            
+            class Test implements List {
+            }
+        """.trimIndent()
+    }
+
+    @Test
+    fun `should add new imports in empty lines`() {
+        TestData.remap("""
+            package test;
+            
+            
+            
+            
+            class Test extends ArrayList implements Closeable {
+            }
+        """.trimIndent()) shouldBe """
+            package test;
+            
+            import java.io.Closeable;
+            import java.util.ArrayList;
+            
+            class Test extends ArrayList implements Closeable {
+            }
+        """.trimIndent()
+    }
+
+    @Test
+    fun `should add new imports in place of removed imports`() {
+        TestData.remap("""
+            package test;
+            
+            import test.Unused1;
+            import test.Unused2;
+            
+            class Test extends ArrayList implements Closeable {
+            }
+        """.trimIndent()) shouldBe """
+            package test;
+            
+            import java.io.Closeable;
+            import java.util.ArrayList;
+            
+            class Test extends ArrayList implements Closeable {
+            }
+        """.trimIndent()
+    }
+
+    @Test
+    fun `preserves star imports`() {
+        TestData.remap("""
+            package test;
+            
+            import test.Unused1;
+            import java.util.*;
+            
+            class Test extends ArrayList implements Closeable {
+            }
+        """.trimIndent()) shouldBe """
+            package test;
+            
+            import java.io.Closeable;
+            import java.util.*;
+            
+            class Test extends ArrayList implements Closeable {
+            }
+        """.trimIndent()
+    }
+
+    @Test
+    fun `should not import classes from the current package`() {
+        TestData.remap("""
+            package b.pkg;
+            
+            class Test extends B implements BInterface {
+            }
+        """.trimIndent()) shouldBe """
+            package b.pkg;
+            
+            class Test extends B implements BInterface {
+            }
+        """.trimIndent()
+    }
+
+    @Test
+    fun `should not import self`() {
+        TestData.remap("test/Test.java", """
+            package test;
+            
+            public class Test {
+                Test inner;
+            }
+        """.trimIndent()) shouldBe """
+            package test;
+            
+            public class Test {
+                Test inner;
+            }
+        """.trimIndent()
+    }
+
+    @Test
+    fun `should not import java-lang package`() {
+        TestData.remap("test/Test.java", """
+            package test;
+            
+            public class Test {
+                Object inner;
+            }
+        """.trimIndent()) shouldBe """
+            package test;
+            
+            public class Test {
+                Object inner;
+            }
+        """.trimIndent()
+    }
+
+    @Test
+    fun `should not import own inner classes`() {
+        TestData.remap("""
+            package test;
+            
+            class Test {
+                TestInner inner;
+                public class TestInner {}
+            }
+        """.trimIndent()) shouldBe """
+            package test;
+            
+            class Test {
+                TestInner inner;
+                public class TestInner {}
+            }
+        """.trimIndent()
+    }
+
+    @Test
+    fun `should not import generic types`() {
+        TestData.remap("""
+            package test;
+            
+            class Test<B> {
+                B inner;
+            }
+        """.trimIndent()) shouldBe """
+            package test;
+            
+            class Test<B> {
+                B inner;
+            }
+        """.trimIndent()
+    }
+
+    @Test
+    fun `should not import variable references`() {
+        TestData.remap("""
+            package test;
+            
+            class Test {
+                Object B;
+                { B = null; }
+                { Object BParent; BParent = B; }
+            }
+        """.trimIndent()) shouldBe """
+            package test;
+            
+            class Test {
+                Object B;
+                { B = null; }
+                { Object BParent; BParent = B; }
+            }
+        """.trimIndent()
+    }
+
+    @Test
+    fun `should not import types that look like fields`() {
+        val content = """
+            package test;
+            
+            import a.pkg.A;
+            
+            class Test extends A {
+                { conflictingField.method(); }
+            }
+        """.trimIndent()
+        TestData.transformer.remap(mapOf(
+            "test/Test.java" to content,
+            "c/conflictingField.java" to "package c; public class conflictingField {}"
+        ))["test/Test.java"]!!.first shouldBe """
+            package test;
+            
+            import b.pkg.B;
+            
+            class Test extends B {
+                { conflictingField.method(); }
+            }
+        """.trimIndent()
+    }
+
+    @Test
+    fun `should not touch static imports (yet)`() {
+        TestData.remap("""
+            package test;
+            
+            import test.Unused1;
+            import test.Unused2;
+            
+            import static test.Unused.method;
+            
+            class Test extends ArrayList implements Closeable {
+            }
+        """.trimIndent()) shouldBe """
+            package test;
+            
+            import java.io.Closeable;
+            import java.util.ArrayList;
+            
+            import static test.Unused.method;
+            
+            class Test extends ArrayList implements Closeable {
+            }
+        """.trimIndent()
+    }
+}
\ No newline at end of file
diff --git a/src/test/kotlin/com/replaymod/gradle/remap/imports/TestJavaAutoImportsFormatting.kt b/src/test/kotlin/com/replaymod/gradle/remap/imports/TestJavaAutoImportsFormatting.kt
new file mode 100644
index 0000000..848eff8
--- /dev/null
+++ b/src/test/kotlin/com/replaymod/gradle/remap/imports/TestJavaAutoImportsFormatting.kt
@@ -0,0 +1,142 @@
+package com.replaymod.gradle.remap.imports
+
+import com.replaymod.gradle.remap.util.TestData
+import io.kotest.matchers.shouldBe
+import org.junit.jupiter.api.Test
+
+class TestJavaAutoImportsFormatting {
+    @Test
+    fun `should separate java(x) from other imports with an empty line if possible`() {
+        TestData.remap("""
+            package test;
+            
+            
+            
+            
+            
+            
+            class Test extends ArrayList implements Closeable, BInterface {
+            }
+        """.trimIndent()) shouldBe """
+            package test;
+            
+            import b.pkg.BInterface;
+            
+            import java.io.Closeable;
+            import java.util.ArrayList;
+            
+            class Test extends ArrayList implements Closeable, BInterface {
+            }
+        """.trimIndent()
+    }
+
+    @Test
+    fun `should put new imports in single line if necessary to preserve original line count`() {
+        TestData.remap("""
+            package test;
+            
+            import test.Unused1;
+            import test.Unused2;
+            
+            class Test extends ArrayList implements Closeable, BInterface {
+            }
+        """.trimIndent()) shouldBe """
+            package test;
+            
+            import b.pkg.BInterface;
+            import java.io.Closeable; import java.util.ArrayList;
+            
+            class Test extends ArrayList implements Closeable, BInterface {
+            }
+        """.trimIndent()
+    }
+
+    @Test
+    fun `should always leave line after imports`() {
+        TestData.remap("""
+            package test;
+            
+            
+            class Test extends ArrayList implements Closeable {
+            }
+        """.trimIndent()) shouldBe """
+            package test;
+            import java.io.Closeable; import java.util.ArrayList;
+            
+            class Test extends ArrayList implements Closeable {
+            }
+        """.trimIndent()
+    }
+
+    @Test
+    fun `should put imports in same line as package if required`() {
+        TestData.remap("""
+            package test;
+            
+            class Test extends ArrayList implements Closeable {
+            }
+        """.trimIndent()) shouldBe """
+            package test; import java.io.Closeable; import java.util.ArrayList;
+            
+            class Test extends ArrayList implements Closeable {
+            }
+        """.trimIndent()
+    }
+
+    @Test
+    fun `should remove unused imports from shared lines`() {
+        TestData.remap("""
+            package test;
+            
+            import java.io.Closeable; import java.util.ArrayList;
+            
+            class Test {
+            }
+        """.trimIndent()) shouldBe """
+            package test;
+            
+            
+            
+            class Test {
+            }
+        """.trimIndent()
+    }
+
+    @Test
+    fun `should remove unused imports from end of shared lines`() {
+        TestData.remap("""
+            package test;
+            
+            import java.io.Closeable; import java.util.ArrayList;
+            
+            class Test implements Closeable {
+            }
+        """.trimIndent()) shouldBe """
+            package test;
+            
+            import java.io.Closeable;
+            
+            class Test implements Closeable {
+            }
+        """.trimIndent()
+    }
+
+    @Test
+    fun `should remove unused imports from start of shared lines`() {
+        TestData.remap("""
+            package test;
+            
+            import java.io.Closeable; import java.util.ArrayList;
+            
+            class Test extends ArrayList {
+            }
+        """.trimIndent()) shouldBe """
+            package test;
+            
+            import java.util.ArrayList;
+            
+            class Test extends ArrayList {
+            }
+        """.trimIndent()
+    }
+}
\ No newline at end of file
diff --git a/src/test/kotlin/com/replaymod/gradle/remap/imports/TestKotlinAutoImports.kt b/src/test/kotlin/com/replaymod/gradle/remap/imports/TestKotlinAutoImports.kt
new file mode 100644
index 0000000..58e2a02
--- /dev/null
+++ b/src/test/kotlin/com/replaymod/gradle/remap/imports/TestKotlinAutoImports.kt
@@ -0,0 +1,26 @@
+package com.replaymod.gradle.remap.imports
+
+import com.replaymod.gradle.remap.util.TestData
+import io.kotest.matchers.shouldBe
+import org.junit.jupiter.api.Test
+
+class TestKotlinAutoImports {
+    @Test
+    fun `should not touch Kotlin files (yet)`() {
+        TestData.remap("test.kt", """
+            package test
+            
+            import java.util.ArrayList
+            
+            class Test {
+            }
+        """.trimIndent()) shouldBe """
+            package test
+            
+            import java.util.ArrayList
+            
+            class Test {
+            }
+        """.trimIndent()
+    }
+}
\ No newline at end of file
diff --git a/src/test/kotlin/com/replaymod/gradle/remap/mapper/mixin/TestMixinAnnotation.kt b/src/test/kotlin/com/replaymod/gradle/remap/mapper/mixin/TestMixinAnnotation.kt
index bc7cd8f..729b00d 100644
--- a/src/test/kotlin/com/replaymod/gradle/remap/mapper/mixin/TestMixinAnnotation.kt
+++ b/src/test/kotlin/com/replaymod/gradle/remap/mapper/mixin/TestMixinAnnotation.kt
@@ -8,9 +8,11 @@ class TestMixinAnnotation {
     @Test
     fun `remaps with class target`() {
         TestData.remap("""
+            import org.spongepowered.asm.mixin.Shadow;
             @org.spongepowered.asm.mixin.Mixin(a.pkg.A.class)
             class MixinA { @Shadow private int aField; }
         """.trimIndent()) shouldBe """
+            import org.spongepowered.asm.mixin.Shadow;
             @org.spongepowered.asm.mixin.Mixin(b.pkg.B.class)
             class MixinA { @Shadow private int bField; }
         """.trimIndent()
@@ -19,9 +21,11 @@ class TestMixinAnnotation {
     @Test
     fun `remaps with string target`() {
         TestData.remap("""
+            import org.spongepowered.asm.mixin.Shadow;
             @org.spongepowered.asm.mixin.Mixin(targets = "a.pkg.A")
             class MixinA { @Shadow private int aField; }
         """.trimIndent()) shouldBe """
+            import org.spongepowered.asm.mixin.Shadow;
             @org.spongepowered.asm.mixin.Mixin(targets = "b.pkg.B")
             class MixinA { @Shadow private int bField; }
         """.trimIndent()
@@ -31,9 +35,11 @@ class TestMixinAnnotation {
     fun `remaps with inner class string target separated by dot`() {
         // FIXME should probably keep the dot?
         TestData.remap("""
+            import org.spongepowered.asm.mixin.Shadow;
             @org.spongepowered.asm.mixin.Mixin(targets = "a.pkg.A.Inner")
             class MixinA { @Shadow private int aField; }
         """.trimIndent()) shouldBe """
+            import org.spongepowered.asm.mixin.Shadow;
             @org.spongepowered.asm.mixin.Mixin(targets = "b.pkg.B${'$'}Inner")
             class MixinA { @Shadow private int bField; }
         """.trimIndent()
@@ -42,9 +48,11 @@ class TestMixinAnnotation {
     @Test
     fun `remaps with inner class string target separated by dollar`() {
         TestData.remap("""
+            import org.spongepowered.asm.mixin.Shadow;
             @org.spongepowered.asm.mixin.Mixin(targets = "a.pkg.A${'$'}Inner")
             class MixinA { @Shadow private int aField; }
         """.trimIndent()) shouldBe """
+            import org.spongepowered.asm.mixin.Shadow;
             @org.spongepowered.asm.mixin.Mixin(targets = "b.pkg.B${'$'}Inner")
             class MixinA { @Shadow private int bField; }
         """.trimIndent()
diff --git a/src/test/kotlin/com/replaymod/gradle/remap/mapper/mixin/TestMixinInjections.kt b/src/test/kotlin/com/replaymod/gradle/remap/mapper/mixin/TestMixinInjections.kt
index da68496..65fbc69 100644
--- a/src/test/kotlin/com/replaymod/gradle/remap/mapper/mixin/TestMixinInjections.kt
+++ b/src/test/kotlin/com/replaymod/gradle/remap/mapper/mixin/TestMixinInjections.kt
@@ -181,16 +181,14 @@ class TestMixinInjections {
     @Test
     fun `remaps @At target`() {
         TestData.remap("""
-            import org.spongepowered.asm.mixin.injection.At;
-            import org.spongepowered.asm.mixin.injection.Inject;
+            import org.spongepowered.asm.mixin.injection.At; import org.spongepowered.asm.mixin.injection.Inject;
             @org.spongepowered.asm.mixin.Mixin(a.pkg.A.class)
             class MixinA {
                 @Inject(method = "aMethod", at = @At(target = "La/pkg/A;aInterfaceMethod()V"))
                 private void test() {}
             }
         """.trimIndent()) shouldBe """
-            import org.spongepowered.asm.mixin.injection.At;
-            import org.spongepowered.asm.mixin.injection.Inject;
+            import org.spongepowered.asm.mixin.injection.At; import org.spongepowered.asm.mixin.injection.Inject;
             @org.spongepowered.asm.mixin.Mixin(b.pkg.B.class)
             class MixinA {
                 @Inject(method = "bMethod", at = @At(target = "Lb/pkg/B;bInterfaceMethod()V"))
@@ -202,16 +200,14 @@ class TestMixinInjections {
     @Test
     fun `remaps @At target without mappings for target`() {
         TestData.remap("""
-            import org.spongepowered.asm.mixin.injection.At;
-            import org.spongepowered.asm.mixin.injection.Inject;
+            import org.spongepowered.asm.mixin.injection.At; import org.spongepowered.asm.mixin.injection.Inject;
             @org.spongepowered.asm.mixin.Mixin(a.pkg.A.class)
             class MixinA {
                 @Inject(method = "aMethod", at = @At(target = "La/pkg/A;unmappedOverloaded(La/pkg/A;)V"))
                 private void test() {}
             }
         """.trimIndent()) shouldBe """
-            import org.spongepowered.asm.mixin.injection.At;
-            import org.spongepowered.asm.mixin.injection.Inject;
+            import org.spongepowered.asm.mixin.injection.At; import org.spongepowered.asm.mixin.injection.Inject;
             @org.spongepowered.asm.mixin.Mixin(b.pkg.B.class)
             class MixinA {
                 @Inject(method = "bMethod", at = @At(target = "Lb/pkg/B;unmappedOverloaded(Lb/pkg/B;)V"))
@@ -223,8 +219,7 @@ class TestMixinInjections {
     @Test
     fun `remaps @At target in constant`() {
         TestData.remap("""
-            import org.spongepowered.asm.mixin.injection.At;
-            import org.spongepowered.asm.mixin.injection.Inject;
+            import org.spongepowered.asm.mixin.injection.At; import org.spongepowered.asm.mixin.injection.Inject;
             @org.spongepowered.asm.mixin.Mixin(a.pkg.A.class)
             class MixinA {
                 private static final String TARGET = "La/pkg/A;aInterfaceMethod()V";
@@ -234,8 +229,7 @@ class TestMixinInjections {
                 private void test2() {}
             }
         """.trimIndent()) shouldBe """
-            import org.spongepowered.asm.mixin.injection.At;
-            import org.spongepowered.asm.mixin.injection.Inject;
+            import org.spongepowered.asm.mixin.injection.At; import org.spongepowered.asm.mixin.injection.Inject;
             @org.spongepowered.asm.mixin.Mixin(b.pkg.B.class)
             class MixinA {
                 private static final String TARGET = "Lb/pkg/B;bInterfaceMethod()V";
diff --git a/src/test/kotlin/com/replaymod/gradle/remap/util/TestData.kt b/src/test/kotlin/com/replaymod/gradle/remap/util/TestData.kt
index c5b8851..d1ed905 100644
--- a/src/test/kotlin/com/replaymod/gradle/remap/util/TestData.kt
+++ b/src/test/kotlin/com/replaymod/gradle/remap/util/TestData.kt
@@ -48,14 +48,21 @@ object TestData {
             findClasspathEntry("BMarkerKt"),
         )
         patternAnnotation = "remap.Pattern"
+        manageImports = true
     }
 
-    fun remap(content: String, patternsBefore: String = "", patternsAfter: String = ""): String = transformer.remap(mapOf(
-        "test.java" to content,
+    fun remap(content: String): String =
+        remap("test.java", content)
+    fun remap(fileName: String, content: String): String =
+        remap(fileName, content, "", "")
+    fun remap(content: String, patternsBefore: String, patternsAfter: String): String =
+        remap("test.java", content, patternsBefore, patternsAfter)
+    fun remap(fileName: String, content: String, patternsBefore: String, patternsAfter: String): String = transformer.remap(mapOf(
+        fileName to content,
         "pattern.java" to "class Patterns {\n$patternsBefore\n}",
     ), mapOf(
         "pattern.java" to "class Patterns {\n$patternsAfter\n}",
-    ))["test.java"]!!.first
+    ))[fileName]!!.first
     fun remapWithErrors(content: String) = transformer.remap(mapOf("test.java" to content))["test.java"]!!
 
     fun remapKt(content: String): String = transformer.remap(mapOf("test.kt" to content))["test.kt"]!!.first