Swift: Assign indexes to extensions looking at all the extensions of a type

Author: jketema

swift/extractor/mangler/SwiftMangler.cpp

@@ -112,32 +112,21 @@ unsigned SwiftMangler::getExtensionIndex(const swift::ExtensionDecl* decl,
   if (auto found = preloadedExtensionIndexes.extract(decl)) {
     return found.mapped();
   }
-  if (auto parentModule = llvm::dyn_cast<swift::ModuleDecl>(parent)) {
-    llvm::SmallVector<swift::Decl*> siblings;
-    parentModule->getTopLevelDecls(siblings);
-    indexExtensions(siblings);
-  } else if (auto iterableParent = llvm::dyn_cast<swift::IterableDeclContext>(parent)) {
-    indexExtensions(iterableParent->getAllMembers());
-  } else {
-    // TODO use a generic logging handle for Swift entities here, once it's available
-    CODEQL_ASSERT(false, "non-local context must be module or iterable decl context");
+
+  auto index = 0u;
+  for (const auto& extension : decl->getExtendedNominal()->getExtensions()) {
+    if (getParent(extension) == parent) {
+      preloadedExtensionIndexes.emplace(extension, index);
+      index++;
+    }
   }
+
   auto found = preloadedExtensionIndexes.extract(decl);
   // TODO use a generic logging handle for Swift entities here, once it's available
   CODEQL_ASSERT(found, "extension not found within parent");
   return found.mapped();
 }
 
-void SwiftMangler::indexExtensions(llvm::ArrayRef<swift::Decl*> siblings) {
-  auto index = 0u;
-  for (auto sibling : siblings) {
-    if (sibling->getKind() == swift::DeclKind::Extension) {
-      preloadedExtensionIndexes.emplace(sibling, index);
-    }
-    ++index;
-  }
-}
-
 SwiftMangledName SwiftMangler::visitGenericTypeParamDecl(const swift::GenericTypeParamDecl* decl) {
   auto ret = visitValueDecl(decl, /*force=*/true);
   if (decl->isParameterPack()) {

swift/extractor/mangler/SwiftMangler.h

@@ -112,7 +112,6 @@ class SwiftMangler : private swift::TypeVisitor<SwiftMangler, SwiftMangledName>,
   virtual SwiftMangledName fetch(const swift::TypeBase* type) = 0;
   SwiftMangledName fetch(swift::Type type) { return fetch(type.getPointer()); }
 
-  void indexExtensions(llvm::ArrayRef<swift::Decl*> siblings);
   unsigned int getExtensionIndex(const swift::ExtensionDecl* decl, const swift::Decl* parent);
   static SwiftMangledName initMangled(const swift::TypeBase* type);
   SwiftMangledName initMangled(const swift::Decl* decl);