Preserve trailing closure index in OriginalArgumentList

Author: hamishknight

include/swift/Sema/IDETypeChecking.h

@@ -243,7 +243,32 @@ namespace swift {
     SmallVector<SourceLoc, 4> labelLocs;
     SourceLoc lParenLoc;
     SourceLoc rParenLoc;
-    bool hasTrailingClosure = false;
+    Optional<unsigned> unlabeledTrailingClosureIdx;
+
+    /// The number of trailing closures in the argument list.
+    unsigned getNumTrailingClosures() const {
+      if (!unlabeledTrailingClosureIdx)
+        return 0;
+      return args.size() - *unlabeledTrailingClosureIdx;
+    }
+
+    /// Whether any unlabeled or labeled trailing closures are present.
+    bool hasAnyTrailingClosures() const {
+      return unlabeledTrailingClosureIdx.hasValue();
+    }
+
+    /// Whether the given index is for an unlabeled trailing closure.
+    bool isUnlabeledTrailingClosureIdx(unsigned i) const {
+      return unlabeledTrailingClosureIdx && *unlabeledTrailingClosureIdx == i;
+    }
+
+    /// Whether the given index is for a labeled trailing closure in an
+    /// argument list with multiple trailing closures.
+    bool isLabeledTrailingClosureIdx(unsigned i) const {
+      if (!unlabeledTrailingClosureIdx)
+        return false;
+      return i > *unlabeledTrailingClosureIdx && i < args.size();
+    }
   };
 
   /// When applying a solution to a constraint system, the type checker rewrites

lib/Sema/CSGen.cpp

@@ -4232,9 +4232,17 @@ OriginalArgumentList
 swift::getOriginalArgumentList(Expr *expr) {
   OriginalArgumentList result;
 
-  auto add = [&](Expr *arg, Identifier label, SourceLoc labelLoc) {
-    if (isa<DefaultArgumentExpr>(arg)) {
+  auto oldTrailingClosureIdx =
+      expr->getUnlabeledTrailingClosureIndexOfPackedArgument();
+  Optional<unsigned> newTrailingClosureIdx;
+
+  auto add = [&](unsigned i, Expr *arg, Identifier label, SourceLoc labelLoc) {
+    if (isa<DefaultArgumentExpr>(arg))
       return;
+
+    if (oldTrailingClosureIdx && *oldTrailingClosureIdx == i) {
+      assert(!newTrailingClosureIdx);
+      newTrailingClosureIdx = result.args.size();
     }
 
     if (auto *varargExpr = dyn_cast<VarargExpansionExpr>(arg)) {
@@ -4260,24 +4268,25 @@ swift::getOriginalArgumentList(Expr *expr) {
   if (auto *parenExpr = dyn_cast<ParenExpr>(expr)) {
     result.lParenLoc = parenExpr->getLParenLoc();
     result.rParenLoc = parenExpr->getRParenLoc();
-    result.hasTrailingClosure = parenExpr->hasTrailingClosure();
-    add(parenExpr->getSubExpr(), Identifier(), SourceLoc());
+    add(0, parenExpr->getSubExpr(), Identifier(), SourceLoc());
   } else if (auto *tupleExpr = dyn_cast<TupleExpr>(expr)) {
     result.lParenLoc = tupleExpr->getLParenLoc();
     result.rParenLoc = tupleExpr->getRParenLoc();
-    result.hasTrailingClosure = tupleExpr->hasTrailingClosure();
 
     auto args = tupleExpr->getElements();
     auto labels = tupleExpr->getElementNames();
     auto labelLocs = tupleExpr->getElementNameLocs();
     for (unsigned i = 0, e = args.size(); i != e; ++i) {
       // Implicit TupleExprs don't always store label locations.
-      add(args[i], labels[i],
+      add(i, args[i], labels[i],
           labelLocs.empty() ? SourceLoc() : labelLocs[i]);
     }
   } else {
-    add(expr, Identifier(), SourceLoc());
+    add(0, expr, Identifier(), SourceLoc());
   }
+  assert(oldTrailingClosureIdx.hasValue() == newTrailingClosureIdx.hasValue());
+  assert(!newTrailingClosureIdx || *newTrailingClosureIdx < result.args.size());
 
+  result.unlabeledTrailingClosureIdx = newTrailingClosureIdx;
   return result;
 }

lib/Sema/TypeCheckCodeCompletion.cpp

@@ -288,18 +288,18 @@ class SanitizeExpr : public ASTWalker {
         new (C) ParenExpr(argList.lParenLoc,
                           argList.args[0],
                           argList.rParenLoc,
-                          argList.hasTrailingClosure);
+                          argList.hasAnyTrailingClosures());
       result->setImplicit();
       return result;
     }
 
     return TupleExpr::create(C,
                              argList.lParenLoc,
+                             argList.rParenLoc,
                              argList.args,
                              argList.labels,
                              argList.labelLocs,
-                             argList.rParenLoc,
-                             argList.hasTrailingClosure,
+                             argList.unlabeledTrailingClosureIdx,
                              /*implicit=*/true);
   }