Add the notion of @_implementationOnly overrides

When an @_implementationOnly import includes Objective-C categories
for existing types, it's useful to be able to override the members
provided in those categories without exposing them to clients of the
framework being built. Allow this as long as the overriding
declaration is marked as @_implementationOnly itself, with an
additional check that the type of the declaration does not change.
(Normally overrides are allowed to change in covariant ways.)

Part of rdar://50827914

include/swift/AST/Attr.def

@@ -396,7 +396,7 @@ SIMPLE_DECL_ATTR(_alwaysEmitIntoClient, AlwaysEmitIntoClient,
   83)
 
 SIMPLE_DECL_ATTR(_implementationOnly, ImplementationOnly,
-  OnImport | UserInaccessible,
+  OnImport | OnFunc | OnConstructor | OnVar | OnSubscript | UserInaccessible,
   84)
 DECL_ATTR(_custom, Custom,
   OnAnyDecl | UserInaccessible,

include/swift/AST/DiagnosticsSema.def

@@ -2485,6 +2485,18 @@ WARNING(warn_implementation_only_conflict,none,
 NOTE(implementation_only_conflict_here,none,
      "imported as implementation-only here", ())
 
+ERROR(implementation_only_decl_non_override,none,
+      "'@_implementationOnly' can only be used on overrides", ())
+ERROR(implementation_only_override_changed_type,none,
+      "'@_implementationOnly' override must have the same type as the "
+      "declaration it overrides (%0)", (Type))
+ERROR(implementation_only_override_without_attr,none,
+      "override of '@_implementationOnly' %0 should also be declared "
+      "'@_implementationOnly'", (DescriptiveDeclKind))
+ERROR(implementation_only_override_import_without_attr,none,
+      "override of %0 imported as implementation-only must be declared "
+      "'@_implementationOnly'", (DescriptiveDeclKind))
+
 // Derived conformances
 ERROR(cannot_synthesize_init_in_extension_of_nonfinal,none,
       "implementation of %0 for non-final class cannot be automatically "

lib/AST/ASTPrinter.cpp

@@ -126,6 +126,10 @@ PrintOptions PrintOptions::printParseableInterfaceFile(bool preferTypeRepr) {
 
   class ShouldPrintForParseableInterface : public ShouldPrintChecker {
     bool shouldPrint(const Decl *D, const PrintOptions &options) override {
+      // Skip anything that is marked `@_implementationOnly` itself.
+      if (D->getAttrs().hasAttribute<ImplementationOnlyAttr>())
+        return false;
+
       // Skip anything that isn't 'public' or '@usableFromInline'.
       if (auto *VD = dyn_cast<ValueDecl>(D)) {
         if (!isPublicOrUsableFromInline(VD)) {

lib/PrintAsObjC/PrintAsObjC.cpp

@@ -204,7 +204,8 @@ public:
   }
   
   bool shouldInclude(const ValueDecl *VD) {
-    return isVisibleToObjC(VD, minRequiredAccess);
+    return isVisibleToObjC(VD, minRequiredAccess) &&
+           !VD->getAttrs().hasAttribute<ImplementationOnlyAttr>();
   }
 
 private:

lib/Sema/TypeCheckAccess.cpp

@@ -1652,6 +1652,9 @@ public:
   explicit ExportabilityChecker(TypeChecker &TC) : TC(TC) {}
 
   static bool shouldSkipChecking(const ValueDecl *VD) {
+    if (VD->getAttrs().hasAttribute<ImplementationOnlyAttr>())
+      return true;
+
     // Accessors are handled as part of their Var or Subscript, and we don't
     // want to redo extension signature checking for them.
     if (isa<AccessorDecl>(VD))
@@ -1681,13 +1684,46 @@ public:
     return true;
   }
 
+  void checkOverride(const ValueDecl *VD) {
+    const ValueDecl *overridden = VD->getOverriddenDecl();
+    if (!overridden)
+      return;
+
+    auto *SF = VD->getDeclContext()->getParentSourceFile();
+    assert(SF && "checking a non-source declaration?");
+
+    ModuleDecl *M = overridden->getModuleContext();
+    if (SF->isImportedImplementationOnly(M)) {
+      TC.diagnose(VD, diag::implementation_only_override_import_without_attr,
+                  overridden->getDescriptiveKind())
+        .fixItInsert(VD->getAttributeInsertionLoc(false),
+                     "@_implementationOnly ");
+      TC.diagnose(overridden, diag::overridden_here);
+      return;
+    }
+
+    if (overridden->getAttrs().hasAttribute<ImplementationOnlyAttr>()) {
+      TC.diagnose(VD, diag::implementation_only_override_without_attr,
+                  overridden->getDescriptiveKind())
+        .fixItInsert(VD->getAttributeInsertionLoc(false),
+                     "@_implementationOnly ");
+      TC.diagnose(overridden, diag::overridden_here);
+      return;
+    }
+
+    // FIXME: Check storage decls where the setter is in a separate module from
+    // the getter, which is a thing Objective-C can do.
+  }
+
   void visit(Decl *D) {
     if (D->isInvalid() || D->isImplicit())
       return;
 
-    if (auto *VD = dyn_cast<ValueDecl>(D))
+    if (auto *VD = dyn_cast<ValueDecl>(D)) {
       if (shouldSkipChecking(VD))
         return;
+      checkOverride(VD);
+    }
 
     DeclVisitor<ExportabilityChecker>::visit(D);
   }
@@ -1729,13 +1765,16 @@ public:
   void checkNamedPattern(const NamedPattern *NP,
                          const llvm::DenseSet<const VarDecl *> &seenVars) {
     const VarDecl *theVar = NP->getDecl();
-    // Only check individual variables if we didn't check an enclosing
-    // TypedPattern.
-    if (seenVars.count(theVar) || theVar->isInvalid())
-      return;
     if (shouldSkipChecking(theVar))
       return;
 
+    checkOverride(theVar);
+
+    // Only check the type of individual variables if we didn't check an
+    // enclosing TypedPattern.
+    if (seenVars.count(theVar) || theVar->isInvalid())
+      return;
+
     checkType(theVar->getInterfaceType(), /*typeRepr*/nullptr, theVar,
               getDiagnoseCallback(theVar), getDiagnoseCallback(theVar));
   }

lib/Sema/TypeCheckAttr.cpp

@@ -772,7 +772,6 @@ public:
     IGNORED_ATTR(IBDesignable)
     IGNORED_ATTR(IBInspectable)
     IGNORED_ATTR(IBOutlet) // checked early.
-    IGNORED_ATTR(ImplementationOnly)
     IGNORED_ATTR(ImplicitlyUnwrappedOptional)
     IGNORED_ATTR(Indirect)
     IGNORED_ATTR(Inline)
@@ -858,6 +857,8 @@ public:
   void visitCustomAttr(CustomAttr *attr);
   void visitPropertyWrapperAttr(PropertyWrapperAttr *attr);
   void visitFunctionBuilderAttr(FunctionBuilderAttr *attr);
+
+  void visitImplementationOnlyAttr(ImplementationOnlyAttr *attr);
 };
 } // end anonymous namespace
 
@@ -2630,6 +2631,69 @@ void AttributeChecker::visitFunctionBuilderAttr(FunctionBuilderAttr *attr) {
   // Any other validation?
 }
 
+void
+AttributeChecker::visitImplementationOnlyAttr(ImplementationOnlyAttr *attr) {
+  if (isa<ImportDecl>(D)) {
+    // These are handled elsewhere.
+    return;
+  }
+
+  auto *VD = cast<ValueDecl>(D);
+  auto *overridden = VD->getOverriddenDecl();
+  if (!overridden) {
+    diagnoseAndRemoveAttr(attr, diag::implementation_only_decl_non_override);
+    return;
+  }
+
+  // Check if VD has the exact same type as what it overrides.
+  // Note: This is specifically not using `swift::getMemberTypeForComparison`
+  // because that erases more information than we want, like `throws`-ness.
+  auto baseInterfaceTy = overridden->getInterfaceType();
+  auto derivedInterfaceTy = VD->getInterfaceType();
+
+  auto selfInterfaceTy = VD->getDeclContext()->getDeclaredInterfaceType();
+
+  auto overrideInterfaceTy =
+      selfInterfaceTy->adjustSuperclassMemberDeclType(overridden, VD,
+                                                      baseInterfaceTy);
+
+  if (isa<AbstractFunctionDecl>(VD)) {
+    // Drop the 'Self' parameter.
+    // FIXME: The real effect here, though, is dropping the generic signature.
+    // This should be okay because it should already be checked as part of
+    // making an override, but that isn't actually the case as of this writing,
+    // and it's kind of suspect anyway.
+    derivedInterfaceTy =
+        derivedInterfaceTy->castTo<AnyFunctionType>()->getResult();
+    overrideInterfaceTy =
+        overrideInterfaceTy->castTo<AnyFunctionType>()->getResult();
+  } else if (isa<SubscriptDecl>(VD)) {
+    // For subscripts, we don't have a 'Self' type, but turn it
+    // into a monomorphic function type.
+    // FIXME: does this actually make sense, though?
+    auto derivedInterfaceFuncTy = derivedInterfaceTy->castTo<AnyFunctionType>();
+    derivedInterfaceTy =
+        FunctionType::get(derivedInterfaceFuncTy->getParams(),
+                          derivedInterfaceFuncTy->getResult());
+    auto overrideInterfaceFuncTy =
+        overrideInterfaceTy->castTo<AnyFunctionType>();
+    overrideInterfaceTy =
+        FunctionType::get(overrideInterfaceFuncTy->getParams(),
+                          overrideInterfaceFuncTy->getResult());
+  }
+
+  if (!derivedInterfaceTy->isEqual(overrideInterfaceTy)) {
+    TC.diagnose(VD, diag::implementation_only_override_changed_type,
+                overrideInterfaceTy);
+    TC.diagnose(overridden, diag::overridden_here);
+    return;
+  }
+
+  // FIXME: When compiling without library evolution enabled, this should also
+  // check whether VD or any of its accessors need a new vtable entry, even if
+  // it won't necessarily be able to say why.
+}
+
 void TypeChecker::checkParameterAttributes(ParameterList *params) {
   for (auto param: *params) {
     checkDeclAttributes(param);

test/PrintAsObjC/Inputs/custom-modules/MiserablePileOfSecrets.h

@@ -0,0 +1,5 @@
+@import Foundation;
+
+@interface NSObject (Secrets)
+- (void)secretMethod;
+@end

test/PrintAsObjC/Inputs/custom-modules/module.map

@@ -47,3 +47,8 @@ module resilient_objc_class {
   header "resilient_objc_class.h"
   export *
 }
+
+module MiserablePileOfSecrets {
+  header "MiserablePileOfSecrets.h"
+  export *
+}

test/PrintAsObjC/imports.swift

@@ -17,12 +17,16 @@
 // CHECK-NEXT: @import MostlyPrivate1;
 // CHECK-NEXT: @import MostlyPrivate1_Private;
 // CHECK-NEXT: @import MostlyPrivate2_Private;
+// CHECK-NEXT: @import ObjectiveC;
 // CHECK-NEXT: @import ctypes.bits;
 
 // NEGATIVE-NOT: ctypes;
 // NEGATIVE-NOT: ImSub;
 // NEGATIVE-NOT: ImplicitSub;
 // NEGATIVE-NOT: MostlyPrivate2;
+// NEGATIVE-NOT: MiserablePileOfSecrets;
+
+// NEGATIVE-NOT: secretMethod
 
 import ctypes.bits
 import Foundation
@@ -40,6 +44,8 @@ import MostlyPrivate1_Private
 // Deliberately not importing MostlyPrivate2
 import MostlyPrivate2_Private
 
+@_implementationOnly import MiserablePileOfSecrets
+
 @objc class Test {
   @objc let word: DWORD = 0
   @objc let number: TimeInterval = 0.0
@@ -57,3 +63,7 @@ import MostlyPrivate2_Private
 
   @objc let mp2priv: MP2PrivateType = 0
 }
+
+@objc public class TestSubclass: NSObject {
+  @_implementationOnly public override func secretMethod() {}
+}

test/Sema/Inputs/implementation-only-override/FooKit.h

@@ -0,0 +1,10 @@
+@interface Base
+@end
+
+@interface Parent : Base
+- (nonnull instancetype)init __attribute__((objc_designated_initializer));
+@end
+
+@interface GenericParent<T: Base *> : Base
+- (nonnull instancetype)init __attribute__((objc_designated_initializer));
+@end

test/Sema/Inputs/implementation-only-override/FooKit_SECRET.h

@@ -0,0 +1,29 @@
+#import <FooKit.h>
+
+@interface SECRETType : Base
+@end
+
+@interface Parent ()
+- (nonnull instancetype)initWithSECRET:(int)secret __attribute__((objc_designated_initializer, swift_name("init(SECRET:)")));
+
+- (void)methodSECRET;
+- (nullable SECRETType *)methodWithSECRETType;
+
+@property (readonly, strong, nullable) Parent *roPropSECRET;
+@property (readwrite, strong, nullable) Parent *rwPropSECRET;
+
+- (nullable Parent *)objectAtIndexedSubscript:(int)index;
+@end
+
+@protocol MandatorySecrets
+- (nonnull instancetype)initWithRequiredSECRET:(int)secret;
+@end
+
+@interface Parent () <MandatorySecrets>
+- (nonnull instancetype)initWithRequiredSECRET:(int)secret __attribute__((objc_designated_initializer));
+@end
+
+@interface GenericParent<T: Base *> ()
+@property (readonly, strong, nullable) T roPropSECRET;
+- (nullable Parent *)objectAtIndexedSubscript:(int)index;
+@end

test/Sema/Inputs/implementation-only-override/module.modulemap

@@ -0,0 +1,9 @@
+module FooKit {
+  header "FooKit.h"
+  export *
+}
+
+module FooKit_SECRET {
+  header "FooKit_SECRET.h"
+  export *
+}

test/Sema/implementation-only-override.swift

@@ -0,0 +1,133 @@
+// RUN: %empty-directory(%t)
+// RUN: %target-typecheck-verify-swift -I %S/Inputs/implementation-only-override -DERRORS -enable-library-evolution -enable-objc-interop
+
+// RUN: %target-swift-frontend -typecheck -emit-module-interface-path %t/Library.swiftinterface -I %S/Inputs/implementation-only-override -enable-library-evolution -enable-objc-interop %s
+// RUN: %FileCheck %s < %t/Library.swiftinterface
+// RUN: %FileCheck -check-prefix=NEGATIVE %s < %t/Library.swiftinterface
+
+// CHECK: import FooKit
+// NEGATIVE-NOT: SECRET
+// NEGATIVE-NOT: subscript
+
+import FooKit
+@_implementationOnly import FooKit_SECRET
+
+// CHECK-LABEL: class GoodChild : FooKit.Parent {
+//  CHECK-NEXT:   @objc override dynamic public init()
+//  CHECK-NEXT:   @objc deinit
+//  CHECK-NEXT: }
+public class GoodChild: Parent {
+  public override init() {}
+  // FIXME: @_implementationOnly on an initializer doesn't exactly make sense,
+  // since they're not inherited.
+  @_implementationOnly public override init(SECRET: Int32) {}
+  @_implementationOnly public required init(requiredSECRET: Int32) {}
+
+  @_implementationOnly public override func methodSECRET() {} // expected-note {{overridden declaration is here}}
+  @_implementationOnly public override func methodWithSECRETType() -> SECRETType? { nil } // expected-note {{overridden declaration is here}}
+  @_implementationOnly public override var roPropSECRET: Parent? { nil } // expected-note {{overridden declaration is here}}
+  @_implementationOnly public override var rwPropSECRET: Parent? { // expected-note {{overridden declaration is here}}
+    get { nil }
+    set {}
+  }
+  @_implementationOnly public override subscript(_ index: Int32) -> Parent? { nil } // expected-note {{overridden declaration is here}}
+}
+
+// CHECK-LABEL: class QuietChild : FooKit.Parent {
+//  CHECK-NEXT:   @objc deinit
+//  CHECK-NEXT: }
+public class QuietChild: Parent {
+  internal override init() {}
+  internal override init(SECRET: Int32) {}
+  internal required init(requiredSECRET: Int32) {}
+}
+
+// CHECK-LABEL: class GoodGenericChild<Toy> : FooKit.Parent {
+//  CHECK-NEXT:   @objc override dynamic public init()
+//  CHECK-NEXT:   @objc deinit
+//  CHECK-NEXT: }
+public class GoodGenericChild<Toy>: Parent {
+  public override init() {}
+  // FIXME: @_implementationOnly on an initializer doesn't exactly make sense,
+  // since they're not inherited.
+  @_implementationOnly public override init(SECRET: Int32) {}
+  @_implementationOnly public required init(requiredSECRET: Int32) {}
+
+  @_implementationOnly public override func methodSECRET() {}
+  @_implementationOnly public override func methodWithSECRETType() -> SECRETType? { nil }
+  @_implementationOnly public override var roPropSECRET: Parent? { nil }
+  @_implementationOnly public override var rwPropSECRET: Parent? {
+    get { nil }
+    set {}
+  }
+  @_implementationOnly public override subscript(_ index: Int32) -> Parent? { nil }
+  @_implementationOnly public override var redefinedPropSECRET: Parent? {
+    get { nil }
+    set {}
+  }
+}
+
+internal class PrivateChild: Parent {
+  override func methodSECRET() {}
+  override func methodWithSECRETType() -> SECRETType? { nil }
+  override var roPropSECRET: Parent? { nil }
+  override var rwPropSECRET: Parent? {
+    get { nil }
+    set {}
+  }
+  override subscript(_ index: Int32) -> Parent? { nil }
+}
+
+internal class PrivateGrandchild: GoodChild {
+  override func methodSECRET() {}
+  override func methodWithSECRETType() -> SECRETType? { nil }
+  override var roPropSECRET: Parent? { nil }
+  override var rwPropSECRET: Parent? {
+    get { nil }
+    set {}
+  }
+  override subscript(_ index: Int32) -> Parent? { nil }
+}
+
+#if ERRORS
+
+public class NaughtyChild: Parent {
+  @_implementationOnly public func nonOverridingMethod() {} // expected-error {{'@_implementationOnly' can only be used on overrides}} {{3-24=}}
+  @_implementationOnly public var nonOverridingProperty: Int { 0 } // expected-error {{'@_implementationOnly' can only be used on overrides}} {{3-24=}}
+
+  public override init() {} // okay
+  public override init(SECRET: Int32) {} // expected-error {{override of initializer imported as implementation-only must be declared '@_implementationOnly'}} {{3-3=@_implementationOnly }}
+  public required init(requiredSECRET: Int32) {} // expected-error {{override of initializer imported as implementation-only must be declared '@_implementationOnly'}} {{3-3=@_implementationOnly }}
+
+  public override func methodSECRET() {} // expected-error {{override of instance method imported as implementation-only must be declared '@_implementationOnly'}} {{3-3=@_implementationOnly }}
+  public override func methodWithSECRETType() -> SECRETType? { nil } // expected-error {{override of instance method imported as implementation-only must be declared '@_implementationOnly'}} {{3-3=@_implementationOnly }} expected-error {{cannot use class 'SECRETType' here; 'FooKit_SECRET' has been imported as implementation-only}} {{none}}
+  public override var roPropSECRET: Parent? { nil } // expected-error {{override of property imported as implementation-only must be declared '@_implementationOnly'}} {{3-3=@_implementationOnly }}
+  public override var rwPropSECRET: Parent? { // expected-error {{override of property imported as implementation-only must be declared '@_implementationOnly'}} {{3-3=@_implementationOnly }}
+    get { nil }
+    set {}
+  }
+  public override subscript(_ index: Int32) -> Parent? { nil } // expected-error {{override of subscript imported as implementation-only must be declared '@_implementationOnly'}} {{3-3=@_implementationOnly }}
+}
+
+public class NaughtyChildByType: Parent {
+  @_implementationOnly public override var roPropSECRET: NaughtyChildByType? { nil } // expected-error {{'@_implementationOnly' override must have the same type as the declaration it overrides ('Parent?')}} {{none}}
+  @_implementationOnly public override subscript(_ index: Int32) -> NaughtyChildByType? { nil } // expected-error {{'@_implementationOnly' override must have the same type as the declaration it overrides ('(Int32) -> Parent?')}} {{none}}
+}
+
+public class NaughtyConcreteChildByType: GenericParent<Parent> {
+  @_implementationOnly public override var roPropSECRET: NaughtyChildByType? { nil } // expected-error {{'@_implementationOnly' override must have the same type as the declaration it overrides ('Parent?')}} {{none}}
+  @_implementationOnly public override subscript(_ index: Int32) -> NaughtyChildByType? { nil } // expected-error {{'@_implementationOnly' override must have the same type as the declaration it overrides ('(Int32) -> Parent?')}} {{none}}
+}
+
+public class NaughtyGrandchild: GoodChild {
+  public override func methodSECRET() {} // expected-error {{override of '@_implementationOnly' instance method should also be declared '@_implementationOnly'}} {{3-3=@_implementationOnly }}
+  public override func methodWithSECRETType() -> SECRETType? { nil } // expected-error {{override of '@_implementationOnly' instance method should also be declared '@_implementationOnly'}} {{3-3=@_implementationOnly }} expected-error {{cannot use class 'SECRETType' here; 'FooKit_SECRET' has been imported as implementation-only}} {{none}}
+  public override var roPropSECRET: Parent? { nil } // expected-error {{override of '@_implementationOnly' property should also be declared '@_implementationOnly'}} {{3-3=@_implementationOnly }}
+  public override var rwPropSECRET: Parent? { // expected-error {{override of '@_implementationOnly' property should also be declared '@_implementationOnly'}} {{3-3=@_implementationOnly }}
+    get { nil }
+    set {}
+  }
+  public override subscript(_ index: Int32) -> Parent? { nil } // expected-error {{override of '@_implementationOnly' subscript should also be declared '@_implementationOnly'}} {{3-3=@_implementationOnly }}
+}
+
+#endif // ERRORS