feat(rust): fix for insert issue

unit-mesh · Jan 1, 2024 · 9ce7734 · 9ce7734
1 parent f452428
commit 9ce7734
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 9 deletions.
diff --git a/rust/src/main/kotlin/cc/unitmesh/rust/provider/RustLivingDocumentation.kt b/rust/src/main/kotlin/cc/unitmesh/rust/provider/RustLivingDocumentation.kt
@@ -3,14 +3,19 @@ package cc.unitmesh.rust.provider
 import cc.unitmesh.devti.custom.document.LivingDocumentationType
 import cc.unitmesh.devti.provider.LivingDocumentation
 import com.intellij.codeInsight.daemon.impl.CollectHighlightsUtil
+import com.intellij.openapi.command.WriteCommandAction
 import com.intellij.openapi.editor.Editor
 import com.intellij.openapi.editor.SelectionModel
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiNameIdentifierOwner
+import com.intellij.psi.codeStyle.CodeStyleManager
 import com.intellij.psi.util.PsiTreeUtil
+import com.intellij.util.IncorrectOperationException
 import org.rust.lang.core.psi.RsFunction
 import org.rust.lang.core.psi.ext.RsNameIdentifierOwner
 import org.rust.lang.core.psi.ext.RsStructOrEnumItemElement
+import org.rust.lang.doc.psi.RsDocComment
+import org.rust.lang.doc.psi.ext.containingDoc
 
 class RustLivingDocumentation : LivingDocumentation {
     override val forbiddenRules: List<String>
@@ -22,13 +27,39 @@ class RustLivingDocumentation : LivingDocumentation {
 
     override fun updateDoc(target: PsiElement, newDoc: String, type: LivingDocumentationType, editor: Editor) {
         val project = target.project
-        val file = target.containingFile
-        val startOffset = target.textRange.startOffset
-        val newEndOffset = startOffset + newDoc.length
+        val codeStyleManager = CodeStyleManager.getInstance(project)
+        WriteCommandAction.runWriteCommandAction(project, "Living Document", "cc.unitmesh.livingDoc", {
+            val startOffset = target.textRange.startOffset
+            val newEndOffset = startOffset + newDoc.length
 
-        editor.document.insertString(startOffset, newDoc)
-        val codeStyleManager = com.intellij.psi.codeStyle.CodeStyleManager.getInstance(project)
-        codeStyleManager.reformatText(file, startOffset, newEndOffset)
+           when (type) {
+                LivingDocumentationType.COMMENT -> {
+                    val psiElementFactory = org.rust.lang.core.psi.RsPsiFactory(project)
+                    val newDocComment = psiElementFactory.createBlockComment(newDoc)
+
+                    if (target is RsDocComment) {
+                        val oldDocComment = target.containingDoc
+                        if (oldDocComment != null) {
+                            oldDocComment.replace(newDocComment)
+                        } else {
+                            target.addBefore(newDocComment, target.firstChild)
+                        }
+                    } else {
+                        target.addBefore(newDocComment, target.firstChild)
+                    }
+                }
+
+                LivingDocumentationType.ANNOTATED -> {
+                    editor.document.insertString(startOffset, newDoc)
+                    codeStyleManager.reformatText(target.containingFile, startOffset, newEndOffset)
+                }
+
+                LivingDocumentationType.CUSTOM -> {
+                    editor.document.insertString(startOffset, newDoc)
+                    codeStyleManager.reformatText(target.containingFile, startOffset, newEndOffset)
+                }
+            }
+        })
     }
 
     override fun findNearestDocumentationTarget(psiElement: PsiElement): PsiNameIdentifierOwner? {

diff --git a/src/main/kotlin/cc/unitmesh/devti/intentions/action/task/LivingDocPromptBuilder.kt b/src/main/kotlin/cc/unitmesh/devti/intentions/action/task/LivingDocPromptBuilder.kt
@@ -63,12 +63,12 @@ open class LivingDocPromptBuilder(
 
     private fun classInstruction(context: ClassContext): String? {
         if (context.name == null) return null
-        return "Write documentation for given class " + context.name + ". You should not add documentation for methods."
+        return "Write documentation for given ${context.root.language} language class " + context.name + "."
     }
 
     private fun methodInstruction(context: MethodContext): String? {
         if (context.name == null) return null
-        var instruction = "Write documentation for given method " + context.name + "."
+        var instruction = "Write documentation for given ${context.root.language} language method " + context.name + "."
         if (context.paramNames.isNotEmpty()) {
             instruction = """
                 $instruction
@@ -101,7 +101,7 @@ open class LivingDocPromptBuilder(
                 }
 
                 contextInstruction(llmQueryContext)
-            } ?: "Write documentation for given code."
+            } ?: "Write documentation for given ${target.language} language code."
 
             instruction.append(basicInstruction)