From f7fc1ed7b5409a252d3cdceea7f413863b817f2c Mon Sep 17 00:00:00 2001
From: Rintaro Ishizaki <rishizaki@apple.com>
Date: Wed, 15 Jul 2020 09:56:20 -0700
Subject: [PATCH] [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
---
 lib/IDE/ExprContextAnalysis.cpp               | 33 +++++++++++++++++--
 .../IDE/crashers_2_fixed/rdar65556791.swift   | 17 ++++++++++
 2 files changed, 47 insertions(+), 3 deletions(-)
 create mode 100644 validation-test/IDE/crashers_2_fixed/rdar65556791.swift

diff --git a/lib/IDE/ExprContextAnalysis.cpp b/lib/IDE/ExprContextAnalysis.cpp
index 14bb5e1b5ef56..10c4f630f1996 100644
--- a/lib/IDE/ExprContextAnalysis.cpp
+++ b/lib/IDE/ExprContextAnalysis.cpp
@@ -213,16 +213,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.
@@ -246,7 +260,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;
   }
 };
 }
diff --git a/validation-test/IDE/crashers_2_fixed/rdar65556791.swift b/validation-test/IDE/crashers_2_fixed/rdar65556791.swift
new file mode 100644
index 0000000000000..256c1945cef0a
--- /dev/null
+++ b/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^#