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

Author: jrose-apple

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 @@ class ObjCPrinter : private DeclVisitor<ObjCPrinter>,
   }
 
   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 @@ class ExportabilityChecker : public DeclVisitor<ExportabilityChecker> {
   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 @@ class ExportabilityChecker : public DeclVisitor<ExportabilityChecker> {
     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 @@ class ExportabilityChecker : public DeclVisitor<ExportabilityChecker> {
   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 @@ class AttributeChecker : public AttributeVisitor<AttributeChecker> {
     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 @@ class AttributeChecker : public AttributeVisitor<AttributeChecker> {
   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