[java-action] IDEA-356782 Incorrect place of caret into text block with injection after enter

- introduced JavaEnterInInjectedTextBlockHandler

GitOrigin-RevId: 7085ebb8d23177b11f49da17df2eb7df98f39b8a

java/java-impl/src/META-INF/JavaPlugin.xml

@@ -1439,7 +1439,11 @@
                           order="before afterUnmatchedBrace"/>
     <lang.smartEnterProcessor language="JAVA"
                               implementationClass="com.intellij.codeInsight.editorActions.smartEnter.JavaSmartEnterProcessor"/>
-    <enterHandlerDelegate implementation="com.intellij.codeInsight.editorActions.JavaEnterInTextBlockHandler"/>
+    <enterHandlerDelegate implementation="com.intellij.codeInsight.editorActions.JavaEnterInTextBlockHandler" id="JavaEnterInTextBlockHandler"/>
+    <enterHandlerDelegate
+      id="JavaEnterInInjectedTextBlockHandler"
+      order="before EnterBetweenBracesHandler, before JavaEnterInTextBlockHandler, first"
+      implementation="com.intellij.codeInsight.editorActions.JavaEnterInInjectedTextBlockHandler"/>
     <breadcrumbsInfoProvider implementation="com.intellij.lang.java.JavaBreadcrumbsInfoProvider"/>
     <editorFileSwapper implementation="com.intellij.codeEditor.JavaEditorFileSwapper"/>
     <binaryFileSourceProvider implementation="com.intellij.codeEditor.JavaBinaryFileSourceProvider"/>

java/java-impl/src/com/intellij/codeInsight/editorActions/JavaEnterInInjectedTextBlockHandler.java

@@ -0,0 +1,145 @@
+// Copyright 2000-2024 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package com.intellij.codeInsight.editorActions;
+
+import com.intellij.codeInsight.editorActions.enter.EnterHandlerDelegateAdapter;
+import com.intellij.injected.editor.DocumentWindow;
+import com.intellij.lang.java.JavaLanguage;
+import com.intellij.openapi.actionSystem.CommonDataKeys;
+import com.intellij.openapi.actionSystem.DataContext;
+import com.intellij.openapi.editor.Document;
+import com.intellij.openapi.editor.Editor;
+import com.intellij.openapi.editor.RangeMarker;
+import com.intellij.openapi.editor.actionSystem.EditorActionHandler;
+import com.intellij.openapi.editor.ex.EditorEx;
+import com.intellij.openapi.project.Project;
+import com.intellij.openapi.util.Ref;
+import com.intellij.openapi.util.TextRange;
+import com.intellij.openapi.util.text.StringUtil;
+import com.intellij.openapi.vfs.VirtualFile;
+import com.intellij.psi.*;
+import com.intellij.psi.util.PsiLiteralUtil;
+import com.intellij.psi.util.PsiTreeUtil;
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+
+public final class JavaEnterInInjectedTextBlockHandler extends EnterHandlerDelegateAdapter {
+  @Nullable
+  private RangeMarker myRangeMarker;
+  @Nullable
+  private String myPreviousIndent;
+
+  @Override
+  public Result preprocessEnter(@NotNull PsiFile file,
+                                @NotNull Editor editor,
+                                @NotNull Ref<Integer> caretOffset,
+                                @NotNull Ref<Integer> caretAdvance,
+                                @NotNull DataContext dataContext,
+                                EditorActionHandler originalHandler) {
+    Document document = file.getFileDocument();
+    if (!(document instanceof DocumentWindow)) return Result.Continue;
+    HostPosition host = getHost(dataContext);
+    if (host == null || host.originalFile() == file || host.originalEditor() == editor) return Result.Continue;
+    if (host.originalFile().getLanguage() != JavaLanguage.INSTANCE) return Result.Continue;
+    PsiElement psiElement = host.originalFile().findElementAt(host.offset());
+    if (!(psiElement instanceof PsiJavaToken textBlock && textBlock.getTokenType() == JavaTokenType.TEXT_BLOCK_LITERAL)) {
+      return Result.Continue;
+    }
+    PsiLiteralExpression literalExpression = PsiTreeUtil.getParentOfType(textBlock, PsiLiteralExpression.class);
+    if (literalExpression == null) return Result.Continue;
+    int endIndex = host.offset() - psiElement.getTextRange().getStartOffset();
+    String text = psiElement.getText();
+    if (endIndex >= text.length() || endIndex < 0) return Result.Continue;
+    String firstPart = text.substring(0, endIndex);
+    if (!firstPart.contains("\n")) return Result.Continue;
+    Document originalDocument = host.originalEditor().getDocument();
+    int lineNumber = originalDocument.getLineNumber(host.offset());
+    int lineStartOffset = originalDocument.getLineStartOffset(lineNumber);
+    RangeMarker marker =
+      originalDocument.createRangeMarker(lineStartOffset, originalDocument.getLineEndOffset(lineNumber));
+    marker.setGreedyToRight(true);
+    myRangeMarker = marker;
+    String previousIndent = PsiLiteralUtil.getTextBlockIndentString(literalExpression);
+    if (previousIndent != null && previousIndent.length() > host.offset() - lineStartOffset) {
+      previousIndent = previousIndent.substring(0, host.offset() - lineStartOffset);
+    }
+    myPreviousIndent = previousIndent;
+    return Result.Continue;
+  }
+
+  @NotNull
+  private static String collectIndent(@NotNull Document document, int offset) {
+    int lineNumber = document.getLineNumber(offset);
+    int lineStartOffset = document.getLineStartOffset(lineNumber);
+    int lineEndOffset = document.getLineEndOffset(lineNumber);
+    String currentLine = document.getText(new TextRange(lineStartOffset, lineEndOffset));
+    int indent = 0;
+    StringBuilder builder = new StringBuilder();
+    while (indent < currentLine.length() && StringUtil.isWhiteSpace(currentLine.charAt(indent))) {
+      builder.append(currentLine.charAt(indent));
+      indent++;
+    }
+    return builder.toString();
+  }
+
+  private static int getIndent(@NotNull Document document, int offset) {
+    return collectIndent(document, offset).length();
+  }
+
+
+  @Override
+  public Result postProcessEnter(@NotNull PsiFile file, @NotNull Editor editor, @NotNull DataContext dataContext) {
+    try {
+      Document document = file.getFileDocument();
+      if (!(document instanceof DocumentWindow)) return Result.Continue;
+      HostPosition host = getHost(dataContext);
+      if (host == null || host.originalFile() == file || host.originalEditor() == editor) return Result.Continue;
+      if (host.originalFile().getLanguage() != JavaLanguage.INSTANCE) return Result.Continue;
+      if (myRangeMarker == null || myPreviousIndent == null) return Result.Continue;
+      PsiElement psiElement = host.originalFile().findElementAt(host.offset());
+      if (!(psiElement instanceof PsiJavaToken javaToken && javaToken.getTokenType() == JavaTokenType.TEXT_BLOCK_LITERAL)) {
+        return Result.Continue;
+      }
+      Document originalDocument = host.originalEditor().getDocument();
+      PsiDocumentManager.getInstance(host.originalFile().getProject()).commitDocument(file.getFileDocument());
+      PsiDocumentManager.getInstance(host.originalFile().getProject()).commitDocument(originalDocument);
+      RangeMarker rangeMarker = myRangeMarker;
+      String previousIndent = myPreviousIndent;
+      if (!rangeMarker.isValid()) return Result.Continue;
+      TextRange changedRange = rangeMarker.getTextRange();
+      int lineNumber = originalDocument.getLineNumber(host.offset());
+      if (lineNumber - 1 < 0) return Result.Continue;
+      int firstChangedLineNumber = originalDocument.getLineNumber(changedRange.getStartOffset());
+      int lastChangedLineNumber = originalDocument.getLineNumber(changedRange.getEndOffset());
+      boolean hasChanges = false;
+      for (int i = firstChangedLineNumber; i <= lastChangedLineNumber; i++) {
+        int lineStartOffset = originalDocument.getLineStartOffset(i);
+        if (getIndent(originalDocument, lineStartOffset) >= previousIndent.length()) {
+          continue;
+        }
+        hasChanges = true;
+        originalDocument.replaceString(lineStartOffset, lineStartOffset, previousIndent);
+      }
+      return hasChanges ? Result.Stop : Result.Continue;
+    }
+    finally {
+      myPreviousIndent = null;
+      myRangeMarker = null;
+    }
+  }
+
+  public record HostPosition(@NotNull PsiFile originalFile, @NotNull Editor originalEditor, int offset) {
+  }
+
+  @Nullable
+  public static HostPosition getHost(@NotNull DataContext dataContext) {
+    Editor data = CommonDataKeys.HOST_EDITOR.getData(dataContext);
+    if (!(data instanceof EditorEx originalEditor)) return null;
+    Project project = CommonDataKeys.PROJECT.getData(dataContext);
+    if (project == null) return null;
+    VirtualFile virtualFile = originalEditor.getVirtualFile();
+    if (virtualFile == null) return null;
+    PsiFile originalFile = PsiManager.getInstance(project).findFile(virtualFile);
+    if (originalFile == null) return null;
+    return new HostPosition(originalFile, originalEditor, originalEditor.getCaretModel().getOffset());
+  }
+}

java/java-tests/testData/codeInsight/editorActions/stringLiteral/enterInInjectedStringBlockLiteralEnd.java

@@ -0,0 +1,10 @@
+import org.intellij.lang.annotations.Language;
+
+class Test {
+
+  @Language("JAVA")
+  String block = """
+    int a = 1;<caret>
+    """;
+
+}

java/java-tests/testData/codeInsight/editorActions/stringLiteral/enterInInjectedStringBlockLiteralEnd_after.java

@@ -0,0 +1,11 @@
+import org.intellij.lang.annotations.Language;
+
+class Test {
+
+  @Language("JAVA")
+  String block = """
+    int a = 1;
+    
+    """;
+
+}

java/java-tests/testData/codeInsight/editorActions/stringLiteral/enterInInjectedStringBlockLiteralMiddle.java

@@ -0,0 +1,10 @@
+import org.intellij.lang.annotations.Language;
+
+class Test {
+
+  @Language("JAVA")
+  String block = """
+    class A{<caret>}
+    """;
+
+}

java/java-tests/testData/codeInsight/editorActions/stringLiteral/enterInInjectedStringBlockLiteralMiddle_after.java

@@ -0,0 +1,12 @@
+import org.intellij.lang.annotations.Language;
+
+class Test {
+
+  @Language("JAVA")
+  String block = """
+    class A{
+    
+    }
+    """;
+
+}

java/java-tests/testData/codeInsight/editorActions/stringLiteral/enterInInjectedStringBlockLiteralStart.java

@@ -0,0 +1,10 @@
+import org.intellij.lang.annotations.Language;
+
+class Test {
+
+  @Language("JAVA")
+  String block = """
+    <caret>int a = 1;
+    """;
+
+}

java/java-tests/testData/codeInsight/editorActions/stringLiteral/enterInInjectedStringBlockLiteralStart_after.java

@@ -0,0 +1,11 @@
+import org.intellij.lang.annotations.Language;
+
+class Test {
+
+  @Language("JAVA")
+  String block = """
+    
+    int a = 1;
+    """;
+
+}

java/java-tests/testSrc/com/intellij/java/codeInsight/editorActions/JavaEnterInStringLiteralTest.java

@@ -1,22 +1,9 @@
-/*
- * Copyright 2000-2017 JetBrains s.r.o.
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
+// Copyright 2000-2024 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
 package com.intellij.java.codeInsight.editorActions;
 
 import com.intellij.application.options.CodeStyle;
 import com.intellij.lang.java.JavaLanguage;
+import com.intellij.pom.java.LanguageLevel;
 import com.intellij.psi.codeStyle.CommonCodeStyleSettings;
 import com.intellij.testFramework.EditorTestUtil;
 import com.intellij.testFramework.LightJavaCodeInsightTestCase;
@@ -24,6 +11,11 @@ import com.intellij.testFramework.LightJavaCodeInsightTestCase;
 public class JavaEnterInStringLiteralTest extends LightJavaCodeInsightTestCase {
   private static final String BASE_PATH = "/codeInsight/editorActions/stringLiteral/";
 
+  @Override
+  protected LanguageLevel getLanguageLevel() {
+    return LanguageLevel.JDK_21;
+  }
+
   public void testNonIndentedTextBlockContent() {
     doTest();
   }
@@ -48,6 +40,18 @@ public class JavaEnterInStringLiteralTest extends LightJavaCodeInsightTestCase {
     doTest();
   }
 
+  public void testEnterInInjectedStringBlockLiteralStart() {
+    doTest();
+  }
+
+  public void testEnterInInjectedStringBlockLiteralMiddle() {
+    doTest();
+  }
+
+  public void testEnterInInjectedStringBlockLiteralEnd() {
+    doTest();
+  }
+
   public void testEnterOpSignOnNextLine() {
     CommonCodeStyleSettings settings = CodeStyle.getSettings(getProject()).getCommonSettings(JavaLanguage.INSTANCE);
     boolean opSignOnNextLine = settings.BINARY_OPERATION_SIGN_ON_NEXT_LINE;