[Sema] Fixits for assoc. type restatement hints should account for where clauses

Author: AnthonyLatsis

include/swift/AST/Decl.h

@@ -1435,6 +1435,8 @@ class alignas(RequirementRepr) TrailingWhereClause final :
     return SourceRange(WhereLoc,
                        getRequirements().back().getSourceRange().End);
   }
+
+  void print(llvm::raw_ostream &OS, bool printWhereKeyword) const;
 };
 
 // A private class for forcing exact field layout.

lib/AST/ASTDumper.cpp

@@ -164,6 +164,23 @@ void RequirementRepr::print(ASTPrinter &out) const {
   printImpl(out, /*AsWritten=*/true);
 }
 
+static void printTrailingRequirements(ASTPrinter &Printer,
+                                      ArrayRef<RequirementRepr> Reqs,
+                                      bool printWhereKeyword) {
+  if (Reqs.empty()) return;
+
+  if (printWhereKeyword)
+    Printer << " where ";
+  interleave(
+      Reqs,
+      [&](const RequirementRepr &req) {
+        Printer.callPrintStructurePre(PrintStructureKind::GenericRequirement);
+        req.print(Printer);
+        Printer.printStructurePost(PrintStructureKind::GenericRequirement);
+      },
+      [&] { Printer << ", "; });
+}
+
 void GenericParamList::print(llvm::raw_ostream &OS) const {
   OS << '<';
   interleave(*this,
@@ -176,14 +193,9 @@ void GenericParamList::print(llvm::raw_ostream &OS) const {
              },
              [&] { OS << ", "; });
 
-  if (!getRequirements().empty()) {
-    OS << " where ";
-    interleave(getRequirements(),
-               [&](const RequirementRepr &req) {
-                 req.print(OS);
-               },
-               [&] { OS << ", "; });
-  }
+  StreamPrinter Printer(OS);
+  printTrailingRequirements(Printer, getRequirements(),
+                            /*printWhereKeyword*/true);
   OS << '>';
 }
 
@@ -192,6 +204,13 @@ void GenericParamList::dump() const {
   llvm::errs() << '\n';
 }
 
+void TrailingWhereClause::print(llvm::raw_ostream &OS,
+                                bool printWhereKeyword) const {
+  StreamPrinter Printer(OS);
+  printTrailingRequirements(Printer, getRequirements(),
+                            printWhereKeyword);
+}
+
 static void printGenericParameters(raw_ostream &OS, GenericParamList *Params) {
   if (!Params)
     return;
@@ -660,9 +679,7 @@ namespace {
       }
       if (auto whereClause = decl->getTrailingWhereClause()) {
         OS << " where requirements: ";
-        interleave(whereClause->getRequirements(),
-                   [&](const RequirementRepr &req) { req.print(OS); },
-                   [&] { OS << ", "; });
+        whereClause->print(OS, /*printWhereKeyword*/false);
       }
       if (decl->overriddenDeclsComputed()) {
         OS << " overridden=";

lib/AST/GenericSignatureBuilder.cpp

@@ -3934,8 +3934,8 @@ ConstraintResult GenericSignatureBuilder::expandConformanceRequirement(
 
   // Retrieve the set of requirements that a given associated type declaration
   // produces, in the form that would be seen in the where clause.
-  auto getAssociatedTypeReqs = [&](AssociatedTypeDecl *assocType,
-                                   const char *start) {
+  const auto getAssociatedTypeReqs = [&](const AssociatedTypeDecl *assocType,
+                                         const char *start) {
     std::string result;
     {
       llvm::raw_string_ostream out(result);
@@ -3949,6 +3949,13 @@ ConstraintResult GenericSignatureBuilder::expandConformanceRequirement(
       }, [&] {
         out << ", ";
       });
+
+      if (const auto whereClause = assocType->getTrailingWhereClause()) {
+        if (!assocType->getInherited().empty())
+          out << ", ";
+
+        whereClause->print(out, /*printWhereKeyword*/false);
+      }
     }
     return result;
   };

test/Generics/associated_type_where_clause_hints.swift

@@ -9,6 +9,9 @@ struct X1 { }
 
 protocol P1 {
   associatedtype A // expected-note 2{{'A' declared here}}
+  associatedtype A2 // expected-note {{'A2' declared here}}
+  associatedtype A3 // expected-note {{'A3' declared here}}
+  associatedtype A4 // expected-note {{'A4' declared here}}
 }
 
 // A typealias in a subprotocol should be written as a same-type
@@ -21,6 +24,13 @@ protocol P2 : P1 {
 // should be written via a where clause.
 protocol P3a : P1 {
   associatedtype A: P0, P0b // expected-warning{{redeclaration of associated type 'A' from protocol 'P1' is better expressed as a 'where' clause on the protocol}}{{18-18= where A: P0, A: P0b}}{{3-29=}}
+  associatedtype A2: P0, P0b where A2.A == Never, A2: P1 // expected-warning{{redeclaration of associated type 'A2' from protocol 'P1' is better expressed as a 'where' clause on the protocol}}{{18-18= where A2: P0, A2: P0b, A2.A == Never, A2 : P1}}{{3-58=}}
+  associatedtype A3 where A3: P0 // expected-warning{{redeclaration of associated type 'A3' from protocol 'P1' is better expressed as a 'where' clause on the protocol}}{{18-18= where A3 : P0}}{{3-34=}}
+
+  // expected-warning@+1 {{redeclaration of associated type 'A4' from protocol 'P1' is better expressed as a 'where' clause on the protocol}}{{18-18= where A4: P0, A4 : Collection, A4.Element == A4.Index, A4.SubSequence == A4}}{{3-52=}}
+  associatedtype A4: P0 where A4: Collection,
+                              A4.Element == A4.Index,
+                              A4.SubSequence == A4 // {{3-52=}} is this line, so it is correct.
 }
 
 // ... unless it has adds a default type witness
@@ -30,10 +40,12 @@ protocol P3b : P1 {
 
 protocol P4: P1 {
   associatedtype B // expected-note{{'B' declared here}}
+  associatedtype B2 // expected-note{{'B2' declared here}}
 }
 
 protocol P5: P4 where A: P0 {
   typealias B = X1 // expected-warning{{typealias overriding associated type 'B' from protocol 'P4' is better expressed as same-type constraint on the protocol}}{{28-28=, B == X1}}{{3-20=}}
+  associatedtype B2: P5 // expected-warning{{redeclaration of associated type 'B2' from protocol 'P4' is better expressed as a 'where' clause on the protocol}}{{28-28=, B2: P5}} {{3-25=}}
 }