[mod-command] ModTemplateBuilder: support end position; use in JavaWithCastSurrounder

GitOrigin-RevId: 1fa5adbd8cac92062303baec88441a4489138b90

java/java-impl/src/com/intellij/codeInsight/generation/surroundWith/JavaWithCastSurrounder.java

@@ -5,87 +5,63 @@ import com.intellij.codeInsight.CodeInsightBundle;
 import com.intellij.codeInsight.guess.GuessManager;
 import com.intellij.codeInsight.lookup.LookupElement;
 import com.intellij.codeInsight.lookup.PsiTypeLookupItem;
-import com.intellij.codeInsight.template.*;
+import com.intellij.codeInsight.template.Expression;
 import com.intellij.codeInsight.template.impl.ConstantNode;
-import com.intellij.openapi.actionSystem.ex.ActionUtil;
-import com.intellij.openapi.application.WriteAction;
-import com.intellij.openapi.editor.Editor;
+import com.intellij.modcommand.ActionContext;
+import com.intellij.modcommand.ModPsiUpdater;
+import com.intellij.openapi.editor.Document;
 import com.intellij.openapi.editor.RangeMarker;
-import com.intellij.openapi.editor.ScrollType;
+import com.intellij.openapi.project.DumbService;
 import com.intellij.openapi.project.Project;
 import com.intellij.openapi.util.TextRange;
 import com.intellij.psi.*;
+import com.intellij.psi.codeStyle.CodeStyleManager;
+import com.intellij.psi.util.PsiTreeUtil;
 import com.intellij.refactoring.introduceField.ElementToWorkOn;
-import com.intellij.util.IncorrectOperationException;
-import org.jetbrains.annotations.NonNls;
+import org.jetbrains.annotations.NotNull;
 
 import java.util.Collections;
 import java.util.LinkedHashSet;
+import java.util.Objects;
 import java.util.Set;
 
-public class JavaWithCastSurrounder extends JavaExpressionSurrounder {
-  private static final @NonNls String TYPE_TEMPLATE_VARIABLE = "type";
-
-  @Override
-  public boolean startInWriteAction() {
-    return false;
-  }
-
+public class JavaWithCastSurrounder extends JavaExpressionModCommandSurrounder {
   @Override
   public boolean isApplicable(PsiExpression expr) {
     return !PsiTypes.voidType().equals(expr.getType());
   }
 
   @Override
-  public TextRange surroundExpression(final Project project, final Editor editor, PsiExpression expr) throws IncorrectOperationException {
-    assert expr.isValid();
-    PsiType[] types = ActionUtil.underModalProgress(
-      project,
-      CodeInsightBundle.message("surround.with.cast.modal.title"),
-      () ->  GuessManager.getInstance(project).guessTypeToCast(expr)
-    );
+  protected void surroundExpression(@NotNull ActionContext context, @NotNull PsiExpression expr, @NotNull ModPsiUpdater updater) {
+    Project project = context.project();
+    PsiType[] types =
+      DumbService.getInstance(project).computeWithAlternativeResolveEnabled(() -> GuessManager.getInstance(project).guessTypeToCast(expr));
     final boolean parenthesesNeeded = expr instanceof PsiPolyadicExpression ||
                                       expr instanceof PsiConditionalExpression ||
                                       expr instanceof PsiAssignmentExpression;
     String exprText = parenthesesNeeded ? "(" + expr.getText() + ")" : expr.getText();
 
-    TextRange range;
-    if (expr.isPhysical()) {
-      range = expr.getTextRange();
-    } else {
-      final RangeMarker rangeMarker = expr.getUserData(ElementToWorkOn.TEXT_RANGE);
-      if (rangeMarker == null) return null;
-      range = rangeMarker.getTextRange();
-    }
-    WriteAction.run(() -> {
-      final Template template = generateTemplate(project, exprText, types);
-      editor.getDocument().deleteString(range.getStartOffset(), range.getEndOffset());
-      editor.getCaretModel().moveToOffset(range.getStartOffset());
-      editor.getScrollingModel().scrollToCaret(ScrollType.RELATIVE);
-      TemplateManager.getInstance(project).startTemplate(editor, template);
-    });
-    return null;
-  }
-
-  private static Template generateTemplate(Project project, String exprText, final PsiType[] suggestedTypes) {
-    final TemplateManager templateManager = TemplateManager.getInstance(project);
-    final Template template = templateManager.createTemplate("", "");
-    template.setToReformat(true);
-
+    RangeMarker rangeMarker = expr.getUserData(ElementToWorkOn.TEXT_RANGE);
+    TextRange range = rangeMarker == null ? expr.getTextRange() : rangeMarker.getTextRange();
+    PsiFile file = updater.getWritable(context.file()); // cannot use expr.getContainingFile(), as expr could be detached (e.g., part of polyadic)
+    Document document = file.getFileDocument();
+    document.replaceString(range.getStartOffset(), range.getEndOffset(), "((x)" + exprText + ")");
+    PsiDocumentManager.getInstance(project).commitDocument(document);
+    PsiTypeCastExpression cast = PsiTreeUtil.getParentOfType(file.findElementAt(range.getStartOffset() + 1), PsiTypeCastExpression.class);
+    if (cast == null) return;
+    cast = (PsiTypeCastExpression)CodeStyleManager.getInstance(project).reformat(cast);
+    updater.moveCaretTo(cast.getParent().getTextRange().getEndOffset());
+    PsiTypeElement castType = Objects.requireNonNull(cast.getCastType());
     Set<LookupElement> itemSet = new LinkedHashSet<>();
-    for (PsiType type : suggestedTypes) {
+    for (PsiType type : types) {
       itemSet.add(PsiTypeLookupItem.createLookupItem(type, null));
     }
 
-    final Result result = suggestedTypes.length > 0 ? new PsiTypeResult(suggestedTypes[0], project) : null;
-
-    Expression expr = new ConstantNode(result).withLookupItems(itemSet.size() > 1 ? itemSet : Collections.emptySet());
-    template.addTextSegment("((");
-    template.addVariable(TYPE_TEMPLATE_VARIABLE, expr, expr, true);
-    template.addTextSegment(")" + exprText + ")");
-    template.addEndVariable();
-
-    return template;
+    String result = types.length > 0 ? types[0].getPresentableText() : "";
+    Expression typeExpr = new ConstantNode(result).withLookupItems(itemSet.size() > 1 ? itemSet : Collections.emptySet());
+    updater.templateBuilder()
+      .field(castType, typeExpr)
+      .finishAt(updater.getCaretOffset());
   }
 
   @Override

java/java-tests/testData/codeInsight/template/postfix/templates/cast/chainCall_after.java

@@ -1,5 +1,5 @@
 public class Foo {
     void m(Object o) {
-        (() "test")<caret>.length()
+      (() "test")<caret>.length()
     }
 }

platform/analysis-api/api-dump-unreviewed.txt

@@ -2825,6 +2825,14 @@ f:com.intellij.modcommand.ModStartTemplate$DependantVariableField
 - range():com.intellij.openapi.util.TextRange
 - varName():java.lang.String
 - withRange(com.intellij.openapi.util.TextRange):com.intellij.modcommand.ModStartTemplate$TemplateField
+f:com.intellij.modcommand.ModStartTemplate$EndField
+- java.lang.Record
+- com.intellij.modcommand.ModStartTemplate$TemplateField
+- <init>(com.intellij.openapi.util.TextRange):V
+- f:equals(java.lang.Object):Z
+- f:hashCode():I
+- range():com.intellij.openapi.util.TextRange
+- withRange(com.intellij.openapi.util.TextRange):com.intellij.modcommand.ModStartTemplate$TemplateField
 f:com.intellij.modcommand.ModStartTemplate$ExpressionField
 - java.lang.Record
 - com.intellij.modcommand.ModStartTemplate$TemplateField
@@ -2843,6 +2851,7 @@ com.intellij.modcommand.ModTemplateBuilder
 - field(com.intellij.psi.PsiElement,java.lang.String):com.intellij.modcommand.ModTemplateBuilder
 - a:field(com.intellij.psi.PsiElement,java.lang.String,com.intellij.codeInsight.template.Expression):com.intellij.modcommand.ModTemplateBuilder
 - a:field(com.intellij.psi.PsiElement,java.lang.String,java.lang.String,Z):com.intellij.modcommand.ModTemplateBuilder
+- a:finishAt(I):com.intellij.modcommand.ModTemplateBuilder
 f:com.intellij.modcommand.ModUpdateFileText
 - java.lang.Record
 - com.intellij.modcommand.ModCommand

platform/analysis-api/src/com/intellij/modcommand/ModStartTemplate.java

@@ -68,4 +68,16 @@ public record ModStartTemplate(@NotNull VirtualFile file, @NotNull List<@NotNull
       return new DependantVariableField(range, varName, dependantVariableName, alwaysStopAt);
     }
   }
+
+  /**
+   * A field to designate the end offset (where caret should be moved after the template is finished)
+   * 
+   * @param range left bound of the range designates the end position, right bound is ignored
+   */
+  public record EndField(@NotNull TextRange range) implements TemplateField {
+    @Override
+    public @NotNull TemplateField withRange(@NotNull TextRange range) {
+      return new EndField(range);
+    }
+  }
 }

platform/analysis-api/src/com/intellij/modcommand/ModTemplateBuilder.java

@@ -50,4 +50,12 @@ public interface ModTemplateBuilder {
    */
   @NotNull ModTemplateBuilder field(@NotNull PsiElement element, @NotNull String varName, @NotNull String dependantVariableName, 
                                     boolean alwaysStopAt);
+
+  /**
+   * Add a finish position to the template. The caret will be moved to a given position after the template is finished
+   * 
+   * @param offset finish position (offset within the file)
+   * @return this builder
+   */
+  @NotNull ModTemplateBuilder finishAt(int offset);
 }

platform/analysis-impl/src/com/intellij/lang/impl/modcommand/PsiUpdateImpl.java

@@ -439,6 +439,13 @@ final class PsiUpdateImpl {
           myTemplateFields.add(new ModStartTemplate.DependantVariableField(range, varName, dependantVariableName, alwaysStopAt));
           return this;
         }
+
+        @Override
+        public @NotNull ModTemplateBuilder finishAt(int offset) {
+          TextRange range = mapRange(TextRange.create(offset, offset));
+          myTemplateFields.add(new ModStartTemplate.EndField(range));
+          return this;
+        }
       };
     }
 

platform/lang-impl/src/com/intellij/lang/impl/modcommand/ModCommandExecutorImpl.java

@@ -283,6 +283,12 @@ public class ModCommandExecutorImpl extends ModCommandBatchExecutorImpl {
           builder.replaceElement(psiFile, variableField.range(), variableField.varName(),
                                  variableField.dependantVariableName(), variableField.alwaysStopAt());
         }
+        else if (field instanceof ModStartTemplate.EndField endField) {
+          PsiElement leaf = psiFile.findElementAt(endField.range().getStartOffset());
+          if (leaf != null) {
+            builder.setEndVariableBefore(leaf);
+          }
+        }
       }
 
       final Template tmpl = builder.buildInlineTemplate();