[ConstraintSystem] NFC: Extract logic related to forming locator for argument info retrieval

Author: xedin

lib/Sema/CSGen.cpp

@@ -3227,18 +3227,11 @@ namespace {
     void associateArgumentLabels(Expr *expr,
                                  ConstraintSystem::ArgumentInfo info,
                                  bool labelsArePermanent = true) {
-      ConstraintLocator *locator = nullptr;
-      if (auto *apply = dyn_cast<ApplyExpr>(expr)) {
-        auto *fnExpr = getArgumentLabelTargetExpr(apply->getFn());
-        locator = CS.getConstraintLocator(fnExpr);
-      } else {
-        locator = CS.getCalleeLocator(expr);
-      }
-
+      assert(expr);
       // Record the labels.
       if (!labelsArePermanent)
         info.Labels = CS.allocateCopy(info.Labels);
-      CS.ArgumentInfos[locator] = info;
+      CS.ArgumentInfos[CS.getArgumentInfoLocator(expr)] = info;
     }
   };
 

lib/Sema/ConstraintSystem.cpp

@@ -2780,23 +2780,25 @@ void ConstraintSystem::generateConstraints(
   }
 }
 
-Optional<ConstraintSystem::ArgumentInfo>
-ConstraintSystem::getArgumentInfo(ConstraintLocator *locator) {
-  Expr *anchor = locator->getAnchor();
+ConstraintLocator *ConstraintSystem::getArgumentInfoLocator(Expr *anchor) {
   if (!anchor)
-    return None;
+    return nullptr;
 
   if (auto *apply = dyn_cast<ApplyExpr>(anchor)) {
     auto *fnExpr = getArgumentLabelTargetExpr(apply->getFn());
-    locator = getConstraintLocator(fnExpr);
-  } else {
-    locator = getCalleeLocator(anchor);
+    return getConstraintLocator(fnExpr);
   }
 
-  auto known = ArgumentInfos.find(locator);
-  if (known != ArgumentInfos.end())
-    return known->second;
+  return getCalleeLocator(anchor);
+}
 
+Optional<ConstraintSystem::ArgumentInfo>
+ConstraintSystem::getArgumentInfo(ConstraintLocator *locator) {
+  if (auto *infoLocator = getArgumentInfoLocator(locator->getAnchor())) {
+    auto known = ArgumentInfos.find(infoLocator);
+    if (known != ArgumentInfos.end())
+      return known->second;
+  }
   return None;
 }
 

lib/Sema/ConstraintSystem.h

@@ -1559,6 +1559,10 @@ class ConstraintSystem {
   /// that locator.
   llvm::DenseMap<ConstraintLocator *, ArgumentInfo> ArgumentInfos;
 
+  /// Form a locator with given anchor which then could be used
+  /// to retrieve argument information cached in the constraint system.
+  ConstraintLocator *getArgumentInfoLocator(Expr *anchor);
+
   /// Retrieve the argument info that is associated with a member
   /// reference at the given locator.
   Optional<ArgumentInfo> getArgumentInfo(ConstraintLocator *locator);