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

Author: jketema

swift/extractor/mangler/SwiftMangler.cpp

@@ -98,46 +98,31 @@ SwiftMangledName SwiftMangler::visitExtensionDecl(const swift::ExtensionDecl* de
     return {};
   }
 
-  auto parent = getParent(decl);
   auto target = decl->getExtendedType();
-  return initMangled(decl) << fetch(target) << getExtensionIndex(decl, parent);
+  return initMangled(decl) << fetch(target) << getExtensionIndex(decl);
 }
 
-unsigned SwiftMangler::getExtensionIndex(const swift::ExtensionDecl* decl,
-                                         const swift::Decl* parent) {
+unsigned SwiftMangler::getExtensionIndex(const swift::ExtensionDecl* decl) {
   // to avoid iterating multiple times on the parent of multiple extensions, we preload extension
   // indexes once for each encountered parent into the `preloadedExtensionIndexes` mapping.
   // Because we mangle declarations only once in a given trap/dispatcher context, we can safely
   // discard preloaded indexes on use
   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");
+  for (const auto& sibling : decl->getExtendedNominal()->getExtensions()) {
+    auto index = 0u;
+    if (sibling->getDeclContext() == decl->getDeclContext()) {
+      preloadedExtensionIndexes.emplace(sibling, 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,8 +112,7 @@ 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);
+  unsigned int getExtensionIndex(const swift::ExtensionDecl* decl);
   static SwiftMangledName initMangled(const swift::TypeBase* type);
   SwiftMangledName initMangled(const swift::Decl* decl);
   SwiftMangledName visitTypeDiscriminatedValueDecl(const swift::ValueDecl* decl);