[SourceKit/CodeFormat] Indent lines in multi-line strings

When the multi-line string is unterminated, the indentation of a line
will be the same as the first previous line that had contents, unless
that line is the line containing the start quotes. In that case the
indentation will be:
  - the same as the start quote indentation, if the quotes are the only
    contents on the line, or
  - an extra indentation level to the start quote line otherwise

Lines within a terminated multi-line string or lines with content will
only ever be indented if their current indentation is invalid (ie. their
indentation level is less than that of their end quotes).

This rule is to prevent any signficant (and possibly unintended)
whitespace being added to existing strings during a whole file/range
format - Xcode does not remove whitespace from whitespace-only lines by
default. This could be improved if the reformat was sent the actual
range rather than a line at a time. Different indentation could then be
chosen if the range was in fact a single line.

Resolves rdar://32181422

Author: bnbarham

lib/IDE/Formatting.cpp

@@ -61,6 +61,13 @@ getLocForContentStartOnSameLine(SourceManager &SM, SourceLoc Loc) {
   return LineStart.getAdvancedLoc(Indentation.size());
 }
 
+/// \returns true if the line at \c Loc is either empty or only contains whitespace
+static bool isLineAtLocEmpty(SourceManager &SM, SourceLoc Loc) {
+  SourceLoc LineStart = Lexer::getLocForStartOfLine(SM, Loc);
+  SourceLoc Next = Lexer::getTokenAtLocation(SM, LineStart, CommentRetentionMode::ReturnAsTokens).getLoc();
+  return Next.isInvalid() || !isOnSameLine(SM, Loc, Next);
+}
+
 /// \returns the first token after the token at \c Loc.
 static Optional<Token>
 getTokenAfter(SourceManager &SM, SourceLoc Loc, bool SkipComments = true) {
@@ -136,6 +143,21 @@ static ClosureExpr *findTrailingClosureFromArgument(Expr *arg) {
   return nullptr;
 }
 
+static size_t calcVisibleWhitespacePrefix(StringRef Line,
+                                          CodeFormatOptions Options) {
+  size_t Indent = 0;
+  for (auto Char : Line) {
+    if (Char == '\t') {
+      Indent += Options.TabWidth;
+    } else if (Char == ' ' || Char == '\v' || Char == '\f') {
+      Indent++;
+    } else {
+      break;
+    }
+  }
+  return Indent;
+}
+
 /// An indentation context of the target location
 struct IndentContext {
   enum ContextKind { Exact, LineStart };
@@ -258,16 +280,14 @@ class FormatContext {
   Optional<IndentContext> InnermostCtx;
   bool InDocCommentBlock;
   bool InCommentLine;
-  bool InStringLiteral;
 
 public:
   FormatContext(SourceManager &SM,
                 Optional<IndentContext> IndentCtx,
                 bool InDocCommentBlock = false,
-                bool InCommentLine = false,
-                bool InStringLiteral = false)
+                bool InCommentLine = false)
     :SM(SM), InnermostCtx(IndentCtx), InDocCommentBlock(InDocCommentBlock),
-     InCommentLine(InCommentLine), InStringLiteral(InStringLiteral) { }
+     InCommentLine(InCommentLine) { }
 
   bool IsInDocCommentBlock() {
     return InDocCommentBlock;
@@ -277,10 +297,6 @@ class FormatContext {
     return InCommentLine;
   }
 
-  bool IsInStringLiteral() const {
-    return InStringLiteral;
-  }
-
   void padToExactColumn(StringBuilder &Builder,
                         const CodeFormatOptions &FmtOptions) {
     assert(isExact() && "Context is not exact?");
@@ -1208,8 +1224,9 @@ class FormatWalker : public ASTWalker {
   bool InDocCommentBlock = false;
   /// Whether the target location appears within a line comment.
   bool InCommentLine = false;
-  /// Whether the target location appears within a string literal.
-  bool InStringLiteral = false;
+
+  /// The range of the string literal the target is inside of (if any, invalid otherwise).
+  CharSourceRange StringLiteralRange;
 
 public:
   explicit FormatWalker(SourceFile &SF, SourceManager &SM, CodeFormatOptions &Options)
@@ -1224,7 +1241,8 @@ class FormatWalker : public ASTWalker {
     CtxOverride.clear();
     TargetLocation = Loc;
     TargetLineLoc = Lexer::getLocForStartOfLine(SM, TargetLocation);
-    InDocCommentBlock = InCommentLine = InStringLiteral = false;
+    InDocCommentBlock = InCommentLine = false;
+    StringLiteralRange = CharSourceRange();
     NodesToSkip.clear();
     CurrentTokIt = TokenList.begin();
 
@@ -1234,14 +1252,99 @@ class FormatWalker : public ASTWalker {
     if (InnermostCtx)
       CtxOverride.applyIfNeeded(SM, *InnermostCtx);
 
+    if (StringLiteralRange.isValid()) {
+      assert(!InDocCommentBlock && !InCommentLine &&
+             "Target is in both a string and comment");
+      InnermostCtx = indentWithinStringLiteral();
+    } else {
+      assert(!InDocCommentBlock || !InCommentLine &&
+             "Target is in both a doc comment block and comment line");
+    }
+
     return FormatContext(SM, InnermostCtx, InDocCommentBlock,
-                         InCommentLine, InStringLiteral);
+                         InCommentLine);
   }
 
+private:
+
+  Optional<IndentContext> indentWithinStringLiteral() {
+    assert(StringLiteralRange.isValid() && "Target is not within a string literal");
+
+    // This isn't ideal since if the user types """""" and then an enter
+    // inside after, we won't indent the end quotes. But indenting the end
+    // quotes could lead to an error in the rest of the string, so best to
+    // avoid it entirely for now.
+    if (isOnSameLine(SM, TargetLineLoc, StringLiteralRange.getEnd()))
+      return IndentContext {TargetLocation, false, IndentContext::Exact};
+
+    // If there's contents before the end quotes then it's likely the quotes
+    // are actually the start quotes of the next string in the file. Pretend
+    // they don't exist so their indent doesn't affect the indenting.
+    SourceLoc EndLineContentLoc =
+        getLocForContentStartOnSameLine(SM, StringLiteralRange.getEnd());
+    bool HaveEndQuotes = CharSourceRange(SM, EndLineContentLoc,
+                                         StringLiteralRange.getEnd())
+        .str().equals(StringRef("\"\"\""));
+
+    if (!HaveEndQuotes) {
+      // Indent to the same indentation level as the first non-empty line
+      // before the target. If that line is the start line then either use
+      // the same indentation of the start quotes if they are on their own
+      // line, or an extra indentation otherwise.
+      //
+      // This will indent lines with content on it as well, which should be
+      // fine since it is quite unlikely anyone would format a range that
+      // includes an unterminated string.
+
+      SourceLoc AlignLoc = TargetLineLoc;
+      while (SM.isBeforeInBuffer(StringLiteralRange.getStart(), AlignLoc)) {
+        AlignLoc = Lexer::getLocForStartOfLine(SM, AlignLoc.getAdvancedLoc(-1));
+        if (!isLineAtLocEmpty(SM, AlignLoc)) {
+          AlignLoc = getLocForContentStartOnSameLine(SM, AlignLoc);
+          break;
+        }
+      }
+
+      if (isOnSameLine(SM, AlignLoc, StringLiteralRange.getStart())) {
+        SourceLoc StartLineContentLoc =
+            getLocForContentStartOnSameLine(SM, StringLiteralRange.getStart());
+        bool StartLineOnlyQuotes = CharSourceRange(SM, StartLineContentLoc,
+                                                   StringLiteralRange.getEnd())
+            .str().startswith(StringRef("\"\"\""));
+        if (!StartLineOnlyQuotes)
+          return IndentContext {StringLiteralRange.getStart(), true};
+
+        AlignLoc = StringLiteralRange.getStart();
+      }
+
+      return IndentContext {AlignLoc, false, IndentContext::Exact};
+    }
+
+    // If there are end quotes, only enforce a minimum indentation. We don't
+    // want to add any other indentation since that could add unintended
+    // whitespace to existing strings. Could change this if the full range
+    // was passed rather than a single line - in that case we *would* indent
+    // if the range was a single empty line.
+
+    CharSourceRange TargetIndentRange =
+        CharSourceRange(SM, Lexer::getLocForStartOfLine(SM, TargetLocation),
+                        TargetLocation);
+    CharSourceRange EndIndentRange =
+        CharSourceRange(SM, Lexer::getLocForStartOfLine(SM, EndLineContentLoc),
+                        EndLineContentLoc);
+
+    size_t TargetIndent =
+        calcVisibleWhitespacePrefix(TargetIndentRange.str(), FmtOptions);
+    size_t EndIndent =
+        calcVisibleWhitespacePrefix(EndIndentRange.str(), FmtOptions);
+    if (TargetIndent >= EndIndent)
+      return IndentContext {TargetLocation, false, IndentContext::Exact};
+
+    return IndentContext {EndLineContentLoc, false, IndentContext::Exact};
+  }
 
 #pragma mark ASTWalker overrides and helpers
 
-private:
   bool walkCustomAttributes(Decl *D) {
     // CustomAttrs of non-param VarDecls are handled when this method is called
     // on their containing PatternBindingDecls (below).
@@ -1330,7 +1433,7 @@ class FormatWalker : public ASTWalker {
         SM.isBeforeInBuffer(E->getStartLoc(), TargetLocation) &&
         SM.isBeforeInBuffer(TargetLocation,
                             Lexer::getLocForEndOfToken(SM, E->getEndLoc()))) {
-      InStringLiteral = true;
+      StringLiteralRange = Lexer::getCharSourceRangeFromSourceRange(SM, E->getSourceRange());
     }
 
     // Create a default indent context for all top-level expressions
@@ -1352,7 +1455,7 @@ class FormatWalker : public ASTWalker {
 
     // Don't visit the child expressions of interpolated strings directly -
     // visit only the argument of each appendInterpolation call instead, and
-    // update InStringLiteral for each segment.
+    // set StringLiteralRange if needed for each segment.
     if (auto *ISL = dyn_cast<InterpolatedStringLiteralExpr>(E)) {
       if (Action.shouldVisitChildren()) {
         llvm::SaveAndRestore<ASTWalker::ParentTy>(Parent, ISL);
@@ -1364,7 +1467,8 @@ class FormatWalker : public ASTWalker {
               // Handle the preceeding string segment.
               CharSourceRange StringRange(SM, PrevStringStart, CE->getStartLoc());
               if (StringRange.contains(TargetLocation)) {
-                InStringLiteral = true;
+                StringLiteralRange =
+                    Lexer::getCharSourceRangeFromSourceRange(SM, E->getSourceRange());
                 return;
               }
               // Walk into the interpolation segment.
@@ -1378,7 +1482,8 @@ class FormatWalker : public ASTWalker {
         SourceLoc End = Lexer::getLocForEndOfToken(SM, ISL->getStartLoc());
         CharSourceRange StringRange(SM, PrevStringStart, End);
         if (StringRange.contains(TargetLocation))
-          InStringLiteral = true;
+          StringLiteralRange =
+              Lexer::getCharSourceRangeFromSourceRange(SM, E->getSourceRange());
 
         return {false, E};
       }
@@ -1467,7 +1572,7 @@ class FormatWalker : public ASTWalker {
   }
 
   void scanTokensUntil(SourceLoc Loc) {
-    if (InDocCommentBlock || InCommentLine)
+    if (InDocCommentBlock || InCommentLine || StringLiteralRange.isValid())
       return;
     for (auto Invalid = Loc.isInvalid(); CurrentTokIt != TokenList.end() &&
          (Invalid || SM.isBeforeInBuffer(CurrentTokIt->getLoc(), Loc));
@@ -1477,16 +1582,16 @@ class FormatWalker : public ASTWalker {
         SourceLoc StartLineLoc = Lexer::getLocForStartOfLine(
             SM, CommentRange.getStart());
 
-        // The -1 is needed in case the past-the-end position is a newline
-        // character. In that case getLocForStartOfLine returns the start of
-        // the next line.
         SourceLoc EndLineLoc = Lexer::getLocForStartOfLine(
-            SM, CommentRange.getEnd().getAdvancedLoc(-1));
+            SM, CommentRange.getEnd());
         auto TokenStr = CurrentTokIt->getRange().str();
         InDocCommentBlock |= SM.isBeforeInBuffer(StartLineLoc, TargetLineLoc) && !SM.isBeforeInBuffer(EndLineLoc, TargetLineLoc) &&
             TokenStr.startswith("/*");
         InCommentLine |= StartLineLoc == TargetLineLoc &&
             TokenStr.startswith("//");
+      } else if (CurrentTokIt->getKind() == tok::unknown &&
+                 CurrentTokIt->getRange().str().startswith("\"\"\"")) {
+        StringLiteralRange = CurrentTokIt->getRange();
       }
     }
   }
@@ -2797,12 +2902,6 @@ class CodeFormatter {
   std::pair<LineRange, std::string> indent(unsigned LineIndex,
                                            FormatContext &FC,
                                            StringRef Text) {
-    if (FC.IsInStringLiteral()) {
-      return std::make_pair(
-          LineRange(LineIndex, 1),
-          swift::ide::getTextForLine(LineIndex, Text, /*Trim*/ false).str());
-    }
-
     if (FC.isExact()) {
       StringRef Line = swift::ide::getTextForLine(LineIndex, Text, /*Trim*/true);
       StringBuilder Builder;

test/SourceKit/CodeFormat/indent-multiline-string-interp.swift

@@ -0,0 +1,57 @@
+let s1 = """
+  \(
+    45
+)
+  """
+
+let s2 = """
+  \(
+    45
+  )
+"""
+
+let s3 = """
+\(
+  45
+)
+  """
+
+let s4 = """
+  foo \( 45 /*
+comment content
+*/) bar
+  """
+
+// RUN: %sourcekitd-test -req=format -line=2 %s >%t.response
+// RUN: %sourcekitd-test -req=format -line=3 %s >>%t.response
+// RUN: %sourcekitd-test -req=format -line=4 %s >>%t.response
+
+// RUN: %sourcekitd-test -req=format -line=8 %s >>%t.response
+// RUN: %sourcekitd-test -req=format -line=9 %s >>%t.response
+// RUN: %sourcekitd-test -req=format -line=10 %s >>%t.response
+
+// RUN: %sourcekitd-test -req=format -line=14 %s >>%t.response
+// RUN: %sourcekitd-test -req=format -line=15 %s >>%t.response
+// RUN: %sourcekitd-test -req=format -line=16 %s >>%t.response
+
+// RUN: %sourcekitd-test -req=format -line=20 %s >>%t.response
+// RUN: %sourcekitd-test -req=format -line=21 %s >>%t.response
+// RUN: %sourcekitd-test -req=format -line=22 %s >>%t.response
+
+// RUN: %FileCheck --strict-whitespace %s <%t.response
+
+// CHECK: key.sourcetext: "  \\("
+// CHECK: key.sourcetext: "    45"
+// CHECK: key.sourcetext: "  )"
+
+// CHECK: key.sourcetext: "  \\("
+// CHECK: key.sourcetext: "    45"
+// CHECK: key.sourcetext: "  )"
+
+// CHECK: key.sourcetext: "  \\("
+// CHECK: key.sourcetext: "    45"
+// CHECK: key.sourcetext: ")"
+
+// CHECK: key.sourcetext: "  foo \\( 45 /*"
+// CHECK: key.sourcetext: "         comment content"
+// CHECK: key.sourcetext: "         */) bar"

test/SourceKit/CodeFormat/indent-multiline-string-nested.swift

@@ -0,0 +1,25 @@
+func foo() {
+    let s1 = """
+        this is line1 in outer string \("""
+  nested string in interpolation
+            """)
+
+    this is line2 in outer string
+        """
+}
+
+// RUN: %sourcekitd-test -req=format -line=3 %s >%t.response
+// RUN: %sourcekitd-test -req=format -line=4 %s >>%t.response
+// RUN: %sourcekitd-test -req=format -line=5 %s >>%t.response
+// RUN: %sourcekitd-test -req=format -line=6 %s >>%t.response
+// RUN: %sourcekitd-test -req=format -line=7 %s >>%t.response
+// RUN: %sourcekitd-test -req=format -line=8 %s >>%t.response
+
+// RUN: %FileCheck --strict-whitespace %s <%t.response
+
+// CHECK: key.sourcetext: "        this is line1 in outer string \\(\"\"\""
+// CHECK: key.sourcetext: "            nested string in interpolation"
+// CHECK: key.sourcetext: "            \"\"\")"
+// CHECK: key.sourcetext: "        "
+// CHECK: key.sourcetext: "        this is line2 in outer string"
+// CHECK: key.sourcetext: "        \"\"\""

test/SourceKit/CodeFormat/indent-multiline-string-trailing-interp.swift

@@ -0,0 +1,31 @@
+// RUN: %sourcekitd-test -req=format -line=24 %s >%t.response
+// RUN: %sourcekitd-test -req=format -line=25 %s >>%t.response
+// RUN: %sourcekitd-test -req=format -line=26 %s >>%t.response
+// RUN: %sourcekitd-test -req=format -line=27 %s >>%t.response
+// RUN: %sourcekitd-test -req=format -line=28 %s >>%t.response
+// RUN: %sourcekitd-test -req=format -line=29 %s >>%t.response
+// RUN: %sourcekitd-test -req=format -line=30 %s >>%t.response
+// RUN: %sourcekitd-test -req=format -line=31 %s >>%t.response
+
+// RUN: %FileCheck --strict-whitespace %s <%t.response
+
+// The end quotes are ignored as it would be impossible to know whether
+// they are part of the interpolation or another string, etc.
+// CHECK: key.sourcetext: "    "
+// CHECK: key.sourcetext: "    this is line1"
+// CHECK: key.sourcetext: ""
+// CHECK: key.sourcetext: "this is line2"
+// CHECK: key.sourcetext: "    "
+// CHECK: key.sourcetext: "    and this is a line with trailing interpolation \\(1 +"
+// CHECK: key.sourcetext: " "
+// CHECK: key.sourcetext: " \"\"\""
+
+let s1 = """
+
+this is line1
+
+    this is line2
+
+ and this is a line with trailing interpolation \(1 +
+
+  """