[ConstraintSystem] Make sure that @autoclosure argument detection works for subscripts/members

Currently logic in `matchCallArguments` could only detect argument
being an @autoclosure parameter for normal calls and operators.

This patch extends it to support subscripts and unresolved member calls.

Author: xedin

lib/Sema/CSSimplify.cpp

@@ -848,18 +848,7 @@ ConstraintSystem::TypeMatchResult constraints::matchCallArguments(
   auto isAutoClosureArg = [&](Expr *anchor, unsigned argIdx) -> bool {
     assert(anchor);
 
-    auto *call = dyn_cast<ApplyExpr>(anchor);
-    if (!call)
-      return false;
-
-    Expr *argExpr = nullptr;
-    if (auto *PE = dyn_cast<ParenExpr>(call->getArg())) {
-      assert(argsWithLabels.size() == 1);
-      argExpr = PE->getSubExpr();
-    } else if (auto *TE = dyn_cast<TupleExpr>(call->getArg())) {
-      argExpr = TE->getElement(argIdx);
-    }
-
+    auto *argExpr = getArgumentExpr(anchor, argIdx);
     if (!argExpr)
       return false;
 

lib/Sema/ConstraintSystem.cpp

@@ -2356,3 +2356,23 @@ Expr *constraints::simplifyLocatorToAnchor(ConstraintSystem &cs,
 
   return locator->getAnchor();
 }
+
+Expr *constraints::getArgumentExpr(Expr *expr, unsigned index) {
+  Expr *argExpr = nullptr;
+  if (auto *AE = dyn_cast<ApplyExpr>(expr))
+    argExpr = AE->getArg();
+  else if (auto *UME = dyn_cast<UnresolvedMemberExpr>(expr))
+    argExpr = UME->getArgument();
+  else if (auto *SE = dyn_cast<SubscriptExpr>(expr))
+    argExpr = SE->getIndex();
+  else
+    return nullptr;
+
+  if (auto *PE = dyn_cast<ParenExpr>(argExpr)) {
+    assert(index == 0);
+    return PE->getSubExpr();
+  }
+
+  assert(isa<TupleExpr>(argExpr));
+  return cast<TupleExpr>(argExpr)->getElement(index);
+}

lib/Sema/ConstraintSystem.h

@@ -3458,6 +3458,13 @@ void simplifyLocator(Expr *&anchor,
 /// null otherwise.
 Expr *simplifyLocatorToAnchor(ConstraintSystem &cs, ConstraintLocator *locator);
 
+/// Retrieve argument at specified index from given expression.
+/// The expression could be "application", "subscript" or "member" call.
+///
+/// \returns argument expression or `nullptr` if given "base" expression
+/// wasn't of one of the kinds listed above.
+Expr *getArgumentExpr(Expr *expr, unsigned index);
+
 class DisjunctionChoice {
   unsigned Index;
   Constraint *Choice;

test/Compatibility/attr_autoclosure.swift

@@ -36,3 +36,12 @@ do {
     foo(c)
   }
 }
+
+func passAutoClosureToSubscript(_ fn: @autoclosure () -> Int) {
+  struct S {
+    subscript(_ fn: @autoclosure () -> Int) -> Int { return fn() }
+  }
+
+  let s = S()
+  let _ = s[fn] // Ok
+}