Support K2 mode in the IJ plugin (#5138)

This PR makes the IJ plugin run in K2 mode with the new Analysis APIs, instead of the old K1 APIs. It also includes cleaning up a bunch of compiler warnings, and some mistakes I saw when migrating to K2. --------- Co-authored-by: Victor Kropp <victor.kropp@jetbrains.com>
2 weeks ago · 1b877ddab1
9 changed files with 247 additions and 126 deletions
--- a/idea-plugin/build.gradle.kts
+++ b/idea-plugin/build.gradle.kts
@ -32,13 +32,7 @@ dependencies {
 }

 intellijPlatform {
-    pluginConfiguration {
-        name = "Compose Multiplatform IDE Support"
-        ideaVersion {
-            sinceBuild = "231.*"
-            untilBuild = "243.*"
-        }
-    }
+    pluginConfiguration { name = "Compose Multiplatform IDE Support" }
    buildSearchableOptions = false
    autoReload = false

@ -56,6 +50,14 @@ tasks {
        targetCompatibility = "21"
    }
    withType<KotlinJvmCompile> { compilerOptions.jvmTarget.set(JvmTarget.JVM_21) }
+
+    runIde {
+        systemProperty("idea.is.internal", true)
+        systemProperty("idea.kotlin.plugin.use.k2", true)
+        jvmArgumentProviders += CommandLineArgumentProvider {
+            listOf("-agentlib:jdwp=transport=dt_socket,server=y,suspend=n,address=5005")
+        }
+    }
 }

 class ProjectProperties(private val project: Project) {
--- a/idea-plugin/src/main/kotlin/org/jetbrains/compose/desktop/ide/preview/ConfigurePreviewTaskNameCache.kt
+++ b/idea-plugin/src/main/kotlin/org/jetbrains/compose/desktop/ide/preview/ConfigurePreviewTaskNameCache.kt
@ -15,6 +15,7 @@ import com.intellij.openapi.project.Project
 import com.intellij.util.concurrency.annotations.RequiresReadLock
 import org.jetbrains.plugins.gradle.settings.GradleSettings
 import org.jetbrains.plugins.gradle.util.GradleConstants
+import java.util.Locale

 internal val DEFAULT_CONFIGURE_PREVIEW_TASK_NAME = "configureDesktopPreview"

@ -38,8 +39,11 @@ internal class ConfigurePreviewTaskNameProviderImpl : ConfigurePreviewTaskNamePr
        return null
    }

-    private fun previewTaskName(targetName: String = "") =
-        "$DEFAULT_CONFIGURE_PREVIEW_TASK_NAME${targetName.capitalize()}"
+    private fun previewTaskName(targetName: String = ""): String {
+        val capitalizedTargetName =
+            targetName.replaceFirstChar { if (it.isLowerCase()) it.titlecase(Locale.getDefault()) else it.toString() }
+        return "$DEFAULT_CONFIGURE_PREVIEW_TASK_NAME$capitalizedTargetName"
+    }

    private fun moduleDataNodeOrNull(project: Project, modulePath: String): DataNode<ModuleData>? {
        val projectDataManager = ProjectDataManager.getInstance()
@ -87,4 +91,4 @@ internal class ConfigurePreviewTaskNameCache(
            cachedTaskName = null
        }
    }
-}
+}
--- a/idea-plugin/src/main/kotlin/org/jetbrains/compose/desktop/ide/preview/PreviewLocation.kt
+++ b/idea-plugin/src/main/kotlin/org/jetbrains/compose/desktop/ide/preview/PreviewLocation.kt
@ -5,6 +5,7 @@

 package org.jetbrains.compose.desktop.ide.preview

+import com.intellij.openapi.components.service
 import com.intellij.openapi.externalSystem.util.ExternalSystemApiUtil
 import com.intellij.openapi.roots.ProjectFileIndex
 import com.intellij.util.concurrency.annotations.RequiresReadLock
@ -20,7 +21,7 @@ internal fun KtNamedFunction.asPreviewFunctionOrNull(): PreviewLocation? {
    val module = ProjectFileIndex.getInstance(project).getModuleForFile(containingFile.virtualFile)
    if (module == null || module.isDisposed) return null

-    val service = project.getService(PreviewStateService::class.java)
+    val service = project.service<PreviewStateService>()
    val previewTaskName = service.configurePreviewTaskNameOrNull(module) ?: DEFAULT_CONFIGURE_PREVIEW_TASK_NAME
    val modulePath = ExternalSystemApiUtil.getExternalProjectPath(module) ?: return null
    return PreviewLocation(fqName = fqName, modulePath = modulePath, taskName = previewTaskName)
--- a/idea-plugin/src/main/kotlin/org/jetbrains/compose/desktop/ide/preview/PreviewStateService.kt
+++ b/idea-plugin/src/main/kotlin/org/jetbrains/compose/desktop/ide/preview/PreviewStateService.kt
@ -7,11 +7,9 @@ package org.jetbrains.compose.desktop.ide.preview

 import com.intellij.openapi.Disposable
 import com.intellij.openapi.components.Service
-import com.intellij.openapi.diagnostic.Logger
 import com.intellij.openapi.externalSystem.model.task.*
 import com.intellij.openapi.externalSystem.service.notification.ExternalSystemProgressNotificationManager
 import com.intellij.openapi.module.Module
-import com.intellij.openapi.project.Project
 import com.intellij.openapi.util.Disposer
 import com.intellij.ui.components.JBLoadingPanel
 import com.intellij.util.concurrency.annotations.RequiresReadLock
@ -22,9 +20,8 @@ import javax.swing.JComponent
 import javax.swing.event.AncestorEvent
 import javax.swing.event.AncestorListener

-@Service
-class PreviewStateService(private val myProject: Project) : Disposable {
-    private val idePreviewLogger = Logger.getInstance("org.jetbrains.compose.desktop.ide.preview")
+@Service(Service.Level.PROJECT)
+class PreviewStateService : Disposable {
    private val previewListener = CompositePreviewListener()
    private val previewManager: PreviewManager = PreviewManagerImpl(previewListener)
    val gradleCallbackPort: Int
@ -35,7 +32,7 @@ class PreviewStateService(private val myProject: Project) : Disposable {
    init {
        val projectRefreshListener = ConfigurePreviewTaskNameCacheInvalidator(configurePreviewTaskNameCache)
        ExternalSystemProgressNotificationManager.getInstance()
-            .addNotificationListener(projectRefreshListener, myProject)
+            .addNotificationListener(projectRefreshListener, this)
    }

    @RequiresReadLock
@ -80,7 +77,6 @@ private class PreviewResizeListener(private val previewManager: PreviewManager)

    override fun ancestorAdded(event: AncestorEvent) {
        updateFrameSize(event.component)
-
    }

    override fun ancestorRemoved(event: AncestorEvent) {
@ -136,4 +132,4 @@ private class ConfigurePreviewTaskNameCacheInvalidator(
            configurePreviewTaskNameCache.invalidate()
        }
    }
-}
+}
--- a/idea-plugin/src/main/kotlin/org/jetbrains/compose/desktop/ide/preview/PreviewToolWindow.kt
+++ b/idea-plugin/src/main/kotlin/org/jetbrains/compose/desktop/ide/preview/PreviewToolWindow.kt
@ -11,23 +11,23 @@ import com.intellij.openapi.project.Project
 import com.intellij.openapi.wm.ToolWindow
 import com.intellij.openapi.wm.ToolWindowFactory
 import com.intellij.ui.components.JBLoadingPanel
-import org.jetbrains.compose.desktop.ide.preview.ui.PreviewPanel
 import java.awt.BorderLayout
+import org.jetbrains.compose.desktop.ide.preview.ui.PreviewPanel

 class PreviewToolWindow : ToolWindowFactory, DumbAware {
-    override fun isApplicable(project: Project): Boolean =
-        isPreviewCompatible(project)
+    @Deprecated("Use isApplicableAsync")
+    override fun isApplicable(project: Project): Boolean = isPreviewCompatible(project)
+
+    override suspend fun isApplicableAsync(project: Project): Boolean = isPreviewCompatible(project)

    override fun init(toolWindow: ToolWindow) {
-        ApplicationManager.getApplication().invokeLater {
-            toolWindow.setIcon(PreviewIcons.COMPOSE)
-        }
+        ApplicationManager.getApplication().invokeLater { toolWindow.setIcon(PreviewIcons.COMPOSE) }
    }

    override fun createToolWindowContent(project: Project, toolWindow: ToolWindow) {
        toolWindow.contentManager.let { content ->
            val panel = PreviewPanel(project)
-            val loadingPanel = JBLoadingPanel(BorderLayout(), project)
+            val loadingPanel = JBLoadingPanel(BorderLayout(), toolWindow.disposable)
            loadingPanel.add(panel, BorderLayout.CENTER)
            content.addContent(content.factory.createContent(loadingPanel, null, false))
            project.service<PreviewStateService>().registerPreviewPanels(panel, loadingPanel)
@ -35,6 +35,5 @@ class PreviewToolWindow : ToolWindowFactory, DumbAware {
    }

    // don't show the toolwindow until a preview is requested
-    override fun shouldBeAvailable(project: Project): Boolean =
-        false
-}
+    override fun shouldBeAvailable(project: Project): Boolean = false
+}
--- a/idea-plugin/src/main/kotlin/org/jetbrains/compose/desktop/ide/preview/locationUtils.kt
+++ b/idea-plugin/src/main/kotlin/org/jetbrains/compose/desktop/ide/preview/locationUtils.kt
@ -21,32 +21,39 @@ import com.intellij.psi.util.CachedValueProvider
 import com.intellij.psi.util.CachedValuesManager
 import com.intellij.psi.util.parentOfType
 import com.intellij.util.concurrency.annotations.RequiresReadLock
+import org.jetbrains.kotlin.analysis.api.analyze
+import org.jetbrains.kotlin.analysis.api.symbols.KaClassLikeSymbol
 import org.jetbrains.kotlin.asJava.findFacadeClass
-import org.jetbrains.kotlin.builtins.KotlinBuiltIns
-import org.jetbrains.kotlin.descriptors.ClassKind
-import org.jetbrains.kotlin.idea.caches.resolve.analyze
-import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.psi.KtClass
+import org.jetbrains.kotlin.psi.KtFile
+import org.jetbrains.kotlin.psi.KtNamedFunction
+import org.jetbrains.kotlin.psi.allConstructors
 import org.jetbrains.kotlin.psi.psiUtil.containingClass
-import org.jetbrains.kotlin.resolve.BindingContext
-import org.jetbrains.kotlin.resolve.descriptorUtil.fqNameSafe
-import org.jetbrains.kotlin.resolve.lazy.BodyResolveMode

-internal const val DESKTOP_PREVIEW_ANNOTATION_FQN = "androidx.compose.desktop.ui.tooling.preview.Preview"
+internal const val DESKTOP_PREVIEW_ANNOTATION_FQN =
+    "androidx.compose.desktop.ui.tooling.preview.Preview"
 internal const val COMPOSABLE_FQ_NAME = "androidx.compose.runtime.Composable"

+private val ComposableAnnotationClassId = ClassId.topLevel(FqName(COMPOSABLE_FQ_NAME))
+private val DesktopPreviewAnnotationClassId =
+    ClassId.topLevel(FqName(DESKTOP_PREVIEW_ANNOTATION_FQN))
+
 /**
 * Utils based on functions from AOSP, taken from
 * tools/adt/idea/compose-designer/src/com/android/tools/idea/compose/preview/util/PreviewElement.kt
 */

 /**
- * Returns whether a `@Composable` [PREVIEW_ANNOTATION_FQN] is defined in a valid location, which can be either:
+ * Returns whether a `@Composable` [DESKTOP_PREVIEW_ANNOTATION_FQN] is defined in a valid location,
+ * which can be either:
 * 1. Top-level functions
- * 2. Non-nested functions defined in top-level classes that have a default (no parameter) constructor
- *
+ * 2. Non-nested functions defined in top-level classes that have a default (no parameter)
+ *    constructor
 */
 private fun KtNamedFunction.isValidPreviewLocation(): Boolean {
-    if (valueParameters.size > 0) return false
+    if (valueParameters.isNotEmpty()) return false
    if (receiverTypeReference != null) return false

    if (isTopLevel) return true
@ -55,7 +62,8 @@ private fun KtNamedFunction.isValidPreviewLocation(): Boolean {
        // This is not a nested method
        val containingClass = containingClass()
        if (containingClass != null) {
-            // We allow functions that are not top level defined in top level classes that have a default (no parameter) constructor.
+            // We allow functions that are not top level defined in top level classes that have a
+            // default (no parameter) constructor.
            if (containingClass.isTopLevel() && containingClass.hasDefaultConstructor()) {
                return true
            }
@ -64,84 +72,67 @@ private fun KtNamedFunction.isValidPreviewLocation(): Boolean {
    return false
 }

-
 /**
 * Computes the qualified name of the class containing this [KtNamedFunction].
 *
- * For functions defined within a Kotlin class, returns the qualified name of that class. For top-level functions, returns the JVM name of
- * the Java facade class generated instead.
- *
+ * For functions defined within a Kotlin class, returns the qualified name of that class. For
+ * top-level functions, returns the JVM name of the Java facade class generated instead.
 */
 internal fun KtNamedFunction.getClassName(): String? =
-    if (isTopLevel) ((parent as? KtFile)?.findFacadeClass())?.qualifiedName else parentOfType<KtClass>()?.getQualifiedName()
-
+    if (isTopLevel) ((parent as? KtFile)?.findFacadeClass())?.qualifiedName
+    else parentOfType<KtClass>()?.getQualifiedName()

-/** Computes the qualified name for a Kotlin Class. Returns null if the class is a kotlin built-in. */
-private fun KtClass.getQualifiedName(): String? {
-    val classDescriptor = analyze(BodyResolveMode.PARTIAL).get(BindingContext.CLASS, this) ?: return null
-    return if (KotlinBuiltIns.isUnderKotlinPackage(classDescriptor) || classDescriptor.kind != ClassKind.CLASS) {
-        null
-    } else {
-        classDescriptor.fqNameSafe.asString()
+/**
+ * Computes the qualified name for a Kotlin Class. Returns null if the class is a kotlin built-in.
+ */
+private fun KtClass.getQualifiedName(): String? =
+    analyze(this) {
+        val classSymbol = symbol
+        return when {
+            classSymbol !is KaClassLikeSymbol -> null
+            classSymbol.classId.isKotlinPackage() -> null
+            else -> classSymbol.classId?.asFqNameString()
+        }
    }
-}
+
+private fun ClassId?.isKotlinPackage() =
+    this != null && startsWith(org.jetbrains.kotlin.builtins.StandardNames.BUILT_INS_PACKAGE_NAME)

 private fun KtClass.hasDefaultConstructor() =
    allConstructors.isEmpty().or(allConstructors.any { it.valueParameters.isEmpty() })

-/**
- * Determines whether this [KtAnnotationEntry] has the specified qualified name.
- * Careful: this does *not* currently take into account Kotlin type aliases (https://kotlinlang.org/docs/reference/type-aliases.html).
- *   Fortunately, type aliases are extremely uncommon for simple annotation types.
- */
-private fun KtAnnotationEntry.fqNameMatches(fqName: String): Boolean {
-    // For inspiration, see IDELightClassGenerationSupport.KtUltraLightSupportImpl.findAnnotation in the Kotlin plugin.
-    val shortName = shortName?.asString() ?: return false
-    return fqName.endsWith(shortName) && fqName == getQualifiedName()
-}
-
-/**
- * Computes the qualified name of this [KtAnnotationEntry].
- * Prefer to use [fqNameMatches], which checks the short name first and thus has better performance.
- */
-private fun KtAnnotationEntry.getQualifiedName(): String? =
-    analyze(BodyResolveMode.PARTIAL).get(BindingContext.ANNOTATION, this)?.fqName?.asString()
-
 internal fun KtNamedFunction.composePreviewFunctionFqn() = "${getClassName()}.${name}"

@RequiresReadLock
 internal fun KtNamedFunction.isValidComposablePreviewFunction(): Boolean {
-    fun isValidComposablePreviewImpl(): Boolean {
-        if (!isValidPreviewLocation()) return false
-
-        var hasComposableAnnotation = false
-        var hasPreviewAnnotation = false
-        val annotationIt = annotationEntries.iterator()
-        while (annotationIt.hasNext() && !(hasComposableAnnotation && hasPreviewAnnotation)) {
-            val annotation = annotationIt.next()
-            hasComposableAnnotation = hasComposableAnnotation || annotation.fqNameMatches(COMPOSABLE_FQ_NAME)
-            hasPreviewAnnotation = hasPreviewAnnotation || annotation.fqNameMatches(DESKTOP_PREVIEW_ANNOTATION_FQN)
-        }
+    fun isValidComposablePreviewImpl(): Boolean =
+        analyze(this) {
+            if (!isValidPreviewLocation()) return false

-        return hasComposableAnnotation && hasPreviewAnnotation
-    }
+            val mySymbol = symbol
+            val hasComposableAnnotation = mySymbol.annotations.contains(ComposableAnnotationClassId)
+            val hasPreviewAnnotation =
+                mySymbol.annotations.contains(DesktopPreviewAnnotationClassId)

-    return CachedValuesManager.getCachedValue(this) {
-        cachedResult(isValidComposablePreviewImpl())
-    }
+            return hasComposableAnnotation && hasPreviewAnnotation
+        }
+
+    return CachedValuesManager.getCachedValue(this) { cachedResult(isValidComposablePreviewImpl()) }
 }

 // based on AndroidComposePsiUtils.kt from AOSP
-internal fun KtNamedFunction.isComposableFunction(): Boolean {
-    return CachedValuesManager.getCachedValue(this) {
-        cachedResult(annotationEntries.any { it.fqNameMatches(COMPOSABLE_FQ_NAME) })
+internal fun KtNamedFunction.isComposableFunction(): Boolean =
+    CachedValuesManager.getCachedValue(this) {
+        val hasComposableAnnotation =
+            analyze(this) { symbol.annotations.contains(ComposableAnnotationClassId) }
+
+        cachedResult(hasComposableAnnotation)
    }
-}

 private fun <T> KtNamedFunction.cachedResult(value: T) =
    CachedValueProvider.Result.create(
        // TODO: see if we can handle alias imports without ruining performance.
        value,
        this.containingKtFile,
-        ProjectRootModificationTracker.getInstance(project)
-    )
+        ProjectRootModificationTracker.getInstance(project),
+    )
--- a/idea-plugin/src/main/kotlin/org/jetbrains/compose/web/ide/run/WebRunLineMarkerContributor.kt
+++ b/idea-plugin/src/main/kotlin/org/jetbrains/compose/web/ide/run/WebRunLineMarkerContributor.kt
@ -15,9 +15,9 @@ class WebRunLineMarkerContributor : RunLineMarkerContributor() {
    override fun getInfo(element: PsiElement): Info? {
        if (element !is LeafPsiElement) return null
        if (element.node.elementType != KtTokens.IDENTIFIER) return null
+        if (element.parent.getAsJsMainFunctionOrNull() == null) return null

-        val jsMain = element.parent.getAsJsMainFunctionOrNull() ?: return null
        val icon = AllIcons.RunConfigurations.TestState.Run
-        return Info(icon, null, ExecutorAction.getActions()[0])
+        return Info(icon, arrayOf(ExecutorAction.getActions()[0]))
    }
 }
--- a/idea-plugin/src/main/kotlin/org/jetbrains/compose/web/ide/run/psiUtils.kt
+++ b/idea-plugin/src/main/kotlin/org/jetbrains/compose/web/ide/run/psiUtils.kt
@ -6,45 +6,169 @@
 package org.jetbrains.compose.web.ide.run

 import com.intellij.psi.PsiElement
-import org.jetbrains.kotlin.builtins.KotlinBuiltIns
-import org.jetbrains.kotlin.descriptors.FunctionDescriptor
-import org.jetbrains.kotlin.idea.caches.resolve.resolveToDescriptorIfAny
-import org.jetbrains.kotlin.idea.project.platform
-import org.jetbrains.kotlin.idea.util.module
+import com.intellij.util.concurrency.annotations.RequiresReadLock
+import org.jetbrains.kotlin.analysis.api.KaSession
+import org.jetbrains.kotlin.analysis.api.analyze
+import org.jetbrains.kotlin.analysis.api.symbols.KaNamedFunctionSymbol
+import org.jetbrains.kotlin.analysis.api.types.KaClassType
+import org.jetbrains.kotlin.analysis.api.types.KaType
+import org.jetbrains.kotlin.analysis.api.types.KaTypeNullability
+import org.jetbrains.kotlin.config.LanguageFeature
+import org.jetbrains.kotlin.idea.base.facet.platform.platform
+import org.jetbrains.kotlin.idea.base.projectStructure.languageVersionSettings
+import org.jetbrains.kotlin.idea.base.psi.KotlinPsiHeuristics
+import org.jetbrains.kotlin.idea.base.util.module
+import org.jetbrains.kotlin.name.StandardClassIds
 import org.jetbrains.kotlin.platform.js.JsPlatforms
 import org.jetbrains.kotlin.psi.KtNamedFunction
-import org.jetbrains.kotlin.resolve.lazy.BodyResolveMode
-import org.jetbrains.kotlin.types.KotlinType
+import org.jetbrains.kotlin.types.Variance

 internal fun PsiElement.getAsJsMainFunctionOrNull(): KtNamedFunction? =
    (this as? KtNamedFunction)?.takeIf { it.isValidJsMain() }

-internal fun KtNamedFunction.isValidJsMain(): Boolean =
-    isTopLevel && isJsPlatform() && isMainFun()
+internal fun KtNamedFunction.isValidJsMain(): Boolean = isTopLevel && isJsPlatform() && isMainFun()

 internal fun KtNamedFunction.isJsPlatform(): Boolean =
-    module?.platform?.let { platform ->
-        platform in JsPlatforms.allJsPlatforms
-    } ?: false
+    module?.platform?.let { platform -> platform in JsPlatforms.allJsPlatforms } == true

-internal fun KtNamedFunction.isMainFun(): Boolean {
-    if (name != "main") return false
+internal fun KtNamedFunction.isMainFun(): Boolean =
+    isMainFromPsiOnly(this) && isMainFromAnalysis(this)

-    val parameters = valueParameters.toList()
-    if (parameters.size > 1) return false
+//////////////////////////////////////////////////////////////////
+// Copied and simplified from PsiOnlyKotlinMainFunctionDetector //
+//////////////////////////////////////////////////////////////////
+@RequiresReadLock
+private fun isMainFromPsiOnly(function: KtNamedFunction): Boolean {
+    if (function.isLocal || function.typeParameters.isNotEmpty()) {
+        return false
+    }

-    val descriptor = resolveToDescriptorIfAny(BodyResolveMode.PARTIAL_NO_ADDITIONAL)
-    return descriptor is FunctionDescriptor
-            && isUnit(descriptor.returnType)
-            && (parameters.isEmpty() || descriptor.hasSingleArrayOfStringsParameter())
+    val isTopLevel = function.isTopLevel
+    val parameterCount =
+        function.valueParameters.size + (if (function.receiverTypeReference != null) 1 else 0)
+
+    if (parameterCount == 0) {
+        if (!isTopLevel) {
+            return false
+        }
+
+        if (
+            !function.languageVersionSettings.supportsFeature(
+                LanguageFeature.ExtendedMainConvention
+            )
+        ) {
+            return false
+        }
+    } else if (parameterCount == 1 && !isMainCheckParameter(function)) {
+        return false
+    } else {
+        return false
+    }
+
+    if ((KotlinPsiHeuristics.findJvmName(function) ?: function.name) != "main") {
+        return false
+    }
+
+    if (!isTopLevel && !KotlinPsiHeuristics.hasJvmStaticAnnotation(function)) {
+        return false
+    }
+
+    val returnTypeReference = function.typeReference
+    return !(returnTypeReference != null &&
+        !KotlinPsiHeuristics.typeMatches(returnTypeReference, StandardClassIds.Unit))
+}
+
+private fun isMainCheckParameter(function: KtNamedFunction): Boolean {
+    val receiverTypeReference = function.receiverTypeReference
+    if (receiverTypeReference != null) {
+        return KotlinPsiHeuristics.typeMatches(
+            receiverTypeReference,
+            StandardClassIds.Array,
+            StandardClassIds.String,
+        )
+    }
+
+    val parameter = function.valueParameters.singleOrNull() ?: return false
+    val parameterTypeReference = parameter.typeReference ?: return false
+
+    return when {
+        parameter.isVarArg ->
+            KotlinPsiHeuristics.typeMatches(parameterTypeReference, StandardClassIds.String)
+        else ->
+            KotlinPsiHeuristics.typeMatches(
+                parameterTypeReference,
+                StandardClassIds.Array,
+                StandardClassIds.String,
+            )
+    }
+}
+
+//////////////////////////////////////////////////////////////////////
+// Copied and simplified from SymbolBasedKotlinMainFunctionDetector //
+//////////////////////////////////////////////////////////////////////
+private fun isMainFromAnalysis(function: KtNamedFunction): Boolean {
+    if (function.isLocal || function.typeParameters.isNotEmpty()) {
+        return false
+    }
+
+    val supportsExtendedMainConvention =
+        function.languageVersionSettings.supportsFeature(LanguageFeature.ExtendedMainConvention)
+
+    val isTopLevel = function.isTopLevel
+    val parameterCount =
+        function.valueParameters.size + (if (function.receiverTypeReference != null) 1 else 0)
+
+    if (parameterCount == 0) {
+        if (!isTopLevel || !supportsExtendedMainConvention) {
+            return false
+        }
+    } else if (parameterCount > 1) {
+        return false
+    }
+
+    analyze(function) {
+        if (parameterCount == 1) {
+            val parameterTypeReference =
+                function.receiverTypeReference
+                    ?: function.valueParameters[0].typeReference
+                    ?: return false
+
+            val parameterType = parameterTypeReference.type
+            if (
+                !parameterType.isResolvedClassType() ||
+                    !parameterType.isSubtypeOf(buildMainParameterType())
+            ) {
+                return false
+            }
+        }
+
+        val functionSymbol = function.symbol
+        if (functionSymbol !is KaNamedFunctionSymbol) {
+            return false
+        }
+
+        if (functionSymbol.name.identifier != "main") {
+            return false
+        }
+
+        if (!function.returnType.isUnitType) {
+            return false
+        }
+    }
+    return true
 }

-private fun isUnit(type: KotlinType?): Boolean =
-    type != null && KotlinBuiltIns.isUnit(type)
+private fun KaSession.buildMainParameterType(): KaType =
+    buildClassType(StandardClassIds.Array) {
+        val argumentType =
+            buildClassType(StandardClassIds.String) { nullability = KaTypeNullability.NON_NULLABLE }
+
+        argument(argumentType, Variance.OUT_VARIANCE)
+        nullability = KaTypeNullability.NULLABLE
+    }

-private fun FunctionDescriptor.hasSingleArrayOfStringsParameter(): Boolean {
-    val parameter = valueParameters.singleOrNull() ?: return false
-    val type = parameter.type
-    val typeArgument = type.arguments.singleOrNull()?.type
-    return KotlinBuiltIns.isArray(type) && KotlinBuiltIns.isString(typeArgument)
-}
+private fun KaType.isResolvedClassType(): Boolean =
+    when (this) {
+        is KaClassType -> typeArguments.mapNotNull { it.type }.all { it.isResolvedClassType() }
+        else -> false
+    }
--- a/idea-plugin/src/main/resources/META-INF/plugin.xml
+++ b/idea-plugin/src/main/resources/META-INF/plugin.xml
@ -23,6 +23,10 @@
    <depends>com.intellij.gradle</depends>
    <depends>org.jetbrains.kotlin</depends>

+    <extensions defaultExtensionNs="org.jetbrains.kotlin">
+        <supportsKotlinPluginMode supportsK2="true" />
+    </extensions>
+
    <extensions defaultExtensionNs="com.intellij">
        <runLineMarkerContributor
                language="kotlin"