refactor(rust): add ReadAction to improve performance

This commit adds the `ReadAction` class from the IntelliJ platform to improve performance in the RustTestService.kt file. The `ReadAction` class is used to wrap code that needs to be executed in a read action, ensuring that it runs on the UI thread and has read access to the PSI tree. This change improves the performance of finding or creating a test file in the RustTestService class. Additionally, the commit also includes some code refactoring and formatting improvements in the ClassContext.kt, RustVariableContextBuilder.kt, and RustCodeModifier.kt files.
unit-mesh · Jan 19, 2024 · 9312338 · 9312338
1 parent 5dfa594
commit 9312338
Show file tree

Hide file tree

Showing 4 changed files with 25 additions and 27 deletions.
diff --git a/rust/src/main/kotlin/cc/unitmesh/rust/context/RustVariableContextBuilder.kt b/rust/src/main/kotlin/cc/unitmesh/rust/context/RustVariableContextBuilder.kt
@@ -7,6 +7,7 @@ import com.intellij.psi.util.PsiTreeUtil
 import org.rust.lang.core.psi.RsFunction
 import org.rust.lang.core.psi.RsImplItem
 import org.rust.lang.core.psi.RsLetDecl
+import org.rust.lang.core.psi.RsNamedFieldDecl
 
 class RustVariableContextBuilder : VariableContextBuilder {
     override fun getVariableContext(
@@ -15,7 +16,7 @@ class RustVariableContextBuilder : VariableContextBuilder {
         withClassContext: Boolean,
         gatherUsages: Boolean
     ): VariableContext? {
-        if (psiElement !is RsLetDecl) return null
+        if (psiElement !is RsNamedFieldDecl) return null
 
         val text = psiElement.text
         val parentOfType = PsiTreeUtil.getParentOfType(psiElement, RsFunction::class.java, true)

diff --git a/rust/src/main/kotlin/cc/unitmesh/rust/provider/RustCodeModifier.kt b/rust/src/main/kotlin/cc/unitmesh/rust/provider/RustCodeModifier.kt
@@ -42,8 +42,8 @@ class RustCodeModifier : CodeModifier {
     }
 
     override fun insertClass(sourceFile: VirtualFile, project: Project, code: String): Boolean {
-        val psiFile = PsiManager.getInstance(project).findFile(sourceFile) ?: return false
         return WriteCommandAction.runWriteCommandAction<Boolean>(project) {
+            val psiFile = PsiManager.getInstance(project).findFile(sourceFile) ?: return@runWriteCommandAction false
             val document = psiFile.viewProvider.document!!
             document.insertString(document.textLength, code)
 

diff --git a/rust/src/main/kotlin/cc/unitmesh/rust/provider/RustTestService.kt b/rust/src/main/kotlin/cc/unitmesh/rust/provider/RustTestService.kt
@@ -6,6 +6,7 @@ import cc.unitmesh.devti.provider.context.TestFileContext
 import cc.unitmesh.rust.context.RustClassContextBuilder
 import cc.unitmesh.rust.context.RustMethodContextBuilder
 import com.intellij.execution.configurations.RunProfile
+import com.intellij.openapi.application.ReadAction
 import com.intellij.openapi.application.runReadAction
 import com.intellij.openapi.project.Project
 import com.intellij.psi.PsiElement
@@ -25,38 +26,34 @@ class RustTestService : WriteTestService() {
 
     override fun findOrCreateTestFile(sourceFile: PsiFile, project: Project, psiElement: PsiElement): TestFileContext? {
         val testable = psiElement is RsImplItem || psiElement is RsFunction
-        val element = if (!testable) {
-            when (val parent = psiElement.parent) {
-                is RsFunction -> parent
-                is RsImplItem -> parent
-                is RsStructItem -> parent
-                is RsEnumItem -> parent
-                else -> {
-                    runReadAction {
+        val element = ReadAction.compute<PsiElement, Throwable> {
+            if (!testable) {
+                when (val parent = psiElement.parent) {
+                    is RsFunction -> parent
+                    is RsImplItem -> parent
+                    is RsStructItem -> parent
+                    is RsEnumItem -> parent
+                    else -> {
                         PsiTreeUtil.getParentOfType(psiElement, RsFunction::class.java, RsImplItem::class.java)
                     }
                 }
-            }
-        } else {
-            psiElement
-        } ?: psiElement
-
-        val currentObject = when (element) {
-            is RsFunction -> {
-                runReadAction {
+            } else {
+                psiElement
+            } ?: psiElement
+        } ?: return null
+
+        val currentObject = ReadAction.compute<String, Throwable> {
+            return@compute when (element) {
+                is RsFunction -> {
                     RustMethodContextBuilder().getMethodContext(element, true, false)?.format()
                 }
-            }
 
-            is RsImplItem -> {
-                runReadAction {
-                    val classContext =
-                        RustClassContextBuilder().getClassContext(element, false) ?: return@runReadAction null
-                    classContext.format()
+                is RsImplItem -> {
+                    RustClassContextBuilder().getClassContext(element, false)?.format() ?: ""
                 }
-            }
 
-            else -> null
+                else -> null
+            }
         }
 
         val imports = PsiTreeUtil.getChildrenOfTypeAsList(sourceFile, RsUseItem::class.java).map {

diff --git a/src/main/kotlin/cc/unitmesh/devti/context/ClassContext.kt b/src/main/kotlin/cc/unitmesh/devti/context/ClassContext.kt
@@ -18,7 +18,7 @@ class ClassContext(
 ) : NamedElementContext(root, text, name) {
     private fun getFieldNames(): List<String> = fields.mapNotNull {
         val variableContextProvider = VariableContextProvider(false, false, false)
-        variableContextProvider.from(it).name
+        variableContextProvider.from(it).format()
     }
 
     private fun getMethodSignatures(): List<String> = methods.mapNotNull {