3 files changed, 29 insertions, 21 deletions

diff --git a/clang/lib/Format/TokenAnnotator.cpp b/clang/lib/Format/TokenAnnotator.cpp
index 38a1c1a5e1e..74db4969867 100644
--- a/clang/lib/Format/TokenAnnotator.cpp
+++ b/clang/lib/Format/TokenAnnotator.cpp
@@ -624,7 +624,9 @@ private:
         Current.Type = determineIncrementUsage(Current);
       } else if (Current.is(tok::exclaim)) {
         Current.Type = TT_UnaryOperator;
-      } else if (Current.isBinaryOperator()) {
+      } else if (Current.isBinaryOperator() &&
+                 (!Current.Previous ||
+                  Current.Previous->isNot(tok::l_square))) {
         Current.Type = TT_BinaryOperator;
       } else if (Current.is(tok::comment)) {
         if (Current.TokenText.startswith("//"))
@@ -1188,6 +1190,8 @@ bool TokenAnnotator::spaceRequiredBetween(const AnnotatedLine &Line,
     return false;
   if (Right.is(tok::ellipsis))
     return false;
+  if (Left.is(tok::l_square) && Right.is(tok::amp))
+    return false;
   if (Right.Type == TT_PointerOrReference)
     return Left.Tok.isLiteral() ||
            ((Left.Type != TT_PointerOrReference) && Left.isNot(tok::l_paren) &&
@@ -1236,7 +1240,8 @@ bool TokenAnnotator::spaceRequiredBetween(const AnnotatedLine &Line,
     return !Left.isOneOf(tok::l_paren, tok::l_square, tok::at) &&
            (Left.isNot(tok::colon) || Left.Type != TT_ObjCMethodExpr);
   if (Left.isOneOf(tok::identifier, tok::greater, tok::r_square) &&
-      Right.is(tok::l_brace) && Right.getNextNonComment())
+      Right.is(tok::l_brace) && Right.getNextNonComment() &&
+      Right.BlockKind != BK_Block)
     return false;
   if (Left.is(tok::period) || Right.is(tok::period))
     return false;
diff --git a/clang/lib/Format/UnwrappedLineParser.cpp b/clang/lib/Format/UnwrappedLineParser.cpp
index 253dbf97fdc..f58ee823032 100644
--- a/clang/lib/Format/UnwrappedLineParser.cpp
+++ b/clang/lib/Format/UnwrappedLineParser.cpp
@@ -278,13 +278,11 @@ void UnwrappedLineParser::calculateBraceTypes() {
       if (!LBraceStack.empty()) {
         if (LBraceStack.back()->BlockKind == BK_Unknown) {
           // If there is a comma, semicolon or right paren after the closing
-          // brace, we assume this is a braced initializer list.
-
-          // FIXME: Note that this currently works only because we do not
-          // use the brace information while inside a braced init list.
-          // Thus, if the parent is a braced init list, we consider all
-          // brace blocks inside it braced init list. That works good enough
-          // for now, but we will need to fix it to correctly handle lambdas.
+          // brace, we assume this is a braced initializer list.  Note that
+          // regardless how we mark inner braces here, we will overwrite the
+          // BlockKind later if we parse a braced list (where all blocks inside
+          // are by default braced lists), or when we explicitly detect blocks
+          // (for example while parsing lambdas).
           //
           // We exclude + and - as they can be ObjC visibility modifiers.
           if (NextTok->isOneOf(tok::comma, tok::semi, tok::r_paren,
@@ -315,12 +313,13 @@ void UnwrappedLineParser::calculateBraceTypes() {
     }
     Tok = NextTok;
     Position += ReadTokens;
-  } while (Tok->Tok.isNot(tok::eof));
+  } while (Tok->Tok.isNot(tok::eof) && !LBraceStack.empty());
   // Assume other blocks for all unclosed opening braces.
   for (unsigned i = 0, e = LBraceStack.size(); i != e; ++i) {
     if (LBraceStack[i]->BlockKind == BK_Unknown)
       LBraceStack[i]->BlockKind = BK_Block;
   }
+
   FormatTok = Tokens->setPosition(StoredPosition);
 }
 
@@ -675,6 +674,7 @@ void UnwrappedLineParser::tryToParseLambda() {
         break;
     }
   }
+  FormatTok->BlockKind = BK_Block;
   nextToken();
   {
     ScopedLineState LineState(*this);
@@ -745,6 +745,9 @@ void UnwrappedLineParser::parseBracedList() {
         tryToParseLambda();
         break;
     case tok::l_brace:
+      // Assume there are no blocks inside a braced init list apart
+      // from the ones we explicitly parse out (like lambdas).
+      FormatTok->BlockKind = BK_BracedInit;
       parseBracedList();
       break;
     case tok::r_brace:
diff --git a/clang/unittests/Format/FormatTest.cpp b/clang/unittests/Format/FormatTest.cpp
index e9bf2c54fd3..3d2cf13c5f3 100644
--- a/clang/unittests/Format/FormatTest.cpp
+++ b/clang/unittests/Format/FormatTest.cpp
@@ -6262,37 +6262,37 @@ TEST_F(FormatTest, FormatsLambdas) {
   // parsing of the unwrapped lines doesn't regress.
   verifyFormat(
       "int c = [b]() mutable {\n"
-      "  return [&b]{\n"
+      "  return [&b] {\n"
       "    return b++;\n"
       "  }();\n"
       "}();\n");
   verifyFormat(
-      "int c = [&]{\n"
-      "  [ = ]{\n"
+      "int c = [&] {\n"
+      "  [=] {\n"
       "    return b++;\n"
       "  }();\n"
       "}();\n");
   verifyFormat(
-      "int c = [ &, &a, a]{\n"
-      "  [ =, c, &d]{\n"
+      "int c = [&, &a, a] {\n"
+      "  [=, c, &d] {\n"
       "    return b++;\n"
       "  }();\n"
       "}();\n");
   verifyFormat(
-      "int c = [&a, &a, a]{\n"
-      "  [ =, a, b, &c]{\n"
+      "int c = [&a, &a, a] {\n"
+      "  [=, a, b, &c] {\n"
       "    return b++;\n"
       "  }();\n"
       "}();\n");
   verifyFormat(
-      "auto c = {[&a, &a, a]{\n"
-      "  [ =, a, b, &c]{\n"
+      "auto c = {[&a, &a, a] {\n"
+      "  [=, a, b, &c] {\n"
       "    return b++;\n"
       "  }();\n"
       "} }\n");
   verifyFormat(
-      "auto c = {[&a, &a, a]{\n"
-      "  [ =, a, b, &c]{\n"
+      "auto c = {[&a, &a, a] {\n"
+      "  [=, a, b, &c] {\n"
       "  }();\n"
       "} }\n");
 }