ADFA-4614: K2-LSP code action — Organize imports

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

@@ -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/KotlinLanguageServer.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/actions/OrganizeImportsAction.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/utils/ImportOrganizer.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/ImportUsageCollector.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
+	}