[ConstraintSystem] Extract logic that identifies application based on its argument node

Author: xedin

include/swift/Sema/ConstraintSystem.h

@@ -4994,6 +4994,13 @@ class ConstraintSystem {
   /// diagnostic.
   void maybeProduceFallbackDiagnostic(SolutionApplicationTarget target) const;
 
+  /// Check whether given AST node represents an argument of an application
+  /// of some sort (call, operator invocation, subscript etc.)
+  /// and return AST node representing and argument index. E.g. for regular
+  /// calls `test(42)` passing `42` should return node representing
+  /// entire call and index `0`.
+  Optional<std::pair<Expr *, unsigned>> isArgumentExpr(Expr *expr);
+
   SWIFT_DEBUG_DUMP;
   SWIFT_DEBUG_DUMPER(dump(Expr *));
 

lib/Sema/ConstraintSystem.cpp

@@ -3770,12 +3770,14 @@ bool ConstraintSystem::diagnoseAmbiguityWithFixes(
     const auto &solution = *entry.first;
     const auto *fix = entry.second;
 
-    if (fix->getLocator()->isForContextualType()) {
+    auto *locator = fix->getLocator();
+
+    if (locator->isForContextualType()) {
       contextualFixes.push_back({&solution, fix});
       continue;
     }
 
-    auto *calleeLocator = solution.getCalleeLocator(fix->getLocator());
+    auto *calleeLocator = solution.getCalleeLocator(locator);
     fixesByCallee[calleeLocator].push_back({&solution, fix});
   }
 
@@ -4533,46 +4535,13 @@ Solution::getFunctionArgApplyInfo(ConstraintLocator *locator) const {
   // It's only valid to use `&` in argument positions, but we need
   // to figure out exactly where it was used.
   if (auto *argExpr = getAsExpr<InOutExpr>(locator->getAnchor())) {
-    auto *argList = cs.getParentExpr(argExpr);
-    assert(argList);
-
-    // `inout` expression might be wrapped in a number of
-    // parens e.g. `test(((&x)))`.
-    if (isa<ParenExpr>(argList)) {
-      for (;;) {
-        auto nextParent = cs.getParentExpr(argList);
-        assert(nextParent && "Incorrect use of `inout` expression");
-
-        // e.g. `test((&x), x: ...)`
-        if (isa<TupleExpr>(nextParent)) {
-          argList = nextParent;
-          break;
-        }
-
-        // e.g. `test(((&x)))`
-        if (isa<ParenExpr>(nextParent)) {
-          argList = nextParent;
-          continue;
-        }
-
-        break;
-      }
-    }
+    auto argInfo = cs.isArgumentExpr(argExpr);
+    assert(argInfo && "Incorrect use of `inout` expression");
 
-    unsigned argIdx = 0;
-    if (auto *tuple = dyn_cast<TupleExpr>(argList)) {
-      auto arguments = tuple->getElements();
+    Expr *call;
+    unsigned argIdx;
 
-      for (auto idx : indices(arguments)) {
-        if (arguments[idx]->getSemanticsProvidingExpr() == argExpr) {
-          argIdx = idx;
-          break;
-        }
-      }
-    }
-
-    auto *call = cs.getParentExpr(argList);
-    assert(call);
+    std::tie(call, argIdx) = *argInfo;
 
     ParameterTypeFlags flags;
     locator = cs.getConstraintLocator(
@@ -4746,6 +4715,58 @@ bool constraints::isStandardComparisonOperator(ASTNode node) {
   return false;
 }
 
+Optional<std::pair<Expr *, unsigned>>
+ConstraintSystem::isArgumentExpr(Expr *expr) {
+  auto *argList = getParentExpr(expr);
+
+  if (isa<ParenExpr>(argList)) {
+    for (;;) {
+      auto *parent = getParentExpr(argList);
+      if (!parent)
+        return None;
+
+      if (isa<TupleExpr>(parent)) {
+        argList = parent;
+        break;
+      }
+
+      // Drop all of the semantically insignificant parens
+      // that might be wrapping an argument e.g. `test(((42)))`
+      if (isa<ParenExpr>(parent)) {
+        argList = parent;
+        continue;
+      }
+
+      break;
+    }
+  }
+
+  if (!(isa<ParenExpr>(argList) || isa<TupleExpr>(argList)))
+    return None;
+
+  auto *application = getParentExpr(argList);
+  if (!application)
+    return None;
+
+  if (!(isa<ApplyExpr>(application) || isa<SubscriptExpr>(application) ||
+        isa<ObjectLiteralExpr>(application)))
+    return None;
+
+  unsigned argIdx = 0;
+  if (auto *tuple = dyn_cast<TupleExpr>(argList)) {
+    auto arguments = tuple->getElements();
+
+    for (auto idx : indices(arguments)) {
+      if (arguments[idx]->getSemanticsProvidingExpr() == expr) {
+        argIdx = idx;
+        break;
+      }
+    }
+  }
+
+  return std::make_pair(application, argIdx);
+}
+
 bool constraints::isOperatorArgument(ConstraintLocator *locator,
                                      StringRef expectedOperator) {
   if (!locator->findLast<LocatorPathElt::ApplyArgToParam>())