ADFA-3754: add AddImport action for Kotlin

actions/src/main/java/com/itsaky/androidide/actions/actionUtils.kt

@@ -24,17 +24,23 @@ import io.github.rosemoe.sora.widget.CodeEditor
 import java.io.File
 import java.nio.file.Path
 
-fun ActionData.getContext(): Context? {
- return get(Context::class.java)
-}
+inline fun <reified T> ActionData.get(): T?=
+	get(T::class.java)
 
-fun ActionData.requireContext(): Context {
- return getContext() ?: throw IllegalArgumentException("No context instance provided")
-}
+inline fun <reified T> ActionData.require(): T =
+	checkNotNull(get<T>())
 
-fun ActionData.requireFile(): File {
- return get(File::class.java) ?: throw IllegalArgumentException("No file instance provided")
-}
+inline fun <reified T> ActionData.has(): Boolean =
+	get<T>() != null
+
+fun ActionData.getContext(): Context? =
+	get<Context>()
+
+fun ActionData.requireContext(): Context =
+	require<Context>()
+
+fun ActionData.requireFile(): File =
+	require<File>()
 
 fun ActionData.requirePath(): Path {
  return requireFile().toPath()

editor/src/main/java/com/itsaky/androidide/editor/ui/EditorActionsMenu.kt

@@ -46,8 +46,10 @@ import com.itsaky.androidide.editor.adapters.IdeEditorAdapter
 import com.itsaky.androidide.editor.databinding.LayoutPopupMenuItemBinding
 import com.itsaky.androidide.editor.ui.EditorActionsMenu.ActionsListAdapter.VH
 import com.itsaky.androidide.idetooltips.TooltipManager
+import com.itsaky.androidide.lsp.api.ILanguageClient
 import com.itsaky.androidide.lsp.api.ILanguageServerRegistry
 import com.itsaky.androidide.lsp.java.JavaLanguageServer
+import com.itsaky.androidide.lsp.kotlin.KotlinLanguageServer
 import com.itsaky.androidide.lsp.models.DiagnosticItem
 import com.itsaky.androidide.lsp.xml.XMLLanguageServer
 import com.itsaky.androidide.resources.R
@@ -298,7 +300,9 @@ open class EditorActionsMenu(val editor: IDEEditor) :
  protected open fun onGetActionLocation() = location
 
  protected open fun onCreateActionData(): ActionData {
+ val languageServerRegistry = ILanguageServerRegistry.default
  val data = ActionData.create(editor.context)
+
  data.put(IDEEditor::class.java, this.editor)
  data.put(
  CodeEditor::class.java,
@@ -307,14 +311,18 @@ open class EditorActionsMenu(val editor: IDEEditor) :
  data.put(File::class.java, editor.file)
  data.put(DiagnosticItem::class.java, getDiagnosticAtCursor())
  data.put(com.itsaky.androidide.models.Range::class.java, editor.cursorLSPRange)
-data.put(
+data.put(
  JavaLanguageServer::class.java,
-ILanguageServerRegistry.default.getServer(JavaLanguageServer.SERVER_ID)
+languageServerRegistry.getServer(JavaLanguageServer.SERVER_ID)
  as? JavaLanguageServer?
  )
+ data.put(KotlinLanguageServer::class.java,
+ languageServerRegistry
+ .getServer(KotlinLanguageServer.SERVER_ID)
+ as? KotlinLanguageServer?)
  data.put(
  XMLLanguageServer::class.java,
-ILanguageServerRegistry.default.getServer(XMLLanguageServer.SERVER_ID)
+languageServerRegistry.getServer(XMLLanguageServer.SERVER_ID)
  as? XMLLanguageServer?
  )
  data.put(TextTarget::class.java, IdeEditorAdapter(this.editor))

lsp/jvm-symbol-index/src/main/kotlin/org/appdevforall/codeonthego/indexing/jvm/JvmSymbolIndex.kt

@@ -126,6 +126,12 @@ class JvmSymbolIndex(
  this.limit = limit
 	})
 
+	fun findBySimpleName(name: String, limit: Int = 200) =
+ query(indexQuery {
+ eq(KEY_NAME, name)
+ this.limit = limit
+ })
+
 	suspend fun findByKey(key: String): JvmSymbol? = get(key)
 
 	fun allPackages(): Sequence<String> = distinctValues(KEY_PACKAGE)

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

@@ -4,6 +4,7 @@ import com.itsaky.androidide.actions.ActionItem
 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
 
 object KotlinCodeActionsMenu : IActionsMenuProvider {
 
@@ -14,6 +15,7 @@ object KotlinCodeActionsMenu : IActionsMenuProvider {
 	override val actions: List<ActionItem> =
  listOf(
  CommentLineAction(KT_LANG, KT_EXTS, KT_LINE_COMMENT_TOKEN),
- UncommentLineAction(KT_LANG, KT_EXTS, KT_LINE_COMMENT_TOKEN)
+ UncommentLineAction(KT_LANG, KT_EXTS, KT_LINE_COMMENT_TOKEN),
+ AddImportAction(),
  )
 }

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

@@ -0,0 +1,134 @@
+package com.itsaky.androidide.lsp.kotlin.actions
+
+import com.itsaky.androidide.actions.ActionData
+import com.itsaky.androidide.actions.has
+import com.itsaky.androidide.actions.markInvisible
+import com.itsaky.androidide.actions.newDialogBuilder
+import com.itsaky.androidide.actions.require
+import com.itsaky.androidide.actions.requireFile
+import com.itsaky.androidide.idetooltips.TooltipTag
+import com.itsaky.androidide.lsp.kotlin.compiler.index.findSymbolBySimpleName
+import com.itsaky.androidide.lsp.kotlin.diagnostic.KotlinDiagnosticExtra
+import com.itsaky.androidide.lsp.kotlin.utils.insertImport
+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.DiagnosticItem
+import com.itsaky.androidide.lsp.models.DocumentChange
+import com.itsaky.androidide.lsp.models.TextEdit
+import com.itsaky.androidide.resources.R
+import org.appdevforall.codeonthego.indexing.jvm.JvmSymbol
+import org.jetbrains.kotlin.analysis.api.fir.diagnostics.KaFirDiagnostic
+import org.slf4j.LoggerFactory
+
+class AddImportAction : BaseKotlinCodeAction() {
+	override var titleTextRes: Int = R.string.action_import_classes
+	override var tooltipTag: String = TooltipTag.EDITOR_CODE_ACTIONS_FIX_IMPORTS
+
+	override val id: String = "ide.editor.lsp.kt.diagnostics.addImport"
+	override var label: String = ""
+
+	companion object {
+ private val logger = LoggerFactory.getLogger(AddImportAction::class.java)
+	}
+
+	override fun prepare(data: ActionData) {
+ super.prepare(data)
+
+ if (!visible || !data.has<DiagnosticItem>()) {
+ markInvisible()
+ return
+ }
+
+ val extra = data.require<DiagnosticItem>().extra as? KotlinDiagnosticExtra
+ if (extra == null) {
+ markInvisible()
+ return
+ }
+
+ val diagnostic = extra.diagnostic as? KaFirDiagnostic.UnresolvedReference?
+ if (diagnostic == null) {
+ markInvisible()
+ return
+ }
+
+ val env = extra.compilationEnv
+ val reference = diagnostic.reference
+ val hasImportableSymbols = env.ktSymbolIndex
+ .findSymbolBySimpleName(reference, limit = 0)
+ .any { it.kind.isClassifier }
+
+ if (!hasImportableSymbols) {
+ markInvisible()
+ return
+ }
+	}
+
+	override suspend fun execAction(data: ActionData): Map<JvmSymbol, List<TextEdit>> {
+ val (diagnostic, env) = data.require<DiagnosticItem>().extra as? KotlinDiagnosticExtra
+ ?: return emptyMap()
+
+ diagnostic as KaFirDiagnostic.UnresolvedReference
+
+ val file = data.requireFile()
+ val nioPath = file.toPath()
+ val ktFile = env.ktSymbolIndex
+ .getOpenedKtFile(nioPath)
+ ?: return emptyMap()
+
+ return env.ktSymbolIndex
+ .findSymbolBySimpleName(diagnostic.reference, limit = 0)
+ .filter { it.kind.isClassifier }
+ .associateWith { symbol -> insertImport(ktFile, symbol.fqName) }
+	}
+
+	override fun postExec(data: ActionData, result: Any) {
+ super.postExec(data, result)
+
+ if (result !is Map<*, *>) {
+ return
+ }
+
+ @Suppress("UNCHECKED_CAST")
+ result as Map<JvmSymbol, List<TextEdit>>
+
+ if (result.isEmpty()) {
+ logger.warn("No classifiers to import.")
+ return
+ }
+
+ val client = data.languageClient
+ ?: run {
+ logger.warn("No language client set. Cannot complete action.")
+ return
+ }
+
+ val file = data.requireFile()
+ val nioPath = file.toPath()
+ val actions =
+ result
+ .map { (symbol, edits) ->
+ CodeActionItem(
+ title = symbol.fqName,
+ changes = listOf(DocumentChange(file = nioPath, edits = edits)),
+ kind = CodeActionKind.QuickFix,
+ command = Command.CMD_FORMAT_CODE,
+ )
+ }
+
+ when (actions.size) {
+ 0 -> logger.error("No code actions found. Cannot completion action.")
+ 1 -> client.performCodeAction(actions[0])
+ else -> newDialogBuilder(data)
+ .setTitle(label)
+ .setItems(actions.map { it.title }.toTypedArray()) { dialog, which ->
+ dialog.dismiss()
+ actions.getOrNull(which)?.also { client.performCodeAction(it) }
+ ?: run {
+ logger.error("Index $which is out of bounds for actions of size ${actions.size}")
+ }
+ }
+ .show()
+ }
+	}
+}

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

@@ -0,0 +1,60 @@
+package com.itsaky.androidide.lsp.kotlin.actions
+
+import android.content.Context
+import android.graphics.drawable.Drawable
+import androidx.annotation.StringRes
+import com.itsaky.androidide.actions.ActionData
+import com.itsaky.androidide.actions.ActionItem
+import com.itsaky.androidide.actions.EditorActionItem
+import com.itsaky.androidide.actions.get
+import com.itsaky.androidide.actions.hasRequiredData
+import com.itsaky.androidide.actions.markInvisible
+import com.itsaky.androidide.actions.requireContext
+import com.itsaky.androidide.actions.requireFile
+import com.itsaky.androidide.lsp.api.ILanguageClient
+import com.itsaky.androidide.lsp.kotlin.KotlinLanguageServer
+import com.itsaky.androidide.utils.DocumentUtils
+import org.slf4j.LoggerFactory
+import java.io.File
+
+abstract class BaseKotlinCodeAction : EditorActionItem {
+
+	override var visible: Boolean = true
+	override var enabled: Boolean = true
+	override var icon: Drawable? = null
+	override var requiresUIThread: Boolean = false
+	override var location: ActionItem.Location = ActionItem.Location.EDITOR_CODE_ACTIONS
+
+	@get:StringRes
+	protected abstract var titleTextRes: Int
+
+	protected val logger = LoggerFactory.getLogger(BaseKotlinCodeAction::class.java)
+
+	override fun prepare(data: ActionData) {
+ super.prepare(data)
+ if (!data.hasRequiredData(
+ Context::class.java,
+ KotlinLanguageServer::class.java,
+ File::class.java
+ )
+ ) {
+ markInvisible()
+ return
+ }
+
+ val context = data.requireContext()
+ val file = data.requireFile()
+ val isKtFile = DocumentUtils.isKotlinFile(file.toPath())
+
+ if (titleTextRes != -1) {
+ label = context.getString(titleTextRes)
+ }
+
+ visible = isKtFile
+ enabled = isKtFile
+	}
+
+	protected val ActionData.languageClient: ILanguageClient?
+ get() = get<KotlinLanguageServer>()
+ ?.client
+}

lsp/kotlin/src/main/java/com/itsaky/androidide/lsp/kotlin/compiler/index/KtSymbolIndex.kt

@@ -138,4 +138,8 @@ internal fun KtSymbolIndex.filesForPackage(packageFqn: String) =
 	fileIndex.getFilesForPackage(packageFqn)
 
 internal fun KtSymbolIndex.subpackageNames(packageFqn: String) =
-	fileIndex.getSubpackageNames(packageFqn)
+	fileIndex.getSubpackageNames(packageFqn)
+
+internal fun KtSymbolIndex.findSymbolBySimpleName(name: String, limit: Int) =
+	(sourceIndex.findBySimpleName(name, 0) + libraryIndex.findBySimpleName(name, 0))
+ .take(limit)

lsp/kotlin/src/main/java/com/itsaky/androidide/lsp/kotlin/diagnostic/KotlinDiagnosticProvider.kt

@@ -18,10 +18,14 @@ import org.jetbrains.kotlin.com.intellij.psi.PsiFile
 import org.jetbrains.kotlin.com.intellij.psi.util.PsiTreeUtil
 import org.slf4j.LoggerFactory
 import java.nio.file.Path
-import kotlin.math.log
 
 private val logger = LoggerFactory.getLogger("KotlinDiagnosticProvider")
 
+internal data class KotlinDiagnosticExtra(
+	val diagnostic: KaDiagnosticWithPsi<*>,
+	val compilationEnv: CompilationEnvironment,
+)
+
 internal fun CompilationEnvironment.collectDiagnosticsFor(file: Path): DiagnosticResult = try {
 	logger.info("Analyzing file: {}", file)
 	return doAnalyze(file)
@@ -63,9 +67,12 @@ private fun CompilationEnvironment.doAnalyze(file: Path): DiagnosticResult {
 
  analyze(ktFile) {
  ktFile.collectDiagnostics(KaDiagnosticCheckerFilter.EXTENDED_AND_COMMON_CHECKERS)
- .forEach { add(it.toDiagnosticItem()) }
+ .forEach { diagnostic ->
+ add(diagnostic.toDiagnosticItem().apply {
+ extra = KotlinDiagnosticExtra(diagnostic, this@doAnalyze)
+ })
+ }
  }
-
  }
 	}