Enable running preview for any Compose Desktop module (#951)

Previously preview only worked in projects,
that define compose.desktop.application {} DSL block

Resolves #908

gradle-plugins/compose/src/main/kotlin/org/jetbrains/compose/desktop/application/internal/configureApplication.kt

@@ -9,7 +9,6 @@ import org.gradle.api.*
 import org.gradle.api.file.Directory
 import org.gradle.api.file.DuplicatesStrategy
 import org.gradle.api.file.FileCollection
-import org.gradle.api.plugins.JavaPluginConvention
 import org.gradle.api.provider.Provider
 import org.gradle.api.tasks.*
 import org.gradle.jvm.tasks.Jar
@@ -17,9 +16,10 @@ import org.jetbrains.compose.desktop.application.dsl.Application
 import org.jetbrains.compose.desktop.application.dsl.TargetFormat
 import org.jetbrains.compose.desktop.application.internal.validation.validatePackageVersions
 import org.jetbrains.compose.desktop.application.tasks.*
-import org.jetbrains.compose.desktop.preview.internal.configureConfigureDesktopPreviewTask
-import org.jetbrains.compose.desktop.preview.tasks.AbstractConfigureDesktopPreviewTask
-import org.jetbrains.kotlin.gradle.dsl.KotlinMultiplatformExtension
+import org.jetbrains.compose.internal.KOTLIN_JVM_PLUGIN_ID
+import org.jetbrains.compose.internal.KOTLIN_MPP_PLUGIN_ID
+import org.jetbrains.compose.internal.javaExt
+import org.jetbrains.compose.internal.mppExt
 import org.jetbrains.kotlin.gradle.plugin.KotlinPlatformType
 import java.io.File
 import java.util.*
@@ -31,10 +31,10 @@ private val defaultJvmArgs = listOf("-Dcompose.application.configure.swing.globa
 // todo: use workers
 fun configureApplicationImpl(project: Project, app: Application) {
     if (app._isDefaultConfigurationEnabled) {
-        if (project.plugins.hasPlugin("org.jetbrains.kotlin.multiplatform")) {
+        if (project.plugins.hasPlugin(KOTLIN_MPP_PLUGIN_ID)) {
             project.configureFromMppPlugin(app)
-        } else if (project.plugins.hasPlugin("org.jetbrains.kotlin.jvm")) {
-            val mainSourceSet = project.convention.getPlugin(JavaPluginConvention::class.java).sourceSets.getByName("main")
+        } else if (project.plugins.hasPlugin(KOTLIN_JVM_PLUGIN_ID)) {
+            val mainSourceSet = project.javaExt.sourceSets.getByName("main")
             app.from(mainSourceSet)
         }
     }
@@ -44,9 +44,8 @@ fun configureApplicationImpl(project: Project, app: Application) {
 }
 
 internal fun Project.configureFromMppPlugin(mainApplication: Application) {
-    val kotlinExt = extensions.getByType(KotlinMultiplatformExtension::class.java)
     var isJvmTargetConfigured = false
-    kotlinExt.targets.all { target ->
+    mppExt.targets.all { target ->
         if (target.platformType == KotlinPlatformType.jvm) {
             if (!isJvmTargetConfigured) {
                 mainApplication.from(target)
@@ -159,10 +158,6 @@ internal fun Project.configurePackagingTasks(apps: Collection<Application>) {
         val run = project.tasks.composeTask<JavaExec>(taskName("run", app)) {
             configureRunTask(app)
         }
-
-        val configureDesktopPreviewTask = project.tasks.composeTask<AbstractConfigureDesktopPreviewTask>("configureDesktopPreview") {
-            configureConfigureDesktopPreviewTask(app)
-        }
     }
 }
 

gradle-plugins/compose/src/main/kotlin/org/jetbrains/compose/desktop/preview/internal/configurePreview.kt

@@ -1,19 +1,39 @@
 package org.jetbrains.compose.desktop.preview.internal
 
 import org.gradle.api.Project
-import org.jetbrains.compose.desktop.application.dsl.Application
-import org.jetbrains.compose.desktop.application.internal.javaHomeOrDefault
-import org.jetbrains.compose.desktop.application.internal.provider
+import org.jetbrains.compose.desktop.application.dsl.ConfigurationSource
 import org.jetbrains.compose.desktop.preview.tasks.AbstractConfigureDesktopPreviewTask
+import org.jetbrains.compose.internal.KOTLIN_JVM_PLUGIN_ID
+import org.jetbrains.compose.internal.KOTLIN_MPP_PLUGIN_ID
+import org.jetbrains.compose.internal.javaExt
+import org.jetbrains.compose.internal.mppExt
+import org.jetbrains.kotlin.gradle.plugin.KotlinPlatformType
+import org.jetbrains.kotlin.gradle.targets.jvm.KotlinJvmTarget
 
 fun Project.initializePreview() {
+    plugins.withId(KOTLIN_MPP_PLUGIN_ID) {
+        mppExt.targets.all { target ->
+            if (target.platformType == KotlinPlatformType.jvm) {
+                val config = ConfigurationSource.KotlinMppTarget(target as KotlinJvmTarget)
+                registerConfigurePreviewTask(project, config, targetName = target.name)
+            }
+        }
+    }
+    plugins.withId(KOTLIN_JVM_PLUGIN_ID) {
+        val config = ConfigurationSource.GradleSourceSet(project.javaExt.sourceSets.getByName("main"))
+        registerConfigurePreviewTask(project, config)
+    }
 }
 
-internal fun AbstractConfigureDesktopPreviewTask.configureConfigureDesktopPreviewTask(app: Application) {
-    app._configurationSource?.let { configSource ->
-        dependsOn(configSource.jarTaskName)
-        previewClasspath = configSource.runtimeClasspath(project)
-        javaHome.set(provider { app.javaHomeOrDefault() })
-        jvmArgs.set(provider { app.jvmArgs })
+private fun registerConfigurePreviewTask(project: Project, config: ConfigurationSource, targetName: String = "") {
+    project.tasks.register(
+        previewTaskName(targetName),
+        AbstractConfigureDesktopPreviewTask::class.java
+    ) { previewTask ->
+        previewTask.dependsOn(config.jarTask(project))
+        previewTask.previewClasspath = config.runtimeClasspath(project)
     }
 }
+
+private fun previewTaskName(targetName: String) =
+    "configureDesktopPreview${targetName.capitalize()}"

gradle-plugins/compose/src/main/kotlin/org/jetbrains/compose/internal/constants.kt

@@ -0,0 +1,9 @@
+/*
+ * Copyright 2020-2021 JetBrains s.r.o. and respective authors and developers.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE.txt file.
+ */
+
+package org.jetbrains.compose.internal
+
+internal const val KOTLIN_MPP_PLUGIN_ID = "org.jetbrains.kotlin.multiplatform"
+internal const val KOTLIN_JVM_PLUGIN_ID = "org.jetbrains.kotlin.jvm"

gradle-plugins/compose/src/main/kotlin/org/jetbrains/compose/internal/projectExtensions.kt

@@ -6,6 +6,7 @@
 package org.jetbrains.compose.internal
 
 import org.gradle.api.Project
+import org.gradle.api.plugins.JavaPluginExtension
 import org.jetbrains.compose.ComposeExtension
 import org.jetbrains.compose.web.WebExtension
 import org.jetbrains.kotlin.gradle.dsl.KotlinMultiplatformExtension
@@ -16,5 +17,14 @@ internal val Project.composeExt: ComposeExtension?
 internal val Project.webExt: WebExtension?
     get() = composeExt?.extensions?.findByType(WebExtension::class.java)
 
-internal val Project.mppExt: KotlinMultiplatformExtension?
+internal val Project.mppExt: KotlinMultiplatformExtension
+    get() = mppExtOrNull ?: error("Could not find KotlinMultiplatformExtension ($project)")
+
+internal val Project.mppExtOrNull: KotlinMultiplatformExtension?
     get() = extensions.findByType(KotlinMultiplatformExtension::class.java)
+
+internal val Project.javaExt: JavaPluginExtension
+    get() = javaExtOrNull ?: error("Could not find JavaPluginExtension ($project)")
+
+internal val Project.javaExtOrNull: JavaPluginExtension?
+    get() = extensions.findByType(JavaPluginExtension::class.java)

gradle-plugins/compose/src/main/kotlin/org/jetbrains/compose/web/WebExtension.kt

@@ -8,6 +8,7 @@ package org.jetbrains.compose.web
 import org.gradle.api.Project
 import org.gradle.api.plugins.ExtensionAware
 import org.jetbrains.compose.internal.mppExt
+import org.jetbrains.compose.internal.mppExtOrNull
 import org.jetbrains.kotlin.gradle.plugin.KotlinTarget
 import org.jetbrains.kotlin.gradle.targets.js.ir.KotlinJsIrTarget
 
@@ -45,7 +46,7 @@ abstract class WebExtension : ExtensionAware {
     }
 
     private fun defaultJsTargetsToConfigure(project: Project): Set<KotlinJsIrTarget> {
-        val mppTargets = project.mppExt?.targets?.asMap?.values ?: emptySet()
+        val mppTargets = project.mppExtOrNull?.targets?.asMap?.values ?: emptySet()
         val jsIRTargets = mppTargets.filterIsInstanceTo(LinkedHashSet<KotlinJsIrTarget>())
 
         return if (jsIRTargets.size > 1) {

gradle-plugins/compose/src/test/kotlin/org/jetbrains/compose/gradle/GradlePluginTest.kt

@@ -6,11 +6,20 @@
 package org.jetbrains.compose.gradle
 
 import org.gradle.testkit.runner.TaskOutcome
+import org.jetbrains.compose.desktop.ui.tooling.preview.rpc.PreviewLogger
+import org.jetbrains.compose.desktop.ui.tooling.preview.rpc.RemoteConnection
+import org.jetbrains.compose.desktop.ui.tooling.preview.rpc.receiveConfigFromGradle
 import org.jetbrains.compose.test.GradlePluginTestBase
 import org.jetbrains.compose.test.TestKotlinVersion
 import org.jetbrains.compose.test.TestProjects
 import org.jetbrains.compose.test.checks
 import org.junit.jupiter.api.Test
+import java.net.ServerSocket
+import java.net.Socket
+import java.net.SocketTimeoutException
+import java.util.concurrent.atomic.AtomicBoolean
+import java.util.concurrent.atomic.AtomicInteger
+import kotlin.concurrent.thread
 
 class GradlePluginTest : GradlePluginTestBase() {
     @Test
@@ -25,4 +34,103 @@ class GradlePluginTest : GradlePluginTestBase() {
                 check.taskOutcome(":compileKotlinJs", TaskOutcome.SUCCESS)
             }
         }
+
+    @Test
+    fun configurePreview() {
+        val isAlive = AtomicBoolean(true)
+        val receivedConfigCount = AtomicInteger(0)
+        val port = AtomicInteger(-1)
+        val connectionThread = thread {
+            val serverSocket = ServerSocket(0).apply {
+                soTimeout = 10_000
+            }
+            port.set(serverSocket.localPort)
+            try {
+                while (isAlive.get()) {
+                    try {
+                        val socket = serverSocket.accept()
+                        val connection = RemoteConnectionImpl(socket, TestPreviewLogger("SERVER"))
+                        val previewConfig = connection.receiveConfigFromGradle()
+                        if (previewConfig != null) {
+                            receivedConfigCount.incrementAndGet()
+                        }
+                    } catch (e: Exception) {
+                        if (!isAlive.get()) break
+
+                        if (e !is SocketTimeoutException) {
+                            e.printStackTrace()
+                            throw e
+                        }
+                    }
+                }
+
+            } finally {
+                serverSocket.close()
+            }
+        }
+
+        val startTimeNs = System.nanoTime()
+        while (port.get() <= 0) {
+            val elapsedTimeNs = System.nanoTime() - startTimeNs
+            val elapsedTimeMs = elapsedTimeNs / 1_000_000L
+            if (elapsedTimeMs > 10_000) {
+                error("Server socket initialization timeout!")
+            }
+            Thread.sleep(200)
+        }
+
+        try {
+            testConfigureDesktopPreivewImpl(port.get())
+        } finally {
+            isAlive.set(false)
+            connectionThread.interrupt()
+            connectionThread.join(5000)
+        }
+
+        val expectedReceivedConfigCount = 2
+        val actualReceivedConfigCount = receivedConfigCount.get()
+        check(actualReceivedConfigCount == 2) {
+            "Expected to receive $expectedReceivedConfigCount preview configs, got $actualReceivedConfigCount"
+        }
+    }
+
+    private fun testConfigureDesktopPreivewImpl(port: Int) {
+        check(port > 0) { "Invalid port: $port" }
+        with(testProject(TestProjects.jvmPreview)) {
+            val portProperty = "-Pcompose.desktop.preview.ide.port=$port"
+            val previewTargetProperty = "-Pcompose.desktop.preview.target=PreviewKt.ExamplePreview"
+            val jvmTask = ":jvm:configureDesktopPreview"
+            gradle(jvmTask, portProperty, previewTargetProperty)
+                .build()
+                .checks { check ->
+                    check.taskOutcome(jvmTask, TaskOutcome.SUCCESS)
+                }
+
+            val mppTask = ":mpp:configureDesktopPreviewDesktop"
+            gradle(mppTask, portProperty, previewTargetProperty)
+                .build()
+                .checks { check ->
+                    check.taskOutcome(mppTask, TaskOutcome.SUCCESS)
+                }
+        }
+    }
+
+    private class TestPreviewLogger(private val prefix: String) : PreviewLogger() {
+        override val isEnabled: Boolean
+            get() = true
+
+        override fun log(s: String) {
+            println("$prefix: $s")
+        }
+    }
+
+    private fun RemoteConnectionImpl(
+        socket: Socket, logger: PreviewLogger
+    ): RemoteConnection {
+        val connectionClass = Class.forName("org.jetbrains.compose.desktop.ui.tooling.preview.rpc.RemoteConnectionImpl")
+        val constructor = connectionClass.constructors.first {
+            it.parameterCount == 3
+        }
+        return constructor.newInstance(socket, logger, {}) as RemoteConnection
+    }
 }

gradle-plugins/compose/src/test/kotlin/org/jetbrains/compose/test/TestProjects.kt

@@ -18,4 +18,5 @@ object TestProjects {
     const val defaultArgsOverride = "application/defaultArgsOverride"
     const val unpackSkiko = "application/unpackSkiko"
     const val jsMpp = "misc/jsMpp"
+    const val jvmPreview = "misc/jvmPreview"
 }

gradle-plugins/compose/src/test/test-projects/misc/jvmPreview/build.gradle

@@ -0,0 +1,9 @@
+subprojects {
+    repositories {
+        mavenLocal()
+        mavenCentral()
+        maven {
+            url 'https://maven.pkg.jetbrains.space/public/p/compose/dev'
+        }
+    }
+}

gradle-plugins/compose/src/test/test-projects/misc/jvmPreview/jvm/build.gradle

@@ -0,0 +1,10 @@
+plugins {
+    id 'org.jetbrains.kotlin.jvm'
+    id 'org.jetbrains.compose'
+}
+
+dependencies {
+    implementation 'org.jetbrains.kotlin:kotlin-stdlib'
+    implementation compose.uiTooling
+    implementation compose.desktop.currentOs
+}

gradle-plugins/compose/src/test/test-projects/misc/jvmPreview/mpp/build.gradle

@@ -0,0 +1,20 @@
+plugins {
+    id 'org.jetbrains.kotlin.multiplatform'
+    id 'org.jetbrains.compose'
+}
+
+kotlin {
+    jvm('desktop') {}
+
+    sourceSets {
+        commonMain.dependencies {
+            api compose.runtime
+            api compose.foundation
+            api compose.material
+            api compose.uiTooling
+        }
+        desktopMain.dependencies {
+            implementation compose.desktop.currentOs
+        }
+    }
+}

gradle-plugins/compose/src/test/test-projects/misc/jvmPreview/mpp/src/commonMain/kotlin/composable.kt

@@ -0,0 +1,21 @@
+/*
+ * Copyright 2020-2021 JetBrains s.r.o. and respective authors and developers.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE.txt file.
+ */
+
+import androidx.compose.material.Text
+import androidx.compose.material.Button
+import androidx.compose.runtime.*
+
+@Composable
+fun ExampleComposable() {
+    var text by remember { mutableStateOf("Hello, World!") }
+
+    Button(onClick = {
+        text = "Hello, $platformName!"
+    }) {
+        Text(text)
+    }
+}
+
+expect val platformName: String

gradle-plugins/compose/src/test/test-projects/misc/jvmPreview/mpp/src/desktopMain/kotlin/preview.kt

@@ -0,0 +1,11 @@
+import androidx.compose.runtime.Composable
+import androidx.compose.desktop.ui.tooling.preview.Preview
+
+@Preview
+@Composable
+fun ExamplePreview() {
+    ExampleComposable()
+}
+
+actual val platformName: String
+    get() = "Desktop"

gradle-plugins/compose/src/test/test-projects/misc/jvmPreview/settings.gradle

@@ -0,0 +1,16 @@
+pluginManagement {
+    plugins {
+        id 'org.jetbrains.kotlin.multiplatform' version 'KOTLIN_VERSION_PLACEHOLDER'
+        id 'org.jetbrains.kotlin.jvm' version 'KOTLIN_VERSION_PLACEHOLDER'
+        id 'org.jetbrains.compose' version 'COMPOSE_VERSION_PLACEHOLDER'
+    }
+    repositories {
+        mavenLocal()
+        gradlePluginPortal()
+        maven {
+            url 'https://maven.pkg.jetbrains.space/public/p/compose/dev'
+        }
+    }
+}
+rootProject.name = 'jvmPreview'
+include(':jvm', ':mpp')

gradle-plugins/preview-rpc/src/main/kotlin/org/jetbrains/compose/desktop/ui/tooling/preview/rpc/RemoteConnection.kt

@@ -39,6 +39,7 @@ abstract class RemoteConnection : AutoCloseable {
     }
 }
 
+// Constructor is also used in GradlePluginTest#configurePreview via reflection
 internal class RemoteConnectionImpl(
     private val socket: Socket,
     private val log: PreviewLogger,

gradle-plugins/preview-rpc/src/main/kotlin/org/jetbrains/compose/desktop/ui/tooling/preview/rpc/commands.kt

@@ -63,7 +63,7 @@ data class ConfigFromGradle(
     val previewHostConfig: PreviewHostConfig
 )
 
-internal fun RemoteConnection.receiveConfigFromGradle(): ConfigFromGradle? {
+fun RemoteConnection.receiveConfigFromGradle(): ConfigFromGradle? {
     var previewClasspath: String? = null
     var previewFqName: String? = null
     var previewHostConfig: PreviewHostConfig? = null

idea-plugin/examples/desktop-project/README.md

@@ -1,5 +0,0 @@
-1. Run from `idea-plugin`:
-```
-./gradlew runIde
-```
-2. Open `idea-plugin/examples/desktop-project` with the test IDE.

idea-plugin/examples/desktop-project/build.gradle.kts

@@ -1,24 +0,0 @@
-import org.jetbrains.compose.compose
-
-plugins {
-    // __KOTLIN_COMPOSE_VERSION__
-    kotlin("jvm") version "1.5.21"
-    // __LATEST_COMPOSE_RELEASE_VERSION__
-    id("org.jetbrains.compose") version "0.5.0-build262"
-}
-
-repositories {
-    mavenCentral()
-    maven("https://maven.pkg.jetbrains.space/public/p/compose/dev")
-}
-
-dependencies {
-    implementation(compose.uiTooling)
-    implementation(compose.desktop.currentOs)
-}
-
-compose.desktop {
-    application {
-        mainClass = "MainKt"
-    }
-}

idea-plugin/examples/desktop-project/gradle.properties

@@ -1,4 +0,0 @@
-org.gradle.jvmargs=-Xmx2048m -Dfile.encoding=UTF-8
-kotlin.code.style=official
-#org.gradle.unsafe.configuration-cache=true
-#org.gradle.unsafe.configuration-cache-problems=warn

idea-plugin/examples/desktop-project/settings.gradle.kts

@@ -1,6 +0,0 @@
-pluginManagement {
-    repositories {
-        gradlePluginPortal()
-        maven("https://maven.pkg.jetbrains.space/public/p/compose/dev")
-    }
-}

idea-plugin/examples/simple-preview-example/README.md

@@ -0,0 +1,5 @@
+1. Run from `idea-plugin`:
+```
+./gradlew runIde
+```
+2. Open the project with the test IDE.

idea-plugin/examples/simple-preview-example/build.gradle.kts

@@ -0,0 +1,21 @@
+buildscript {
+    repositories {
+        mavenLocal()
+        mavenCentral()
+        maven("https://maven.pkg.jetbrains.space/public/p/compose/dev")
+    }
+    dependencies {
+        // __LATEST_COMPOSE_RELEASE_VERSION__
+        classpath("org.jetbrains.compose:compose-gradle-plugin:0.5.0-build262")
+        // __KOTLIN_COMPOSE_VERSION__
+        classpath(kotlin("gradle-plugin", version = "1.5.21"))
+    }
+}
+
+subprojects {
+    repositories {
+        mavenLocal()
+        mavenCentral()
+        maven("https://maven.pkg.jetbrains.space/public/p/compose/dev")
+    }
+}

idea-plugin/examples/simple-preview-example/gradle.properties

@@ -0,0 +1,3 @@
+kotlin.code.style=official
+#org.gradle.unsafe.configuration-cache=true
+#org.gradle.unsafe.configuration-cache-problems=warn

idea-plugin/examples/simple-preview-example/mpp-jvm/build.gradle.kts

@@ -0,0 +1,26 @@
+import org.jetbrains.compose.compose
+
+plugins {
+    kotlin("multiplatform")
+    id("org.jetbrains.compose")
+}
+
+kotlin {
+    jvm("desktop")
+
+    sourceSets {
+        named("commonMain") {
+            dependencies {
+                api(compose.runtime)
+                api(compose.foundation)
+                api(compose.material)
+                api(compose.uiTooling)
+            }
+        }
+        named("desktopMain") {
+            dependencies {
+                implementation(compose.desktop.currentOs)
+            }
+        }
+    }
+}

idea-plugin/examples/simple-preview-example/mpp-jvm/src/commonMain/kotlin/App.kt

@@ -0,0 +1,15 @@
+import androidx.compose.material.Button
+import androidx.compose.material.MaterialTheme
+import androidx.compose.material.Text
+import androidx.compose.runtime.*
+
+@Composable
+fun App() {
+    MaterialTheme {
+        Button(onClick = {}) {
+            Text("Hello, ${getPlatformName()}!")
+        }
+    }
+}
+
+expect fun getPlatformName(): String

idea-plugin/examples/simple-preview-example/mpp-jvm/src/desktopMain/kotlin/DesktopApp.kt

@@ -0,0 +1,10 @@
+import androidx.compose.runtime.*
+import androidx.compose.desktop.ui.tooling.preview.Preview
+
+actual fun getPlatformName(): String = "Desktop"
+
+@Preview
+@Composable
+fun DesktopAppPreview() {
+    App()
+}

idea-plugin/examples/simple-preview-example/pure-jvm/build.gradle.kts

@@ -0,0 +1,12 @@
+import org.jetbrains.compose.compose
+
+plugins {
+    kotlin("jvm")
+    id("org.jetbrains.compose")
+}
+
+dependencies {
+    implementation(compose.desktop.currentOs)
+    // todo: remove after update
+    implementation(compose.uiTooling)
+}

idea-plugin/examples/simple-preview-example/pure-jvm/src/main/kotlin/preview.kt

@@ -0,0 +1,16 @@
+import androidx.compose.material.Text
+import androidx.compose.material.Button
+import androidx.compose.runtime.*
+import androidx.compose.desktop.ui.tooling.preview.Preview
+
+@Preview
+@Composable
+fun ExamplePreview() {
+    var text by remember { mutableStateOf("Hello, World!") }
+
+    Button(onClick = {
+        text = "Hello, Desktop!"
+    }) {
+        Text(text)
+    }
+}

idea-plugin/examples/simple-preview-example/settings.gradle.kts

@@ -0,0 +1 @@
+include(":mpp-jvm", ":pure-jvm")

idea-plugin/src/main/kotlin/org/jetbrains/compose/desktop/ide/preview/ConfigurePreviewTaskNameCache.kt

@@ -0,0 +1,91 @@
+/*
+ * Copyright 2020-2021 JetBrains s.r.o. and respective authors and developers.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE.txt file.
+ */
+
+package org.jetbrains.compose.desktop.ide.preview
+
+import com.intellij.openapi.externalSystem.model.DataNode
+import com.intellij.openapi.externalSystem.model.ProjectKeys
+import com.intellij.openapi.externalSystem.model.project.ModuleData
+import com.intellij.openapi.externalSystem.service.project.ProjectDataManager
+import com.intellij.openapi.externalSystem.util.ExternalSystemApiUtil
+import com.intellij.openapi.module.Module
+import com.intellij.openapi.project.Project
+import com.intellij.util.concurrency.annotations.RequiresReadLock
+import org.jetbrains.kotlin.idea.configuration.KotlinTargetData
+import org.jetbrains.plugins.gradle.settings.GradleSettings
+import org.jetbrains.plugins.gradle.util.GradleConstants
+
+internal val DEFAULT_CONFIGURE_PREVIEW_TASK_NAME = "configureDesktopPreview"
+
+internal interface ConfigurePreviewTaskNameProvider {
+    @RequiresReadLock
+    fun configurePreviewTaskNameOrNull(module: Module):  String?
+}
+
+internal class ConfigurePreviewTaskNameProviderImpl : ConfigurePreviewTaskNameProvider {
+    @RequiresReadLock
+    override fun configurePreviewTaskNameOrNull(module: Module): String? {
+        val modulePath = ExternalSystemApiUtil.getExternalProjectPath(module) ?: return null
+        val moduleNode = moduleDataNodeOrNull(module.project, modulePath)
+        if (moduleNode != null) {
+            val target = ExternalSystemApiUtil.getChildren(moduleNode, KotlinTargetData.KEY).singleOrNull()
+            if (target != null) {
+                return previewTaskName(target.data.externalName)
+            }
+        }
+
+        return null
+    }
+
+    private fun previewTaskName(targetName: String = "") =
+        "$DEFAULT_CONFIGURE_PREVIEW_TASK_NAME${targetName.capitalize()}"
+
+    private fun moduleDataNodeOrNull(project: Project, modulePath: String): DataNode<ModuleData>? {
+        val projectDataManager = ProjectDataManager.getInstance()
+        for (settings in GradleSettings.getInstance(project).linkedProjectsSettings) {
+            val projectInfo = projectDataManager.getExternalProjectData(project, GradleConstants.SYSTEM_ID, settings.externalProjectPath)
+            val projectNode = projectInfo?.externalProjectStructure ?: continue
+            val moduleNodes = ExternalSystemApiUtil.getChildren(projectNode, ProjectKeys.MODULE)
+            for (moduleNode in moduleNodes) {
+                val externalProjectPath = moduleNode.data.linkedExternalProjectPath
+                if (externalProjectPath == modulePath) {
+                    return moduleNode
+                }
+            }
+        }
+        return null
+    }
+}
+
+internal class ConfigurePreviewTaskNameCache(
+    private val provider: ConfigurePreviewTaskNameProvider
+) : ConfigurePreviewTaskNameProvider {
+    private var cachedModuleId: String? = null
+    private var cachedTaskName: String? = null
+
+    @RequiresReadLock
+    override fun configurePreviewTaskNameOrNull(module: Module): String? {
+        val externalProjectPath = ExternalSystemApiUtil.getExternalProjectPath(module) ?: return null
+        val moduleId = "$externalProjectPath#${module.name}"
+
+        synchronized(this) {
+            if (moduleId == cachedModuleId) return cachedTaskName
+        }
+
+        val taskName = provider.configurePreviewTaskNameOrNull(module)
+        synchronized(this) {
+            cachedTaskName = taskName
+            cachedModuleId = moduleId
+        }
+        return taskName
+    }
+
+    fun invalidate() {
+        synchronized(this) {
+            cachedModuleId = null
+            cachedTaskName = null
+        }
+    }
+}

idea-plugin/src/main/kotlin/org/jetbrains/compose/desktop/ide/preview/PreviewActions.kt

@@ -57,7 +57,7 @@ private fun buildPreviewViaGradle(project: Project, previewLocation: PreviewLoca
     val settings = ExternalSystemTaskExecutionSettings()
     settings.executionName = "Preview: ${previewLocation.fqName}"
     settings.externalProjectPath = previewLocation.modulePath
-    settings.taskNames = listOf("configureDesktopPreview")
+    settings.taskNames = listOf(previewLocation.taskName)
     settings.vmOptions = gradleVmOptions
     settings.externalSystemIdString = GradleConstants.SYSTEM_ID.id
     val previewService = project.service<PreviewStateService>()

idea-plugin/src/main/kotlin/org/jetbrains/compose/desktop/ide/preview/PreviewLocation.kt

@@ -6,22 +6,23 @@
 package org.jetbrains.compose.desktop.ide.preview
 
 import com.intellij.openapi.externalSystem.util.ExternalSystemApiUtil
+import com.intellij.openapi.roots.ProjectFileIndex
 import com.intellij.util.concurrency.annotations.RequiresReadLock
-import org.jetbrains.kotlin.idea.util.projectStructure.module
+import org.jetbrains.kotlin.idea.debugger.getService
 import org.jetbrains.kotlin.psi.KtNamedFunction
 
-data class PreviewLocation(val fqName: String, val modulePath: String)
+data class PreviewLocation(val fqName: String, val modulePath: String, val taskName: String)
 
 @RequiresReadLock
 internal fun KtNamedFunction.asPreviewFunctionOrNull(): PreviewLocation? {
-    if (isValidComposablePreviewFunction()) {
-        val fqName = composePreviewFunctionFqn()
-        val module = module?.let { ExternalSystemApiUtil.getExternalProjectPath(it) }
-        if (module != null) {
-            return PreviewLocation(fqName = fqName, modulePath = module)
-        }
-    }
+    if (!isValidComposablePreviewFunction()) return null
 
-    return null
+    val fqName = composePreviewFunctionFqn()
+    val module = ProjectFileIndex.getInstance(project).getModuleForFile(containingFile.virtualFile)
+    if (module == null || module.isDisposed) return null
 
+    val service = project.getService<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

@@ -9,10 +9,15 @@ import com.intellij.notification.NotificationGroupManager
 import com.intellij.notification.NotificationType
 import com.intellij.openapi.Disposable
 import com.intellij.openapi.components.Service
+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
 import org.jetbrains.compose.desktop.ui.tooling.preview.rpc.*
+import org.jetbrains.kotlin.idea.framework.GRADLE_SYSTEM_ID
 import javax.swing.JComponent
 import javax.swing.event.AncestorEvent
 import javax.swing.event.AncestorListener
@@ -23,9 +28,22 @@ class PreviewStateService(private val myProject: Project) : Disposable {
     private val previewManager: PreviewManager = PreviewManagerImpl(previewListener)
     val gradleCallbackPort: Int
         get() = previewManager.gradleCallbackPort
+    private val configurePreviewTaskNameCache =
+        ConfigurePreviewTaskNameCache(ConfigurePreviewTaskNameProviderImpl())
+
+    init {
+        val projectRefreshListener = ConfigurePreviewTaskNameCacheInvalidator(configurePreviewTaskNameCache)
+        ExternalSystemProgressNotificationManager.getInstance()
+            .addNotificationListener(projectRefreshListener, myProject)
+    }
+
+    @RequiresReadLock
+    internal fun configurePreviewTaskNameOrNull(module: Module): String? =
+        configurePreviewTaskNameCache.configurePreviewTaskNameOrNull(module)
 
     override fun dispose() {
         previewManager.close()
+        configurePreviewTaskNameCache.invalidate()
     }
 
     internal fun registerPreviewPanels(
@@ -112,3 +130,20 @@ private class LoadingPanelUpdater(private val panel: JBLoadingPanel) : PreviewLi
         panel.stopLoading()
     }
 }
+
+// ExternalSystemTaskNotificationListenerAdapter is used,
+// because ExternalSystemTaskNotificationListener interface's API
+// was changed between 2020.3 and 2021.1, so a direct implementation
+// would not work with both 2020.3 and 2021.1
+private class ConfigurePreviewTaskNameCacheInvalidator(
+    private val configurePreviewTaskNameCache: ConfigurePreviewTaskNameCache
+) : ExternalSystemTaskNotificationListenerAdapter(null) {
+    override fun onStart(id: ExternalSystemTaskId, workingDir: String?) {
+        if (
+            id.projectSystemId == GRADLE_SYSTEM_ID &&
+            id.type == ExternalSystemTaskType.RESOLVE_PROJECT
+        ) {
+            configurePreviewTaskNameCache.invalidate()
+        }
+    }
+}