[ConstraintSystem] Add a new locator element - argument attribute

lib/Sema/CSDiagnostics.cpp

@@ -1239,6 +1239,14 @@ bool RValueTreatedAsLValueFailure::diagnoseAsError() {
   if (auto callExpr = dyn_cast<ApplyExpr>(diagExpr)) {
     Expr *argExpr = callExpr->getArg();
     loc = callExpr->getFn()->getLoc();
+    auto *locator = getLocator();
+
+    // `argument attribute` is used for identification purposes
+    // only, so it could be looked through in this situation.
+    if (locator->isLastElement<LocatorPathElt::ArgumentAttribute>()) {
+      auto path = locator->getPath();
+      locator = getConstraintLocator(getRawAnchor(), path.drop_back());
+    }
 
     if (isa<PrefixUnaryExpr>(callExpr) || isa<PostfixUnaryExpr>(callExpr)) {
       subElementDiagID = diag::cannot_apply_lvalue_unop_to_subelement;
@@ -1247,16 +1255,14 @@ bool RValueTreatedAsLValueFailure::diagnoseAsError() {
     } else if (isa<BinaryExpr>(callExpr)) {
       subElementDiagID = diag::cannot_apply_lvalue_binop_to_subelement;
       rvalueDiagID = diag::cannot_apply_lvalue_binop_to_rvalue;
-      auto argTuple = dyn_cast<TupleExpr>(argExpr);
-      diagExpr = argTuple->getElement(0);
-    } else if (getLocator()->getPath().size() > 0) {
-      auto argElt =
-          getLocator()->castLastElementTo<LocatorPathElt::ApplyArgToParam>();
-
+      diagExpr = castToExpr(simplifyLocatorToAnchor(locator));
+    } else if (auto argElt =
+                   locator
+                       ->getLastElementAs<LocatorPathElt::ApplyArgToParam>()) {
       subElementDiagID = diag::cannot_pass_rvalue_inout_subelement;
       rvalueDiagID = diag::cannot_pass_rvalue_inout;
       if (auto argTuple = dyn_cast<TupleExpr>(argExpr))
-        diagExpr = argTuple->getElement(argElt.getArgIdx());
+        diagExpr = argTuple->getElement(argElt->getArgIdx());
       else if (auto parens = dyn_cast<ParenExpr>(argExpr))
         diagExpr = parens->getSubExpr();
     } else {

lib/Sema/CSFix.cpp

@@ -119,6 +119,10 @@ bool TreatRValueAsLValue::diagnose(const Solution &solution,
 
 TreatRValueAsLValue *TreatRValueAsLValue::create(ConstraintSystem &cs,
                                    ConstraintLocator *locator) {
+  if (locator->isLastElement<LocatorPathElt::ApplyArgToParam>())
+    locator = cs.getConstraintLocator(
+        locator, LocatorPathElt::ArgumentAttribute::forInOut());
+
   return new (cs.getAllocator()) TreatRValueAsLValue(cs, locator);
 }
 
@@ -168,6 +172,10 @@ bool MarkExplicitlyEscaping::diagnose(const Solution &solution,
 MarkExplicitlyEscaping *
 MarkExplicitlyEscaping::create(ConstraintSystem &cs, Type lhs, Type rhs,
                                ConstraintLocator *locator) {
+  if (locator->isLastElement<LocatorPathElt::ApplyArgToParam>())
+    locator = cs.getConstraintLocator(
+        locator, LocatorPathElt::ArgumentAttribute::forEscaping());
+
   return new (cs.getAllocator()) MarkExplicitlyEscaping(cs, lhs, rhs, locator);
 }
 

lib/Sema/ConstraintLocator.cpp

@@ -57,6 +57,7 @@ void ConstraintLocator::Profile(llvm::FoldingSetNodeID &id, ASTNode anchor,
       id.AddPointer(elt.castTo<LocatorPathElt::PatternMatch>().getPattern());
       break;
 
+    case ArgumentAttribute:
     case GenericArgument:
     case NamedTupleElement:
     case TupleElement:
@@ -125,6 +126,7 @@ unsigned LocatorPathElt::getNewSummaryFlags() const {
   case ConstraintLocator::ImplicitCallAsFunction:
   case ConstraintLocator::TernaryBranch:
   case ConstraintLocator::PatternMatch:
+  case ConstraintLocator::ArgumentAttribute:
     return 0;
 
   case ConstraintLocator::FunctionArgument:
@@ -500,6 +502,25 @@ void ConstraintLocator::dump(SourceManager *sm, raw_ostream &out) const {
     case PatternMatch:
       out << "pattern match";
       break;
+
+    case ArgumentAttribute: {
+      using AttrLoc = LocatorPathElt::ArgumentAttribute;
+
+      auto attrElt = elt.castTo<AttrLoc>();
+      out << "argument attribute: ";
+
+      switch (attrElt.getAttr()) {
+      case AttrLoc::Attribute::InOut:
+        out << "inout";
+        break;
+
+      case AttrLoc::Attribute::Escaping:
+        out << "@escaping";
+        break;
+      }
+
+      break;
+    }
     }
   }
   out << ']';

lib/Sema/ConstraintLocator.h

@@ -86,6 +86,7 @@ class ConstraintLocator : public llvm::FoldingSetNode {
     case SynthesizedArgument:
     case KeyPathDynamicMember:
     case TernaryBranch:
+    case ArgumentAttribute:
       return 1;
 
     case TypeParameterRequirement:
@@ -847,6 +848,31 @@ class LocatorPathElt::PatternMatch final : public LocatorPathElt {
   }
 };
 
+class LocatorPathElt::ArgumentAttribute final : public LocatorPathElt {
+public:
+  enum Attribute : uint8_t { InOut, Escaping };
+
+private:
+  ArgumentAttribute(Attribute attr)
+      : LocatorPathElt(ConstraintLocator::ArgumentAttribute,
+                       static_cast<uint8_t>(attr)) {}
+
+public:
+  Attribute getAttr() const { return static_cast<Attribute>(getValue(0)); }
+
+  static ArgumentAttribute forInOut() {
+    return ArgumentAttribute(Attribute::InOut);
+  }
+
+  static ArgumentAttribute forEscaping() {
+    return ArgumentAttribute(Attribute::Escaping);
+  }
+
+  static bool classof(const LocatorPathElt *elt) {
+    return elt->getKind() == ConstraintLocator::ArgumentAttribute;
+  }
+};
+
 /// A simple stack-only builder object that constructs a
 /// constraint locator without allocating memory.
 ///

lib/Sema/ConstraintLocatorPathElts.def

@@ -181,6 +181,14 @@ CUSTOM_LOCATOR_PATH_ELT(TernaryBranch)
 /// Performing a pattern patch.
 CUSTOM_LOCATOR_PATH_ELT(PatternMatch)
 
+/// Points to a particular attribute associated with one of
+/// the arguments e.g. `inout` or its type e.g. `@escaping`.
+///
+/// This is very useful when dealing with argument-to-parameter
+/// failures because it allows to express in the locator kind
+/// of a problem.
+CUSTOM_LOCATOR_PATH_ELT(ArgumentAttribute)
+
 #undef LOCATOR_PATH_ELT
 #undef CUSTOM_LOCATOR_PATH_ELT
 #undef SIMPLE_LOCATOR_PATH_ELT

lib/Sema/ConstraintSystem.cpp

@@ -462,6 +462,10 @@ ConstraintLocator *ConstraintSystem::getCalleeLocator(
     return getConstraintLocator(anchor, LocatorPathElt::ApplyFunction());
   }
 
+  if (locator->isLastElement<LocatorPathElt::ArgumentAttribute>()) {
+    return getConstraintLocator(anchor, path.drop_back());
+  }
+
   // If we have a locator that starts with a key path component element, we
   // may have a callee given by a property or subscript component.
   if (auto componentElt =
@@ -3527,6 +3531,13 @@ void constraints::simplifyLocator(ASTNode &anchor,
       continue;
     }
 
+    case ConstraintLocator::ArgumentAttribute: {
+      // At this point we should have already found argument expression
+      // this attribute belogs to, so we can leave this element in place
+      // because it points out exact location useful for diagnostics.
+      break;
+    }
+
     default:
       // FIXME: Lots of other cases to handle.
       break;