feat: String enter handler

CHANGELOG.md

@@ -17,6 +17,11 @@ Changelog structure reference:
 
 ## [Unreleased]
 
+### Added
+
+- Zig
+  - Enter key handling in strings and multi line strings
+
 ### Fixed
 
 - Zig

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

@@ -32,6 +32,13 @@
 
         <lang.formatter language="Zig" implementationClass="com.falsepattern.zigbrains.zig.formatter.ZigFormattingModelBuilder"/>
 
+        <enterHandlerDelegate
+                id="ZigEnterInTextBlockHandler"
+                implementation="com.falsepattern.zigbrains.zig.editing.ZigEnterInTextBlockHandler" />
+        <enterHandlerDelegate
+                id="ZigEnterInQuotedStringHandler"
+                implementation="com.falsepattern.zigbrains.zig.editing.ZigEnterInQuotedStringHandler" />
+
         <postStartupActivity implementation="com.falsepattern.zigbrains.zig.lsp.ZLSStartupActivity"/>
 
         <!-- LSP textDocument/signatureHelp -->

modules/zig/src/main/grammar/ZigString.flex

@@ -30,6 +30,11 @@ import static com.intellij.psi.StringEscapesTokenTypes.*;
 %implements FlexLexer
 %function advance
 %type IElementType
+%{
+    public ZigStringLexer() {
+
+    }
+%}
 
 hex=[0-9a-fA-F]
 

modules/zig/src/main/java/com/falsepattern/zigbrains/zig/editing/ZigEnterInQuotedStringHandler.java

@@ -0,0 +1,63 @@
+package com.falsepattern.zigbrains.zig.editing;
+
+import com.falsepattern.zigbrains.zig.parser.ZigFile;
+import com.falsepattern.zigbrains.zig.psi.ZigStringLiteral;
+import com.falsepattern.zigbrains.zig.psi.ZigTypes;
+import com.falsepattern.zigbrains.zig.stringlexer.ZigStringLexer;
+import com.falsepattern.zigbrains.zig.util.PsiTextUtil;
+import com.falsepattern.zigbrains.zig.util.ZigStringUtil;
+import com.intellij.application.options.CodeStyle;
+import com.intellij.codeInsight.editorActions.JavaLikeQuoteHandler;
+import com.intellij.codeInsight.editorActions.enter.EnterHandlerDelegateAdapter;
+import com.intellij.lang.ASTNode;
+import com.intellij.lexer.StringLiteralLexer;
+import com.intellij.openapi.actionSystem.DataContext;
+import com.intellij.openapi.editor.Document;
+import com.intellij.openapi.editor.Editor;
+import com.intellij.openapi.editor.actionSystem.EditorActionHandler;
+import com.intellij.openapi.util.Ref;
+import com.intellij.openapi.util.TextRange;
+import com.intellij.openapi.util.text.StringUtil;
+import com.intellij.psi.PsiDocumentManager;
+import com.intellij.psi.PsiElement;
+import com.intellij.psi.PsiFile;
+import com.intellij.psi.StringEscapesTokenTypes;
+import com.intellij.psi.impl.source.tree.LeafPsiElement;
+import com.intellij.psi.tree.IElementType;
+import lombok.RequiredArgsConstructor;
+import lombok.SneakyThrows;
+import lombok.val;
+import org.jetbrains.annotations.NotNull;
+
+import java.io.StringReader;
+import java.util.ArrayList;
+
+@RequiredArgsConstructor
+public class ZigEnterInQuotedStringHandler extends EnterHandlerDelegateAdapter {
+    @Override
+    public Result preprocessEnter(@NotNull PsiFile file,
+                                  @NotNull Editor editor,
+                                  @NotNull Ref<Integer> caretOffsetRef,
+                                  @NotNull Ref<Integer> caretAdvanceRef,
+                                  @NotNull DataContext dataContext,
+                                  EditorActionHandler originalHandler) {
+        if (!(file instanceof ZigFile)) {
+            return Result.Continue;
+        }
+
+        val caretOffset = (int)caretOffsetRef.get();
+        var psiAtOffset = file.findElementAt(caretOffset);
+        if (psiAtOffset instanceof LeafPsiElement leaf) {
+            if ( leaf.getTokenType() == ZigTypes.STRING_LITERAL_SINGLE) {
+                psiAtOffset = leaf.getParent();
+            }
+        }
+        if (psiAtOffset instanceof ZigStringLiteral str &&
+            !str.isMultiLine() &&
+            str.getTextOffset() < caretOffset) {
+            PsiTextUtil.splitString(editor, str, caretOffset, true);
+            return Result.Stop;
+        }
+        return Result.Continue;
+    }
+}

modules/zig/src/main/java/com/falsepattern/zigbrains/zig/editing/ZigEnterInTextBlockHandler.java

@@ -0,0 +1,97 @@
+package com.falsepattern.zigbrains.zig.editing;
+
+import com.falsepattern.zigbrains.zig.parser.ZigFile;
+import com.intellij.codeInsight.editorActions.enter.EnterHandlerDelegateAdapter;
+import com.intellij.openapi.actionSystem.DataContext;
+import com.intellij.openapi.editor.Editor;
+import com.intellij.openapi.editor.actionSystem.EditorActionHandler;
+import com.intellij.openapi.util.Ref;
+import com.intellij.openapi.util.TextRange;
+import com.intellij.openapi.util.text.StringUtil;
+import com.intellij.psi.PsiDocumentManager;
+import com.intellij.psi.PsiElement;
+import com.intellij.psi.PsiFile;
+import lombok.RequiredArgsConstructor;
+import lombok.val;
+import org.jetbrains.annotations.NotNull;
+
+import java.util.ArrayList;
+
+@RequiredArgsConstructor
+public class ZigEnterInTextBlockHandler extends EnterHandlerDelegateAdapter {
+    @Override
+    public Result preprocessEnter(@NotNull PsiFile file,
+                                  @NotNull Editor editor,
+                                  @NotNull Ref<Integer> caretOffsetRef,
+                                  @NotNull Ref<Integer> caretAdvanceRef,
+                                  @NotNull DataContext dataContext,
+                                  EditorActionHandler originalHandler) {
+        if (!(file instanceof ZigFile)) {
+            return Result.Continue;
+        }
+        for (val assistant: ZigMultiLineAssistant.Assistants.ASSISTANTS) {
+            val result = preprocessEnterWithAssistant(file, editor, assistant);
+            if (result != null)
+                return result;
+        }
+        return Result.Continue;
+    }
+
+    private static <T extends PsiElement> Result preprocessEnterWithAssistant(@NotNull PsiFile file,
+                                                                              @NotNull Editor editor,
+                                                                              ZigMultiLineAssistant<T> assistant) {
+        val offset = editor.getCaretModel().getOffset();
+        val textBlock = getTextBlockAt(file, offset, assistant);
+        if (textBlock == null) {
+            return null;
+        }
+        val textBlockOffset = textBlock.getTextOffset();
+        val document = editor.getDocument();
+        val textBlockLine = document.getLineNumber(textBlockOffset);
+        val textBlockLineStart = document.getLineStartOffset(textBlockLine);
+        val indentPre = textBlockOffset - textBlockLineStart;
+        val project = textBlock.getProject();
+        val lineNumber = document.getLineNumber(offset);
+        val lineStartOffset = document.getLineStartOffset(lineNumber);
+        val text = document.getText(new TextRange(lineStartOffset, offset + 1));
+        val parts = new ArrayList<>(StringUtil.split(text, assistant.prefix));
+        if (parts.size() <= 1)
+            return Result.Continue;
+        if (parts.size() > 2) {
+            val sb = new StringBuilder();
+            sb.append(parts.get(1));
+            while (parts.size() > 2) {
+                sb.append(assistant.prefix);
+                sb.append(parts.remove(2));
+            }
+            parts.set(1, sb.toString());
+        }
+        val indentPost = measureSpaces(parts.get(1));
+        val newLine = '\n' + StringUtil.repeatSymbol(' ', indentPre) + assistant.prefix + StringUtil.repeatSymbol(' ', indentPost);
+        document.insertString(offset, newLine);
+        PsiDocumentManager.getInstance(project).commitDocument(document);
+        editor.getCaretModel().moveToOffset(offset + newLine.length());
+        return Result.Stop;
+    }
+
+    private static int measureSpaces(String str) {
+        for (int i = 0; i < str.length(); i++) {
+            val c = str.charAt(i);
+            switch (c) {
+                case ' ':
+                case '\t':
+                    continue;
+                default:
+                    return i;
+            }
+        }
+        return str.length();
+    }
+
+    private static <T extends PsiElement> T getTextBlockAt(PsiFile file, int offset, ZigMultiLineAssistant<T> assistant) {
+        val psiAtOffset = file.findElementAt(offset);
+        if (psiAtOffset == null)
+            return null;
+        return assistant.acceptPSI(psiAtOffset);
+    }
+}

modules/zig/src/main/java/com/falsepattern/zigbrains/zig/editing/ZigMultiLineAssistant.java

@@ -0,0 +1,61 @@
+package com.falsepattern.zigbrains.zig.editing;
+
+import com.falsepattern.zigbrains.zig.psi.ZigStringLiteral;
+import com.falsepattern.zigbrains.zig.psi.ZigTypes;
+import com.intellij.psi.PsiComment;
+import com.intellij.psi.PsiElement;
+import com.intellij.psi.impl.source.tree.LeafPsiElement;
+import com.intellij.psi.tree.IElementType;
+import lombok.RequiredArgsConstructor;
+import org.jetbrains.annotations.NotNull;
+
+import java.util.List;
+
+@RequiredArgsConstructor
+public abstract class ZigMultiLineAssistant<T extends PsiElement> {
+    public static class Assistants {
+        private Assistants() {}
+
+        public static final List<ZigMultiLineAssistant<?>> ASSISTANTS = List.of(new StringAssistant(),
+                                                                                new CommentAssistant("//", ZigTypes.LINE_COMMENT),
+                                                                                new CommentAssistant("///", ZigTypes.DOC_COMMENT),
+                                                                                new CommentAssistant("//!", ZigTypes.CONTAINER_DOC_COMMENT));
+    }
+    public final String prefix;
+
+    public abstract T acceptPSI(@NotNull PsiElement element);
+
+    public static class StringAssistant extends ZigMultiLineAssistant<ZigStringLiteral> {
+        public StringAssistant() {
+            super("\\\\");
+        }
+
+        @Override
+        public ZigStringLiteral acceptPSI(final @NotNull PsiElement element) {
+            final PsiElement candidate;
+            if (element instanceof LeafPsiElement leaf) {
+                if (leaf.getTokenType() == ZigTypes.STRING_LITERAL_MULTI) {
+                    candidate = leaf.getParent();
+                } else {
+                    return null;
+                }
+            } else {
+                candidate = element;
+            }
+            return candidate instanceof ZigStringLiteral str && str.isMultiLine() ? str : null;
+        }
+    }
+
+    public static class CommentAssistant extends ZigMultiLineAssistant<PsiComment> {
+        private final IElementType tokenType;
+        public CommentAssistant(String prefix, IElementType type) {
+            super(prefix);
+            tokenType = type;
+        }
+
+        @Override
+        public PsiComment acceptPSI(@NotNull PsiElement element) {
+            return element instanceof PsiComment comment && comment.getTokenType() == tokenType ? comment : null;
+        }
+    }
+}

modules/zig/src/main/java/com/falsepattern/zigbrains/zig/psi/ZigStringElementManipulator.java

@@ -25,7 +25,7 @@ public class ZigStringElementManipulator extends AbstractElementManipulator<ZigS
             throws IncorrectOperationException {
         val originalContext = element.getText();
         val isMulti = element.isMultiLine();
-        final String replacement;
+        final CharSequence replacement;
         if (isMulti) {
             val contentRanges = element.getContentRanges();
             val contentBuilder = new StringBuilder();
@@ -51,9 +51,7 @@ public class ZigStringElementManipulator extends AbstractElementManipulator<ZigS
                 }
             }
             val content = contentBuilder.toString();
-            val pfx = PsiTextUtil.getIndentString(element) + "\\\\";
+            replacement = ZigStringUtil.prefixWithTextBlockEscape(PsiTextUtil.getIndentSize(element), "\\\\", content, false, true);
-            replacement = Arrays.stream(content.split("(\\r\\n|\\r|\\n)")).map(line -> pfx + line).collect(
-                    Collectors.joining("\n"));
         } else {
             val elementRange = getRangeInElement(element);
             replacement = "\"" +

modules/zig/src/main/java/com/falsepattern/zigbrains/zig/util/PsiTextUtil.java

@@ -1,11 +1,23 @@
 package com.falsepattern.zigbrains.zig.util;
 
+import com.falsepattern.zigbrains.zig.stringlexer.ZigStringLexer;
+import com.intellij.codeInsight.editorActions.enter.EnterHandlerDelegate;
+import com.intellij.lang.ASTNode;
+import com.intellij.lexer.FlexAdapter;
+import com.intellij.lexer.FlexLexer;
+import com.intellij.lexer.Lexer;
+import com.intellij.openapi.editor.Editor;
 import com.intellij.openapi.util.TextRange;
 import com.intellij.openapi.util.text.StringUtil;
 import com.intellij.psi.PsiElement;
+import com.intellij.psi.StringEscapesTokenTypes;
+import com.intellij.psi.tree.IElementType;
+import com.intellij.util.MathUtil;
+import lombok.SneakyThrows;
 import lombok.val;
 import org.jetbrains.annotations.NotNull;
 
+import java.io.StringReader;
 import java.util.ArrayList;
 import java.util.List;
 
@@ -64,4 +76,56 @@ public class PsiTextUtil {
         val indent = getIndentSize(element);
         return " ".repeat(Math.max(0, indent));
     }
+
+    public static void splitString(@NotNull Editor editor,
+                                   @NotNull PsiElement psiAtOffset,
+                                   int caretOffset,
+                                   boolean insertNewlineAtCaret) {
+        val document = editor.getDocument();
+        ASTNode token = psiAtOffset.getNode();
+        val text = document.getText();
+
+        TextRange range = token.getTextRange();
+        val lexer = new FlexAdapter(new ZigStringLexer());
+        lexer.start(text, range.getStartOffset(), range.getEndOffset());
+        caretOffset = skipStringLiteralEscapes(caretOffset, lexer);
+        caretOffset = MathUtil.clamp(caretOffset, range.getStartOffset() + 1, range.getEndOffset() - 1);
+        val unescapedPrefix = ZigStringUtil.unescape(text.substring(range.getStartOffset() + 1, caretOffset), false);
+        val unescapedSuffix = ZigStringUtil.unescape(text.substring(caretOffset, range.getEndOffset() - 1), false);
+        val stringRange = document.createRangeMarker(range.getStartOffset(), range.getEndOffset());
+        stringRange.setGreedyToRight(true);
+        val lineNumber = document.getLineNumber(caretOffset);
+        val lineOffset = document.getLineStartOffset(lineNumber);
+        val indent = stringRange.getStartOffset() - lineOffset;
+        document.deleteString(stringRange.getStartOffset(), stringRange.getEndOffset());
+        document.insertString(stringRange.getStartOffset(),
+                              ZigStringUtil.prefixWithTextBlockEscape(indent,
+                                                                      "\\\\",
+                                                                      insertNewlineAtCaret ? unescapedPrefix + "\n" : unescapedPrefix,
+                                                                      false,
+                                                                      true));
+        caretOffset = stringRange.getEndOffset();
+        document.insertString(caretOffset,
+                              ZigStringUtil.prefixWithTextBlockEscape(indent,
+                                                                      "\\\\",
+                                                                      unescapedSuffix,
+                                                                      false,
+                                                                      false));
+        document.insertString(stringRange.getEndOffset(), "\n" + " ".repeat(indent));
+        stringRange.dispose();
+        editor.getCaretModel().moveToOffset(caretOffset);
+    }
+    @SneakyThrows
+    protected static int skipStringLiteralEscapes(int caretOffset, Lexer lexer) {
+        while (lexer.getTokenType() != null) {
+            if (lexer.getTokenStart() < caretOffset && caretOffset < lexer.getTokenEnd()) {
+                if (StringEscapesTokenTypes.STRING_LITERAL_ESCAPES.contains(lexer.getTokenType())) {
+                    caretOffset = lexer.getTokenEnd();
+                }
+                break;
+            }
+            lexer.advance();
+        }
+        return caretOffset;
+    }
 }

modules/zig/src/main/java/com/falsepattern/zigbrains/zig/util/ZigStringUtil.java

@@ -9,7 +9,9 @@ import lombok.val;
 import org.jetbrains.annotations.NotNull;
 
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.List;
+import java.util.regex.Pattern;
 import java.util.stream.Collectors;
 
 public class ZigStringUtil {
@@ -68,6 +70,35 @@ public class ZigStringUtil {
         return result.toString();
     }
 
+    private static final Pattern NL_MATCHER = Pattern.compile("(\\r\\n|\\r|\\n)");
+    private static final String[] COMMON_INDENTS;
+    static {
+        val count = 32;
+        val sb = new StringBuilder(count);
+        COMMON_INDENTS = new String[count];
+        for (int i = 0; i < count; i++) {
+            COMMON_INDENTS[i] = sb.toString();
+            sb.append(" ");
+        }
+    }
+
+    public static CharSequence prefixWithTextBlockEscape(int indent, CharSequence marker, CharSequence content, boolean indentFirst, boolean prefixFirst) {
+        val indentStr = indent >= 0 ? indent < COMMON_INDENTS.length ? COMMON_INDENTS[indent] : " ".repeat(indent) : "";
+        val parts = Arrays.asList(NL_MATCHER.split(content, -1));
+        val result = new StringBuilder(content.length() + marker.length() * parts.size() + indentStr.length() * parts.size());
+        if (indentFirst) {
+            result.append(indentStr);
+        }
+        if (prefixFirst) {
+            result.append(marker);
+        }
+        result.append(parts.getFirst());
+        for (val part: parts.subList(1, parts.size())) {
+            result.append("\n").append(indentStr).append(marker).append(part);
+        }
+        return result;
+    }
+
     @UtilityClass
     private static class Escaper {
         private static final Int2IntMap ESC_TO_CODE = new Int2IntOpenHashMap();