Merge pull request #219 from xonixx/Incorrect-auto-insertion-of-}-on-…

…Enter-#216

Incorrect auto insertion of } on enter #216

src/main/java/intellij_awk/AwkEnterAfterUnmatchedBraceHandler.java

@@ -0,0 +1,65 @@
+package intellij_awk;
+
+import com.intellij.codeInsight.editorActions.enter.EnterAfterUnmatchedBraceHandler;
+import com.intellij.openapi.util.Pair;
+import com.intellij.psi.*;
+import com.intellij.util.text.CharArrayUtil;
+import intellij_awk.psi.AwkFile;
+import intellij_awk.psi.AwkStatement;
+import org.jetbrains.annotations.NotNull;
+
+public class AwkEnterAfterUnmatchedBraceHandler extends EnterAfterUnmatchedBraceHandler {
+  @Override
+  public boolean isApplicable(@NotNull PsiFile file, int caretOffset) {
+    return file instanceof AwkFile;
+  }
+
+  @Override
+  protected Pair<PsiElement, Integer> calculateOffsetToInsertClosingBrace(
+      @NotNull PsiFile file, @NotNull CharSequence text, int offset) {
+
+    String rest = text.subSequence(offset, text.length()).toString();
+
+    // The idea is that to understand where to put the closing `}` we need to recognize the longest
+    // statement that goes right after `{`. Since the AWK PSI tree can be broken due to incomplete
+    // code, the only reliable way to do so is to (re-)parse the rest of the file (after `{`) and
+    // take the first statement.
+    String code = "BEGIN{" + rest + "}";
+
+    AwkFile awkFile =
+        (AwkFile)
+            PsiFileFactory.getInstance(file.getProject())
+                .createFileFromText("dummy.awk", AwkFileType.INSTANCE, code);
+
+    AwkStatement awkStatement =
+        (AwkStatement) AwkUtil.findFirstMatchedDeep(awkFile, AwkStatement.class::isInstance);
+
+    if (awkStatement == null) {
+      return Pair.create(null, CharArrayUtil.shiftForwardUntil(text, offset, "\n"));
+    }
+
+    //    System.out.println("awkStatement="+awkStatement.getTextLength());
+    //    System.out.println("awkStatement="+awkStatement.getText());
+
+    int pos;
+    PsiElement possiblyComment = AwkUtil.getNextNotWhitespace(awkStatement);
+    if (possiblyComment instanceof PsiComment) {
+      // {<caret>print 123 # comment
+      // here    ^_______^ is statement
+      pos =
+          offset
+              + (possiblyComment.getTextRange().getEndOffset()
+                  - awkStatement.getTextRange().getStartOffset());
+    } else {
+      // why not just `offset + awkStatement.getTextLength();` ?
+      // it's because due to the peculiarity of grammar the recognized `if` statement can have the
+      // final newline as part of it
+      pos = offset + (awkStatement.getText().trim()).length();
+    }
+
+    //    System.out.println(
+    //        "res=" + text.subSequence(0, pos) + "<here>" + text.subSequence(pos, text.length()));
+
+    return Pair.create(null, pos);
+  }
+}

src/main/resources/META-INF/plugin.xml

@@ -108,6 +108,9 @@
         <codeInsight.parameterInfo language="AWK"
                                    implementationClass="intellij_awk.AwkParameterInfoHandler"/>
 
+        <enterHandlerDelegate implementation="intellij_awk.AwkEnterAfterUnmatchedBraceHandler" order="before afterUnmatchedBrace"/>
+
+
     </extensions>
 
     <actions>

src/test/java/intellij_awk/AwkBracesAndQuoteAutoCloseTests.java

@@ -73,6 +73,75 @@ public void testQuote1() {
         "BEGIN {\n  print \"<caret>\"\n  print \"\"}");
   }
 
+  public void testEnterCurlyBrace1() {
+    doTest(
+        '\n',
+        "function f() {\n    {<caret>print 123\n}",
+        "function f() {\n    {\n        <caret>print 123\n    }\n}");
+  }
+
+  public void testEnterCurlyBrace1_1() {
+    doTest(
+        '\n',
+        "function f() {\n    {<caret>print 123   # comment\n}",
+        "function f() {\n    {\n        <caret>print 123   # comment\n    }\n}");
+  }
+
+  public void testEnterCurlyBrace2() {
+    doTest(
+        '\n',
+        "function f() {\n    if(1){<caret>print 123\n}",
+        "function f() {\n    if(1){\n        <caret>print 123\n    }\n}");
+  }
+
+  public void testEnterCurlyBrace3() {
+    doTest(
+        '\n',
+        "function f() {\n    while(1){<caret>print 123\n}",
+        "function f() {\n    while(1){\n        <caret>print 123\n    }\n}");
+  }
+
+  public void testEnterCurlyBrace4() {
+    doTest(
+        '\n',
+        "function f() {\n    for (;;) {<caret>f(1)\n}",
+            "function f() {\n    for (;;) {\n        <caret>f(1)\n    }\n}");
+  }
+  public void testEnterCurlyBrace5() {
+    doTest('\n', "{<caret>print 123", "{\n    <caret>print 123\n}");
+  }
+  public void testEnterCurlyBrace6() {
+    doTest('\n', "BEGIN{<caret>print 123", "BEGIN{\n    <caret>print 123\n}");
+  }
+  public void testEnterCurlyBrace7() {
+    doTest('\n', "NF {<caret>print 123", "NF {\n    <caret>print 123\n}");
+  }
+  public void testEnterCurlyBrace8() {
+    doTest('\n', "function f()\n{<caret>print 123", "function f()\n{\n    <caret>print 123\n}");
+  }
+  public void testEnterCurlyBrace9_1() {
+    doTestEnterCurlyBraceComplex("if (2)");
+  }
+  public void testEnterCurlyBrace9_2() {
+    doTestEnterCurlyBraceComplex("while (2)");
+  }
+  public void testEnterCurlyBrace9_3() {
+    doTestEnterCurlyBraceComplex("for(;;)");
+  }
+  public void testEnterCurlyBrace10() {
+    doTest(
+        '\n',
+        "function f() {\n    if(1){<caret>\n}",
+        "function f() {\n    if(1){\n        <caret>\n    }\n}");
+  }
+
+  private void doTestEnterCurlyBraceComplex(String opener) {
+    doTest(
+        '\n',
+        "function f() {\n    if (1) {<caret>" + opener + " {\n        print 123}\n}",
+        "function f() {\n    if (1) {\n        "+opener+" {\n            print 123}\n    }\n}");
+  }
+
   private void doTest(char brace, String code, String expectedCode) {
     myFixture.configureByText("a.awk", code);
     myFixture.type(brace);