ADFA-4614: K2-LSP code action — Organize imports (#1502)

Author: itsaky-adfa

docs/superpowers/plans/2026-06-25-enhance-ai-assistant-plugin-ui.md

@@ -5,9 +5,9 @@ import com.itsaky.androidide.lsp.actions.CommentLineAction
 import com.itsaky.androidide.lsp.actions.IActionsMenuProvider
 import com.itsaky.androidide.lsp.actions.UncommentLineAction
 import com.itsaky.androidide.lsp.kotlin.actions.AddImportAction
+import com.itsaky.androidide.lsp.kotlin.actions.OrganizeImportsAction
 
 object KotlinCodeActionsMenu : IActionsMenuProvider {
-
 	private const val KT_LANG = "kt"
 	private val KT_EXTS = listOf("kt", "kts")
 	private const val KT_LINE_COMMENT_TOKEN = "//"
@@ -17,5 +17,6 @@ object KotlinCodeActionsMenu : IActionsMenuProvider {
 			CommentLineAction(KT_LANG, KT_EXTS, KT_LINE_COMMENT_TOKEN),
 			UncommentLineAction(KT_LANG, KT_EXTS, KT_LINE_COMMENT_TOKEN),
 			AddImportAction(),
+			OrganizeImportsAction(),
 		)
-}
+}

lsp/kotlin/src/main/java/com/itsaky/androidide/lsp/kotlin/KotlinCodeActionsMenu.kt

@@ -30,6 +30,7 @@ import com.itsaky.androidide.eventbus.events.file.FileRenameEvent
 import com.itsaky.androidide.lsp.api.ILanguageClient
 import com.itsaky.androidide.lsp.api.ILanguageServer
 import com.itsaky.androidide.lsp.api.IServerSettings
+import com.itsaky.androidide.lsp.kotlin.compiler.CompilationEnvironment
 import com.itsaky.androidide.lsp.kotlin.compiler.Compiler
 import com.itsaky.androidide.lsp.kotlin.compiler.KotlinProjectModel
 import com.itsaky.androidide.lsp.kotlin.compiler.index.KT_SOURCE_FILE_INDEX_KEY
@@ -118,6 +119,9 @@ class KotlinLanguageServer : ILanguageServer {
 		this.compiler?.updateLanguageClient(client)
 	}
 
+	/** Returns the [CompilationEnvironment] responsible for [file], or null if the compiler is not ready. */
+	internal fun compilationEnvironmentFor(file: Path): CompilationEnvironment? = compiler?.compilationEnvironmentFor(file)
+
 	override fun applySettings(settings: IServerSettings?) {
 		this._settings = settings
 	}

lsp/kotlin/src/main/java/com/itsaky/androidide/lsp/kotlin/KotlinLanguageServer.kt

@@ -0,0 +1,93 @@
+package com.itsaky.androidide.lsp.kotlin.actions
+
+import com.itsaky.androidide.actions.ActionData
+import com.itsaky.androidide.actions.get
+import com.itsaky.androidide.actions.requireFile
+import com.itsaky.androidide.lsp.kotlin.KotlinLanguageServer
+import com.itsaky.androidide.lsp.kotlin.compiler.AbstractCompilationEnvironment
+import com.itsaky.androidide.lsp.kotlin.compiler.modules.analyzeMaybeDangling
+import com.itsaky.androidide.lsp.kotlin.compiler.read
+import com.itsaky.androidide.lsp.kotlin.utils.collectImportUsage
+import com.itsaky.androidide.lsp.kotlin.utils.organizedImportBlock
+import com.itsaky.androidide.lsp.kotlin.utils.toRange
+import com.itsaky.androidide.lsp.models.CodeActionItem
+import com.itsaky.androidide.lsp.models.CodeActionKind
+import com.itsaky.androidide.lsp.models.Command
+import com.itsaky.androidide.lsp.models.DocumentChange
+import com.itsaky.androidide.lsp.models.TextEdit
+import com.itsaky.androidide.models.Range
+import com.itsaky.androidide.resources.R
+import org.slf4j.LoggerFactory
+import java.nio.file.Path
+
+class OrganizeImportsAction : BaseKotlinCodeAction() {
+	override var titleTextRes: Int = R.string.action_organize_imports
+	override val id: String = "ide.editor.lsp.kt.organizeImports"
+	override var label: String = ""
+
+	companion object {
+		private val logger = LoggerFactory.getLogger(OrganizeImportsAction::class.java)
+	}
+
+	override suspend fun execAction(data: ActionData): List<TextEdit> {
+		val server = data.get<KotlinLanguageServer>() ?: return emptyList()
+		val nioPath = data.requireFile().toPath()
+		val env = server.compilationEnvironmentFor(nioPath) ?: return emptyList()
+		return computeOrganizeEdit(env, nioPath)
+	}
+
+	/**
+	 * Computes the text edits that organize the imports of the file at [nioPath] within [env].
+	 * The current [org.jetbrains.kotlin.psi.KtFile] is fetched BEFORE entering [read] (deadlock
+	 * rule: never block on `getCurrentKtFile(...).get()` inside `project.read`). Returns an empty
+	 * list when there is nothing to do (no imports, already organized, or no usable range) *and*
+	 * whenever anything in this pipeline (the `.get()`, analysis, or PSI access) throws: the action
+	 * framework only catches [IllegalArgumentException] and this runs on a coroutine scope with no
+	 * exception handler, so an uncaught throw here would crash the app. Degrading to zero edits is
+	 * always safe -- it just leaves the imports as-is, never produces a partial/incorrect rewrite.
+	 */
+	internal fun computeOrganizeEdit(
+		env: AbstractCompilationEnvironment,
+		nioPath: Path,
+	): List<TextEdit> =
+		runCatching {
+			val ktFile = env.ktSymbolIndex.getCurrentKtFile(nioPath).get() ?: return emptyList()
+			if (ktFile.importDirectives.isEmpty()) return emptyList()
+			env.project.read {
+				val usage = analyzeMaybeDangling(ktFile) { collectImportUsage(ktFile) }
+				val newText = organizedImportBlock(ktFile, usage) ?: return@read emptyList()
+				val range = ktFile.importList?.textRange?.toRange(ktFile) ?: return@read emptyList()
+				if (range == Range.NONE) return@read emptyList()
+				listOf(TextEdit(range, newText))
+			}
+		}.getOrElse { e ->
+			logger.warn("Failed to organize imports", e)
+			emptyList()
+		}
+
+	override fun postExec(
+		data: ActionData,
+		result: Any,
+	) {
+		super.postExec(data, result)
+		if (result !is List<*> || result.isEmpty()) return
+
+		@Suppress("UNCHECKED_CAST")
+		result as List<TextEdit>
+
+		val client =
+			data.languageClient ?: run {
+				logger.warn("No language client set. Cannot organize imports.")
+				return
+			}
+		val file = data.requireFile()
+		client.performCodeAction(
+			CodeActionItem(
+				title = label,
+				changes = listOf(DocumentChange(file = file.toPath(), edits = result)),
+				kind = CodeActionKind.QuickFix,
+				command = Command("", ""), // no post-action command (no CMD_FORMAT_CODE)
+			),
+		)
+	}
+}

lsp/kotlin/src/main/java/com/itsaky/androidide/lsp/kotlin/actions/OrganizeImportsAction.kt

@@ -0,0 +1,102 @@
+package com.itsaky.androidide.lsp.kotlin.utils
+
+import org.jetbrains.kotlin.com.intellij.psi.util.PsiTreeUtil
+import org.jetbrains.kotlin.kdoc.psi.api.KDoc
+import org.jetbrains.kotlin.psi.KtFile
+import org.jetbrains.kotlin.psi.KtImportDirective
+
+/**
+ * What a file's body actually uses, expressed as plain strings so nothing crosses an `analyze`
+ * lifetime boundary. Produced by [collectImportUsage].
+ *
+ * @property usedFqNames importable fully-qualified names referenced by the body.
+ * @property usedPackages parent packages/objects of used symbols (for wildcard matching).
+ * @property unresolvedNames short names of body references that failed to resolve; an import
+ *   matching one of these is kept, since a resolution failure can't prove the import unused.
+ */
+internal data class ImportUsage(
+	val usedFqNames: Set<String>,
+	val usedPackages: Set<String>,
+	val unresolvedNames: Set<String> = emptySet(),
+)
+
+/** JVM packages that Kotlin imports with a wildcard by default; explicit named imports from these are redundant. */
+internal val DEFAULT_STAR_PACKAGES: Set<String> =
+	setOf(
+		"kotlin",
+		"kotlin.annotation",
+		"kotlin.collections",
+		"kotlin.comparisons",
+		"kotlin.io",
+		"kotlin.ranges",
+		"kotlin.sequences",
+		"kotlin.text",
+		"kotlin.jvm",
+		"java.lang",
+	)
+
+private val KDOC_LINK = Regex("""\[([^\]\s]+)]""")
+
+/**
+ * Computes the canonical import block for [ktFile] given [usage]: unused/redundant imports removed,
+ * survivors deduped and lexicographically sorted. Returns null when the imports are already in that
+ * exact form (no edit needed). The returned text has no surrounding newlines.
+ */
+internal fun organizedImportBlock(
+	ktFile: KtFile,
+	usage: ImportUsage,
+): String? {
+	val directives = ktFile.importDirectives
+	if (directives.isEmpty()) return null
+
+	val filePackage = ktFile.packageFqName.asString()
+	val kdocNames = collectKDocLinkNames(ktFile)
+
+	val newLines =
+		directives
+			.filter { keepImport(it, usage, filePackage, kdocNames) }
+			.mapNotNull { it.importPath?.let { path -> "import $path" } }
+			.distinct()
+			.sorted()
+
+	val currentLines = directives.mapNotNull { it.importPath?.let { path -> "import $path" } }
+
+	if (newLines == currentLines) return null
+	return newLines.joinToString(System.lineSeparator())
+}
+
+private fun keepImport(
+	directive: KtImportDirective,
+	usage: ImportUsage,
+	filePackage: String,
+	kdocNames: Set<String>,
+): Boolean {
+	val fqName = directive.importedFqName ?: return true // malformed import -> keep (conservative)
+	val fqNameStr = fqName.asString()
+	val alias = directive.aliasName
+	val shortName = alias ?: fqName.shortName().asString()
+
+	// Conservative: keep anything referenced by short name/alias in a KDoc link.
+	if (shortName in kdocNames) return true
+
+	// Conservative: an unresolved body reference by this short name can't prove the import dead.
+	if (shortName in usage.unresolvedNames) return true
+
+	if (directive.isAllUnder) {
+		// Wildcard: keep iff some used symbol lives in this package/object.
+		return fqNameStr in usage.usedPackages
+	}
+
+	val parentPackage = fqName.parent().asString()
+	// Redundant named imports (only when not aliased — an alias is meaningful).
+	if (alias == null && parentPackage in DEFAULT_STAR_PACKAGES) return false
+	if (alias == null && parentPackage == filePackage) return false
+
+	return fqNameStr in usage.usedFqNames
+}
+
+private fun collectKDocLinkNames(ktFile: KtFile): Set<String> =
+	PsiTreeUtil
+		.collectElementsOfType(ktFile, KDoc::class.java)
+		.flatMap { kdoc -> KDOC_LINK.findAll(kdoc.text).map { it.groupValues[1].substringAfterLast('.') } }
+		.toSet()

lsp/kotlin/src/main/java/com/itsaky/androidide/lsp/kotlin/utils/ImportOrganizer.kt

@@ -0,0 +1,105 @@
+package com.itsaky.androidide.lsp.kotlin.utils
+
+import org.jetbrains.kotlin.analysis.api.KaSession
+import org.jetbrains.kotlin.analysis.api.resolution.successfulFunctionCallOrNull
+import org.jetbrains.kotlin.analysis.api.resolution.symbol
+import org.jetbrains.kotlin.analysis.api.symbols.KaCallableSymbol
+import org.jetbrains.kotlin.analysis.api.symbols.KaClassLikeSymbol
+import org.jetbrains.kotlin.analysis.api.symbols.KaConstructorSymbol
+import org.jetbrains.kotlin.analysis.api.symbols.KaSymbol
+import org.jetbrains.kotlin.idea.references.mainReference
+import org.jetbrains.kotlin.psi.KtArrayAccessExpression
+import org.jetbrains.kotlin.psi.KtCallExpression
+import org.jetbrains.kotlin.psi.KtDestructuringDeclaration
+import org.jetbrains.kotlin.psi.KtDestructuringDeclarationEntry
+import org.jetbrains.kotlin.psi.KtElement
+import org.jetbrains.kotlin.psi.KtFile
+import org.jetbrains.kotlin.psi.KtForExpression
+import org.jetbrains.kotlin.psi.KtImportList
+import org.jetbrains.kotlin.psi.KtNameReferenceExpression
+import org.jetbrains.kotlin.psi.KtOperationReferenceExpression
+import org.jetbrains.kotlin.psi.KtPropertyDelegate
+import org.jetbrains.kotlin.psi.psiUtil.collectDescendantsOfType
+import org.jetbrains.kotlin.psi.psiUtil.getParentOfType
+
+/**
+ * Collects the importable fq-names (and their packages) referenced by [ktFile]'s body. MUST be
+ * called inside [analyzeMaybeDangling]. Returns only plain strings, so nothing escapes the analyze
+ * lifetime. A reference that fails to resolve doesn't join the used set; instead its short name is
+ * recorded in [ImportUsage.unresolvedNames] so its import is kept. Both paths are safe: they lead to
+ * keeping an import, never removing a used one.
+ */
+internal fun KaSession.collectImportUsage(ktFile: KtFile): ImportUsage {
+	val usedFqNames = HashSet<String>()
+	val usedPackages = HashSet<String>()
+	val unresolvedNames = HashSet<String>()
+
+	fun record(symbol: KaSymbol?) {
+		val fq = symbol?.importableFqNameString() ?: return
+		usedFqNames += fq
+		val pkg = fq.substringBeforeLast('.', missingDelimiterValue = "")
+		if (pkg.isNotEmpty()) usedPackages += pkg
+	}
+
+	fun recordAll(symbols: Collection<KaSymbol>?) {
+		symbols?.forEach(::record)
+	}
+
+	// 1) Plain name / type references (excluding the import list itself). A null (or thrown)
+	// resolution is treated as unresolved and its short name kept, so a used-but-unresolvable
+	// reference never drops its import. A non-null, non-importable symbol (local, param) is a
+	// clean resolve: it records nothing and is not unresolved.
+	ktFile.collectDescendantsOfType<KtNameReferenceExpression>().forEach { ref ->
+		if (ref.getParentOfType<KtImportList>(strict = false) != null) return@forEach
+		val symbol = runCatching { ref.mainReference.resolveToSymbol() }.getOrNull()
+		if (symbol != null) record(symbol) else unresolvedNames += ref.getReferencedName()
+	}
+
+	// 1b) Implicit-convention references that carry more than one resolution target and so don't
+	// resolve through `resolveToSymbol()` (singular; returns null when ambiguous) but do resolve
+	// through `resolveToSymbols()` (plural). Confirmed empirically:
+	//  - KtForExpression: resolves to [iterator(), hasNext(), next()] -- iterator is the
+	//    user-importable one for a `for (x in foo)` loop.
+	//  - KtDestructuringDeclarationEntry (one per destructured variable): resolves to that
+	//    variable's own componentN() symbol.
+	//  - KtPropertyDelegate: resolves to the delegate's getValue()/setValue() symbol(s).
+	// Recording every returned symbol is safe: extra (e.g. stdlib Iterator.next) symbols only ever
+	// keep an import, never drop a used one.
+	ktFile.collectDescendantsOfType<KtForExpression>().forEach { forExpr ->
+		runCatching { recordAll(forExpr.mainReference?.resolveToSymbols()) }
+	}
+	ktFile.collectDescendantsOfType<KtDestructuringDeclarationEntry>().forEach { entry ->
+		runCatching { recordAll(entry.mainReference?.resolveToSymbols()) }
+	}
+	ktFile.collectDescendantsOfType<KtPropertyDelegate>().forEach { delegate ->
+		runCatching { recordAll(delegate.mainReference?.resolveToSymbols()) }
+	}
+
+	// 2) Convention / operator call sites (no textual name reference).
+	ktFile.collectDescendantsOfType<KtElement>().forEach { element ->
+		val isConvention =
+			element is KtOperationReferenceExpression ||
+				element is KtArrayAccessExpression ||
+				element is KtCallExpression ||
+				element is KtForExpression ||
+				element is KtDestructuringDeclaration ||
+				element is KtPropertyDelegate
+		if (!isConvention) return@forEach
+		runCatching {
+			record(element.resolveToCall()?.successfulFunctionCallOrNull()?.symbol)
+		}
+	}
+
+	return ImportUsage(usedFqNames, usedPackages, unresolvedNames)
+}
+
+private fun KaSymbol.importableFqNameString(): String? =
+	when (this) {
+		// A constructor's own callableId is null, so it must map to its containing class -- the name
+		// that's actually imported. Covers `Foo()` calls and `@Foo` annotations (both resolve to the
+		// constructor). Must precede the KaCallableSymbol branch, which a constructor also matches.
+		is KaConstructorSymbol -> containingClassId?.asSingleFqName()?.asString()
+		is KaClassLikeSymbol -> classId?.asSingleFqName()?.asString()
+		is KaCallableSymbol -> callableId?.asSingleFqName()?.asString()
+		else -> null
+	}