Allow stored properties in @_objcImpls

beccadax · beccadax · commit 3a2f12a0998d · 2022-10-18T17:21:56.000-07:00
Stored properties are only allowed in the extension implementing the class's main interface, not its categories. This also means banning `@objc final`, which is unenforceable anyway when ObjC subclasses are allowed, and therefore allowing `@objc let` and `@objc static` properties to be overridden if they're declared in objcImplementations.
diff --git a/include/swift/AST/Decl.h b/include/swift/AST/Decl.h
@@ -1501,6 +1501,11 @@ class ExtensionDecl final : public GenericContext, public Decl,
   /// behaviors for its members.
   bool isObjCImplementation() const;
 
+  /// Returns the name of the category specified by the \c \@_objcImplementation
+  /// attribute, or \c None if the name is invalid. Do not call unless
+  /// \c isObjCImplementation() returns \c true.
+  Optional<Identifier> getCategoryNameForObjCImplementation() const;
+
   /// Returns the \c clang::ObjCCategoryDecl or \c clang::ObjCInterfaceDecl
   /// implemented by this extension.
   ///
diff --git a/include/swift/AST/DeclContext.h b/include/swift/AST/DeclContext.h
@@ -824,6 +824,11 @@ class IterableDeclContext {
   /// abstractions on top of member loading, such as a name lookup table.
   DeclRange getCurrentMembersWithoutLoading() const;
 
+  /// Return the context that contains the actual implemented members. This
+  /// is \em usually just \c this, but if \c this is an imported class or
+  /// category, it may be a Swift extension instead.
+  IterableDeclContext *getImplementationContext();
+
   /// Add a member to this context.
   ///
   /// If the hint decl is specified, the new decl is inserted immediately
diff --git a/include/swift/AST/DiagnosticsSema.def b/include/swift/AST/DiagnosticsSema.def
@@ -1511,6 +1511,10 @@ ERROR(attr_objc_implementation_category_not_found,none,
 NOTE(attr_objc_implementation_fixit_remove_category_name,none,
      "remove arguments to implement the main '@interface' for this class",
      ())
+ERROR(attr_objc_implementation_no_objc_final,none,
+      "%0 %1 cannot be 'final' because Objective-C subclasses of %2 can "
+      "override it",
+      (DescriptiveDeclKind, ValueDecl *, ValueDecl *))
 
 ERROR(member_of_objc_implementation_not_objc_or_final,none,
       "%0 %1 does not match any %0 declared in the headers for %2; did you use "
diff --git a/lib/AST/Decl.cpp b/lib/AST/Decl.cpp
@@ -1471,17 +1471,15 @@ bool ExtensionDecl::isObjCImplementation() const {
 
 const clang::ObjCContainerDecl *
 ExtensionDecl::getInterfaceForObjCImplementation() const {
-  assert(isObjCImplementation());
-
-  auto attr = getAttrs().getAttribute<ObjCImplementationAttr>();
-  if (attr->isCategoryNameInvalid())
+  auto categoryName = getCategoryNameForObjCImplementation();
+  if (!categoryName)
     return nullptr;
 
   auto CD = dyn_cast_or_null<ClassDecl>(getExtendedNominal());
   if (!CD)
     return nullptr;
 
-  auto importedDecl = CD->getImportedObjCCategory(attr->CategoryName);
+  auto importedDecl = CD->getImportedObjCCategory(*categoryName);
   if (!importedDecl)
     return nullptr;
 
@@ -1490,6 +1488,17 @@ ExtensionDecl::getInterfaceForObjCImplementation() const {
   return cast<clang::ObjCContainerDecl>(importedDecl->getClangDecl());
 }
 
+Optional<Identifier>
+ExtensionDecl::getCategoryNameForObjCImplementation() const {
+  assert(isObjCImplementation());
+
+  auto attr = getAttrs().getAttribute<ObjCImplementationAttr>();
+  if (attr->isCategoryNameInvalid())
+    return None;
+
+  return attr->CategoryName;
+}
+
 PatternBindingDecl::PatternBindingDecl(SourceLoc StaticLoc,
                                        StaticSpellingKind StaticSpelling,
                                        SourceLoc VarLoc,
diff --git a/lib/ClangImporter/ClangImporter.cpp b/lib/ClangImporter/ClangImporter.cpp
@@ -5032,6 +5032,31 @@ TinyPtrVector<ValueDecl *> ClangRecordMemberLookup::evaluate(
   return result;
 }
 
+IterableDeclContext *IterableDeclContext::getImplementationContext() {
+  // Only ClangNodes have @_objcImplementation extensions.
+  if (!getDecl()->hasClangNode())
+    return this;
+
+  ClassDecl *classDecl = getAsGenericContext()->getSelfClassDecl();
+  if (!classDecl)
+    return this;
+
+  // Figure out the category name of the extension that will match this decl.
+  Identifier categoryName{};
+  if (getIterableContextKind() == IterableDeclContextKind::ExtensionDecl) {
+    auto category = cast<clang::ObjCCategoryDecl>(getDecl()->getClangDecl());
+    categoryName = getASTContext().getIdentifier(category->getName());
+  }
+
+  for (ExtensionDecl *ext : classDecl->getExtensions()) {
+    if (ext->isObjCImplementation()
+        && ext->getCategoryNameForObjCImplementation() == categoryName)
+      return ext;
+  }
+
+  return this;
+}
+
 Decl *ClangCategoryLookupRequest::
 evaluate(Evaluator &evaluator, ClangCategoryLookupDescriptor desc) const {
   const ClassDecl *CD = desc.classDecl;
diff --git a/lib/Sema/TypeCheckDecl.cpp b/lib/Sema/TypeCheckDecl.cpp
@@ -288,15 +288,16 @@ enum class ImplicitlyFinalReason : unsigned {
 }
 
 static bool inferFinalAndDiagnoseIfNeeded(ValueDecl *D, ClassDecl *cls,
+                                          FinalAttr *explicitFinalAttr,
                                           StaticSpellingKind staticSpelling) {
   // Are there any reasons to infer 'final'? Prefer 'static' over the class
   // being final for the purposes of diagnostics.
   Optional<ImplicitlyFinalReason> reason;
   if (staticSpelling == StaticSpellingKind::KeywordStatic) {
     reason = ImplicitlyFinalReason::Static;
 
-    if (auto finalAttr = D->getAttrs().getAttribute<FinalAttr>()) {
-      auto finalRange = finalAttr->getRange();
+    if (explicitFinalAttr) {
+      auto finalRange = explicitFinalAttr->getRange();
       if (finalRange.isValid()) {
         auto &context = D->getASTContext();
         context.Diags.diagnose(finalRange.Start, diag::static_decl_already_final)
@@ -802,13 +803,30 @@ PrimaryAssociatedTypesRequest::evaluate(Evaluator &evaluator,
 
 bool
 IsFinalRequest::evaluate(Evaluator &evaluator, ValueDecl *decl) const {
+  auto explicitFinalAttr = decl->getAttrs().getAttribute<FinalAttr>();
   if (isa<ClassDecl>(decl))
-    return decl->getAttrs().hasAttribute<FinalAttr>();
+    return explicitFinalAttr;
 
   auto cls = decl->getDeclContext()->getSelfClassDecl();
   if (!cls)
     return false;
 
+  // Objective-C doesn't have a way to prevent overriding, so if this member
+  // is accessible from ObjC and it's possible to subclass its parent from ObjC,
+  // `final` doesn't make sense even when inferred.
+  //
+  // FIXME: This is technically true iff `!cls->hasKnownSwiftImplementation()`,
+  //        but modeling that would be source-breaking, so instead we only
+  //        enforce it in `@_objcImplementation extension`s.
+  if (auto ext = dyn_cast<ExtensionDecl>(decl->getDeclContext()))
+    if (ext->isObjCImplementation() && decl->isObjC()) {
+      if (explicitFinalAttr)
+        diagnoseAndRemoveAttr(decl, explicitFinalAttr,
+                              diag::attr_objc_implementation_no_objc_final,
+                              decl->getDescriptiveKind(), decl, cls);
+      return false;
+    }
+
   switch (decl->getKind()) {
     case DeclKind::Var: {
       // Properties are final if they are declared 'static' or a 'let'
@@ -831,7 +849,8 @@ IsFinalRequest::evaluate(Evaluator &evaluator, ValueDecl *decl) const {
         // If this variable is a class member, mark it final if the
         // class is final, or if it was declared with 'let'.
         auto *PBD = VD->getParentPatternBinding();
-        if (PBD && inferFinalAndDiagnoseIfNeeded(decl, cls, PBD->getStaticSpelling()))
+        if (PBD && inferFinalAndDiagnoseIfNeeded(decl, cls, explicitFinalAttr,
+                                                 PBD->getStaticSpelling()))
           return true;
 
         if (VD->isLet()) {
@@ -856,7 +875,8 @@ IsFinalRequest::evaluate(Evaluator &evaluator, ValueDecl *decl) const {
     case DeclKind::Func: {
       // Methods declared 'static' are final.
       auto staticSpelling = cast<FuncDecl>(decl)->getStaticSpelling();
-      if (inferFinalAndDiagnoseIfNeeded(decl, cls, staticSpelling))
+      if (inferFinalAndDiagnoseIfNeeded(decl, cls, explicitFinalAttr,
+                                        staticSpelling))
         return true;
       break;
     }
@@ -889,7 +909,8 @@ IsFinalRequest::evaluate(Evaluator &evaluator, ValueDecl *decl) const {
     case DeclKind::Subscript: {
       // Member subscripts.
       auto staticSpelling = cast<SubscriptDecl>(decl)->getStaticSpelling();
-      if (inferFinalAndDiagnoseIfNeeded(decl, cls, staticSpelling))
+      if (inferFinalAndDiagnoseIfNeeded(decl, cls, explicitFinalAttr,
+                                        staticSpelling))
         return true;
       break;
     }
@@ -898,10 +919,7 @@ IsFinalRequest::evaluate(Evaluator &evaluator, ValueDecl *decl) const {
       break;
   }
 
-  if (decl->getAttrs().hasAttribute<FinalAttr>())
-    return true;
-
-  return false;
+  return explicitFinalAttr;
 }
 
 bool IsMoveOnlyRequest::evaluate(Evaluator &evaluator, ValueDecl *decl) const {
diff --git a/lib/Sema/TypeCheckDeclPrimary.cpp b/lib/Sema/TypeCheckDeclPrimary.cpp
@@ -1217,10 +1217,24 @@ static void checkDynamicSelfType(ValueDecl *decl, Type type) {
 /// extension, it is either `final` or `@objc` (which may have been inferred by
 /// checking whether it shadows an imported declaration).
 static void checkObjCImplementationMemberAvoidsVTable(ValueDecl *VD) {
+  // We check the properties instead of their accessors.
+  if (isa<AccessorDecl>(VD))
+    return;
+
+  // Are we in an @_objcImplementation extension?
   auto ED = dyn_cast<ExtensionDecl>(VD->getDeclContext());
-  if (!ED || !ED->isObjCImplementation()) return;
+  if (!ED || !ED->isObjCImplementation())
+    return;
 
-  if (VD->isSemanticallyFinal() || VD->isObjC()) return;
+  assert(ED->getSelfClassDecl() &&
+         !ED->getSelfClassDecl()->hasKnownSwiftImplementation() &&
+         "@_objcImplementation on non-class or Swift class?");
+
+  if (VD->isSemanticallyFinal() || VD->isObjC()) {
+    assert(!VD->isObjC() || VD->isDynamic() &&
+           "@objc decls in @_objcImplementations should be dynamic!");
+    return;
+  }
 
   auto &diags = VD->getASTContext().Diags;
   diags.diagnose(VD, diag::member_of_objc_implementation_not_objc_or_final,
diff --git a/lib/Sema/TypeCheckStorage.cpp b/lib/Sema/TypeCheckStorage.cpp
@@ -93,24 +93,27 @@ static bool contextAllowsPatternBindingWithoutVariables(DeclContext *dc) {
   return true;
 }
 
-static bool hasStoredProperties(NominalTypeDecl *decl) {
+static bool hasStoredProperties(NominalTypeDecl *decl,
+                                IterableDeclContext *implDecl) {
   bool isForeignReferenceTy =
       isa<ClassDecl>(decl) && cast<ClassDecl>(decl)->isForeignReferenceType();
 
   return (isa<StructDecl>(decl) ||
           (isa<ClassDecl>(decl) &&
-           (!decl->hasClangNode() || isForeignReferenceTy)));
+           (!decl->hasClangNode() || isForeignReferenceTy
+                || (decl != implDecl))));
 }
 
-static void computeLoweredStoredProperties(NominalTypeDecl *decl) {
+static void computeLoweredStoredProperties(NominalTypeDecl *decl,
+                                           IterableDeclContext *implDecl) {
   // If declaration has a type wrapper, make sure that
   // `$_storage` property is synthesized.
   if (decl->hasTypeWrapper())
     (void)decl->getTypeWrapperProperty();
 
   // Just walk over the members of the type, forcing backing storage
-  // for lazy properties, property and type wrappers to be synthesized.
-  for (auto *member : decl->getMembers()) {
+  // for lazy properties and property wrappers to be synthesized.
+  for (auto *member : implDecl->getMembers()) {
     auto *var = dyn_cast<VarDecl>(member);
     if (!var || var->isStatic())
       continue;
@@ -160,6 +163,7 @@ static void computeLoweredStoredProperties(NominalTypeDecl *decl) {
 /// in a deterministic order.
 static void enumerateStoredPropertiesAndMissing(
     NominalTypeDecl *decl,
+    IterableDeclContext *implDecl,
     llvm::function_ref<void(VarDecl *)> addStoredProperty,
     llvm::function_ref<void(MissingMemberDecl *)> addMissing) {
   // If we have a distributed actor, find the id and actorSystem
@@ -169,7 +173,7 @@ static void enumerateStoredPropertiesAndMissing(
   VarDecl *distributedActorSystem = nullptr;
   if (decl->isDistributedActor()) {
     ASTContext &ctx = decl->getASTContext();
-    for (auto *member : decl->getMembers()) {
+    for (auto *member : implDecl->getMembers()) {
       if (auto *var = dyn_cast<VarDecl>(member)) {
         if (!var->isStatic() && var->hasStorage()) {
           if (var->getName() == ctx.Id_id) {
@@ -190,7 +194,7 @@ static void enumerateStoredPropertiesAndMissing(
       addStoredProperty(distributedActorSystem);
   }
 
-  for (auto *member : decl->getMembers()) {
+  for (auto *member : implDecl->getMembers()) {
     if (auto *var = dyn_cast<VarDecl>(member)) {
       if (!var->isStatic() && var->hasStorage()) {
         // Skip any properties that we already emitted explicitly
@@ -209,20 +213,29 @@ static void enumerateStoredPropertiesAndMissing(
   }
 }
 
+static bool isInSourceFile(IterableDeclContext *idc) {
+  const DeclContext *dc = idc->getAsGenericContext();
+  return isa<SourceFile>(dc->getModuleScopeContext());
+}
+
 ArrayRef<VarDecl *>
 StoredPropertiesRequest::evaluate(Evaluator &evaluator,
                                   NominalTypeDecl *decl) const {
-  if (!hasStoredProperties(decl))
+  // If this is an imported class with an @_objcImplementation extension, get
+  // members from the extension instead.
+  IterableDeclContext *implDecl = decl->getImplementationContext();
+
+  if (!hasStoredProperties(decl, implDecl))
     return ArrayRef<VarDecl *>();
 
   SmallVector<VarDecl *, 4> results;
 
   // Unless we're in a source file we don't have to do anything
   // special to lower lazy properties and property wrappers.
-  if (isa<SourceFile>(decl->getModuleScopeContext()))
-    computeLoweredStoredProperties(decl);
+  if (isInSourceFile(implDecl))
+    computeLoweredStoredProperties(decl, implDecl);
 
-  enumerateStoredPropertiesAndMissing(decl,
+  enumerateStoredPropertiesAndMissing(decl, implDecl,
     [&](VarDecl *var) {
       results.push_back(var);
     },
@@ -234,17 +247,21 @@ StoredPropertiesRequest::evaluate(Evaluator &evaluator,
 ArrayRef<Decl *>
 StoredPropertiesAndMissingMembersRequest::evaluate(Evaluator &evaluator,
                                                    NominalTypeDecl *decl) const {
-  if (!hasStoredProperties(decl))
+  // If this is an imported class with an @_objcImplementation extension, get
+  // members from the extension instead.
+  IterableDeclContext *implDecl = decl->getImplementationContext();
+
+  if (!hasStoredProperties(decl, implDecl))
     return ArrayRef<Decl *>();
 
   SmallVector<Decl *, 4> results;
 
   // Unless we're in a source file we don't have to do anything
   // special to lower lazy properties and property wrappers.
-  if (isa<SourceFile>(decl->getModuleScopeContext()))
-    computeLoweredStoredProperties(decl);
+  if (isInSourceFile(implDecl))
+    computeLoweredStoredProperties(decl, implDecl);
 
-  enumerateStoredPropertiesAndMissing(decl,
+  enumerateStoredPropertiesAndMissing(decl, implDecl,
     [&](VarDecl *var) {
       results.push_back(var);
     },
@@ -3295,8 +3312,17 @@ static void finishStorageImplInfo(AbstractStorageDecl *storage,
     if (isa<EnumDecl>(dc)) {
       storage->diagnose(diag::enum_stored_property);
       info = StorageImplInfo::getMutableComputed();
-    } else if (isa<ExtensionDecl>(dc) &&
-              !storage->getAttrs().getAttribute<DynamicReplacementAttr>()) {
+    } else if (auto ext = dyn_cast<ExtensionDecl>(dc)) {
+      // Extensions can dynamically replace a stored property.
+      if (storage->getAttrs().getAttribute<DynamicReplacementAttr>())
+        return;
+
+      // @_objcImplementation extensions on a non-category can declare stored
+      // properties; StoredPropertiesRequest knows to look for them there.
+      if (ext->isObjCImplementation() &&
+          ext->getCategoryNameForObjCImplementation() == Identifier())
+        return;
+
       storage->diagnose(diag::extension_stored_property);
 
       info = (info.supportsMutation()
diff --git a/test/decl/ext/Inputs/objc_implementation.h b/test/decl/ext/Inputs/objc_implementation.h
diff --git a/test/decl/ext/objc_implementation.swift b/test/decl/ext/objc_implementation.swift
