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>

idea-plugin/build.gradle.kts

@@ -32,13 +32,7 @@ dependencies {
 }
 
 intellijPlatform {
-    pluginConfiguration {
+    pluginConfiguration { name = "Compose Multiplatform IDE Support" }
-        name = "Compose Multiplatform IDE Support"
-        ideaVersion {
-            sinceBuild = "231.*"
-            untilBuild = "243.*"
-        }
-    }
     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) {

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 = "") =
+    private fun previewTaskName(targetName: String = ""): String {
-        "$DEFAULT_CONFIGURE_PREVIEW_TASK_NAME${targetName.capitalize()}"
+        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()

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)

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
+@Service(Service.Level.PROJECT)
-class PreviewStateService(private val myProject: Project) : Disposable {
+class PreviewStateService : Disposable {
-    private val idePreviewLogger = Logger.getInstance("org.jetbrains.compose.desktop.ide.preview")
     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) {

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 =
+    @Deprecated("Use isApplicableAsync")
-        isPreviewCompatible(project)
+    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 {
+        ApplicationManager.getApplication().invokeLater { toolWindow.setIcon(PreviewIcons.COMPOSE) }
-            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 =
+    override fun shouldBeAvailable(project: Project): Boolean = false
-        false
 }

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.name.ClassId
-import org.jetbrains.kotlin.descriptors.ClassKind
+import org.jetbrains.kotlin.name.FqName
-import org.jetbrains.kotlin.idea.caches.resolve.analyze
+import org.jetbrains.kotlin.psi.KtClass
-import org.jetbrains.kotlin.psi.*
+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
+ * For functions defined within a Kotlin class, returns the qualified name of that class. For
- * the Java facade class generated instead.
+ * 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? {
+ * Computes the qualified name for a Kotlin Class. Returns null if the class is a kotlin built-in.
-    val classDescriptor = analyze(BodyResolveMode.PARTIAL).get(BindingContext.CLASS, this) ?: return null
+ */
-    return if (KotlinBuiltIns.isUnderKotlinPackage(classDescriptor) || classDescriptor.kind != ClassKind.CLASS) {
+private fun KtClass.getQualifiedName(): String? =
-        null
+    analyze(this) {
-    } else {
+        val classSymbol = symbol
-        classDescriptor.fqNameSafe.asString()
+        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 {
+    fun isValidComposablePreviewImpl(): Boolean =
-        if (!isValidPreviewLocation()) return false
+        analyze(this) {
-
+            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)
-        }
 
-        return hasComposableAnnotation && hasPreviewAnnotation
+            val mySymbol = symbol
-    }
+            val hasComposableAnnotation = mySymbol.annotations.contains(ComposableAnnotationClassId)
+            val hasPreviewAnnotation =
+                mySymbol.annotations.contains(DesktopPreviewAnnotationClassId)
 
-    return CachedValuesManager.getCachedValue(this) {
+            return hasComposableAnnotation && hasPreviewAnnotation
-        cachedResult(isValidComposablePreviewImpl())
+        }
-    }
+
+    return CachedValuesManager.getCachedValue(this) { cachedResult(isValidComposablePreviewImpl()) }
 }
 
 // based on AndroidComposePsiUtils.kt from AOSP
-internal fun KtNamedFunction.isComposableFunction(): Boolean {
+internal fun KtNamedFunction.isComposableFunction(): Boolean =
-    return CachedValuesManager.getCachedValue(this) {
+    CachedValuesManager.getCachedValue(this) {
-        cachedResult(annotationEntries.any { it.fqNameMatches(COMPOSABLE_FQ_NAME) })
+        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),
     )

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]))
     }
 }

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 com.intellij.util.concurrency.annotations.RequiresReadLock
-import org.jetbrains.kotlin.descriptors.FunctionDescriptor
+import org.jetbrains.kotlin.analysis.api.KaSession
-import org.jetbrains.kotlin.idea.caches.resolve.resolveToDescriptorIfAny
+import org.jetbrains.kotlin.analysis.api.analyze
-import org.jetbrains.kotlin.idea.project.platform
+import org.jetbrains.kotlin.analysis.api.symbols.KaNamedFunctionSymbol
-import org.jetbrains.kotlin.idea.util.module
+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.Variance
-import org.jetbrains.kotlin.types.KotlinType
 
 internal fun PsiElement.getAsJsMainFunctionOrNull(): KtNamedFunction? =
     (this as? KtNamedFunction)?.takeIf { it.isValidJsMain() }
 
-internal fun KtNamedFunction.isValidJsMain(): Boolean =
+internal fun KtNamedFunction.isValidJsMain(): Boolean = isTopLevel && isJsPlatform() && isMainFun()
-    isTopLevel && isJsPlatform() && isMainFun()
 
 internal fun KtNamedFunction.isJsPlatform(): Boolean =
-    module?.platform?.let { platform ->
+    module?.platform?.let { platform -> platform in JsPlatforms.allJsPlatforms } == true
-        platform in JsPlatforms.allJsPlatforms
-    } ?: false
 
-internal fun KtNamedFunction.isMainFun(): Boolean {
+internal fun KtNamedFunction.isMainFun(): Boolean =
-    if (name != "main") return false
+    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)
+    val isTopLevel = function.isTopLevel
-    return descriptor is FunctionDescriptor
+    val parameterCount =
-            && isUnit(descriptor.returnType)
+        function.valueParameters.size + (if (function.receiverTypeReference != null) 1 else 0)
-            && (parameters.isEmpty() || descriptor.hasSingleArrayOfStringsParameter())
+
+    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 isUnit(type: KotlinType?): Boolean =
+private fun isMainCheckParameter(function: KtNamedFunction): Boolean {
-    type != null && KotlinBuiltIns.isUnit(type)
+    val receiverTypeReference = function.receiverTypeReference
+    if (receiverTypeReference != null) {
+        return KotlinPsiHeuristics.typeMatches(
+            receiverTypeReference,
+            StandardClassIds.Array,
+            StandardClassIds.String,
+        )
+    }
 
-private fun FunctionDescriptor.hasSingleArrayOfStringsParameter(): Boolean {
+    val parameter = function.valueParameters.singleOrNull() ?: return false
-    val parameter = valueParameters.singleOrNull() ?: return false
+    val parameterTypeReference = parameter.typeReference ?: return false
-    val type = parameter.type
+
-    val typeArgument = type.arguments.singleOrNull()?.type
+    return when {
-    return KotlinBuiltIns.isArray(type) && KotlinBuiltIns.isString(typeArgument)
+        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 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 KaType.isResolvedClassType(): Boolean =
+    when (this) {
+        is KaClassType -> typeArguments.mapNotNull { it.type }.all { it.isResolvedClassType() }
+        else -> false
+    }

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"