Merge pull request swiftlang#33397 from bnbarham/benb/multiline-indentation-32181422

Author: swift-ci

include/swift/Parse/Lexer.h

@@ -352,7 +352,7 @@ class Lexer {
   static SourceLoc getLocForStartOfLine(SourceManager &SM, SourceLoc Loc);
 
   /// Retrieve the source location for the end of the line containing the
-  /// given token, which is the location of the start of the next line.
+  /// given location, which is the location of the start of the next line.
   static SourceLoc getLocForEndOfLine(SourceManager &SM, SourceLoc Loc);
 
   /// Retrieve the string used to indent the line that contains the given

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;

lib/Parse/Lexer.cpp

@@ -2706,12 +2706,12 @@ static SourceLoc getLocForStartOfTokenInBuf(SourceManager &SM,
 // Find the start of the given line.
 static const char *findStartOfLine(const char *bufStart, const char *current) {
   while (current != bufStart) {
-    if (current[0] == '\n' || current[0] == '\r') {
+    --current;
+
+    if (current[0] == '\n') {
       ++current;
       break;
     }
-
-    --current;
   }
 
   return current;
@@ -2779,19 +2779,16 @@ SourceLoc Lexer::getLocForEndOfLine(SourceManager &SM, SourceLoc Loc) {
   if (BufferID < 0)
     return SourceLoc();
 
-  // Use fake language options; language options only affect validity
-  // and the exact token produced.
-  LangOptions FakeLangOpts;
+  CharSourceRange entireRange = SM.getRangeForBuffer(BufferID);
+  StringRef Buffer = SM.extractText(entireRange);
 
-  // Here we return comments as tokens because either the caller skipped
-  // comments and normally we won't be at the beginning of a comment token
-  // (making this option irrelevant), or the caller lexed comments and
-  // we need to lex just the comment token.
-  Lexer L(FakeLangOpts, SM, BufferID, nullptr, LexerMode::Swift,
-          HashbangMode::Allowed, CommentRetentionMode::ReturnAsTokens);
-  L.restoreState(State(Loc));
-  L.skipToEndOfLine(/*EatNewline=*/true);
-  return getSourceLoc(L.CurPtr);
+  // Windows line endings are \r\n. Since we want the start of the next
+  // line, just look for \n so the \r is skipped through.
+  size_t Offset = SM.getLocOffsetInBuffer(Loc, BufferID);
+  Offset = Buffer.find('\n', Offset);
+  if (Offset == StringRef::npos)
+    return SourceLoc();
+  return getSourceLoc(Buffer.data() + Offset + 1);
 }
 
 StringRef Lexer::getIndentationForLine(SourceManager &SM, SourceLoc Loc,

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"