fix: slow file searches (fixes #652)

carlrobertoh · Aug 13, 2024 · 325dfc5 · 325dfc5
1 parent 1d7f3e8
commit 325dfc5
Show file tree

Hide file tree

Showing 6 changed files with 45 additions and 40 deletions.
diff --git a/src/main/kotlin/ee/carlrobert/codegpt/ui/textarea/CustomTextPaneKeyAdapter.kt b/src/main/kotlin/ee/carlrobert/codegpt/ui/textarea/CustomTextPaneKeyAdapter.kt
@@ -1,15 +1,12 @@
 package ee.carlrobert.codegpt.ui.textarea
 
-import com.intellij.openapi.application.runInEdt
 import com.intellij.openapi.components.service
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.util.TextRange
 import com.jetbrains.rd.util.AtomicReference
 import ee.carlrobert.codegpt.CodeGPTKeys
 import ee.carlrobert.codegpt.ui.textarea.suggestion.SuggestionsPopupManager
-import kotlinx.coroutines.CoroutineScope
-import kotlinx.coroutines.Dispatchers
-import kotlinx.coroutines.launch
+import kotlinx.coroutines.*
 import java.awt.event.KeyAdapter
 import java.awt.event.KeyEvent
 import javax.swing.text.StyleContext
@@ -22,6 +19,8 @@ class CustomTextPaneKeyAdapter(
     onWebSearchIncluded: () -> Unit
 ) : KeyAdapter() {
 
+    private var updateSuggestionsJob: Job? = null
+    private val scope = CoroutineScope(Dispatchers.Default + SupervisorJob())
     private val suggestionsPopupManager =
         SuggestionsPopupManager(project, textPane, onWebSearchIncluded)
     private val popupOpenedAtRange: AtomicReference<TextRange?> = AtomicReference(null)
@@ -100,15 +99,22 @@ class CustomTextPaneKeyAdapter(
     }
 
     private fun updateSuggestions() {
-        CoroutineScope(Dispatchers.Default).launch {
-            runInEdt {
-                val lastAtIndex = textPane.text.lastIndexOf('@')
+        updateSuggestionsJob?.cancel()
+        updateSuggestionsJob = scope.launch {
+            withContext(Dispatchers.Main) {
+                val text = textPane.text
+                val lastAtIndex = text.lastIndexOf('@')
                 if (lastAtIndex != -1) {
-                    val lastAtSearchIndex = textPane.text.lastIndexOf(':')
+                    val lastAtSearchIndex = text.lastIndexOf(':')
                     if (lastAtSearchIndex != -1) {
-                        val searchText = textPane.text.substring(lastAtSearchIndex + 1)
+                        val searchText = text.substring(lastAtSearchIndex + 1)
                         if (searchText.isNotEmpty()) {
-                            suggestionsPopupManager.updateSuggestions(searchText)
+                            launch {
+                                suggestionsPopupManager.updateSuggestions(
+                                    searchText,
+                                    updateSuggestionsJob
+                                )
+                            }
                         }
                     }
                 } else {

diff --git a/src/main/kotlin/ee/carlrobert/codegpt/ui/textarea/FileSearchService.kt b/src/main/kotlin/ee/carlrobert/codegpt/ui/textarea/FileSearchService.kt
@@ -6,19 +6,10 @@ import com.intellij.openapi.vfs.VirtualFile
 import ee.carlrobert.codegpt.CodeGPTKeys
 import ee.carlrobert.codegpt.ReferencedFile
 import ee.carlrobert.codegpt.actions.IncludeFilesInContextNotifier
-import ee.carlrobert.codegpt.util.file.FileUtil
-import kotlinx.coroutines.*
 import java.io.File
 
 @Service
 class FileSearchService private constructor(val project: Project) {
-    private val scope = CoroutineScope(Dispatchers.Default + SupervisorJob())
-
-    fun searchFiles(searchText: String): List<VirtualFile> = runBlocking {
-        withContext(scope.coroutineContext) {
-            FileUtil.searchProjectFiles(project, searchText)
-        }
-    }
 
     fun addFileToSession(file: VirtualFile) {
         val filesIncluded =

diff --git a/src/main/kotlin/ee/carlrobert/codegpt/ui/textarea/suggestion/SuggestionsPopupManager.kt b/src/main/kotlin/ee/carlrobert/codegpt/ui/textarea/suggestion/SuggestionsPopupManager.kt
@@ -1,14 +1,11 @@
 package ee.carlrobert.codegpt.ui.textarea.suggestion
 
-import com.intellij.openapi.application.runInEdt
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.ui.popup.JBPopup
 import com.intellij.vcsUtil.showAbove
 import ee.carlrobert.codegpt.ui.textarea.CustomTextPane
 import ee.carlrobert.codegpt.ui.textarea.suggestion.item.*
-import kotlinx.coroutines.Dispatchers
-import kotlinx.coroutines.runBlocking
-import kotlinx.coroutines.withContext
+import kotlinx.coroutines.*
 import java.awt.Dimension
 import java.awt.Point
 import javax.swing.DefaultListModel
@@ -21,6 +18,7 @@ class SuggestionsPopupManager(
     private val textPane: CustomTextPane,
     onWebSearchIncluded: () -> Unit,
 ) {
+    private val scope = CoroutineScope(Dispatchers.Default + SupervisorJob())
 
     private var selectedActionGroup: SuggestionGroupItem? = null
     private var popup: JBPopup? = null
@@ -75,13 +73,12 @@ class SuggestionsPopupManager(
         list.selectNext()
     }
 
-    fun updateSuggestions(searchText: String? = null) {
-        val suggestions = runBlocking {
-            withContext(Dispatchers.Default) {
-                selectedActionGroup?.getSuggestions(searchText) ?: emptyList()
-            }
+    suspend fun updateSuggestions(searchText: String? = null, updateSuggestionsJob: Job? = null) {
+        val suggestions = withContext(Dispatchers.Default) {
+            selectedActionGroup?.getSuggestions(searchText, updateSuggestionsJob) ?: emptyList()
         }
-        runInEdt {
+
+        withContext(Dispatchers.Main) {
             listModel.clear()
             listModel.addAll(suggestions)
             list.revalidate()
@@ -107,7 +104,9 @@ class SuggestionsPopupManager(
 
             is SuggestionGroupItem -> {
                 selectedActionGroup = item
-                updateSuggestions()
+                scope.launch {
+                    updateSuggestions()
+                }
                 textPane.appendHighlightedText(item.groupPrefix, withWhitespace = false)
                 textPane.requestFocus()
             }

diff --git a/src/main/kotlin/ee/carlrobert/codegpt/ui/textarea/suggestion/item/SuggestionGroupItems.kt b/src/main/kotlin/ee/carlrobert/codegpt/ui/textarea/suggestion/item/SuggestionGroupItems.kt
@@ -12,9 +12,10 @@ import ee.carlrobert.codegpt.settings.GeneralSettings
 import ee.carlrobert.codegpt.settings.documentation.DocumentationSettings
 import ee.carlrobert.codegpt.settings.service.ServiceType
 import ee.carlrobert.codegpt.ui.DocumentationDetails
-import ee.carlrobert.codegpt.ui.textarea.FileSearchService
 import ee.carlrobert.codegpt.util.ResourceUtil.getDefaultPersonas
+import ee.carlrobert.codegpt.util.file.FileUtil
 import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.Job
 import kotlinx.coroutines.withContext
 import java.time.Instant
 import java.time.format.DateTimeParseException
@@ -24,7 +25,7 @@ class FileSuggestionGroupItem(private val project: Project) : SuggestionGroupIte
     override val groupPrefix = "file:"
     override val icon = AllIcons.FileTypes.Any_type
 
-    override suspend fun getSuggestions(searchText: String?): List<SuggestionActionItem> {
+    override suspend fun getSuggestions(searchText: String?, updateSuggestionsJob: Job?): List<SuggestionActionItem> {
         if (searchText == null) {
             val projectFileIndex = project.service<ProjectFileIndex>()
             return readAction {
@@ -33,9 +34,7 @@ class FileSuggestionGroupItem(private val project: Project) : SuggestionGroupIte
                     .toFileSuggestions()
             }
         }
-        return project.service<FileSearchService>()
-            .searchFiles(searchText)
-            .toFileSuggestions()
+        return FileUtil.searchProjectFiles(project, searchText, job = updateSuggestionsJob).toFileSuggestions()
     }
 
     private fun Iterable<VirtualFile>.toFileSuggestions() = take(10).map { FileActionItem(it) }
@@ -48,7 +47,7 @@ class FolderSuggestionGroupItem(private val project: Project) : SuggestionGroupI
     override val groupPrefix = "folder:"
     override val icon = AllIcons.Nodes.Folder
 
-    override suspend fun getSuggestions(searchText: String?): List<SuggestionActionItem> {
+    override suspend fun getSuggestions(searchText: String?, updateSuggestionsJob: Job?): List<SuggestionActionItem> {
         if (searchText == null) {
             return getProjectFolders(project).toFolderSuggestions()
         }
@@ -83,7 +82,7 @@ class PersonaSuggestionGroupItem : SuggestionGroupItem {
     override val groupPrefix = "persona:"
     override val icon = AllIcons.General.User
 
-    override suspend fun getSuggestions(searchText: String?): List<SuggestionActionItem> =
+    override suspend fun getSuggestions(searchText: String?, updateSuggestionsJob: Job?): List<SuggestionActionItem> =
         getDefaultPersonas()
             .filter {
                 if (searchText.isNullOrEmpty()) {
@@ -102,7 +101,7 @@ class DocumentationSuggestionGroupItem : SuggestionGroupItem {
     override val icon = AllIcons.Toolwindows.Documentation
     override val enabled = GeneralSettings.getSelectedService() == ServiceType.CODEGPT
 
-    override suspend fun getSuggestions(searchText: String?): List<SuggestionActionItem> =
+    override suspend fun getSuggestions(searchText: String?, updateSuggestionsJob: Job?): List<SuggestionActionItem> =
         service<DocumentationSettings>().state.documentations
             .sortedByDescending { parseDateTime(it.lastUsedDateTime) }
             .filter {

diff --git a/src/main/kotlin/ee/carlrobert/codegpt/ui/textarea/suggestion/item/SuggestionItem.kt b/src/main/kotlin/ee/carlrobert/codegpt/ui/textarea/suggestion/item/SuggestionItem.kt
@@ -2,6 +2,7 @@ package ee.carlrobert.codegpt.ui.textarea.suggestion.item
 
 import com.intellij.openapi.project.Project
 import ee.carlrobert.codegpt.ui.textarea.CustomTextPane
+import kotlinx.coroutines.Job
 import javax.swing.Icon
 
 interface SuggestionItem {
@@ -18,5 +19,8 @@ interface SuggestionActionItem : SuggestionItem {
 interface SuggestionGroupItem : SuggestionItem {
     val groupPrefix: String
 
-    suspend fun getSuggestions(searchText: String? = null): List<SuggestionItem>
+    suspend fun getSuggestions(
+        searchText: String? = null,
+        updateSuggestionsJob: Job?
+    ): List<SuggestionItem>
 }
diff --git a/src/main/kotlin/ee/carlrobert/codegpt/util/file/FileUtil.kt b/src/main/kotlin/ee/carlrobert/codegpt/util/file/FileUtil.kt
@@ -15,6 +15,7 @@ import com.intellij.openapi.util.text.StringUtil
 import com.intellij.openapi.vfs.VirtualFile
 import com.intellij.openapi.vfs.VirtualFileFilter
 import ee.carlrobert.codegpt.settings.service.llama.LlamaSettings.getLlamaModelsPath
+import kotlinx.coroutines.Job
 import java.io.File
 import java.io.FileOutputStream
 import java.io.IOException
@@ -228,11 +229,16 @@ object FileUtil {
         project: Project,
         query: String,
         maxResults: Int = 6,
+        job: Job? = null
     ): List<VirtualFile> {
-        val results = mutableListOf<SearchResult>()
+        val results = TreeSet<SearchResult>(compareByDescending { it.score })
         val fileIndex = project.service<ProjectFileIndex>()
 
         fileIndex.iterateContent({ file ->
+            if (job != null && job.isCancelled) {
+                return@iterateContent false
+            }
+
             val score = calculateScore(file, query)
             if (score > 0) {
                 results.add(SearchResult(file, score))