1 files changed, 63 insertions, 1 deletions

diff --git a/clang/lib/Format/Format.cpp b/clang/lib/Format/Format.cpp
index 7085aaf8b77..3bedc6a4cd5 100644
--- a/clang/lib/Format/Format.cpp
+++ b/clang/lib/Format/Format.cpp
@@ -751,6 +751,67 @@ private:
   std::vector<TokenAnnotation> Annotations;
 };
 
+class LexerBasedFormatTokenSource : public FormatTokenSource {
+public:
+  LexerBasedFormatTokenSource(Lexer &Lex, SourceManager &SourceMgr)
+      : GreaterStashed(false),
+        Lex(Lex),
+        SourceMgr(SourceMgr),
+        IdentTable(Lex.getLangOpts()) {
+    Lex.SetKeepWhitespaceMode(true);
+  }
+
+  virtual FormatToken getNextToken() {
+    if (GreaterStashed) {
+      FormatTok.NewlinesBefore = 0;
+      FormatTok.WhiteSpaceStart =
+          FormatTok.Tok.getLocation().getLocWithOffset(1);
+      FormatTok.WhiteSpaceLength = 0;
+      GreaterStashed = false;
+      return FormatTok;
+    }
+
+    FormatTok = FormatToken();
+    Lex.LexFromRawLexer(FormatTok.Tok);
+    FormatTok.WhiteSpaceStart = FormatTok.Tok.getLocation();
+
+    // Consume and record whitespace until we find a significant token.
+    while (FormatTok.Tok.is(tok::unknown)) {
+      FormatTok.NewlinesBefore += tokenText(FormatTok.Tok).count('\n');
+      FormatTok.WhiteSpaceLength += FormatTok.Tok.getLength();
+
+      if (FormatTok.Tok.is(tok::eof))
+        return FormatTok;
+      Lex.LexFromRawLexer(FormatTok.Tok);
+    }
+
+    if (FormatTok.Tok.is(tok::raw_identifier)) {
+      const IdentifierInfo &Info = IdentTable.get(tokenText(FormatTok.Tok));
+      FormatTok.Tok.setKind(Info.getTokenID());
+    }
+
+    if (FormatTok.Tok.is(tok::greatergreater)) {
+      FormatTok.Tok.setKind(tok::greater);
+      GreaterStashed = true;
+    }
+
+    return FormatTok;
+  }
+
+private:
+  FormatToken FormatTok;
+  bool GreaterStashed;
+  Lexer &Lex;
+  SourceManager &SourceMgr;
+  IdentifierTable IdentTable;
+
+  /// Returns the text of \c FormatTok.
+  StringRef tokenText(Token &Tok) {
+    return StringRef(SourceMgr.getCharacterData(Tok.getLocation()),
+                     Tok.getLength());
+  }
+};
+
 class Formatter : public UnwrappedLineConsumer {
 public:
   Formatter(const FormatStyle &Style, Lexer &Lex, SourceManager &SourceMgr,
@@ -766,7 +827,8 @@ public:
   }
 
   tooling::Replacements format() {
-    UnwrappedLineParser Parser(Style, Lex, SourceMgr, *this);
+    LexerBasedFormatTokenSource Tokens(Lex, SourceMgr);
+    UnwrappedLineParser Parser(Style, Tokens, *this);
     StructuralError = Parser.parse();
     for (std::vector<UnwrappedLine>::iterator I = UnwrappedLines.begin(),
                                               E = UnwrappedLines.end();