ClangImporter: Record the error wrapper's struct 'Code' member enum in a side table

Author: slavapestov

include/swift/ClangImporter/ClangImporter.h

@@ -51,11 +51,13 @@ class ClangModuleUnit;
 class ClangNode;
 class Decl;
 class DeclContext;
+class EnumDecl;
 class ImportDecl;
 class IRGenOptions;
 class LazyResolver;
 class ModuleDecl;
 class NominalTypeDecl;
+class StructDecl;
 class TypeDecl;
 class VisibleDeclConsumer;
 enum class SelectorSplitKind;
@@ -213,6 +215,9 @@ class ClangImporter final : public ClangModuleLoader {
                       StringRef relatedEntityKind,
                       llvm::function_ref<void(TypeDecl *)> receiver) override;
 
+  /// Look up the nested 'Code' enum for an error wrapper struct.
+  EnumDecl *lookupErrorCodeEnum(const StructDecl *errorWrapper) const;
+
   /// Look for textually included declarations from the bridging header.
   ///
   /// \param filter returns true if the given clang decl/macro should be

lib/ClangImporter/ClangImporter.cpp

@@ -2171,21 +2171,10 @@ static bool isVisibleFromModule(const ClangModuleUnit *ModuleFilter,
     case ClangImporterSynthesizedTypeAttr::Kind::NSErrorWrapper:
     case ClangImporterSynthesizedTypeAttr::Kind::NSErrorWrapperAnon: {
       ASTContext &Ctx = ContainingUnit->getASTContext();
-      auto LookupFlags =
-          NominalTypeDecl::LookupDirectFlags::IgnoreNewExtensions;
+      auto *Importer = static_cast<ClangImporter *>(Ctx.getClangModuleLoader());
       auto WrapperStruct = cast<StructDecl>(VD);
-      TinyPtrVector<ValueDecl *> LookupResults =
-          WrapperStruct->lookupDirect(Ctx.Id_Code, LookupFlags);
-      assert(!LookupResults.empty() && "imported error enum without Code");
-
-      auto CodeEnumIter = llvm::find_if(LookupResults,
-                                        [&](ValueDecl *Member) -> bool {
-        return Member->getDeclContext() == WrapperStruct;
-      });
-      assert(CodeEnumIter != LookupResults.end() &&
-             "could not find Code enum in wrapper struct");
-      assert((*CodeEnumIter)->hasClangNode());
-      ClangNode = (*CodeEnumIter)->getClangNode();
+      auto *CodeEnum = Importer->lookupErrorCodeEnum(WrapperStruct);
+      ClangNode = CodeEnum->getClangNode();
       break;
     }
     }
@@ -2496,6 +2485,10 @@ void ClangImporter::lookupValue(DeclName name, VisibleDeclConsumer &consumer) {
   });
 }
 
+EnumDecl *ClangImporter::lookupErrorCodeEnum(const StructDecl *errorWrapper) const {
+  return Impl.lookupErrorCodeEnum(errorWrapper);
+}
+
 void ClangImporter::lookupTypeDecl(
     StringRef rawName, ClangTypeKind kind,
     llvm::function_ref<void(TypeDecl *)> receiver) {
@@ -2794,20 +2787,13 @@ ClangModuleUnit::lookupNestedType(Identifier name,
                                   const NominalTypeDecl *baseType) const {
   // Special case for error code enums: try looking directly into the struct
   // first. But only if it looks like a synthesized error wrapped struct.
-  if (name == getASTContext().Id_Code && !baseType->hasClangNode() &&
-      isa<StructDecl>(baseType) && !baseType->hasLazyMembers() &&
-      baseType->isChildContextOf(this)) {
-    auto *mutableBase = const_cast<NominalTypeDecl *>(baseType);
-    auto flags = OptionSet<NominalTypeDecl::LookupDirectFlags>();
-    flags |= NominalTypeDecl::LookupDirectFlags::IgnoreNewExtensions;
-    auto codeEnum = mutableBase->lookupDirect(name, flags);
-    // Double-check that we actually have a good result. It's possible what we
-    // found is /not/ a synthesized error struct, but just something that looks
-    // like it. But if we still found a good result we should return that.
-    if (codeEnum.size() == 1 && isa<TypeDecl>(codeEnum.front()))
-      return cast<TypeDecl>(codeEnum.front());
-    if (codeEnum.size() > 1)
-      return nullptr;
+  if (name == getASTContext().Id_Code &&
+      !baseType->hasClangNode() &&
+      isa<StructDecl>(baseType)) {
+    auto *wrapperStruct = cast<StructDecl>(baseType);
+    if (auto *codeEnum = owner.lookupErrorCodeEnum(wrapperStruct))
+      return codeEnum;
+
     // Otherwise, fall back and try via lookup table.
   }
 

lib/ClangImporter/ImportDecl.cpp

@@ -2948,6 +2948,9 @@ namespace {
           // Add the 'Code' enum to the error wrapper.
           errorWrapper->addMember(enumDecl);
           Impl.addAlternateDecl(enumDecl, errorWrapper);
+
+          // Stash the 'Code' enum so we can find it later.
+          Impl.ErrorCodeEnums[errorWrapper] = enumDecl;
         }
 
         // The enumerators go into this enumeration.

lib/ClangImporter/ImportType.cpp

@@ -821,23 +821,15 @@ namespace {
 
     /// Retrieve the 'Code' type for a bridged NSError, or nullptr if
     /// this is not a bridged NSError type.
-    static TypeDecl *getBridgedNSErrorCode(TypeDecl *decl) {
-      auto nominal = dyn_cast<NominalTypeDecl>(decl);
-      if (!nominal) return nullptr;
+    TypeDecl *getBridgedNSErrorCode(TypeDecl *decl) {
+      auto errorWrapper = dyn_cast<StructDecl>(decl);
+      if (!errorWrapper) return nullptr;
 
-      const DeclAttributes &allAttrs = decl->getAttrs();
+      const DeclAttributes &allAttrs = errorWrapper->getAttrs();
       for (auto attr : allAttrs.getAttributes<SynthesizedProtocolAttr>()) {
         if (attr->getProtocolKind() ==
             KnownProtocolKind::BridgedStoredNSError) {
-          auto &ctx = nominal->getASTContext();
-          auto flags = OptionSet<NominalTypeDecl::LookupDirectFlags>();
-          flags |= NominalTypeDecl::LookupDirectFlags::IgnoreNewExtensions;
-          auto lookup = nominal->lookupDirect(ctx.Id_Code, flags);
-          for (auto found : lookup) {
-            if (auto codeDecl = dyn_cast<TypeDecl>(found))
-              return codeDecl;
-          }
-          llvm_unreachable("couldn't find 'Code' within bridged error type");
+          return Impl.lookupErrorCodeEnum(errorWrapper);
         }
       }
 

lib/ClangImporter/ImporterImpl.h

@@ -498,6 +498,10 @@ class LLVM_LIBRARY_VISIBILITY ClangImporter::Implementation
   llvm::DenseMap<const NominalTypeDecl *, TinyPtrVector<ConstructorDecl *>>
       ConstructorsForNominal;
 
+  /// Keep track of the nested 'Code' enum for imported error wrapper
+  /// structs.
+  llvm::DenseMap<const StructDecl *, EnumDecl *> ErrorCodeEnums;
+
   /// Retrieve the alternative declaration for the given imported
   /// Swift declaration.
   ArrayRef<ValueDecl *> getAlternateDecls(Decl *decl) {
@@ -1220,6 +1224,13 @@ class LLVM_LIBRARY_VISIBILITY ClangImporter::Implementation
     return result;
   }
 
+  EnumDecl *lookupErrorCodeEnum(const StructDecl *errorWrapper) {
+    auto found = ErrorCodeEnums.find(errorWrapper);
+    if (found == ErrorCodeEnums.end())
+      return nullptr;
+    return found->second;
+  }
+
   virtual void
   loadAllMembers(Decl *D, uint64_t unused) override;
 

lib/IDE/Utils.cpp

@@ -959,15 +959,13 @@ ClangNode swift::ide::getEffectiveClangNode(const Decl *decl) {
     return clangNode;
 
   // Or via the nested "Code" enum.
-  if (auto nominal =
-      const_cast<NominalTypeDecl *>(dyn_cast<NominalTypeDecl>(decl))) {
-    auto &ctx = nominal->getASTContext();
-    auto flags = OptionSet<NominalTypeDecl::LookupDirectFlags>();
-    flags |= NominalTypeDecl::LookupDirectFlags::IgnoreNewExtensions;
-    for (auto code : nominal->lookupDirect(ctx.Id_Code, flags)) {
-      if (auto clangDecl = code->getClangDecl())
-        return clangDecl;
-    }
+  if (auto *errorWrapper = dyn_cast<StructDecl>(decl)) {
+    auto &ctx = errorWrapper->getASTContext();
+    auto *importer = static_cast<ClangImporter *>(ctx.getClangModuleLoader());
+    if (importer)
+      if (auto *code = importer->lookupErrorCodeEnum(errorWrapper))
+        if (auto clangDecl = code->getClangDecl())
+          return clangDecl;
   }
 
   return ClangNode();