From e263419088198ef2b2753dfb0852c410aa295b98 Mon Sep 17 00:00:00 2001
From: Yaroslav Chernyshev <yaroslav.chernyshev@jetbrains.com>
Date: Thu, 23 Jul 2020 19:42:48 +0300
Subject: [PATCH] Populate Kotlin source set roots with KotlinSourceRootTypes

By default Kotlin source root is populated with JavaSourceRootType and
changed during import only if real source folder exists.
Now Kotlin source roots has valid source type right after creation

#KT-37106 Fixed
---
 .../KotlinSourceSetDataService.kt             |  2 ++
 .../kotlin/idea/roots/projectRootUtils.kt     | 33 +++++++++++++++++++
 2 files changed, 35 insertions(+)

diff --git a/idea/idea-gradle/src/org/jetbrains/kotlin/idea/configuration/KotlinSourceSetDataService.kt b/idea/idea-gradle/src/org/jetbrains/kotlin/idea/configuration/KotlinSourceSetDataService.kt
index a0d5119756d..ac5fa767a3a 100644
--- a/idea/idea-gradle/src/org/jetbrains/kotlin/idea/configuration/KotlinSourceSetDataService.kt
+++ b/idea/idea-gradle/src/org/jetbrains/kotlin/idea/configuration/KotlinSourceSetDataService.kt
@@ -30,6 +30,7 @@ import org.jetbrains.kotlin.idea.inspections.gradle.findAll
 import org.jetbrains.kotlin.idea.inspections.gradle.findKotlinPluginVersion
 import org.jetbrains.kotlin.idea.platform.IdePlatformKindTooling
 import org.jetbrains.kotlin.idea.roots.migrateNonJvmSourceFolders
+import org.jetbrains.kotlin.idea.roots.populateNonJvmSourceRootTypes
 import org.jetbrains.kotlin.platform.TargetPlatform
 import org.jetbrains.kotlin.platform.impl.JvmIdePlatformKind
 import org.jetbrains.kotlin.platform.impl.NativeIdePlatformKind
@@ -61,6 +62,7 @@ class KotlinSourceSetDataService : AbstractProjectDataService<GradleSourceSetDat
 
             if (platform.platforms.any { it != KotlinPlatform.JVM && it != KotlinPlatform.ANDROID }) {
                 migrateNonJvmSourceFolders(rootModel)
+                populateNonJvmSourceRootTypes(nodeToImport, ideModule)
             }
 
             configureFacet(sourceSetData, kotlinSourceSet, mainModuleData, ideModule, modelsProvider)?.let { facet ->
diff --git a/idea/src/org/jetbrains/kotlin/idea/roots/projectRootUtils.kt b/idea/src/org/jetbrains/kotlin/idea/roots/projectRootUtils.kt
index 8ac32cd8b7a..221d492e0bd 100644
--- a/idea/src/org/jetbrains/kotlin/idea/roots/projectRootUtils.kt
+++ b/idea/src/org/jetbrains/kotlin/idea/roots/projectRootUtils.kt
@@ -5,6 +5,12 @@
 
 package org.jetbrains.kotlin.idea.roots
 
+import com.intellij.openapi.externalSystem.model.DataNode
+import com.intellij.openapi.externalSystem.model.ProjectKeys
+import com.intellij.openapi.externalSystem.model.project.ContentRootData
+import com.intellij.openapi.externalSystem.model.project.ExternalSystemSourceType
+import com.intellij.openapi.externalSystem.service.project.manage.SourceFolderManager
+import com.intellij.openapi.externalSystem.util.ExternalSystemApiUtil
 import com.intellij.openapi.module.Module
 import com.intellij.openapi.module.ModuleManager
 import com.intellij.openapi.project.Project
@@ -13,11 +19,14 @@ import com.intellij.openapi.roots.ModifiableRootModel
 import com.intellij.openapi.roots.ModuleRootManager
 import com.intellij.openapi.roots.ProjectRootManager
 import com.intellij.openapi.roots.SourceFolder
+import com.intellij.openapi.util.io.FileUtil
 import com.intellij.openapi.vfs.NonPhysicalFileSystem
 import com.intellij.openapi.vfs.VirtualFile
+import com.intellij.openapi.vfs.VirtualFileManager
 import com.intellij.psi.PsiCodeFragment
 import com.intellij.psi.PsiFile
 import com.intellij.util.containers.ContainerUtil
+import com.intellij.util.io.URLUtil
 import org.jetbrains.jps.model.JpsElement
 import org.jetbrains.jps.model.ex.JpsElementBase
 import org.jetbrains.jps.model.java.JavaModuleSourceRootTypes
@@ -28,6 +37,7 @@ import org.jetbrains.jps.model.module.JpsModuleSourceRoot
 import org.jetbrains.jps.model.module.JpsModuleSourceRootType
 import org.jetbrains.kotlin.config.*
 import org.jetbrains.kotlin.idea.framework.KotlinSdkType
+import org.jetbrains.plugins.gradle.model.data.GradleSourceSetData
 import java.util.*
 
 private fun JpsModuleSourceRoot.getOrCreateProperties() =
@@ -59,6 +69,29 @@ fun migrateNonJvmSourceFolders(modifiableRootModel: ModifiableRootModel) {
     KotlinSdkType.setUpIfNeeded()
 }
 
+private val ContentRootData.SourceRoot.pathAsUrl
+    get() = VirtualFileManager.constructUrl(URLUtil.FILE_PROTOCOL, FileUtil.toSystemIndependentName(path))
+
+fun populateNonJvmSourceRootTypes(sourceSetNode: DataNode<GradleSourceSetData>, module: Module) {
+    val sourceFolderManager = SourceFolderManager.getInstance(module.project)
+    val contentRootDataNodes = ExternalSystemApiUtil.findAll(sourceSetNode, ProjectKeys.CONTENT_ROOT)
+    val contentRootDataList = contentRootDataNodes.mapNotNull { it.data }
+    if (contentRootDataList.isEmpty()) return
+
+    val externalToKotlinSourceTypes = mapOf(
+        ExternalSystemSourceType.SOURCE to SourceKotlinRootType,
+        ExternalSystemSourceType.TEST to TestSourceKotlinRootType
+    )
+    externalToKotlinSourceTypes.forEach { (externalType, kotlinType) ->
+        val sourcesRoots = contentRootDataList.flatMap { it.getPaths(externalType) }
+        sourcesRoots.forEach {
+            if (!FileUtil.exists(it.path)) {
+                sourceFolderManager.addSourceFolder(module, it.pathAsUrl, kotlinType)
+            }
+        }
+    }
+}
+
 fun getKotlinAwareDestinationSourceRoots(project: Project): List<VirtualFile> {
     return ModuleManager.getInstance(project).modules.flatMap { it.collectKotlinAwareDestinationSourceRoots() }
 }