[AutoDiff] [Sema] Fix '@differentiable' witness matching regression. (#34533)

When checking protocol conformances with `@differentiable` requirements, the type checker is supposed to accept omissions of `@differentiable` attributes when there exsits an attribute that covers a superset of the differentiation configuration. This was accidentally regressed in #33776 which made the following test case fail to compile. This is fixed by adjusting the witness matching conditions.

```swift
// rdar://70348904 reproducer:
public protocol P: Differentiable {
    @differentiable(wrt: self)
    @differentiable(wrt: (self, x))
    func foo(_ x: Float) -> Float
}

public struct S: P {}

extension S {
    // This had worked until #33776.
    @differentiable(wrt: (self, x))
    public func foo(_ x: Float) -> Float { x }
}
```

Also fix some suboptimal diagnostics where more information could be shown.

Resolves rdar://70348904.

Author: rxwei

include/swift/AST/DiagnosticsSema.def

@@ -3121,8 +3121,6 @@ ERROR(differentiable_attr_layout_req_unsupported,none,
       ())
 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_invalid_context,none,
      "candidate is missing explicit '%0' attribute to satisfy requirement %1 "
      "(in protocol %3); explicit attribute is necessary because candidate is "

lib/Sema/TypeCheckProtocol.cpp

@@ -383,20 +383,27 @@ matchWitnessDifferentiableAttr(DeclContext *dc, ValueDecl *req,
               reqDiffAttr->getParameterIndices()))
         supersetConfig = witnessConfig;
     }
+
+    // If no exact witness derivative configuration was found, check conditions
+    // for creating an implicit witness `@differentiable` attribute with the
+    // exact derivative configuration.
     if (!foundExactConfig) {
-      // If no exact witness derivative configuration was found, check
-      // conditions for creating an implicit witness `@differentiable` attribute
-      // with the exact derivative configuration.
-
-      // 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() ==
+      auto witnessInDifferentFile =
+          dc->getParentSourceFile() !=
+          witness->getDeclContext()->getParentSourceFile();
+      auto witnessInDifferentTypeContext =
+          dc->getInnermostTypeContext() !=
           witness->getDeclContext()->getInnermostTypeContext();
-      auto sameModule = dc->getModuleScopeContext() ==
-                        witness->getDeclContext()->getModuleScopeContext();
-      if (!sameTypeContext || !sameModule) {
+      // Produce an error instead of creating an implicit `@differentiable`
+      // attribute if any of the following conditions are met:
+      // - The witness is in a different file than the conformance
+      //   declaration.
+      // - The witness is in a different type context (i.e. extension) than
+      //   the conformance declaration, and there is no existing
+      //   `@differentiable` attribute that covers the required differentiation
+      //   parameters.
+      if (witnessInDifferentFile ||
+          (witnessInDifferentTypeContext && !supersetConfig)) {
         // FIXME(TF-1014): `@differentiable` attribute diagnostic does not
         // appear if associated type inference is involved.
         if (auto *vdWitness = dyn_cast<VarDecl>(witness)) {
@@ -2492,30 +2499,14 @@ 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 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_invalid_context,
-              reqDiffAttrString, req->getName(), conformance->getType(),
-              conformance->getProtocol()->getDeclaredInterfaceType())
-          .fixItInsert(match.Witness->getStartLoc(), reqDiffAttrString + ' ');
-    }
-    // Otherwise, emit a general "missing attribute" diagnostic.
-    else {
-      diags
-          .diagnose(witness, diag::protocol_witness_missing_differentiable_attr,
-                    reqDiffAttrString)
-          .fixItInsert(witness->getStartLoc(), reqDiffAttrString + ' ');
-    }
+    diags
+        .diagnose(
+            witness,
+            diag::
+                protocol_witness_missing_differentiable_attr_invalid_context,
+            reqDiffAttrString, req->getName(), conformance->getType(),
+            conformance->getProtocol()->getDeclaredInterfaceType())
+        .fixItInsert(match.Witness->getStartLoc(), reqDiffAttrString + ' ');
     break;
   }
   case MatchKind::EnumCaseWithAssociatedValues:

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

@@ -16,7 +16,8 @@ protocol Protocol1: Differentiable {
   func internalMethod3(_ x: Float) -> Float
 }
 
-protocol Protocol2: Differentiable {
+public protocol Protocol2: Differentiable {
+  @differentiable(wrt: self)
   @differentiable(wrt: (self, x))
   func internalMethod4(_ x: Float) -> Float
 }

test/AutoDiff/Sema/ImplicitDifferentiableAttributeCrossFile/main.swift

@@ -23,6 +23,8 @@
 // RUN: %target-swift-frontend -c %s -primary-file %S/Inputs/other_file.swift
 // RUN: not %target-build-swift %s %S/Inputs/other_file.swift
 
+import _Differentiation
+
 // Error: conformance is in different file than witnesses.
 // expected-error @+1 {{type 'ConformingStruct' does not conform to protocol 'Protocol1'}}
 extension ConformingStruct: Protocol1 {}
@@ -33,3 +35,16 @@ extension ConformingStruct: Protocol2 {
     x
   }
 }
+
+public final class ConformingStructWithSupersetAttr: Protocol2 {}
+
+// rdar://70348904: Witness mismatch failure when a matching witness with a *superset* `@differentiable`
+// attribute is specified.
+// 
+// Note that public witnesses are required to explicitly specify `@differentiable` attributes except
+// those w.r.t. parameters that have already been covered by an existing `@differentiable` attribute.
+extension ConformingStructWithSupersetAttr {
+  // @differentiable(wrt: self) // Omitting this is okay.
+  @differentiable
+  public func internalMethod4(_ x: Float) -> Float { x }
+}

test/AutoDiff/Sema/differentiable_attr_type_checking.swift

@@ -291,39 +291,39 @@ public struct PublicDiffAttrConformance: ProtocolRequirements {
   var y: Float
 
   // FIXME(TF-284): Fix unexpected diagnostic.
-  // expected-note @+2 {{candidate is missing attribute '@differentiable'}} {{10-10=@differentiable }}
+  // expected-note @+2 {{candidate is missing explicit '@differentiable' attribute to satisfy requirement}} {{10-10=@differentiable }}
   // expected-note @+1 {{candidate has non-matching type '(x: Float, y: Float)'}}
   public init(x: Float, y: Float) {
     self.x = x
     self.y = y
   }
 
   // FIXME(TF-284): Fix unexpected diagnostic.
-  // expected-note @+2 {{candidate is missing attribute '@differentiable'}} {{10-10=@differentiable }}
+  // expected-note @+2 {{candidate is missing explicit '@differentiable' attribute to satisfy requirement}} {{10-10=@differentiable }}
   // expected-note @+1 {{candidate has non-matching type '(x: Float, y: Int)'}}
   public init(x: Float, y: Int) {
     self.x = x
     self.y = Float(y)
   }
 
-  // expected-note @+2 {{candidate is missing attribute '@differentiable'}} {{10-10=@differentiable }}
+  // expected-note @+2 {{candidate is missing explicit '@differentiable' attribute to satisfy requirement}} {{10-10=@differentiable }}
   // expected-note @+1 {{candidate has non-matching type '(Float, Float) -> Float'}}
   public func amb(x: Float, y: Float) -> Float {
     return x
   }
 
-  // expected-note @+2 {{candidate is missing attribute '@differentiable(wrt: x)'}} {{10-10=@differentiable(wrt: x) }}
+  // expected-note @+2 {{candidate is missing explicit '@differentiable(wrt: x)' attribute to satisfy requirement}} {{10-10=@differentiable(wrt: x) }}
   // expected-note @+1 {{candidate has non-matching type '(Float, Int) -> Float'}}
   public func amb(x: Float, y: Int) -> Float {
     return x
   }
 
-  // expected-note @+1 {{candidate is missing attribute '@differentiable'}}
+  // expected-note @+1 {{candidate is missing explicit '@differentiable' attribute to satisfy requirement}}
   public func f1(_ x: Float) -> Float {
     return x
   }
 
-  // expected-note @+2 {{candidate is missing attribute '@differentiable'}}
+  // expected-note @+2 {{candidate is missing explicit '@differentiable' attribute to satisfy requirement}}
   @differentiable(wrt: (self, x))
   public func f2(_ x: Float, _ y: Float) -> Float {
     return x + y
@@ -558,7 +558,7 @@ public struct AttemptsToSatisfyRequirement: HasRequirement {
   // This `@differentiable` attribute does not satisfy the requirement because
   // it is mroe constrained than the requirement's `@differentiable` attribute.
   @differentiable(where T: CustomStringConvertible)
-  // expected-note @+1 {{candidate is missing attribute '@differentiable(wrt: (x, y))'}}
+  // expected-note @+1 {{candidate is missing explicit '@differentiable(wrt: (x, y))' attribute to satisfy requirement}}
   public func requirement<T: Differentiable>(_ x: T, _ y: T) -> T { x }
 }