Convert single-line raw strings into multi-line on pressing <enter>

src/EditorFeatures/CSharp/RawStringLiteral/RawStringLiteralCommandHandler_Return.cs

@@ -5,6 +5,7 @@
 using System.Linq;
 using System.Threading;
 using Microsoft.CodeAnalysis.CSharp;
+using Microsoft.CodeAnalysis.CSharp.Syntax;
 using Microsoft.CodeAnalysis.Editor.Shared.Extensions;
 using Microsoft.CodeAnalysis.Editor.Shared.Utilities;
 using Microsoft.CodeAnalysis.Indentation;
@@ -50,74 +51,197 @@ public bool ExecuteCommand(ReturnKeyCommandArgs args, CommandExecutionContext co
         if (position >= currentSnapshot.Length)
             return false;
 
-        if (currentSnapshot[position] != '"')
+        var cancellationToken = context.OperationContext.UserCancellationToken;
+
+        var document = currentSnapshot.GetOpenDocumentInCurrentContextWithChanges();
+        if (document == null)
             return false;
 
-        var quotesBefore = 0;
-        var quotesAfter = 0;
+        return currentSnapshot[position] == '"'
+            ? ExecuteReturnCommandBeforeQuoteCharacter()
+            : ExecuteReturnCommandNotBeforeQuoteCharacter();
 
-        for (int i = position, n = currentSnapshot.Length; i < n; i++)
+        bool ExecuteReturnCommandBeforeQuoteCharacter()
         {
-            if (currentSnapshot[i] != '"')
-                break;
-
-            quotesAfter++;
+            var quotesBefore = 0;
+            var quotesAfter = 0;
+
+            // Ensure we're in between a balanced set of quotes, with at least 3 quotes on each side.
+
+            var currentSnapshot = subjectBuffer.CurrentSnapshot;
+            for (int i = position, n = currentSnapshot.Length; i < n; i++)
+            {
+                if (currentSnapshot[i] != '"')
+                    break;
+
+                quotesAfter++;
+            }
+
+            for (var i = position - 1; i >= 0; i--)
+            {
+                if (currentSnapshot[i] != '"')
+                    break;
+
+                quotesBefore++;
+            }
+
+            // We support two cases here.  Something simple like `"""$$"""`.  In this case, we have to be hitting enter
+            // inside balanced quotes.  But we also support `"""goo$$"""`.  In this case it's ok if quotes are not
+            // balanced.  We're going to go through the non-empty path involving adding multiple newlines to the final
+            // text.
+
+            var isEmpty = quotesBefore > 0;
+            if (isEmpty && quotesAfter != quotesBefore)
+                return false;
+
+            if (quotesAfter < 3)
+                return false;
+
+            // Looks promising based on text alone.  Now ensure we're actually on a raw string token/expression.
+            var parsedDocument = ParsedDocument.CreateSynchronously(document, cancellationToken);
+
+            var token = parsedDocument.Root.FindToken(position);
+            if (token.Kind() is not (SyntaxKind.SingleLineRawStringLiteralToken or
+                                     SyntaxKind.MultiLineRawStringLiteralToken or
+                                     SyntaxKind.InterpolatedSingleLineRawStringStartToken or
+                                     SyntaxKind.InterpolatedMultiLineRawStringStartToken) ||
+                token.Parent is not ExpressionSyntax expression)
+            {
+                return false;
+            }
+
+            return MakeEdit(parsedDocument, expression, isEmpty);
         }
 
-        for (var i = position - 1; i >= 0; i--)
+        bool ExecuteReturnCommandNotBeforeQuoteCharacter()
         {
-            if (currentSnapshot[i] != '"')
-                break;
-
-            quotesBefore++;
+            // If the caret is not on a quote, we need to find whether we are within the contents of a single-line raw
+            // string literal but not inside an interpolation. If we are inside a raw string literal and the caret is not on
+            // top of a quote, it is part of the literal's text. Here we try to ensure that the literal's closing quotes are
+            // properly placed in their own line We could reach this point after pressing enter within a single-line raw
+            // string
+
+            var parsedDocument = ParsedDocument.CreateSynchronously(document, cancellationToken);
+
+            var token = parsedDocument.Root.FindToken(position);
+            ExpressionSyntax expression;
+            switch (token.Kind())
+            {
+                case SyntaxKind.SingleLineRawStringLiteralToken when token.Parent is ExpressionSyntax parentExpression:
+                    expression = parentExpression;
+                    break;
+
+                case SyntaxKind.InterpolatedStringTextToken:
+                case SyntaxKind.OpenBraceToken:
+                    if (token is not
+                        {
+                            Parent.Parent: InterpolatedStringExpressionSyntax
+                            {
+                                StringStartToken.RawKind: (int)SyntaxKind.InterpolatedSingleLineRawStringStartToken,
+                            } interpolatedStringExpression,
+                        })
+                    {
+                        return false;
+                    }
+
+                    if (token.Kind() is SyntaxKind.OpenBraceToken && position != token.SpanStart)
+                        return false;
+
+                    expression = interpolatedStringExpression;
+                    break;
+
+                default:
+                    return false;
+            }
+
+            return MakeEdit(parsedDocument, expression, isEmpty: false);
         }
 
-        if (quotesAfter != quotesBefore)
-            return false;
-
-        if (quotesAfter < 3)
-            return false;
-
-        return SplitRawString(textView, subjectBuffer, span.Start.Position, CancellationToken.None);
-    }
-
-    private bool SplitRawString(ITextView textView, ITextBuffer subjectBuffer, int position, CancellationToken cancellationToken)
-    {
-        var document = subjectBuffer.CurrentSnapshot.GetOpenDocumentInCurrentContextWithChanges();
-        if (document == null)
-            return false;
-
-        var parsedDocument = ParsedDocument.CreateSynchronously(document, cancellationToken);
-
-        var token = parsedDocument.Root.FindToken(position);
-        if (token.Kind() is not (SyntaxKind.SingleLineRawStringLiteralToken or
-                                 SyntaxKind.MultiLineRawStringLiteralToken or
-                                 SyntaxKind.InterpolatedSingleLineRawStringStartToken or
-                                 SyntaxKind.InterpolatedMultiLineRawStringStartToken))
+        bool MakeEdit(
+            ParsedDocument parsedDocument,
+            ExpressionSyntax expression,
+            bool isEmpty)
         {
-            return false;
+            var project = document.Project;
+            var indentationOptions = subjectBuffer.GetIndentationOptions(_editorOptionsService, project.GetFallbackAnalyzerOptions(), project.Services, explicitFormat: false);
+            var indentation = expression.GetFirstToken().GetPreferredIndentation(parsedDocument, indentationOptions, cancellationToken);
+
+            var newLine = indentationOptions.FormattingOptions.NewLine;
+
+            using var transaction = CaretPreservingEditTransaction.TryCreate(
+                CSharpEditorResources.Split_raw_string, textView, _undoHistoryRegistry, _editorOperationsFactoryService);
+            var edit = subjectBuffer.CreateEdit();
+
+            if (isEmpty)
+            {
+                // If the literal is empty, we just want to help the user transform it into a multiline raw string
+                // literal with the extra empty newline between the delimiters to place the caret at
+                edit.Insert(position, newLine + newLine + indentation);
+
+                var finalSnapshot = edit.Apply();
+
+                // move caret to the right location in virtual space for the blank line we added.
+                var lineInNewSnapshot = finalSnapshot.GetLineFromPosition(position);
+                var nextLine = finalSnapshot.GetLineFromLineNumber(lineInNewSnapshot.LineNumber + 1);
+                textView.Caret.MoveTo(new VirtualSnapshotPoint(nextLine, indentation.Length));
+            }
+            else
+            {
+                // Otherwise, we're starting with a raw string that has content in it.  That's something like:
+                // """GooBar""".  If we hit enter at the `G` we only want to insert a single new line before the caret.
+                // However, if we were to hit enter anywhere after that, we want two new lines inserted.  One after the
+                // `"""` and one at the caret itself.
+                var newLineAndIndentation = newLine + indentation;
+
+                // Add a newline at the position of the end literal
+                var closingStart = GetStartPositionOfClosingDelimiter(expression);
+                edit.Insert(closingStart, newLineAndIndentation);
+
+                // Add a newline at the caret's position, to insert the newline that the user requested
+                edit.Insert(position, newLineAndIndentation);
+
+                // Also add a newline at the start of the text, only if there is text before the caret's position
+                var insertedLinesBeforeCaret = 1;
+                var openingEnd = GetEndPositionOfOpeningDelimiter(expression);
+                if (openingEnd != position)
+                {
+                    insertedLinesBeforeCaret = 2;
+                    edit.Insert(openingEnd, newLineAndIndentation);
+                }
+
+                var finalSnapshot = edit.Apply();
+
+                var lineInNewSnapshot = finalSnapshot.GetLineFromPosition(openingEnd);
+                var nextLine = finalSnapshot.GetLineFromLineNumber(lineInNewSnapshot.LineNumber + insertedLinesBeforeCaret);
+                textView.Caret.MoveTo(new VirtualSnapshotPoint(nextLine, indentation.Length));
+            }
+
+            transaction?.Complete();
+            return true;
         }
 
-        var indentationOptions = subjectBuffer.GetIndentationOptions(_editorOptionsService, document.Project.GetFallbackAnalyzerOptions(), document.Project.Services, explicitFormat: false);
-        var indentation = token.GetPreferredIndentation(parsedDocument, indentationOptions, cancellationToken);
-
-        var newLine = indentationOptions.FormattingOptions.NewLine;
+        int GetStartPositionOfClosingDelimiter(ExpressionSyntax expression)
+        {
+            if (expression is InterpolatedStringExpressionSyntax interpolatedStringExpression)
+                return interpolatedStringExpression.StringEndToken.Span.Start;
 
-        using var transaction = CaretPreservingEditTransaction.TryCreate(
-            CSharpEditorResources.Split_string, textView, _undoHistoryRegistry, _editorOperationsFactoryService);
+            var index = expression.Span.End;
+            while (currentSnapshot[index - 1] == '"')
+                index--;
 
-        var edit = subjectBuffer.CreateEdit();
+            return index;
+        }
 
-        // apply the change:
-        edit.Insert(position, newLine + newLine + indentation);
-        var snapshot = edit.Apply();
+        int GetEndPositionOfOpeningDelimiter(ExpressionSyntax expression)
+        {
+            if (expression is InterpolatedStringExpressionSyntax interpolatedStringExpression)
+                return interpolatedStringExpression.StringStartToken.Span.End;
 
-        // move caret:
-        var lineInNewSnapshot = snapshot.GetLineFromPosition(position);
-        var nextLine = snapshot.GetLineFromLineNumber(lineInNewSnapshot.LineNumber + 1);
-        textView.Caret.MoveTo(new VirtualSnapshotPoint(nextLine, indentation.Length));
+            var index = expression.Span.Start;
+            while (currentSnapshot[index] == '"')
+                index++;
 
-        transaction?.Complete();
-        return true;
+            return index;
+        }
     }
 }