[Type checker] Be more rigorous about extracting argument labels from calls.

Whenever we have a call, retrieve the argument labels from the
argument structurally and associate them with the callee. We were
previously doing this as a separate AST walk (which was unnecessary),
so fold that into constraint generation for a CallExpr.

This is a slightly-pared-back version of
3753d77 that isn't so rigid in its
interpretation of ASTs. I'll tighten up the semantics over time.

Author: DougGregor

include/swift/AST/Expr.h

@@ -1621,6 +1621,9 @@ class ParenExpr : public IdentityExpr {
   /// \brief Whether this expression has a trailing closure as its argument.
   bool hasTrailingClosure() const { return HasTrailingClosure; }
 
+  /// Create a new, implicit parenthesized expression.
+  static ParenExpr *createImplicit(ASTContext &ctx, Expr *expr);
+
   static bool classof(const Expr *E) { return E->getKind() == ExprKind::Paren; }
 };
 
@@ -3195,7 +3198,15 @@ class CallExpr : public ApplyExpr {
     SourceLoc FnLoc = getFn()->getLoc(); 
     return FnLoc.isValid() ? FnLoc : getArg()->getLoc();
   }
-  
+
+  /// Retrieve the expression that direct represents the callee.
+  ///
+  /// The "direct" callee is the expression representing the callee
+  /// after looking through top-level constructs that don't affect the
+  /// identity of the callee, e.g., extra parentheses, optional
+  /// unwrapping (?)/forcing (!), etc.
+  Expr *getDirectCallee() const;
+
   static bool classof(const Expr *E) { return E->getKind() == ExprKind::Call; }
 };
 

lib/AST/Expr.cpp

@@ -922,6 +922,14 @@ SequenceExpr *SequenceExpr::create(ASTContext &ctx, ArrayRef<Expr*> elements) {
   return ::new(Buffer) SequenceExpr(elements);
 }
 
+ParenExpr *ParenExpr::createImplicit(ASTContext &ctx, Expr *expr) {
+  ParenExpr *result = new (ctx) ParenExpr(SourceLoc(), expr, SourceLoc(),
+                                          /*hasTrailingClosure=*/false,
+                                          expr->getType());
+  result->setImplicit();
+  return result;
+}
+
 SourceLoc TupleExpr::getStartLoc() const {
   if (LParenLoc.isValid()) return LParenLoc;
   if (getNumElements() == 0) return SourceLoc();
@@ -1033,6 +1041,25 @@ ValueDecl *ApplyExpr::getCalledValue() const {
   return ::getCalledValue(Fn);
 }
 
+Expr *CallExpr::getDirectCallee() const {
+  auto fn = getFn();
+  while (true) {
+    fn = fn->getSemanticsProvidingExpr();
+
+    if (auto force = dyn_cast<ForceValueExpr>(fn)) {
+      fn = force->getSubExpr();
+      continue;
+    }
+
+    if (auto bind = dyn_cast<BindOptionalExpr>(fn)) {
+      fn = bind->getSubExpr();
+      continue;
+    }
+
+    return fn;
+  }
+}
+
 RebindSelfInConstructorExpr::RebindSelfInConstructorExpr(Expr *SubExpr,
                                                          VarDecl *Self)
   : Expr(ExprKind::RebindSelfInConstructor, /*Implicit=*/true,

lib/Sema/CSGen.cpp

@@ -1597,6 +1597,12 @@ namespace {
 
       // If there is an argument, apply it.
       if (auto arg = expr->getArgument()) {
+        // Identify and record the argument labels for this call.
+        if (auto argumentLabels = getArgumentLabelsForCallArgument(arg)) {
+          CS.CalleeArgumentLabels[CS.getConstraintLocator(expr)] =
+            *argumentLabels;
+        }
+
         // The result type of the function must be convertible to the base type.
         // TODO: we definitely want this to include ImplicitlyUnwrappedOptional; does it
         // need to include everything else in the world?
@@ -2429,11 +2435,52 @@ namespace {
       return expr->getType();
     }
 
+    using ArgumentLabelState = ConstraintSystem::ArgumentLabelState;
+
+    /// Extract argument labels from the given call argument.
+    Optional<ArgumentLabelState> getArgumentLabelsForCallArgument(Expr *arg) {
+      // Parentheses.
+      if (auto paren = dyn_cast<ParenExpr>(arg)) {
+        return ArgumentLabelState{
+                 CS.allocateCopy(llvm::makeArrayRef(Identifier())),
+                 paren->hasTrailingClosure() };
+      }
+
+      // FIXME: Hack because not all CallExprs come in with
+      // ParenExpr/TupleExpr arguments.
+      if (!isa<TupleExpr>(arg)) return None;
+
+      // Tuples.
+      auto tuple = cast<TupleExpr>(arg);
+
+      // If we have element names, use 'em.
+      if (tuple->hasElementNames()) {
+        return ArgumentLabelState{tuple->getElementNames(),
+                                  tuple->hasTrailingClosure()};
+      }
+
+      // Otherwise, there are no argument labels.
+      llvm::SmallVector<Identifier, 4> names(tuple->getNumElements(),
+                                             Identifier());
+      return ArgumentLabelState{CS.allocateCopy(names),
+                                tuple->hasTrailingClosure()};
+    }
+
     Type visitApplyExpr(ApplyExpr *expr) {
       Type outputTy;
 
       auto fnExpr = expr->getFn();
-      
+
+      // Identify and record the argument labels for a call.
+      if (auto call = dyn_cast<CallExpr>(expr)) {
+        if (auto argumentLabels =
+              getArgumentLabelsForCallArgument(expr->getArg())) {
+          auto callee = call->getDirectCallee();
+          CS.CalleeArgumentLabels[CS.getConstraintLocator(callee)] =
+            *argumentLabels;
+        }
+      }
+
       if (isa<DeclRefExpr>(fnExpr)) {
         if (auto fnType = fnExpr->getType()->getAs<AnyFunctionType>()) {
           outputTy = fnType->getResult();
@@ -2536,7 +2583,7 @@ namespace {
       FunctionType::ExtInfo extInfo;
       if (isa<ClosureExpr>(fnExpr->getSemanticsProvidingExpr()))
         extInfo = extInfo.withNoEscape();
-      
+
       auto funcTy = FunctionType::get(expr->getArg()->getType(), outputTy,
                                       extInfo);
 
@@ -2999,87 +3046,12 @@ namespace {
     /// \brief Ignore declarations.
     bool walkToDeclPre(Decl *decl) override { return false; }
   };
-
-  /// AST walker that records the keyword arguments provided at each
-  /// call site.
-  class ArgumentLabelWalker : public ASTWalker {
-    ConstraintSystem &CS;
-    llvm::DenseMap<Expr *, Expr *> ParentMap;
-
-  public:
-    ArgumentLabelWalker(ConstraintSystem &cs, Expr *expr) 
-      : CS(cs), ParentMap(expr->getParentMap()) { }
-
-    using State = ConstraintSystem::ArgumentLabelState;
-
-    void associateArgumentLabels(Expr *arg, State labels,
-                                 bool labelsArePermanent) {
-      // Our parent must be a call.
-      auto call = dyn_cast_or_null<CallExpr>(ParentMap[arg]);
-      if (!call) 
-        return;
-
-      // We must have originated at the call argument.
-      if (arg != call->getArg())
-        return;
-
-      // Dig out the function, looking through, parentheses, ?, and !.
-      auto fn = call->getFn();
-      do {
-        fn = fn->getSemanticsProvidingExpr();
-
-        if (auto force = dyn_cast<ForceValueExpr>(fn)) {
-          fn = force->getSubExpr();
-          continue;
-        }
-
-        if (auto bind = dyn_cast<BindOptionalExpr>(fn)) {
-          fn = bind->getSubExpr();
-          continue;
-        }
-
-        break;
-      } while (true);
-
-      // Record the labels.
-      if (!labelsArePermanent)
-        labels.Labels = CS.allocateCopy(labels.Labels);
-      CS.ArgumentLabels[CS.getConstraintLocator(fn)] = labels;
-    }
-
-    std::pair<bool, Expr *> walkToExprPre(Expr *expr) override {
-      if (auto tuple = dyn_cast<TupleExpr>(expr)) {
-        if (tuple->hasElementNames())
-          associateArgumentLabels(expr,
-                                  { tuple->getElementNames(),
-                                    tuple->hasTrailingClosure() },
-                                  /*labelsArePermanent*/ true);
-        else {
-          llvm::SmallVector<Identifier, 4> names(tuple->getNumElements(),
-                                                 Identifier()); 
-          associateArgumentLabels(expr, { names, tuple->hasTrailingClosure() },
-                                  /*labelsArePermanent*/ false);
-        }
-      } else if (auto paren = dyn_cast<ParenExpr>(expr)) {
-        associateArgumentLabels(paren,
-                                { { Identifier() },
-                                  paren->hasTrailingClosure() },
-                                /*labelsArePermanent*/ false);
-      }
-
-      return { true, expr };
-    }
-  };
-
 } // end anonymous namespace
 
 Expr *ConstraintSystem::generateConstraints(Expr *expr) {
   // Remove implicit conversions from the expression.
   expr = expr->walk(SanitizeExpr(getTypeChecker()));
 
-  // Walk the expression to associate labeled arguments.
-  expr->walk(ArgumentLabelWalker(*this, expr));
-
   // Walk the expression, generating constraints.
   ConstraintGenerator cg(*this);
   ConstraintWalker cw(cg);

lib/Sema/CSSimplify.cpp

@@ -2612,8 +2612,8 @@ getArgumentLabels(ConstraintSystem &cs, ConstraintLocatorBuilder locator) {
   if (!parts.empty())
     return None;
 
-  auto known = cs.ArgumentLabels.find(cs.getConstraintLocator(anchor));
-  if (known == cs.ArgumentLabels.end())
+  auto known = cs.CalleeArgumentLabels.find(cs.getConstraintLocator(anchor));
+  if (known == cs.CalleeArgumentLabels.end())
     return None;
 
   return known->second;

lib/Sema/ConstraintSystem.h

@@ -1042,16 +1042,25 @@ class ConstraintSystem {
   /// we're exploring. 
   SolverState *solverState = nullptr;
 
+  /// Describes argument labels as they occur in a call.
   struct ArgumentLabelState {
+    /// The actual argument labels provided at the call site.
+    ///
+    /// The contents of this array are only guaranteed to live until
+    /// the constraint system is destroyed.
     ArrayRef<Identifier> Labels;
+
+    /// Whether the last of the arguments was written as a trailing
+    /// closure. It will have an empty Identifier().
     bool HasTrailingClosure;
   };
 
-  /// A mapping from the constraint locators for references to various
-  /// names (e.g., member references, normal name references, possible
-  /// constructions) to the argument labels provided in the call to
-  /// that locator.
-  llvm::DenseMap<ConstraintLocator *, ArgumentLabelState> ArgumentLabels;
+  /// The argument labels that are associated with a given callee.
+  ///
+  /// The argument labels for a particular callee are recorded at the
+  /// time of constraint generation based on the syntactic call
+  /// argument.
+  llvm::DenseMap<ConstraintLocator *, ArgumentLabelState> CalleeArgumentLabels;
 
   /// FIXME: This is a workaround for the way we perform protocol
   /// conformance checking for generic requirements, where we re-use

lib/Sema/MiscDiagnostics.cpp

@@ -2830,14 +2830,16 @@ class ObjCSelectorWalker : public ASTWalker {
     bool hadParens = false;
     auto lookThroughParens = [&](Expr *arg, bool outermost) -> Expr * {
       if (auto parenExpr = dyn_cast<ParenExpr>(arg)) {
-        if (!outermost) {
+        if (!outermost && !parenExpr->isImplicit()) {
           hadParens = true;
           return parenExpr->getSubExpr()->getSemanticsProvidingExpr();
         }
 
         arg = parenExpr->getSubExpr();
         if (auto innerParenExpr = dyn_cast<ParenExpr>(arg)) {
-          hadParens = true;
+          if (!innerParenExpr->isImplicit())
+            hadParens = true;
+
           arg = innerParenExpr->getSubExpr();
         }
       }

lib/Sema/TypeCheckConstraints.cpp

@@ -792,10 +792,8 @@ namespace {
               continue;
             }
 
-            // Look through identity, force-value, and 'try' expressions.
-            if (isa<IdentityExpr>(ancestor) ||
-                isa<ForceValueExpr>(ancestor) ||
-                isa<AnyTryExpr>(ancestor)) {
+            // Look through force-value, and 'try' expressions.
+            if (isa<ForceValueExpr>(ancestor) || isa<AnyTryExpr>(ancestor)) {
               if (target)
                 target = ancestor;
               continue;