refactor(rust): modify RustTestService and VariableContext

- Modify RustTestService in order to improve readability and maintainability. - Refactor VariableContext to remove unnecessary comments and improve code structure.
unit-mesh · Jan 19, 2024 · dd14143 · dd14143
1 parent 9312338
commit dd14143
Show file tree

Hide file tree

Showing 4 changed files with 27 additions and 24 deletions.
diff --git a/rust/src/233/main/kotlin/cc/unitmesh/rust/context/RustClassContextBuilder.kt b/rust/src/233/main/kotlin/cc/unitmesh/rust/context/RustClassContextBuilder.kt
@@ -7,6 +7,7 @@ import com.intellij.psi.util.PsiTreeUtil
 import org.rust.lang.core.psi.RsEnumItem
 import org.rust.lang.core.psi.RsFunction
 import org.rust.lang.core.psi.RsImplItem
+import org.rust.lang.core.psi.RsMembers
 import org.rust.lang.core.psi.RsStructItem
 import org.rust.lang.core.psi.ext.RsStructOrEnumItemElement
 import org.rust.lang.core.psi.ext.expandedFields
@@ -38,7 +39,9 @@ class RustClassContextBuilder : ClassContextBuilder {
 
             is RsImplItem -> {
                 val structItem = psiElement.implementingType?.item ?: return null
-                val functions = PsiTreeUtil.getChildrenOfTypeAsList(psiElement, RsFunction::class.java)
+                val functions = PsiTreeUtil.getChildrenOfTypeAsList(psiElement, RsMembers::class.java)
+                    .flatMap { PsiTreeUtil.getChildrenOfTypeAsList(it, RsFunction::class.java) }
+
                 val fields = when (structItem) {
                     is RsStructItem -> structItem.expandedFields
                     is RsEnumItem -> structItem.enumBody?.enumVariantList?.map { it } ?: emptyList()
@@ -48,7 +51,7 @@ class RustClassContextBuilder : ClassContextBuilder {
                 return ClassContext(
                     psiElement,
                     psiElement.text,
-                    psiElement.name,
+                    structItem.name,
                     functions,
                     fields,
                     emptyList(),

diff --git a/rust/src/main/kotlin/cc/unitmesh/rust/provider/RustCodeModifier.kt b/rust/src/main/kotlin/cc/unitmesh/rust/provider/RustCodeModifier.kt
@@ -2,6 +2,7 @@ package cc.unitmesh.rust.provider
 
 import cc.unitmesh.devti.context.builder.CodeModifier
 import com.intellij.lang.Language
+import com.intellij.openapi.application.runReadAction
 import com.intellij.openapi.command.WriteCommandAction
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.vfs.VirtualFile
@@ -19,7 +20,9 @@ class RustCodeModifier : CodeModifier {
         // 1. check mod test exits
         val psiFile = PsiManager.getInstance(project).findFile(sourceFile) ?: return false
         val rsFile = psiFile as? RsFile ?: return false
-        val testMod = PsiTreeUtil.findChildOfType(rsFile, RsModItem::class.java)
+        val testMod = runReadAction {
+            PsiTreeUtil.findChildOfType(rsFile, RsModItem::class.java)
+        }
         if (testMod == null) {
             insertClass(sourceFile, project, code)
         } else {

diff --git a/rust/src/main/kotlin/cc/unitmesh/rust/provider/RustTestService.kt b/rust/src/main/kotlin/cc/unitmesh/rust/provider/RustTestService.kt
@@ -7,15 +7,13 @@ 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
 import com.intellij.psi.PsiFile
 import com.intellij.psi.util.PsiTreeUtil
 import org.rust.cargo.runconfig.command.CargoCommandConfiguration
 import org.rust.lang.RsLanguage
 import org.rust.lang.core.psi.*
-import org.rust.lang.core.psi.ext.implementingType
 
 class RustTestService : WriteTestService() {
     override fun runConfigurationClass(project: Project): Class<out RunProfile> = CargoCommandConfiguration::class.java
@@ -26,9 +24,10 @@ class RustTestService : WriteTestService() {
 
     override fun findOrCreateTestFile(sourceFile: PsiFile, project: Project, psiElement: PsiElement): TestFileContext? {
         val testable = psiElement is RsImplItem || psiElement is RsFunction
+        var elementName = sourceFile.name
         val element = ReadAction.compute<PsiElement, Throwable> {
             if (!testable) {
-                when (val parent = psiElement.parent) {
+                val value = when (val parent = psiElement.parent) {
                     is RsFunction -> parent
                     is RsImplItem -> parent
                     is RsStructItem -> parent
@@ -37,10 +36,20 @@ class RustTestService : WriteTestService() {
                         PsiTreeUtil.getParentOfType(psiElement, RsFunction::class.java, RsImplItem::class.java)
                     }
                 }
+
+                elementName = when (value) {
+                    is RsFunction -> value.name
+                    is RsImplItem -> value.name
+                    is RsStructItem -> value.name
+                    is RsEnumItem -> value.name
+                    else -> null
+                } ?: ""
+
+                return@compute value
             } else {
-                psiElement
-            } ?: psiElement
-        } ?: return null
+                return@compute psiElement
+            }
+        } ?: psiElement
 
         val currentObject = ReadAction.compute<String, Throwable> {
             return@compute when (element) {
@@ -66,7 +75,7 @@ class RustTestService : WriteTestService() {
             false,
             sourceFile.virtualFile,
             relevantClasses,
-            "",
+            elementName,
             RsLanguage,
             currentObject,
             imports

diff --git a/src/main/kotlin/cc/unitmesh/devti/context/VariableContext.kt b/src/main/kotlin/cc/unitmesh/devti/context/VariableContext.kt
@@ -10,7 +10,7 @@ class VariableContext(
     override val text: String,
     override val name: String?,
     val enclosingMethod: PsiElement? = null,
-    val enclosingClass: PsiElement?= null,
+    val enclosingClass: PsiElement? = null,
     val usages: List<PsiReference> = emptyList(),
     val includeMethodContext: Boolean = false,
     val includeClassContext: Boolean = false
@@ -34,21 +34,9 @@ class VariableContext(
         }
     }
 
-    /**
-     * Returns a formatted string representation of the method.
-     *
-     * The returned string includes the following information:
-     * - The name of the method, or "_" if the name is null.
-     * - The name of the method's context, or "_" if the context is null.
-     * - The name of the class's context, or "_" if the context is null.
-     *
-     * @return A formatted string representation of the method.
-     */
     override fun format(): String {
         return """
-            var name: ${name ?: "_"}
-            var method name: ${methodContext?.name ?: "_"}
-            var class name: ${classContext?.name ?: "_"}
+            'variable -> ${root.text}
         """.trimIndent()
     }
 }