[Diagnostics] Unify DiagnosticFailure and ConstraintFix classes for extra & failures.

Make `InvalidUseOfAddressOf` a `ContextualFailure`, make `ReturnAddressOf`
a `ContextualMismatch`, and extend this failure to cover using `&` with a
non-inout argument.

Author: hborla

lib/Sema/CSDiagnostics.cpp

@@ -908,17 +908,6 @@ bool MissingAddressOfFailure::diagnoseAsError() {
   return true;
 }
 
-bool ExtraAddressOfFailure::diagnoseAsError() {
-  if (hasComplexLocator())
-    return false;
-
-  auto anchor = getAnchor();
-  emitDiagnostic(getAnchor()->getLoc(), diag::extra_address_of, getToType())
-      .highlight(anchor->getSourceRange())
-      .fixItRemove(anchor->getStartLoc());
-  return true;
-}
-
 bool MissingExplicitConversionFailure::diagnoseAsError() {
   if (hasComplexLocator())
     return false;
@@ -3619,6 +3608,16 @@ SourceLoc InvalidUseOfAddressOf::getLoc() const {
 }
 
 bool InvalidUseOfAddressOf::diagnoseAsError() {
+  if (auto argApplyInfo = getFunctionArgApplyInfo(getLocator())) {
+    if (!argApplyInfo->getParameterFlags().isInOut()) {
+      auto anchor = getAnchor();
+      emitDiagnostic(anchor->getLoc(), diag::extra_address_of, getToType())
+          .highlight(anchor->getSourceRange())
+          .fixItRemove(anchor->getStartLoc());
+      return true;
+    }
+  }
+
   emitDiagnostic(getLoc(), diag::extraneous_address_of);
   return true;
 }

lib/Sema/CSDiagnostics.h

@@ -843,14 +843,28 @@ class MissingAddressOfFailure final : public ContextualFailure {
   bool diagnoseAsError() override;
 };
 
-/// Diagnose failures when passing an inout argument to a non-inout parameter.
-class ExtraAddressOfFailure final : public ContextualFailure {
+/// Diagnose extraneous use of address of (`&`) which could only be
+/// associated with arguments to inout parameters e.g.
+///
+/// ```swift
+/// struct S {}
+///
+/// var a: S = ...
+/// var b: S = ...
+///
+/// a = &b
+/// ```
+class InvalidUseOfAddressOf final : public ContextualFailure {
 public:
-  ExtraAddressOfFailure(Expr *expr, ConstraintSystem &cs, Type argTy,
-                        Type paramTy, ConstraintLocator *locator)
-      : ContextualFailure(expr, cs, argTy, paramTy, locator) {}
+  InvalidUseOfAddressOf(Expr *root, ConstraintSystem &cs, Type lhs, Type rhs,
+                        ConstraintLocator *locator)
+      : ContextualFailure(root, cs, lhs, rhs, locator) {}
 
   bool diagnoseAsError() override;
+
+protected:
+  /// Compute location of the failure for diagnostic.
+  SourceLoc getLoc() const;
 };
 
 /// Diagnose mismatches relating to tuple destructuring.
@@ -1455,30 +1469,6 @@ class InvalidMethodRefInKeyPath final : public InvalidMemberRefInKeyPath {
   bool diagnoseAsError() override;
 };
 
-/// Diagnose extraneous use of address of (`&`) which could only be
-/// associated with arguments to inout parameters e.g.
-///
-/// ```swift
-/// struct S {}
-///
-/// var a: S = ...
-/// var b: S = ...
-///
-/// a = &b
-/// ```
-class InvalidUseOfAddressOf final : public FailureDiagnostic {
-public:
-  InvalidUseOfAddressOf(Expr *root, ConstraintSystem &cs,
-                        ConstraintLocator *locator)
-      : FailureDiagnostic(root, cs, locator) {}
-
-  bool diagnoseAsError() override;
-
-protected:
-  /// Compute location of the failure for diagnostic.
-  SourceLoc getLoc() const;
-};
-
 /// Diagnose an attempt return something from a function which
 /// doesn't have a return type specified e.g.
 ///

lib/Sema/CSFix.cpp

@@ -114,19 +114,6 @@ AddAddressOf *AddAddressOf::create(ConstraintSystem &cs, Type argTy,
   return new (cs.getAllocator()) AddAddressOf(cs, argTy, paramTy, locator);
 }
 
-bool RemoveAddressOfArg::diagnose(Expr *root, bool asNote) const {
-  auto &cs = getConstraintSystem();
-  ExtraAddressOfFailure failure(root, cs, getFromType(), getToType(),
-                                getLocator());
-  return failure.diagnose(asNote);
-}
-
-RemoveAddressOfArg *
-RemoveAddressOfArg::create(ConstraintSystem &cs, Type argTy, Type paramTy,
-                           ConstraintLocator *locator) {
-  return new (cs.getAllocator()) RemoveAddressOfArg(cs, argTy, paramTy, locator);
-}
-
 bool TreatRValueAsLValue::diagnose(Expr *root, bool asNote) const {
   RValueTreatedAsLValueFailure failure(getConstraintSystem(), getLocator());
   return failure.diagnose(asNote);
@@ -636,13 +623,14 @@ KeyPathContextualMismatch::create(ConstraintSystem &cs, Type lhs, Type rhs,
 }
 
 bool RemoveAddressOf::diagnose(Expr *root, bool asNote) const {
-  InvalidUseOfAddressOf failure(root, getConstraintSystem(), getLocator());
+  InvalidUseOfAddressOf failure(root, getConstraintSystem(), getFromType(),
+                                getToType(), getLocator());
   return failure.diagnose(asNote);
 }
 
-RemoveAddressOf *RemoveAddressOf::create(ConstraintSystem &cs,
+RemoveAddressOf *RemoveAddressOf::create(ConstraintSystem &cs, Type lhs, Type rhs,
                                          ConstraintLocator *locator) {
-  return new (cs.getAllocator()) RemoveAddressOf(cs, locator);
+  return new (cs.getAllocator()) RemoveAddressOf(cs, lhs, rhs, locator);
 }
 
 bool RemoveReturn::diagnose(Expr *root, bool asNote) const {

lib/Sema/CSFix.h

@@ -545,19 +545,20 @@ class AddAddressOf final : public ContextualMismatch {
                               ConstraintLocator *locator);
 };
 
-/// Remove the \c & of an argument for a non-inout parameter.
-class RemoveAddressOfArg final : public ContextualMismatch {
-  RemoveAddressOfArg(ConstraintSystem &cs, Type argTy, Type paramTy,
-                     ConstraintLocator *locator)
-      : ContextualMismatch(cs, argTy, paramTy, locator) {}
+class RemoveAddressOf final : public ContextualMismatch {
+  RemoveAddressOf(ConstraintSystem &cs, Type lhs, Type rhs,
+                  ConstraintLocator *locator)
+      : ContextualMismatch(cs, FixKind::RemoveAddressOf, lhs, rhs, locator) {}
 
 public:
-  std::string getName() const override { return "remove address of argument"; }
+  std::string getName() const override {
+    return "remove extraneous use of `&`";
+  }
 
   bool diagnose(Expr *root, bool asNote = false) const override;
 
-  static RemoveAddressOfArg *create(ConstraintSystem &cs, Type argTy,
-                                    Type paramTy, ConstraintLocator *locator);
+  static RemoveAddressOf *create(ConstraintSystem &cs, Type lhs, Type rhs,
+                                 ConstraintLocator *locator);
 };
 
 /// Detect situations where two type's generic arguments must
@@ -1180,21 +1181,6 @@ class AllowInvalidRefInKeyPath final : public ConstraintFix {
                                           ConstraintLocator *locator);
 };
 
-class RemoveAddressOf final : public ConstraintFix {
-  RemoveAddressOf(ConstraintSystem &cs, ConstraintLocator *locator)
-      : ConstraintFix(cs, FixKind::RemoveAddressOf, locator) {}
-
-public:
-  std::string getName() const override {
-    return "remove extraneous use of `&`";
-  }
-
-  bool diagnose(Expr *root, bool asNote = false) const override;
-
-  static RemoveAddressOf *create(ConstraintSystem &cs,
-                                 ConstraintLocator *locator);
-};
-
 class RemoveReturn final : public ConstraintFix {
   RemoveReturn(ConstraintSystem &cs, ConstraintLocator *locator)
       : ConstraintFix(cs, FixKind::RemoveReturn, locator) {}

lib/Sema/CSSimplify.cpp

@@ -2357,7 +2357,8 @@ bool ConstraintSystem::repairFailures(
 
       if (isa<InOutExpr>(AE->getSrc())) {
         conversionsOrFixes.push_back(
-            RemoveAddressOf::create(*this, getConstraintLocator(locator)));
+            RemoveAddressOf::create(*this, lhs, rhs,
+                                    getConstraintLocator(locator)));
         return true;
       }
 
@@ -2443,7 +2444,7 @@ bool ConstraintSystem::repairFailures(
                                TypeMatchFlags::TMF_ApplyingFix, locator);
 
       if (result.isSuccess())
-        conversionsOrFixes.push_back(RemoveAddressOfArg::create(*this, lhs,
+        conversionsOrFixes.push_back(RemoveAddressOf::create(*this, lhs,
             rhs, getConstraintLocator(locator)));
     }
 
@@ -3005,7 +3006,8 @@ ConstraintSystem::matchTypes(Type type1, Type type2, ConstraintKind kind,
 
       if (kind == ConstraintKind::OperatorArgumentConversion) {
         conversionsOrFixes.push_back(
-            RemoveAddressOf::create(*this, getConstraintLocator(locator)));
+            RemoveAddressOf::create(*this, type1, type2,
+                                    getConstraintLocator(locator)));
         break;
       }