From f598f97fb6bac2c66316bd07b304d0942c0e8fb6 Mon Sep 17 00:00:00 2001
From: Mikhail Mazurkevich <Mikhail.Mazurkevich@jetbrains.com>
Date: Mon, 1 Apr 2024 22:03:41 +0400
Subject: [PATCH] [workspace model] Make separate get folder for test source
 roots and group source roots

GitOrigin-RevId: a6c11b19052def30963454560f1db34e8ec95581
---
 .../src/WorkspaceModelGenerator.kt            | 23 ++++++++++++++++---
 1 file changed, 20 insertions(+), 3 deletions(-)

diff --git a/plugins/devkit/intellij.devkit.workspaceModel/src/WorkspaceModelGenerator.kt b/plugins/devkit/intellij.devkit.workspaceModel/src/WorkspaceModelGenerator.kt
index 3a1cf502cb68..35d71c2d404d 100644
--- a/plugins/devkit/intellij.devkit.workspaceModel/src/WorkspaceModelGenerator.kt
+++ b/plugins/devkit/intellij.devkit.workspaceModel/src/WorkspaceModelGenerator.kt
@@ -55,7 +55,7 @@ class WorkspaceModelGenerator(private val project: Project, private val coroutin
       LOG.info("Acceptable module source roots not found")
       return
     }
-    acceptedSourceRoots.map { sourceRoot ->
+    acceptedSourceRoots.forEach { sourceRoot ->
       withContext(Dispatchers.EDT) {
         DumbService.getInstance(project).completeJustSubmittedTasks() // Waiting for smart mode
         CodeWriter.generate(
@@ -64,16 +64,30 @@ class WorkspaceModelGenerator(private val project: Project, private val coroutin
           explicitApiEnabled = module.explicitApiEnabled,
           isTestModule = module.isTestModule, // TODO(It doesn't work for all modules)
           targetFolderGenerator = { createGeneratedSourceFolder(module, sourceRoot) },
-          existingTargetFolder = { sourceRoot.file?.parent?.children?.find { it.name == GENERATED_FOLDER_NAME } }
+          existingTargetFolder = { calculateExistingGenFolder(sourceRoot) }
         )
       }
     }
   }
 
+  private fun calculateExistingGenFolder(sourceFolder: SourceFolder): VirtualFile? {
+    return sourceFolder.file?.parent?.children?.find {
+      if (sourceFolder.isTestSource) {
+        it.name == TEST_GENERATED_FOLDER_NAME
+      } else {
+        it.name == GENERATED_FOLDER_NAME
+      }
+    }
+  }
+
   private fun createGeneratedSourceFolder(module: Module, sourceFolder: SourceFolder): VirtualFile? {
     // Create gen folder if it doesn't exist
     val generatedFolder = WriteAction.compute<VirtualFile, RuntimeException> {
-      VfsUtil.createDirectoryIfMissing(sourceFolder.file?.parent, GENERATED_FOLDER_NAME)
+      if (sourceFolder.isTestSource) {
+        VfsUtil.createDirectoryIfMissing(sourceFolder.file?.parent, TEST_GENERATED_FOLDER_NAME)
+      } else {
+        VfsUtil.createDirectoryIfMissing(sourceFolder.file?.parent, GENERATED_FOLDER_NAME)
+      }
     }
 
     val moduleRootManager = ModuleRootManager.getInstance(module)
@@ -112,6 +126,7 @@ class WorkspaceModelGenerator(private val project: Project, private val coroutin
     //}
     return contentEntries.flatMap { it.sourceFolders.asIterable() }.filter {
       if (it.file == null) return@filter false
+      if (it.rootType !is JavaSourceRootType)  return@filter false
       val javaSourceRootProperties = it.jpsElement.properties as? JavaSourceRootProperties
       if (javaSourceRootProperties == null) return@filter true
       return@filter !javaSourceRootProperties.isForGeneratedSources
@@ -131,6 +146,8 @@ class WorkspaceModelGenerator(private val project: Project, private val coroutin
   companion object {
     const val GENERATED_FOLDER_NAME = "gen"
 
+    const val TEST_GENERATED_FOLDER_NAME = "testGen"
+
     private const val RIDER_MODULES_PREFIX = "intellij.rider"
 
     val modulesWithAbstractTypes: Set<String> = setOf("intellij.platform.workspace.storage.testEntities")