[CSDiagnostics] Emit fix-its to insert requirement stubs for missing protocols in context (#33628)

* [CSDiagnostics] Emit fix-its to insert requirement stubs in editor mode for missing protocols in context

* [CSDiagnostics] Only include missing requirements in stub fix-it for missing protocols in context

The conforming type may already have declarations that could satisfy a requirement, so we shouldn't include it in the fix-it

Author: theblixguy

include/swift/AST/ProtocolConformance.h

@@ -465,21 +465,22 @@ class NormalProtocolConformance : public RootProtocolConformance,
   uint64_t LoaderContextData;
   friend class ASTContext;
 
+  void resolveLazyInfo() const;
+
+  void differenceAndStoreConditionalRequirements() const;
+
+public:
   NormalProtocolConformance(Type conformingType, ProtocolDecl *protocol,
                             SourceLoc loc, DeclContext *dc,
                             ProtocolConformanceState state)
-    : RootProtocolConformance(ProtocolConformanceKind::Normal, conformingType),
-      ProtocolAndState(protocol, state), Loc(loc), ContextAndInvalid(dc, false)
-  {
+      : RootProtocolConformance(ProtocolConformanceKind::Normal,
+                                conformingType),
+        ProtocolAndState(protocol, state), Loc(loc),
+        ContextAndInvalid(dc, false) {
     assert(!conformingType->hasArchetype() &&
            "ProtocolConformances should store interface types");
   }
 
-  void resolveLazyInfo() const;
-
-  void differenceAndStoreConditionalRequirements() const;
-
-public:
   /// Get the protocol being conformed to.
   ProtocolDecl *getProtocol() const { return ProtocolAndState.getPointer(); }
 

lib/Sema/CSDiagnostics.cpp

@@ -17,8 +17,10 @@
 #include "CSDiagnostics.h"
 #include "ConstraintSystem.h"
 #include "MiscDiagnostics.h"
+#include "TypeCheckProtocol.h"
 #include "TypoCorrection.h"
 #include "swift/AST/ASTContext.h"
+#include "swift/AST/ASTPrinter.h"
 #include "swift/AST/Decl.h"
 #include "swift/AST/ExistentialLayout.h"
 #include "swift/AST/Expr.h"
@@ -2768,9 +2770,11 @@ bool ContextualFailure::tryProtocolConformanceFixIt(
 
   // Let's build a list of protocols that the context does not conform to.
   SmallVector<std::string, 8> missingProtoTypeStrings;
+  SmallVector<ProtocolDecl *, 8> missingProtocols;
   for (auto protocol : layout.getProtocols()) {
     if (!TypeChecker::conformsToProtocol(fromType, protocol->getDecl(), getDC())) {
       missingProtoTypeStrings.push_back(protocol->getString());
+      missingProtocols.push_back(protocol->getDecl());
     }
   }
 
@@ -2792,8 +2796,6 @@ bool ContextualFailure::tryProtocolConformanceFixIt(
 
   // Emit a diagnostic to inform the user that they need to conform to the
   // missing protocols.
-  //
-  // TODO: Maybe also insert the requirement stubs?
   auto conformanceDiag =
       emitDiagnostic(diag::assign_protocol_conformance_fix_it, unwrappedToType,
                      nominal->getDescriptiveKind(), fromType);
@@ -2808,6 +2810,35 @@ bool ContextualFailure::tryProtocolConformanceFixIt(
     conformanceDiag.fixItInsert(nameEndLoc, ": " + protoString);
   }
 
+  // Emit fix-its to insert requirement stubs if we're in editor mode.
+  if (!getASTContext().LangOpts.DiagnosticsEditorMode) {
+    return true;
+  }
+
+  {
+    llvm::SmallString<128> Text;
+    llvm::raw_svector_ostream SS(Text);
+    llvm::SetVector<ValueDecl *> missingWitnesses;
+    for (auto protocol : missingProtocols) {
+      auto conformance = NormalProtocolConformance(
+          nominal->getDeclaredType(), protocol, SourceLoc(), nominal,
+          ProtocolConformanceState::Incomplete);
+      ConformanceChecker checker(getASTContext(), &conformance,
+                                 missingWitnesses);
+      checker.resolveValueWitnesses();
+      checker.resolveTypeWitnesses();
+    }
+
+    for (auto decl : missingWitnesses) {
+      swift::printRequirementStub(decl, nominal, nominal->getDeclaredType(),
+                                  nominal->getStartLoc(), SS);
+    }
+
+    if (!Text.empty()) {
+      conformanceDiag.fixItInsertAfter(nominal->getBraces().Start, Text.str());
+    }
+  }
+
   return true;
 }
 

test/decl/protocol/fixits_missing_protocols_in_context.swift

@@ -0,0 +1,45 @@
+// RUN: %target-swift-frontend -typecheck -diagnostics-editor-mode -verify %s
+
+// Test that we emit fix-its to insert requirement stubs for the missing protocol conformance, in addition to adding the conformance.
+
+protocol P {
+  func method()
+  var property: Int { get }
+}
+
+class C {
+  var p: P?
+
+  func assign() {
+    p = self
+    // expected-error@-1 {{cannot assign value of type 'C' to type 'P?'}}
+    // expected-note@-2 {{add missing conformance to 'P' to class 'C'}} {{8-8=: P}} {{10-10=\n    func method() {\n        <#code#>\n    \}\n\n    var property: Int\n}}
+  }
+}
+
+// Test that we don't emit fix-it to insert a requirement stub if there is already a satisfying witness.
+
+class C1 {
+  var p: P?
+
+  func assign() {
+    p = self
+    // expected-error@-1 {{cannot assign value of type 'C1' to type 'P?'}}
+    // expected-note@-2 {{add missing conformance to 'P' to class 'C1'}} {{9-9=: P}} {{11-11=\n    var property: Int\n}}
+  }
+
+  func method() {}
+}
+
+class C2 {
+  var p: P?
+
+  func assign() {
+    p = self
+    // expected-error@-1 {{cannot assign value of type 'C2' to type 'P?'}}
+    // expected-note@-2 {{add missing conformance to 'P' to class 'C2'}} {{9-9=: P}}
+  }
+
+  func method() {}
+  var property: Int = 0
+}