Incorrect auto-insertion of } on Enter #216

xonixx · Jul 7, 2024 · e521242 · e521242
1 parent 9a48618
commit e521242
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 17 deletions.
diff --git a/src/main/java/intellij_awk/AwkEnterAfterUnmatchedBraceHandler.java b/src/main/java/intellij_awk/AwkEnterAfterUnmatchedBraceHandler.java
@@ -18,35 +18,47 @@ public boolean isApplicable(@NotNull PsiFile file, int caretOffset) {
   protected Pair<PsiElement, Integer> calculateOffsetToInsertClosingBrace(
       @NotNull PsiFile file, @NotNull CharSequence text, int offset) {
 
-    System.out.println("text="+text);
-    System.out.println("offset="+offset);
-
     String rest = text.subSequence(offset, text.length()).toString();
-    System.out.println("rest="+rest);
-
-    String code = "BEGIN{"+rest+"}";
 
-    System.out.println("code="+code);
+    // 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)
+    AwkFile awkFile =
+        (AwkFile)
             PsiFileFactory.getInstance(file.getProject())
-                    .createFileFromText("dummy.awk", AwkFileType.INSTANCE, code);
+                .createFileFromText("dummy.awk", AwkFileType.INSTANCE, code);
 
-    AwkStatement awkStatement = (AwkStatement) AwkUtil.findFirstMatchedDeep(awkFile, AwkStatement.class::isInstance);
+    AwkStatement awkStatement =
+        (AwkStatement) AwkUtil.findFirstMatchedDeep(awkFile, AwkStatement.class::isInstance);
 
     if (awkStatement == null) {
       return Pair.create(null, CharArrayUtil.shiftForwardUntil(text, offset, "\n"));
     }
 
-    System.out.println("zzz:"+CharArrayUtil.shiftForwardUntil(text, offset, "\n")+":"+(offset + awkStatement.getTextLength()+1));
-    System.out.println("awkStatement="+awkStatement.getTextLength());
-    System.out.println("awkStatement="+awkStatement.getText());
+    //    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();
+    }
 
-//    int pos = offset + awkStatement.getTextLength();
-    int pos = offset + (awkStatement.getText().trim()).length();
-
-    System.out.println("res="+text.subSequence(0, pos) + "<here>" + text.subSequence(pos, text.length()));
+    //    System.out.println(
+    //        "res=" + text.subSequence(0, pos) + "<here>" + text.subSequence(pos, text.length()));
 
     return Pair.create(null, pos);
   }

diff --git a/src/test/java/intellij_awk/AwkBracesAndQuoteAutoCloseTests.java b/src/test/java/intellij_awk/AwkBracesAndQuoteAutoCloseTests.java
@@ -80,6 +80,13 @@ public void testEnterCurlyBrace1() {
         "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',