From 3d7239a3dda459128bd2d2258545178a9490c134 Mon Sep 17 00:00:00 2001
From: Carl-Robert Linnupuu <carlrobertoh@gmail.com>
Date: Fri, 26 Jul 2024 14:39:04 +0300
Subject: [PATCH] refactor: improve folder suggestions update strategy

---
 .../ui/textarea/SuggestionUpdateStrategy.kt   | 35 ++++++-------------
 1 file changed, 11 insertions(+), 24 deletions(-)

diff --git a/src/main/kotlin/ee/carlrobert/codegpt/ui/textarea/SuggestionUpdateStrategy.kt b/src/main/kotlin/ee/carlrobert/codegpt/ui/textarea/SuggestionUpdateStrategy.kt
index fc044354..b9877da0 100644
--- a/src/main/kotlin/ee/carlrobert/codegpt/ui/textarea/SuggestionUpdateStrategy.kt
+++ b/src/main/kotlin/ee/carlrobert/codegpt/ui/textarea/SuggestionUpdateStrategy.kt
@@ -1,23 +1,20 @@
 package ee.carlrobert.codegpt.ui.textarea
 
 import com.intellij.openapi.application.readAction
-import com.intellij.openapi.application.runReadAction
 import com.intellij.openapi.components.service
 import com.intellij.openapi.fileEditor.FileEditorManager
 import com.intellij.openapi.project.Project
+import com.intellij.openapi.roots.ContentIterator
 import com.intellij.openapi.roots.ProjectFileIndex
-import com.intellij.openapi.vfs.VirtualFileManager
+import com.intellij.openapi.vfs.VirtualFile
 import ee.carlrobert.codegpt.util.ResourceUtil
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.launch
 import kotlinx.coroutines.withContext
 import java.io.File
-import java.nio.file.Files
 import java.nio.file.Path
 import javax.swing.DefaultListModel
-import kotlin.io.path.absolutePathString
-import kotlin.io.path.isDirectory
 
 interface SuggestionUpdateStrategy {
     fun populateSuggestions(
@@ -96,33 +93,23 @@ class FolderSuggestionActionStrategy : SuggestionUpdateStrategy {
         }
     }
 
-    private suspend fun findProjectFolders(project: Project): List<String> {
-        val projectRoot = project.basePath?.let { Path.of(it) } ?: return emptyList()
-        val projectFileIndex = project.service<ProjectFileIndex>()
-        val virtualFileManager = VirtualFileManager.getInstance()
-
-        return withContext(Dispatchers.IO) {
+    private suspend fun findProjectFolders(project: Project): List<String> =
+        withContext(Dispatchers.IO) {
             val uniqueFolders = mutableSetOf<String>()
-            Files.walk(projectRoot)
-                .filter { path ->
-                    val file = virtualFileManager.findFileByNioPath(path)
-                    val isProjectFile =
-                        file != null && runReadAction {
-                            projectFileIndex.isInSourceContent(file)
-                                    || projectFileIndex.isInTestSourceContent(file)
-                        }
-                    path.isDirectory() && !path.startsWith(".") && isProjectFile
-                }
-                .forEach { folder ->
-                    val folderPath = folder.absolutePathString()
+            val iterator = ContentIterator { file: VirtualFile ->
+                if (file.isDirectory && !file.name.startsWith(".")) {
+                    val folderPath = file.path
                     if (uniqueFolders.none { it.startsWith(folderPath) }) {
                         uniqueFolders.removeAll { it.startsWith(folderPath) }
                         uniqueFolders.add(folderPath)
                     }
                 }
+                true
+            }
+
+            project.service<ProjectFileIndex>().iterateContent(iterator)
             uniqueFolders.toList()
         }
-    }
 }
 
 class PersonaSuggestionActionStrategy : SuggestionUpdateStrategy {