SourceKit/Formatting: avoid indenting for consecutive dot-member calls

rdar://52392291

Author: nkcsgexi

include/swift/Basic/SourceManager.h

@@ -231,9 +231,14 @@ class SourceManager {
   void verifyAllBuffers() const;
 
   /// Translate line and column pair to the offset.
+  /// If the column number is the maximum unsinged int, return the offset of the end of the line.
   llvm::Optional<unsigned> resolveFromLineCol(unsigned BufferId, unsigned Line,
                                               unsigned Col) const;
 
+  /// Translate the end position of the given line to the offset.
+  llvm::Optional<unsigned> resolveOffsetForEndOfLine(unsigned BufferId,
+                                                     unsigned Line) const;
+
   SourceLoc getLocForLineCol(unsigned BufferId, unsigned Line, unsigned Col) const {
     auto Offset = resolveFromLineCol(BufferId, Line, Col);
     return Offset.hasValue() ? getLocForOffset(BufferId, Offset.getValue()) :

lib/Basic/SourceLoc.cpp

@@ -310,12 +310,19 @@ void CharSourceRange::dump(const SourceManager &SM) const {
   print(llvm::errs(), SM);
 }
 
+llvm::Optional<unsigned>
+SourceManager::resolveOffsetForEndOfLine(unsigned BufferId,
+                                         unsigned Line) const {
+  return resolveFromLineCol(BufferId, Line, ~0u);
+}
+
 llvm::Optional<unsigned> SourceManager::resolveFromLineCol(unsigned BufferId,
                                                            unsigned Line,
                                                            unsigned Col) const {
   if (Line == 0 || Col == 0) {
     return None;
   }
+  const bool LineEnd = Col == ~0u;
   auto InputBuf = getLLVMSourceMgr().getMemoryBuffer(BufferId);
   const char *Ptr = InputBuf->getBufferStart();
   const char *End = InputBuf->getBufferEnd();
@@ -331,14 +338,18 @@ llvm::Optional<unsigned> SourceManager::resolveFromLineCol(unsigned BufferId,
     return None;
   }
   Ptr = LineStart;
-
   // The <= here is to allow for non-inclusive range end positions at EOF
-  for (; Ptr <= End; ++Ptr) {
+  for (; ; ++Ptr) {
     --Col;
     if (Col == 0)
       return Ptr - InputBuf->getBufferStart();
-    if (*Ptr == '\n')
-      break;
+    if (*Ptr == '\n' || Ptr == End) {
+      if (LineEnd) {
+        return Ptr - InputBuf->getBufferStart();
+      } else {
+        break;
+      }
+    }
   }
   return None;
 }

lib/IDE/Formatting.cpp

@@ -29,14 +29,35 @@ struct SiblingAlignInfo {
 };
 
 struct TokenInfo {
-  const Token *StartOfLineTarget;
-  const Token *StartOfLineBeforeTarget;
-  TokenInfo(const Token *StartOfLineTarget,
-            const Token *StartOfLineBeforeTarget) :
-    StartOfLineTarget(StartOfLineTarget),
-    StartOfLineBeforeTarget(StartOfLineBeforeTarget) {}
-  TokenInfo() : TokenInfo(nullptr, nullptr) {}
-  operator bool() { return StartOfLineTarget && StartOfLineBeforeTarget; }
+  // The tokens appearing at the start of lines, from the first line to the
+  // line under indentation.
+  std::vector<const Token*> LineStarts;
+  const Token *getLineStarter(unsigned idx = 0) const {
+    auto it = LineStarts.rbegin() + idx;
+    return it < LineStarts.rend() ? *it : nullptr;
+  }
+  operator bool() { return LineStarts.size() > 1; }
+  bool isRBraceDotsPattern() const {
+    for (auto It = LineStarts.rbegin(), end = LineStarts.rend();
+         It + 1 < end; ++It) {
+      auto* CurrentLine = *It;
+      auto* PreviousLine = *(It + 1);
+      if (CurrentLine->getKind() == tok::period &&
+          PreviousLine->getKind() == tok::period) {
+        // If the previous line starts with dot too, move further up.
+        continue;
+      } else if (CurrentLine->getKind() == tok::period &&
+                 PreviousLine->getKind() == tok::r_brace &&
+                 PreviousLine + 1 == CurrentLine) {
+        // Check if the previous line starts with '}' and the period of the
+        // current line is immediately after the '}'
+        return true;
+      } else {
+        return false;
+      }
+    }
+    return false;
+  }
 };
 
 using StringBuilder = llvm::SmallString<64>;
@@ -285,17 +306,16 @@ class FormatContext {
       return false;
 
     if (TInfo) {
-      if (TInfo.StartOfLineTarget->getKind() == tok::l_brace &&
-          isKeywordPossibleDeclStart(*TInfo.StartOfLineBeforeTarget) &&
-          TInfo.StartOfLineBeforeTarget->isKeyword())
+      if (TInfo.getLineStarter()->getKind() == tok::l_brace &&
+          isKeywordPossibleDeclStart(*TInfo.getLineStarter(1)) &&
+          TInfo.getLineStarter(1)->isKeyword())
         return false;
       // VStack {
       //   ...
       // }
       // .onAppear { <---- No indentation here.
-      if (TInfo.StartOfLineTarget->getKind() == tok::period &&
-          TInfo.StartOfLineBeforeTarget->getKind() == tok::r_brace &&
-          TInfo.StartOfLineBeforeTarget + 1 == TInfo.StartOfLineTarget) {
+      // .onAppear1 { <---- No indentation here.
+      if (TInfo.isRBraceDotsPattern()) {
         return false;
       }
     }
@@ -365,7 +385,7 @@ class FormatContext {
       SignatureEnd = SD->getSignatureSourceRange().End;
     }
     if (SignatureEnd.isValid() && TInfo &&
-        TInfo.StartOfLineTarget->getLoc() == SignatureEnd) {
+        TInfo.getLineStarter()->getLoc() == SignatureEnd) {
       return false;
     }
 
@@ -862,7 +882,7 @@ class CodeFormatter {
 
     // If having sibling locs to align with, respect siblings.
     auto isClosingSquare =
-      ToInfo && ToInfo.StartOfLineTarget->getKind() == tok::r_square;
+      ToInfo && ToInfo.getLineStarter()->getKind() == tok::r_square;
     if (!isClosingSquare && FC.HasSibling()) {
       StringRef Line = swift::ide::getTextForLine(LineIndex, Text, /*Trim*/true);
       StringBuilder Builder;
@@ -940,33 +960,37 @@ class TokenInfoCollector {
   SourceManager &SM;
   ArrayRef<Token> Tokens;
   unsigned Line;
-
-  struct Comparator {
-    SourceManager &SM;
-    Comparator(SourceManager &SM) : SM(SM) {}
-    bool operator()(const Token &T, unsigned Line) const {
-      return SM.getLineNumber(T.getLoc()) < Line;
-    }
-    bool operator()(unsigned Line, const Token &T) const {
-      return Line < SM.getLineNumber(T.getLoc());
-    }
-  };
-
+  // The location of the end of the line under indentation, we don't need to
+  // collect tokens after this location.
+  SourceLoc EndLimit;
 public:
-  TokenInfoCollector(SourceManager &SM, ArrayRef<Token> Tokens,
-         unsigned Line) : SM(SM), Tokens(Tokens), Line(Line) {}
+  TokenInfoCollector(SourceManager &SM,
+                     unsigned BufferId,
+                     ArrayRef<Token> Tokens, unsigned Line):
+                     SM(SM), Tokens(Tokens), Line(Line) {
+    if (auto Offset = SM.resolveOffsetForEndOfLine(BufferId, Line)) {
+      EndLimit = SM.getLocForOffset(BufferId, *Offset);
+    }
+  }
 
   TokenInfo collect() {
-    if (Line == 0)
+    if (EndLimit.isInvalid()) {
       return TokenInfo();
-    Comparator Comp(SM);
-    auto LineMatch = [this] (const Token* T, unsigned Line) {
-      return T != Tokens.end() && SM.getLineNumber(T->getLoc()) == Line;
-    };
-    auto TargetIt = std::lower_bound(Tokens.begin(), Tokens.end(), Line, Comp);
-    auto LineBefore = std::lower_bound(Tokens.begin(), TargetIt, Line - 1, Comp);
-    if (LineMatch(TargetIt, Line) && LineMatch(LineBefore, Line - 1))
-      return TokenInfo(TargetIt, LineBefore);
+    }
+    TokenInfo Result;
+    for(auto &T: Tokens) {
+      if (!T.isAtStartOfLine())
+        continue;
+      if (SM.isBeforeInBuffer(EndLimit, T.getLoc())) {
+        if (!Result.LineStarts.empty()) {
+          if (SM.getLineNumber(Result.getLineStarter()->getLoc()) == Line) {
+            return Result;
+          }
+        }
+        return TokenInfo();
+      }
+      Result.LineStarts.push_back(&T);
+    }
     return TokenInfo();
   }
 };
@@ -1051,7 +1075,8 @@ std::pair<LineRange, std::string> swift::ide::reformat(LineRange Range,
   FormatContext FC = walker.walkToLocation(Loc);
   CodeFormatter CF(Options);
   unsigned Line = Range.startLine();
-  return CF.indent(Line, FC, Text, TokenInfoCollector(SM, walker.getTokens(),
+  return CF.indent(Line, FC, Text, TokenInfoCollector(SM, SourceBufferID,
+                                                      walker.getTokens(),
                                                       Line).collect());
 }
 

test/SourceKit/CodeFormat/indent-closure.swift

@@ -74,6 +74,16 @@ func foo11() {
     }
 }
 
+func foo12() {
+    VStack {
+    }
+        .onAppear1()
+        .onAppear2() {}
+        .onAppear3() {
+        }
+        .onAppear4() {}
+}
+
 // RUN: %sourcekitd-test -req=format -line=3 -length=1 %s >%t.response
 // RUN: %sourcekitd-test -req=format -line=4 -length=1 %s >>%t.response
 // RUN: %sourcekitd-test -req=format -line=5 -length=1 %s >>%t.response
@@ -105,6 +115,12 @@ func foo11() {
 // RUN: %sourcekitd-test -req=format -line=67 -length=1 %s >>%t.response
 // RUN: %sourcekitd-test -req=format -line=73 -length=1 %s >>%t.response
 
+// RUN: %sourcekitd-test -req=format -line=80 -length=1 %s >>%t.response
+// RUN: %sourcekitd-test -req=format -line=81 -length=1 %s >>%t.response
+// RUN: %sourcekitd-test -req=format -line=82 -length=1 %s >>%t.response
+// RUN: %sourcekitd-test -req=format -line=83 -length=1 %s >>%t.response
+// RUN: %sourcekitd-test -req=format -line=84 -length=1 %s >>%t.response
+
 // RUN: %FileCheck --strict-whitespace %s <%t.response
 
 // CHECK: key.sourcetext: "        var abc = 1"
@@ -145,3 +161,9 @@ func foo11() {
 // CHECK: key.sourcetext: "    Something() ["
 // CHECK: key.sourcetext: "    ].whatever"
 // CHECK: key.sourcetext: "    .onAppear {"
+
+// CHECK: key.sourcetext: "    .onAppear1()"
+// CHECK: key.sourcetext: "    .onAppear2() {}"
+// CHECK: key.sourcetext: "    .onAppear3() {"
+// CHECK: key.sourcetext: "    }"
+// CHECK: key.sourcetext: "    .onAppear4() {}"