Formatter: Detect ObjC message expressions after 'in' in loop

Before:
  for (id foo in[self getStuffFor : bla]) {
  }

Now:
  for (id foo in [self getStuffFor:bla]) {
  }

"in" is treated as loop keyword if the line starts with "for", and as a
regular identifier else. To check for "in", its IdentifierInfo is handed
through a few layers.



git-svn-id: https://llvm.org/svn/llvm-project/cfe/trunk@174889 91177308-0d34-0410-b5e6-96231b3b80d8

4 files changed, 26 insertions, 6 deletions

diff --git a/lib/Format/Format.cpp b/lib/Format/Format.cpp
index 6dc84a10c6..347e31470d 100644
--- a/lib/Format/Format.cpp
+++ b/lib/Format/Format.cpp
@@ -842,6 +842,8 @@ public:
     return FormatTok;
   }
 
+  IdentifierTable &getIdentTable() { return IdentTable; }
+
 private:
   FormatToken FormatTok;
   bool GreaterStashed;
@@ -909,7 +911,8 @@ public:
     UnwrappedLineParser Parser(Diag, Style, Tokens, *this);
     StructuralError = Parser.parse();
     unsigned PreviousEndOfLineColumn = 0;
-    TokenAnnotator Annotator(Style, SourceMgr, Lex);
+    TokenAnnotator Annotator(Style, SourceMgr, Lex,
+                             Tokens.getIdentTable().get("in"));
     for (unsigned i = 0, e = AnnotatedLines.size(); i != e; ++i) {
       Annotator.annotate(AnnotatedLines[i]);
     }
diff --git a/lib/Format/TokenAnnotator.cpp b/lib/Format/TokenAnnotator.cpp
index 8f2e00ebfd..1e814b2d39 100644
--- a/lib/Format/TokenAnnotator.cpp
+++ b/lib/Format/TokenAnnotator.cpp
@@ -78,9 +78,10 @@ static const AnnotatedToken *getNextToken(const AnnotatedToken &Tok) {
 /// into template parameter lists.
 class AnnotatingParser {
 public:
-  AnnotatingParser(SourceManager &SourceMgr, Lexer &Lex, AnnotatedLine &Line)
+  AnnotatingParser(SourceManager &SourceMgr, Lexer &Lex, AnnotatedLine &Line,
+                   IdentifierInfo &Ident_in)
       : SourceMgr(SourceMgr), Lex(Lex), Line(Line), CurrentToken(&Line.First),
-        KeywordVirtualFound(false) {
+        KeywordVirtualFound(false), Ident_in(Ident_in) {
     Contexts.push_back(Context(1, /*IsExpression=*/ false));
     Contexts.back().LookForFunctionName = Line.MustBeDeclaration;
   }
@@ -195,6 +196,7 @@ public:
         !Parent || Parent->is(tok::colon) || Parent->is(tok::l_square) ||
         Parent->is(tok::l_paren) || Parent->is(tok::kw_return) ||
         Parent->is(tok::kw_throw) || isUnaryOperator(*Parent) ||
+        Parent->Type == TT_ObjCForIn ||
         getBinOpPrecedence(Parent->FormatTok.Tok.getKind(), true, true) >
         prec::Unknown;
     bool StartsObjCArrayLiteral = Parent && Parent->is(tok::at);
@@ -383,6 +385,11 @@ public:
     case tok::kw_template:
       parseTemplateDeclaration();
       break;
+    case tok::identifier:
+      if (Line.First.is(tok::kw_for) &&
+          Tok->FormatTok.Tok.getIdentifierInfo() == &Ident_in)
+        Tok->Type = TT_ObjCForIn;
+      break;
     default:
       break;
     }
@@ -684,6 +691,7 @@ private:
   AnnotatedLine &Line;
   AnnotatedToken *CurrentToken;
   bool KeywordVirtualFound;
+  IdentifierInfo &Ident_in;
 };
 
 /// \brief Parses binary expressions by inserting fake parenthesis based on
@@ -763,7 +771,7 @@ private:
 };
 
 void TokenAnnotator::annotate(AnnotatedLine &Line) {
-  AnnotatingParser Parser(SourceMgr, Lex, Line);
+  AnnotatingParser Parser(SourceMgr, Lex, Line, Ident_in);
   Line.Type = Parser.parseLine();
   if (Line.Type == LT_Invalid)
     return;
diff --git a/lib/Format/TokenAnnotator.h b/lib/Format/TokenAnnotator.h
index c19b486c1c..506f271dc5 100644
--- a/lib/Format/TokenAnnotator.h
+++ b/lib/Format/TokenAnnotator.h
@@ -38,6 +38,7 @@ enum TokenType {
   TT_ObjCArrayLiteral,
   TT_ObjCBlockLParen,
   TT_ObjCDecl,
+  TT_ObjCForIn,
   TT_ObjCMethodExpr,
   TT_ObjCMethodSpecifier,
   TT_ObjCProperty,
@@ -178,8 +179,9 @@ inline prec::Level getPrecedence(const AnnotatedToken &Tok) {
 /// \c UnwrappedLine.
 class TokenAnnotator {
 public:
-  TokenAnnotator(const FormatStyle &Style, SourceManager &SourceMgr, Lexer &Lex)
-      : Style(Style), SourceMgr(SourceMgr), Lex(Lex) {
+  TokenAnnotator(const FormatStyle &Style, SourceManager &SourceMgr, Lexer &Lex,
+                 IdentifierInfo &Ident_in)
+      : Style(Style), SourceMgr(SourceMgr), Lex(Lex), Ident_in(Ident_in) {
   }
 
   void annotate(AnnotatedLine &Line);
@@ -201,6 +203,9 @@ private:
   const FormatStyle &Style;
   SourceManager &SourceMgr;
   Lexer &Lex;
+
+  // Contextual keywords:
+  IdentifierInfo &Ident_in;
 };
 
 } // end namespace format
diff --git a/unittests/Format/FormatTest.cpp b/unittests/Format/FormatTest.cpp
index f9e73bb102..9b9f5c71d2 100644
--- a/unittests/Format/FormatTest.cpp
+++ b/unittests/Format/FormatTest.cpp
@@ -2427,6 +2427,10 @@ TEST_F(FormatTest, FormatObjCMethodExpr) {
   verifyFormat("[foo bar:baz] % [foo bar:baz];");
   // Whew!
 
+  verifyFormat("return in[42];");
+  verifyFormat("for (id foo in [self getStuffFor:bla]) {\n"
+               "}");
+
   verifyFormat("[self stuffWithInt:(4 + 2) float:4.5];");
   verifyFormat("[self stuffWithInt:a ? b : c float:4.5];");
   verifyFormat("[self stuffWithInt:a ? [self foo:bar] : c];");