Update requirements for explicit @differentiable attribute on protocol witnesses.

Require explicit `@differentiable` attribute on witnesses declared in a
different type context or file than the protocol conformance.

Author: dan-zheng

include/swift/AST/DiagnosticsSema.def

@@ -3112,10 +3112,11 @@ ERROR(overriding_decl_missing_differentiable_attr,none,
       "overriding declaration is missing attribute '%0'", (StringRef))
 NOTE(protocol_witness_missing_differentiable_attr,none,
      "candidate is missing attribute '%0'", (StringRef))
-NOTE(protocol_witness_missing_differentiable_attr_other_file,none,
+NOTE(protocol_witness_missing_differentiable_attr_invalid_context,none,
      "candidate is missing explicit '%0' attribute to satisfy requirement %1 "
-     "(in protocol %3) because it is declared in a different file than the "
-     "conformance of %2 to %3", (StringRef, DeclName, Type, Type))
+     "(in protocol %3); explicit attribute is necessary because candidate is "
+     "declared in a different type context or file than the conformance of %2 "
+     "to %3", (StringRef, DeclName, Type, Type))
 
 // @derivative
 ERROR(derivative_attr_expected_result_tuple,none,

lib/Sema/TypeCheckProtocol.cpp

@@ -388,11 +388,15 @@ matchWitnessDifferentiableAttr(DeclContext *dc, ValueDecl *req,
       // conditions for creating an implicit witness `@differentiable` attribute
       // with the exact derivative configuration.
 
-      // If witness is declared in a different file as the conformance, we
-      // should not create an implicit `@differentiable` attribute on the
-      // witness. Produce an error.
-      if (dc->getModuleScopeContext() !=
-          witness->getDeclContext()->getModuleScopeContext()) {
+      // If witness is declared in a different file or type context than the
+      // conformance, we should not create an implicit `@differentiable`
+      // attribute on the witness. Produce an error.
+      auto sameTypeContext =
+          dc->getInnermostTypeContext() ==
+          witness->getDeclContext()->getInnermostTypeContext();
+      auto sameModule = dc->getModuleScopeContext() ==
+                        witness->getDeclContext()->getModuleScopeContext();
+      if (!sameTypeContext || !sameModule) {
         // FIXME(TF-1014): `@differentiable` attribute diagnostic does not
         // appear if associated type inference is involved.
         if (auto *vdWitness = dyn_cast<VarDecl>(witness)) {
@@ -2454,14 +2458,19 @@ diagnoseMatch(ModuleDecl *module, NormalProtocolConformance *conformance,
     llvm::raw_string_ostream os(reqDiffAttrString);
     reqAttr->print(os, req, omitWrtClause);
     os.flush();
-    // If the witness is declared in a different file than the conformance, emit
-    // a specialized diagnostic.
-    if (conformance->getDeclContext()->getModuleScopeContext() !=
-            witness->getDeclContext()->getModuleScopeContext()) {
+    // If the witness is declared in a different file or type context than the
+    // conformance, emit a specialized diagnostic.
+    auto sameModule = conformance->getDeclContext()->getModuleScopeContext() !=
+                      witness->getDeclContext()->getModuleScopeContext();
+    auto sameTypeContext =
+        conformance->getDeclContext()->getInnermostTypeContext() !=
+        witness->getDeclContext()->getInnermostTypeContext();
+    if (sameModule || sameTypeContext) {
       diags
           .diagnose(
               witness,
-              diag::protocol_witness_missing_differentiable_attr_other_file,
+              diag::
+                  protocol_witness_missing_differentiable_attr_invalid_context,
               reqDiffAttrString, req->getName(), conformance->getType(),
               conformance->getProtocol()->getDeclaredInterfaceType())
           .fixItInsert(match.Witness->getStartLoc(), reqDiffAttrString + ' ');

test/AutoDiff/Sema/ImplicitDifferentiableAttributeCrossFile/Inputs/other_file.swift

@@ -32,20 +32,20 @@ protocol Protocol2: Differentiable {
 
 struct ConformingStruct: Differentiable {
   // Error for missing `@differentiable` attribute.
-  // expected-note @+1 {{candidate is missing explicit '@differentiable' attribute to satisfy requirement 'internalMethod1' (in protocol 'Protocol1') because it is declared in a different file than the conformance of 'ConformingStruct' to 'Protocol1'}} {{3-3=@differentiable }}
+  // expected-note @+1 {{candidate is missing explicit '@differentiable' attribute to satisfy requirement 'internalMethod1' (in protocol 'Protocol1'); explicit attribute is necessary because candidate is declared in a different type context or file than the conformance of 'ConformingStruct' to 'Protocol1'}} {{3-3=@differentiable }}
   func internalMethod1(_ x: Float) -> Float {
     x
   }
 
   // Error for missing `@differentiable` superset attribute.
-  // expected-note @+2 {{candidate is missing explicit '@differentiable' attribute to satisfy requirement 'internalMethod2' (in protocol 'Protocol1') because it is declared in a different file than the conformance of 'ConformingStruct' to 'Protocol1'}} {{3-3=@differentiable }}
+  // expected-note @+2 {{candidate is missing explicit '@differentiable' attribute to satisfy requirement 'internalMethod2' (in protocol 'Protocol1'); explicit attribute is necessary because candidate is declared in a different type context or file than the conformance of 'ConformingStruct' to 'Protocol1'}} {{3-3=@differentiable }}
   @differentiable(wrt: x)
   func internalMethod2(_ x: Float) -> Float {
     x
   }
 
   // Error for missing `@differentiable` subset attribute.
-  // expected-note @+2 {{candidate is missing explicit '@differentiable(wrt: x)' attribute to satisfy requirement 'internalMethod3' (in protocol 'Protocol1') because it is declared in a different file than the conformance of 'ConformingStruct' to 'Protocol1'}} {{3-3=@differentiable(wrt: x) }}
+  // expected-note @+2 {{candidate is missing explicit '@differentiable(wrt: x)' attribute to satisfy requirement 'internalMethod3' (in protocol 'Protocol1'); explicit attribute is necessary because candidate is declared in a different type context or file than the conformance of 'ConformingStruct' to 'Protocol1'}} {{3-3=@differentiable(wrt: x) }}
   @differentiable(wrt: (self, x))
   func internalMethod3(_ x: Float) -> Float {
     x

test/AutoDiff/Sema/ImplicitDifferentiableAttributeCrossFile/Inputs/other_file_protocol_default_implementation_witness.swift

@@ -10,7 +10,7 @@ protocol P2: P1 {}
 
 extension P2 {
   @differentiable(wrt: (self, input))
-  // expected-note @+1 {{candidate is missing explicit '@differentiable(wrt: self)' attribute to satisfy requirement 'callAsFunction' (in protocol 'P1') because it is declared in a different file than the conformance of 'ConformingStruct' to 'P1'}}
+  // expected-note @+1 {{candidate is missing explicit '@differentiable(wrt: self)' attribute to satisfy requirement 'callAsFunction' (in protocol 'P1'); explicit attribute is necessary because candidate is declared in a different type context or file than the conformance of 'ConformingStruct' to 'P1'}}
   public func callAsFunction(_ input: Float) -> Float {
     return input
   }