Merge pull request #570 from JetBrains/alex99999/fixes

Fix #569
JetBrains · Oct 23, 2024 · b11d090 · b11d090
2 parents 70d503c + cc45743
commit b11d090
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 10 deletions.
diff --git a/src/main/kotlin/org/arend/codeInsight/ArendTypedHandler.kt b/src/main/kotlin/org/arend/codeInsight/ArendTypedHandler.kt
@@ -9,9 +9,8 @@ import com.intellij.openapi.editor.Editor
 import com.intellij.openapi.project.Project
 import com.intellij.psi.PsiDocumentManager
 import com.intellij.psi.PsiFile
-import com.intellij.psi.util.elementType
-import com.intellij.refactoring.suggested.endOffset
-import com.intellij.refactoring.suggested.startOffset
+import com.intellij.psi.util.endOffset
+import com.intellij.psi.util.startOffset
 import com.intellij.util.text.CharArrayCharSequence
 import org.arend.psi.ArendElementTypes.*
 import org.arend.settings.ArendSettings
@@ -53,24 +52,32 @@ class ArendTypedHandler : TypedHandlerDelegate() {
         )
     }
 
-    private fun addParens(project: Project, editor: Editor, file: PsiFile) {
+    private fun addParensToGoal(project: Project, editor: Editor, file: PsiFile) {
         val element = file.findElementAt(editor.selectionModel.selectionStart) ?: return
         val document = editor.document
         PsiDocumentManager.getInstance(project).commitDocument(document)
         document.insertString(element.startOffset, "(")
         document.insertString(element.endOffset + 1, ")")
     }
 
+    private fun addBrackets(project: Project, editor: Editor) {
+        val document = editor.document
+        PsiDocumentManager.getInstance(project).commitDocument(document)
+        document.insertString(editor.selectionModel.selectionStart, "{")
+        document.insertString(editor.selectionModel.selectionEnd, "}")
+    }
+
     override fun beforeSelectionRemoved(c: Char, project: Project, editor: Editor, file: PsiFile): Result {
         val selectedText = editor.selectionModel.selectedText
         if (c == '(' && selectedText == "{?}") {
-            addParens(project, editor, file)
+            addParensToGoal(project, editor, file)
+            return Result.STOP
+        } else if (BRACKETS.contains(c.toString()) && BRACKETS.contains(selectedText)) {
+            changeCorrespondingBracket(c, project, editor, file)
+            return Result.CONTINUE
+        } else if (c == '{' && selectedText?.isNotEmpty() == true) {
+            addBrackets(project, editor)
             return Result.STOP
-        } else if (BRACKETS.contains(c.toString())) {
-            if (BRACKETS.contains(selectedText)) {
-                changeCorrespondingBracket(c, project, editor, file)
-                return Result.CONTINUE
-            }
         }
         return SelectionQuotingTypedHandler().beforeSelectionRemoved(c, project, editor, file)
     }

diff --git a/src/test/kotlin/org/arend/codeInsight/ArendTypedHandlerTest.kt b/src/test/kotlin/org/arend/codeInsight/ArendTypedHandlerTest.kt
@@ -33,4 +33,6 @@ class ArendTypedHandlerTest : ArendTestBase() {
     fun `test parens goal`() = check("""\func f => {-caret-}{?}""", """\func f => ({?})""", '(')
 
     fun `test parens goal 2`() = check("""\func f => {-caret-}{?}""", """\func f => ({?}""", '(', false)
+
+    fun `test braces 2`() = check("""\class Foo {\n  | foo : {-caret-}Nat\n}""", """\class Foo {\n  | foo : {Nat}\n}""", '{')
 }