[AutoDiff] Require same access level for original/derivative functions.

Require `@derivative` functions and their original functions to have the same
access level. Public original functions may also have internal
`@usableFromInline` derivatives, as a special case. Diagnose access level
mismatches.

This simplifies derivative registration rules, and may enable simplification of
AutoDiff symbol linkage rules.

Resolves TF-1099 and TF-1160.

Author: dan-zheng

include/swift/AST/DiagnosticsSema.def

@@ -3054,6 +3054,22 @@ ERROR(derivative_attr_original_already_has_derivative,none,
       "a derivative already exists for %0", (DeclName))
 NOTE(derivative_attr_duplicate_note,none,
      "other attribute declared here", ())
+ERROR(derivative_attr_access_level_lower_than_original,none,
+      "derivative functions for "
+      "%select{private|fileprivate|internal|public or '@usableFromInline'|open}1 "
+      "original function %0 must also be "
+      "%select{private|fileprivate|internal|public or '@usableFromInline'|open}1, "
+      "but %2 is %select{private|fileprivate|internal|public|open}3",
+      (/*original*/ DeclName, /*original*/ AccessLevel,
+       /*derivative*/ DeclName, /*derivative*/ AccessLevel))
+ERROR(derivative_attr_access_level_higher_than_original,none,
+      "the original function of "
+      "%select{a private|a fileprivate|an internal|a public or '@usableFromInline'|an open}3 "
+      "derivative function must also be "
+      "%select{private|fileprivate|internal|public or '@usableFromInline'|open}3, "
+      "but %0 is %select{private|fileprivate|internal|public|open}1",
+      (/*original*/ DeclName, /*original*/ AccessLevel,
+       /*derivative*/ DeclName, /*derivative*/ AccessLevel))
 
 // @transpose
 ERROR(transpose_attr_invalid_linearity_parameter_or_result,none,

lib/Sema/TypeCheckAttr.cpp

@@ -4540,6 +4540,37 @@ static bool typeCheckDerivativeAttr(ASTContext &Ctx, Decl *D,
   }
   attr->setOriginalFunction(originalAFD);
 
+  // Returns true if:
+  // - Original function and derivative function has the same access level.
+  // - Original function is public and derivative function is internal
+  //   `@usableFromInline`. This is the only special case.
+  auto compatibleAccessLevels = [&]() {
+    if (originalAFD->getFormalAccess() == derivative->getFormalAccess())
+      return true;
+    return originalAFD->getFormalAccess() == AccessLevel::Public &&
+           derivative->getEffectiveAccess() == AccessLevel::Public;
+  };
+
+  // Check access level compatibility for original and derivative functions.
+  if (!compatibleAccessLevels()) {
+    auto diagID = diag::derivative_attr_access_level_higher_than_original;
+    AccessLevel originalAccess;
+    AccessLevel derivativeAccess;
+    if (originalAFD->getEffectiveAccess() < derivative->getEffectiveAccess()) {
+      originalAccess =
+          originalAFD->getFormalAccessScope().accessLevelForDiagnostics();
+      derivativeAccess = derivative->getEffectiveAccess();
+    } else {
+      diagID = diag::derivative_attr_access_level_lower_than_original;
+      originalAccess = originalAFD->getEffectiveAccess();
+      derivativeAccess =
+          derivative->getFormalAccessScope().accessLevelForDiagnostics();
+    }
+    diags.diagnose(originalName.Loc, diagID, originalAFD->getName(),
+                   originalAccess, derivative->getName(), derivativeAccess);
+    return true;
+  }
+
   // Get the resolved differentiability parameter indices.
   auto *resolvedDiffParamIndices = attr->getParameterIndices();
 

test/AutoDiff/SILGen/sil_differentiability_witness.swift

@@ -44,7 +44,7 @@ public func foo_vjp(_ x: Float) -> (value: Float, pullback: (Float) -> Float) {
 func bar<T>(_ x: Float, _ y: T) -> Float { x }
 
 @derivative(of: bar)
-public func bar_jvp<T>(_ x: Float, _ y: T) -> (value: Float, differential: (Float) -> Float) {
+func bar_jvp<T>(_ x: Float, _ y: T) -> (value: Float, differential: (Float) -> Float) {
   (x, { $0 })
 }
 

test/AutoDiff/Sema/derivative_attr_type_checking.swift

@@ -746,6 +746,7 @@ extension InoutParameters {
 // Test cross-file derivative registration.
 
 extension FloatingPoint where Self: Differentiable {
+  @usableFromInline
   @derivative(of: rounded)
   func vjpRounded() -> (
     value: Self,
@@ -802,3 +803,112 @@ extension HasADefaultDerivative {
     (x, { 10 * $0 })
   }
 }
+
+// MARK: - Original function visibility = derivative function visibility
+
+public func public_original_public_derivative(_ x: Float) -> Float { x }
+@derivative(of: public_original_public_derivative)
+public func _public_original_public_derivative(_ x: Float) -> (value: Float, pullback: (Float) -> Float) {
+  fatalError()
+}
+
+public func public_original_usablefrominline_derivative(_ x: Float) -> Float { x }
+@usableFromInline
+@derivative(of: public_original_usablefrominline_derivative)
+func _public_original_usablefrominline_derivative(_ x: Float) -> (value: Float, pullback: (Float) -> Float) {
+  fatalError()
+}
+
+@usableFromInline
+func usablefrominline_original_public_derivative(_ x: Float) -> Float { x }
+// expected-error @+1 {{derivative functions for public or '@usableFromInline' original function 'usablefrominline_original_public_derivative' must also be public or '@usableFromInline', but '_usablefrominline_original_public_derivative' is public}}
+@derivative(of: usablefrominline_original_public_derivative)
+public func _usablefrominline_original_public_derivative(_ x: Float) -> (value: Float, pullback: (Float) -> Float) {
+  fatalError()
+}
+
+func internal_original_internal_derivative(_ x: Float) -> Float { x }
+@derivative(of: internal_original_internal_derivative)
+func _internal_original_internal_derivative(_ x: Float) -> (value: Float, pullback: (Float) -> Float) {
+  fatalError()
+}
+
+private func private_original_private_derivative(_ x: Float) -> Float { x }
+@derivative(of: private_original_private_derivative)
+private func _private_original_private_derivative(_ x: Float) -> (value: Float, pullback: (Float) -> Float) {
+  fatalError()
+}
+
+fileprivate func fileprivate_original_fileprivate_derivative(_ x: Float) -> Float { x }
+@derivative(of: fileprivate_original_fileprivate_derivative)
+fileprivate func _fileprivate_original_fileprivate_derivative(_ x: Float) -> (value: Float, pullback: (Float) -> Float) {
+  fatalError()
+}
+
+fileprivate func fileprivate_original_private_derivative(_ x: Float) -> Float { x }
+// expected-error @+1 {{derivative functions for fileprivate original function 'fileprivate_original_private_derivative' must also be fileprivate, but '_fileprivate_original_private_derivative' is private}}
+@derivative(of: fileprivate_original_private_derivative)
+private func _fileprivate_original_private_derivative(_ x: Float) -> (value: Float, pullback: (Float) -> Float) {
+  fatalError()
+}
+
+private func private_original_fileprivate_derivative(_ x: Float) -> Float { x }
+// expected-error @+1 {{derivative functions for fileprivate original function 'private_original_fileprivate_derivative' must also be fileprivate, but '_private_original_fileprivate_derivative' is fileprivate}}
+@derivative(of: private_original_fileprivate_derivative)
+fileprivate func _private_original_fileprivate_derivative(_ x: Float) -> (value: Float, pullback: (Float) -> Float) {
+  fatalError()
+}
+
+// MARK: - Original function visibility < derivative function visibility
+
+func internal_original_public_derivative(_ x: Float) -> Float { x }
+// expected-error @+1 {{the original function of a public or '@usableFromInline' derivative function must also be public or '@usableFromInline', but 'internal_original_public_derivative' is internal}}
+@derivative(of: internal_original_public_derivative)
+public func _internal_original_public_derivative(_ x: Float) -> (value: Float, pullback: (Float) -> Float) {
+  fatalError()
+}
+
+private func private_original_usablefrominline_derivative(_ x: Float) -> Float { x }
+// expected-error @+1 {{the original function of a public or '@usableFromInline' derivative function must also be public or '@usableFromInline', but 'private_original_usablefrominline_derivative' is private}}
+@derivative(of: private_original_usablefrominline_derivative)
+@usableFromInline
+func _private_original_usablefrominline_derivative(_ x: Float) -> (value: Float, pullback: (Float) -> Float) {
+  fatalError()
+}
+
+private func private_original_public_derivative(_ x: Float) -> Float { x }
+// expected-error @+1 {{the original function of a public or '@usableFromInline' derivative function must also be public or '@usableFromInline', but 'private_original_public_derivative' is private}}
+@derivative(of: private_original_public_derivative)
+public func _private_original_public_derivative(_ x: Float) -> (value: Float, pullback: (Float) -> Float) {
+  fatalError()
+}
+
+private func private_original_internal_derivative(_ x: Float) -> Float { x }
+// expected-error @+1 {{the original function of an internal derivative function must also be internal, but 'private_original_internal_derivative' is private}}
+@derivative(of: private_original_internal_derivative)
+func _private_original_internal_derivative(_ x: Float) -> (value: Float, pullback: (Float) -> Float) {
+  fatalError()
+}
+
+// MARK: - Original function visibility > derivative function visibility
+
+public func public_original_private_derivative(_ x: Float) -> Float { x }
+// expected-error @+1 {{derivative functions for public or '@usableFromInline' original function 'public_original_private_derivative' must also be public or '@usableFromInline', but '_public_original_private_derivative' is fileprivate}}
+@derivative(of: public_original_private_derivative)
+fileprivate func _public_original_private_derivative(_ x: Float) -> (value: Float, pullback: (Float) -> Float) {
+  fatalError()
+}
+
+public func public_original_internal_derivative(_ x: Float) -> Float { x }
+// expected-error @+1 {{derivative functions for public or '@usableFromInline' original function 'public_original_internal_derivative' must also be public or '@usableFromInline', but '_public_original_internal_derivative' is internal}}
+@derivative(of: public_original_internal_derivative)
+func _public_original_internal_derivative(_ x: Float) -> (value: Float, pullback: (Float) -> Float) {
+  fatalError()
+}
+
+func internal_original_fileprivate_derivative(_ x: Float) -> Float { x }
+// expected-error @+1 {{derivative functions for internal original function 'internal_original_fileprivate_derivative' must also be internal, but '_internal_original_fileprivate_derivative' is fileprivate}}
+@derivative(of: internal_original_fileprivate_derivative)
+fileprivate func _internal_original_fileprivate_derivative(_ x: Float) -> (value: Float, pullback: (Float) -> Float) {
+  fatalError()
+}