From cc457439aac61f7ef15bdc6e52e7d2812a27def8 Mon Sep 17 00:00:00 2001
From: Aleksei Luchinin <alex999990009@gmail.com>
Date: Wed, 23 Oct 2024 14:01:55 +0200
Subject: [PATCH] Fix #569

---
 .../arend/codeInsight/ArendTypedHandler.kt    | 27 ++++++++++++-------
 .../codeInsight/ArendTypedHandlerTest.kt      |  2 ++
 2 files changed, 19 insertions(+), 10 deletions(-)

diff --git a/src/main/kotlin/org/arend/codeInsight/ArendTypedHandler.kt b/src/main/kotlin/org/arend/codeInsight/ArendTypedHandler.kt
index 9a07fac87..6893d840b 100644
--- 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,7 +52,7 @@ 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)
@@ -61,16 +60,24 @@ class ArendTypedHandler : TypedHandlerDelegate() {
         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
index eb45eddf7..8e2633449 100644
--- 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}""", '{')
 }