@@ -597,6 +597,7 @@ ModuleDecl::ModuleDecl(Identifier name, ASTContext &ctx,
597597 Bits.ModuleDecl .HasIncrementalInfo = 0 ;
598598 Bits.ModuleDecl .HasHermeticSealAtLink = 0 ;
599599 Bits.ModuleDecl .IsConcurrencyChecked = 0 ;
600+ Bits.ModuleDecl .ObjCNameLookupCachePopulated = 0 ;
600601}
601602
602603ImplicitImportList ModuleDecl::getImplicitImports () const {
@@ -1176,6 +1177,60 @@ void ModuleDecl::getTopLevelDeclsWhereAttributesMatch(
11761177 FORWARD (getTopLevelDeclsWhereAttributesMatch, (Results, matchAttributes));
11771178}
11781179
1180+ void ModuleDecl::lookupTopLevelDeclsByObjCName (SmallVectorImpl<Decl *> &Results,
1181+ DeclName name) {
1182+ if (!isObjCNameLookupCachePopulated ())
1183+ populateObjCNameLookupCache ();
1184+
1185+ // A top level decl can't be special anyways
1186+ if (name.isSpecial ())
1187+ return ;
1188+
1189+ auto resultsForFileUnit = ObjCNameLookupCache.find (name.getBaseIdentifier ());
1190+ if (resultsForFileUnit == ObjCNameLookupCache.end ())
1191+ return ;
1192+
1193+ Results.append (resultsForFileUnit->second .begin (),
1194+ resultsForFileUnit->second .end ());
1195+ }
1196+
1197+ void ModuleDecl::populateObjCNameLookupCache () {
1198+ SmallVector<Decl *> topLevelObjCExposedDeclsInFileUnit;
1199+ auto hasObjCAttrNamePredicate = [](const DeclAttributes &attrs) -> bool {
1200+ return attrs.hasAttribute <ObjCAttr>();
1201+ };
1202+
1203+ for (FileUnit *file : getFiles ()) {
1204+ file->getTopLevelDeclsWhereAttributesMatch (
1205+ topLevelObjCExposedDeclsInFileUnit, hasObjCAttrNamePredicate);
1206+ if (auto *synth = file->getSynthesizedFile ()) {
1207+ synth->getTopLevelDeclsWhereAttributesMatch (
1208+ topLevelObjCExposedDeclsInFileUnit, hasObjCAttrNamePredicate);
1209+ }
1210+ }
1211+
1212+ for (Decl *decl : topLevelObjCExposedDeclsInFileUnit) {
1213+ if (ValueDecl *VD = dyn_cast<ValueDecl>(decl); VD && VD->hasName ()) {
1214+ const auto &declObjCAttribute = VD->getAttrs ().getAttribute <ObjCAttr>();
1215+ // No top level decl (class, protocol, extension etc.) is allowed to have a
1216+ // compound name, @objc provided or otherwise. Global functions are allowed to
1217+ // have compound names, but not allowed to have @objc attributes. Thus we
1218+ // are sure to not hit asserts getting the simple name.
1219+ //
1220+ // Similarly, init, dealloc and subscript (the special names) can't be top
1221+ // level decls, so we won't hit asserts getting the base identifier out of the
1222+ // value decl.
1223+ const Identifier &declObjCName =
1224+ declObjCAttribute->hasName ()
1225+ ? declObjCAttribute->getName ()->getSimpleName ()
1226+ : VD->getName ().getBaseIdentifier ();
1227+ ObjCNameLookupCache[declObjCName].push_back (decl);
1228+ }
1229+ }
1230+
1231+ setIsObjCNameLookupCachePopulated (true );
1232+ }
1233+
11791234void SourceFile::getTopLevelDecls (SmallVectorImpl<Decl*> &Results) const {
11801235 auto decls = getTopLevelDecls ();
11811236 Results.append (decls.begin (), decls.end ());
@@ -3347,6 +3402,9 @@ bool SourceFile::shouldCrossImport() const {
33473402void ModuleDecl::clearLookupCache () {
33483403 getASTContext ().getImportCache ().clear ();
33493404
3405+ setIsObjCNameLookupCachePopulated (false );
3406+ ObjCNameLookupCache.clear ();
3407+
33503408 if (!Cache)
33513409 return ;
33523410
0 commit comments