Remember only one @_objcImpl implementation

Simpliies the implementation dramatically.

Author: beccadax

include/swift/AST/Decl.h

@@ -724,11 +724,11 @@ class alignas(1 << DeclAlignInBits) Decl : public ASTAllocated<Decl> {
   llvm::PointerUnion<DeclContext *, ASTContext *> Context;
 
   /// The imported Clang declaration representing the \c @_objcInterface for
-  /// this declaration, or \c nullptr if there is none.
+  /// this declaration (or vice versa), or \c nullptr if there is none.
   ///
   /// If \c this (an otherwise nonsensical value), the value has not yet been
   /// computed.
-  Decl *CachedObjCInterfaceDecl;
+  Decl *CachedObjCImplementationDecl;
 
   Decl(const Decl&) = delete;
   void operator=(const Decl&) = delete;
@@ -747,7 +747,7 @@ class alignas(1 << DeclAlignInBits) Decl : public ASTAllocated<Decl> {
 protected:
 
   Decl(DeclKind kind, llvm::PointerUnion<DeclContext *, ASTContext *> context)
-    : Context(context), CachedObjCInterfaceDecl(this) {
+    : Context(context), CachedObjCImplementationDecl(this) {
     Bits.OpaqueBits = 0;
     Bits.Decl.Kind = unsigned(kind);
     Bits.Decl.Invalid = false;
@@ -986,34 +986,26 @@ class alignas(1 << DeclAlignInBits) Decl : public ASTAllocated<Decl> {
   /// returns the imported declaration. Otherwise return \c nullptr.
   ///
   /// \seeAlso ExtensionDecl::isObjCInterface()
-  ///
-  /// \note Currently this only links an @_objcImplementation extension to its
-  ///       class or category; in the future it should also work for members.
   Decl *getObjCInterfaceDecl() const;
 
   /// If this is the ObjC interface of a declaration implemented in Swift,
-  /// returns the implementating declarations. Otherwise return \c nullptr.
-  ///
-  /// In valid code, there should only be one implementation per interface, but
-  /// there can be several in invalid code.
+  /// returns the implementating declaration. Otherwise return \c nullptr.
   ///
   /// \seeAlso ExtensionDecl::isObjCInterface()
-  ///
-  /// \note Currently this only links an imported class or category to its
-  ///       extension; in the future it should also work for members.
-  ArrayRef<Decl *> getObjCImplementationDecls() const;
+  Decl *getObjCImplementationDecl() const;
 
-  Optional<Decl *> getCachedObjCInterfaceDecl() const {
-    if (CachedObjCInterfaceDecl == this)
+  Optional<Decl *> getCachedObjCImplementationDecl() const {
+    if (CachedObjCImplementationDecl == this)
       return None;
-    return CachedObjCInterfaceDecl;
+    return CachedObjCImplementationDecl;
   }
 
-  void setCachedObjCInterfaceDecl(Decl *decl) {
-    assert((CachedObjCInterfaceDecl == this)
+  void setCachedObjCImplementationDecl(Decl *decl) {
+    assert((CachedObjCImplementationDecl == this
+              || CachedObjCImplementationDecl == decl)
                && "can't change CachedObjCInterfaceDecl once it's computed");
     assert(decl != this && "can't form circular reference");
-    CachedObjCInterfaceDecl = decl;
+    CachedObjCImplementationDecl = decl;
   }
 
   /// Return the GenericContext if the Decl has one.

include/swift/ClangImporter/ClangImporterRequests.h

@@ -234,32 +234,32 @@ class ClangCategoryLookupRequest
 /// using \c \@_objcImplementation.
 struct ObjCInterfaceAndImplementation final {
   Decl *interfaceDecl;
-  llvm::ArrayRef<Decl *> implementationDecls;
+  Decl *implementationDecl;
 
   ObjCInterfaceAndImplementation(Decl *interfaceDecl,
-                                 llvm::ArrayRef<Decl *> implementationDecls)
-      : interfaceDecl(interfaceDecl), implementationDecls(implementationDecls)
+                                 Decl *implementationDecl)
+      : interfaceDecl(interfaceDecl), implementationDecl(implementationDecl)
   {
     assert(interfaceDecl && "interface is non-null");
-    assert(!implementationDecls.empty() && "implementation is non-empty");
+    assert(implementationDecl && "implementation is non-empty");
   }
 
   ObjCInterfaceAndImplementation()
-      : interfaceDecl(nullptr), implementationDecls(nullptr, nullptr) {}
+      : interfaceDecl(nullptr), implementationDecl(nullptr) {}
 
   operator bool() const {
     return interfaceDecl;
   }
 
   friend llvm::hash_code
   hash_value(const ObjCInterfaceAndImplementation &pair) {
-    return llvm::hash_combine(pair.interfaceDecl, pair.implementationDecls);
+    return llvm::hash_combine(pair.interfaceDecl, pair.implementationDecl);
   }
 
   friend bool operator==(const ObjCInterfaceAndImplementation &lhs,
                          const ObjCInterfaceAndImplementation &rhs) {
     return lhs.interfaceDecl == rhs.interfaceDecl
-               && lhs.implementationDecls == rhs.implementationDecls;
+               && lhs.implementationDecl == rhs.implementationDecl;
   }
 
   friend bool operator!=(const ObjCInterfaceAndImplementation &lhs,

include/swift/IRGen/Linking.h

@@ -1560,7 +1560,7 @@ class LinkEntity {
 private:
   static bool isObjCImplementation(NominalTypeDecl *NTD) {
     if (NTD)
-      return !NTD->getObjCImplementationDecls().empty();
+      return NTD->getObjCImplementationDecl();
     return false;
   }
   static bool isObjCImplementation(CanType ty) {

lib/ClangImporter/ClangImporter.cpp

@@ -5003,9 +5003,8 @@ TinyPtrVector<ValueDecl *> ClangRecordMemberLookup::evaluate(
 }
 
 IterableDeclContext *IterableDeclContext::getImplementationContext() {
-  auto implDecls = getDecl()->getObjCImplementationDecls();
-  if (!implDecls.empty())
-    if (auto implExt = dyn_cast<ExtensionDecl>(implDecls.front()))
+  if (auto implDecl = getDecl()->getObjCImplementationDecl())
+    if (auto implExt = dyn_cast<ExtensionDecl>(implDecl))
       return implExt;
 
   return this;
@@ -5089,19 +5088,11 @@ findInterfaceGivenImpl(ClassDecl *classDecl, ExtensionDecl *ext) {
 }
 
 static ObjCInterfaceAndImplementation
-cacheImplsAndConstructResult(Decl *interface,
-                             llvm::TinyPtrVector<Decl *> impls) {
+constructResult(Decl *interface, llvm::TinyPtrVector<Decl *> impls) {
   if (impls.empty())
     return ObjCInterfaceAndImplementation();
 
-  // Cache the array immediately so that an ArrayRef will have a stable
-  // address for it.
-  auto &importer =
-      ClangImporter::Implementation::get(interface->getASTContext());
-  auto &cachedImpls = importer.ImplementationsForObjCInterfaces[interface];
-
-  cachedImpls = std::make_unique<llvm::TinyPtrVector<Decl *>>(impls);
-  return ObjCInterfaceAndImplementation(interface, *cachedImpls);
+  return ObjCInterfaceAndImplementation(interface, impls.front());
 }
 
 static ObjCInterfaceAndImplementation
@@ -5149,7 +5140,7 @@ findContextInterfaceAndImplementation(DeclContext *dc) {
   // look for extensions implementing it.
 
   auto implDecls = findImplsGivenInterface(classDecl, categoryName);
-  return cacheImplsAndConstructResult(interfaceDecl, implDecls);
+  return constructResult(interfaceDecl, implDecls);
 }
 
 ObjCInterfaceAndImplementation ObjCInterfaceAndImplementationRequest::
@@ -5187,15 +5178,14 @@ void swift::simple_display(llvm::raw_ostream &out,
   out << "clang interface ";
   simple_display(out, pair.interfaceDecl);
   out << " with @_objcImplementation ";
-  simple_display(out, pair.implementationDecls);
-  out << "Looking up @interface for ";
+  simple_display(out, pair.implementationDecl);
 }
 
 SourceLoc
 swift::extractNearestSourceLoc(const ObjCInterfaceAndImplementation &pair) {
-  if (pair.implementationDecls.empty())
+  if (pair.implementationDecl)
     return SourceLoc();
-  return extractNearestSourceLoc(pair.implementationDecls.front());
+  return extractNearestSourceLoc(pair.implementationDecl);
 }
 
 Decl *Decl::getObjCInterfaceDecl() const {
@@ -5208,14 +5198,14 @@ Decl *Decl::getObjCInterfaceDecl() const {
              .interfaceDecl;
 }
 
-ArrayRef<Decl *> Decl::getObjCImplementationDecls() const {
+Decl *Decl::getObjCImplementationDecl() const {
   if (!hasClangNode())
     // This *is* the implementation, if it has one.
     return {};
 
   ObjCInterfaceAndImplementationRequest req{const_cast<Decl *>(this)};
   return evaluateOrDefault(getASTContext().evaluator, req, {})
-             .implementationDecls;
+             .implementationDecl;
 }
 
 IterableDeclContext *ClangCategoryLookupRequest::

lib/ClangImporter/ClangImporterRequests.cpp

@@ -26,78 +26,52 @@ using namespace swift;
 Optional<ObjCInterfaceAndImplementation>
 ObjCInterfaceAndImplementationRequest::getCachedResult() const {
   auto passedDecl = std::get<0>(getStorage());
-  Decl *interface;
+  if (!passedDecl)
+    return {};
 
-  // An imported decl can only be an interface; a native decl can only be an
-  // implementation.
-  if (!passedDecl || passedDecl->hasClangNode()) {
-    interface = passedDecl;
-  } else {
-    if (auto cached = passedDecl->getCachedObjCInterfaceDecl())
-      interface = *cached;
-    else
-      return None;
-  }
+  auto cachedDecl = passedDecl->getCachedObjCImplementationDecl();
 
-  // They passed in either null, or a Swift decl we've previously computed is
-  // *not* an @_objcImplementation.
-  if (!interface)
-    return ObjCInterfaceAndImplementation();
+  // !cachedDecl means that no decl has been cached and we need to evaluate the
+  // request.
+  if (!cachedDecl)
+    return None;
 
-  // Okay, now find the implementations.
-  auto &importer =
-      ClangImporter::Implementation::get(interface->getASTContext());
-  auto impls = importer.ImplementationsForObjCInterfaces.find(interface);
+  // nullptr cachedDecl means that the lack of a decl was cached.
+  else if (!*cachedDecl)
+    return {};
 
-  // Uncached.
-  if (impls == importer.ImplementationsForObjCInterfaces.end())
-    return None;
+  // A decl was cached! Arbitrarily guess that we looked up the implementation
+  // from the interface.
+  ObjCInterfaceAndImplementation result{passedDecl, *cachedDecl};
 
-  // They passed in a Clang decl that we've previously computed has no impls.
-  if (!impls->second) {
-    assert(interface == passedDecl
-           && "cached empty result doesn't contain the decl we looked up from???");
-    return ObjCInterfaceAndImplementation();
-  }
+  // An imported decl can only be an interface; a native decl can only be an
+  // implementation. If `implementationDecl` has a Clang node, we must have
+  // looked up the interface from the implementation.
+  if (result.implementationDecl->hasClangNode())
+    std::swap(result.interfaceDecl, result.implementationDecl);
 
-  assert(interface == passedDecl || llvm::is_contained(*impls->second, passedDecl)
-         && "cached result doesn't contain the decl we looked up from???");
-  return ObjCInterfaceAndImplementation(interface, *impls->second);
+  return result;
 }
 
 void ObjCInterfaceAndImplementationRequest::
 cacheResult(ObjCInterfaceAndImplementation value) const {
   Decl *decl = std::get<0>(getStorage());
-  auto &importer =
-      ClangImporter::Implementation::get(decl->getASTContext());
 
   if (!value) {
     // `decl` is neither an interface nor an implementation.
-    decl->setCachedObjCInterfaceDecl(nullptr);
-    if (decl->hasClangNode())
-      importer.ImplementationsForObjCInterfaces[decl] = nullptr;
-
+    decl->setCachedObjCImplementationDecl(nullptr);
     return;
   }
 
-  assert(value.interfaceDecl == decl ||
-         llvm::is_contained(value.implementationDecls, decl)
-             && "ought to return the input as one of the results!");
-
-  // `evaluate()` already updates `importer.ImplementationsForObjCInterfaces`
-  // so that it can return a persistent `ArrayRef` in the `value`, so we just
-  // check that the array was cached correctly for `getCachedResult()` to find.
-  auto &cachedImpls =
-      importer.ImplementationsForObjCInterfaces[value.interfaceDecl];
-  assert(cachedImpls
-         && cachedImpls->begin() == value.implementationDecls.begin()
-         && cachedImpls->end() == value.implementationDecls.end() &&
-         "did not cache expected result?");
-
   // Cache computed pointers from implementations to interfaces.
-  value.interfaceDecl->setCachedObjCInterfaceDecl(nullptr);
-  for (auto implDecl : value.implementationDecls)
-    implDecl->setCachedObjCInterfaceDecl(value.interfaceDecl);
+  value.interfaceDecl->
+      setCachedObjCImplementationDecl(value.implementationDecl);
+  value.implementationDecl->
+      setCachedObjCImplementationDecl(value.interfaceDecl);
+
+  // If this was a duplicate implementation, cache a null so we don't recompute.
+  if (decl != value.interfaceDecl && decl != value.implementationDecl)
+    decl->setCachedObjCImplementationDecl(nullptr);
 }
 
 // Define request evaluation functions for each of the name lookup requests.

lib/ClangImporter/ImporterImpl.h

@@ -607,15 +607,6 @@ class LLVM_LIBRARY_VISIBILITY ClangImporter::Implementation
   }
 
 public:
-  /// Mapping from imported clang decls to native Swift implementations via
-  /// \c \@_objcImplementation extensions. In valid code, there should be only
-  /// one implementation per interface.
-  ///
-  /// \c std::unique_ptr is used so that we can form a stable \c ArrayRef for
-  /// each value vector.
-  llvm::DenseMap<Decl *, std::unique_ptr<llvm::TinyPtrVector<Decl *>>>
-    ImplementationsForObjCInterfaces;
-
   /// Keep track of subscript declarations based on getter/setter
   /// pairs.
   llvm::DenseMap<std::pair<FuncDecl *, FuncDecl *>, SubscriptDecl *> Subscripts;

lib/IRGen/GenClass.cpp

@@ -161,7 +161,7 @@ namespace {
       auto theClass = classType.getClassOrBoundGenericClass();
       assert(theClass);
 
-      if (!theClass->getObjCImplementationDecls().empty())
+      if (theClass->getObjCImplementationDecl())
         Options |= ClassMetadataFlags::ClassHasObjCImplementation;
 
       if (theClass->isGenericContext() && !theClass->hasClangNode())
@@ -1342,7 +1342,7 @@ namespace {
       auto *theClass = getClass();
 
       if (theClass->hasClangNode() &&
-          theClass->getObjCImplementationDecls().empty())
+          !theClass->getObjCImplementationDecl())
         return IGM.getAddrOfObjCClass(theClass, NotForDefinition);
 
       // Note that getClassMetadataStrategy() will return

lib/IRGen/GenDecl.cpp

@@ -4716,10 +4716,11 @@ llvm::GlobalValue *IRGenModule::defineTypeMetadata(
     SmallVector<std::pair<Size, SILDeclRef>, 8> vtableEntries) {
   assert(init);
 
-  auto isPrespecialized = concreteType->getAnyGeneric() &&
-                          concreteType->getAnyGeneric()->isGenericContext();
-  auto isObjCImpl =
-      !concreteType->getAnyGeneric()->getObjCImplementationDecls().empty();
+  auto concreteTypeDecl = concreteType->getAnyGeneric();
+  auto isPrespecialized = concreteTypeDecl &&
+                          concreteTypeDecl->isGenericContext();
+  bool isObjCImpl = concreteTypeDecl &&
+                    concreteTypeDecl->getObjCImplementationDecl();
 
   if (isPattern) {
     assert(isConstant && "Type metadata patterns must be constant");
@@ -4834,7 +4835,7 @@ IRGenModule::getAddrOfTypeMetadata(CanType concreteType,
   } else if (nominal) {
     // The symbol for native non-generic nominal type metadata is generated at
     // the aliased address point (see defineTypeMetadata() above).
-    if (!nominal->getObjCImplementationDecls().empty()) {
+    if (nominal->getObjCImplementationDecl()) {
       defaultVarTy = ObjCClassStructTy;
     } else {
       assert(!nominal->hasClangNode());
@@ -4871,7 +4872,7 @@ IRGenModule::getAddrOfTypeMetadata(CanType concreteType,
   switch (canonicality) {
   case TypeMetadataCanonicality::Canonical: {
     auto classDecl = concreteType->getClassOrBoundGenericClass();
-    if (classDecl && !classDecl->getObjCImplementationDecls().empty()) {
+    if (classDecl && classDecl->getObjCImplementationDecl()) {
       entity = LinkEntity::forObjCClass(classDecl);
     } else {
       entity = LinkEntity::forTypeMetadata(

lib/IRGen/GenMeta.cpp

@@ -4384,7 +4384,7 @@ namespace {
 static void emitObjCClassSymbol(IRGenModule &IGM,
                                 ClassDecl *classDecl,
                                 llvm::Constant *metadata) {
-  if (!classDecl->getObjCImplementationDecls().empty())
+  if (classDecl->getObjCImplementationDecl())
     // Should already have this symbol.
     return;
 

lib/IRGen/GenReflection.cpp

@@ -1560,7 +1560,7 @@ void IRGenModule::emitFieldDescriptor(const NominalTypeDecl *D) {
   }
 
   if (auto *CD = dyn_cast<ClassDecl>(D)) {
-    if (!CD->getObjCImplementationDecls().empty())
+    if (CD->getObjCImplementationDecl())
       needsFieldDescriptor = false;
   }
 

lib/SIL/IR/SIL.cpp

@@ -36,7 +36,7 @@ FormalLinkage swift::getDeclLinkage(const ValueDecl *D) {
   // Clang declarations are public and can't be assured of having a
   // unique defining location.
   if (isa<ClangModuleUnit>(fileContext) &&
-          D->getObjCImplementationDecls().empty())
+          !D->getObjCImplementationDecl())
     return FormalLinkage::PublicNonUnique;
 
   switch (D->getEffectiveAccess()) {