[CodeCompletion] Fix a crash in CCExprRemover

- Handle cases where getArgumentLabelLocs().size() == 0
- Add some assertions to verify invariants
- Explicit handling of 'llvm::Optional' for 'getUnlabeledTrailingClosureIndex()'
- Avoid walking into nodes after the removing happens

rdar://problem/65556791
(cherry picked from commit f7fc1ed)

Author: rintaro

lib/IDE/ExprContextAnalysis.cpp

@@ -226,16 +226,30 @@ class CCExprRemover: public ASTWalker, public ExprVisitor<CCExprRemover, Expr *>
     } else if (auto tuple = dyn_cast<TupleExpr>(E->getArg())) {
       lParenLoc = tuple->getLParenLoc();
       rParenLoc = tuple->getRParenLoc();
+
+      assert((!E->getUnlabeledTrailingClosureIndex().hasValue() ||
+              (tuple->getNumElements() == E->getArgumentLabels().size() &&
+               tuple->getNumElements() == E->getArgumentLabelLocs().size())) &&
+             "CallExpr with trailing closure must have the same number of "
+             "argument labels");
+      assert(tuple->getNumElements() == E->getArgumentLabels().size());
+      assert(tuple->getNumElements() == E->getArgumentLabelLocs().size() ||
+             E->getArgumentLabelLocs().size() == 0);
+
+      bool hasArgumentLabelLocs = E->getArgumentLabelLocs().size() > 0;
+
       for (unsigned i = 0, e = tuple->getNumElements(); i != e; ++i) {
         if (isa<CodeCompletionExpr>(tuple->getElement(i))) {
           removing = true;
           continue;
         }
 
-        if (i < E->getUnlabeledTrailingClosureIndex()) {
+        if (!E->getUnlabeledTrailingClosureIndex().hasValue() ||
+            i < *E->getUnlabeledTrailingClosureIndex()) {
           // Normal arguments.
           argLabels.push_back(E->getArgumentLabels()[i]);
-          argLabelLocs.push_back(E->getArgumentLabelLocs()[i]);
+          if (hasArgumentLabelLocs)
+            argLabelLocs.push_back(E->getArgumentLabelLocs()[i]);
           args.push_back(tuple->getElement(i));
         } else {
           // Trailing closure arguments.
@@ -259,7 +273,20 @@ class CCExprRemover: public ASTWalker, public ExprVisitor<CCExprRemover, Expr *>
   }
 
   std::pair<bool, Expr *> walkToExprPre(Expr *E) override {
-    return {true, visit(E)};
+    if (Removed)
+      return {false, nullptr};
+    E = visit(E);
+    return {!Removed, E};
+  }
+
+  std::pair<bool, Stmt *> walkToStmtPre(Stmt *S) override {
+    if (Removed)
+      return {false, nullptr};
+    return {true, S};
+  }
+
+  bool walkToDeclPre(Decl *D) override {
+    return !Removed;
   }
 };
 }

validation-test/IDE/crashers_2_fixed/rdar65556791.swift

@@ -0,0 +1,17 @@
+// RUN: %target-swift-ide-test -code-completion -code-completion-token=A -source-filename=%s
+// RUN: %target-swift-ide-test -code-completion -code-completion-token=B -source-filename=%s
+// RUN: %target-swift-ide-test -code-completion -code-completion-token=C -source-filename=%s
+
+func myFunc(_: Int, _: Undefined) {}
+
+undefined {
+  myFunc($0, undefined)
+} #^A^#
+
+undefined(x: 1) {
+  myFunc($0, undefined)
+} #^B^#
+
+undefined(1, 2) {
+  myFunc { 1 }
+} #^C^#