[CFI] Refactor CfiFunctionIndex to externalize GUID calculation

llvm/include/llvm/IR/ModuleSummaryIndex.h

@@ -18,11 +18,13 @@
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/SmallString.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/ADT/StringMap.h"
 #include "llvm/ADT/StringRef.h"
+#include "llvm/ADT/StringSet.h"
 #include "llvm/ADT/iterator_range.h"
 #include "llvm/IR/ConstantRange.h"
 #include "llvm/IR/GlobalValue.h"
@@ -1317,72 +1319,77 @@ struct TypeIdSummary {
   std::map<uint64_t, WholeProgramDevirtResolution> WPDRes;
 };
 
+/// Encapsulate the names of CFI target functions. It interfaces with ThinLTO to
+/// determine efficiently which of the names need to be exported for a
+/// particular module.
 class CfiFunctionIndex {
-  DenseMap<GlobalValue::GUID, std::set<std::string, std::less<>>> Index;
-  using IndexIterator =
-      DenseMap<GlobalValue::GUID,
-               std::set<std::string, std::less<>>>::const_iterator;
-  using NestedIterator = std::set<std::string, std::less<>>::const_iterator;
+  // `Names` is the authoritative source of data. `ThinLTOToNamesIndex` is there
+  // just to efficiently retrieve which names in this index need exporting for
+  // a particular module index. We cannot guarantee the ThinLTO GUIDs are
+  // collision - free, so we associate a collection to a guid. Functions with
+  // the same name may have different GUIDs, too. So we index a list of names
+  // with the same GUID under that GUID key. We don't need the reverse because
+  // the queries from ThinLTO use GUIDs as key.
+  // Note that StringSet rehashing doesn't move keys, so we can safely store the
+  // StringRef value inserted in `Names` in ThinLTOToNamesIndex, and avoid
+  // copies.
+  // Design note: we could do away with Names and use ThinLTOToNamesIndex as
+  // index and data source, but opted against, for a small heap penalty, to
+  // avoid confusion wrt the role GUIDs play in this case: they are an artifact
+  // of the need to interface with ThinLTO, not otherwise necessary to CFI.
+  StringSet<> Names;
+
+  using InternalIndexGroup = SetVector<StringRef>;
+  DenseMap<GlobalValue::GUID, InternalIndexGroup> ThinLTOToNamesIndex;
+
+  using NestedIterator = InternalIndexGroup::const_iterator;
 
 public:
-  // Iterates keys of the DenseMap.
-  class GUIDIterator : public iterator_adaptor_base<GUIDIterator, IndexIterator,
-                                                    std::forward_iterator_tag,
-                                                    GlobalValue::GUID> {
-    using base = GUIDIterator::iterator_adaptor_base;
-
-  public:
-    GUIDIterator() = default;
-    explicit GUIDIterator(IndexIterator I) : base(I) {}
-
-    GlobalValue::GUID operator*() const { return this->wrapped()->first; }
-  };
-
   CfiFunctionIndex() = default;
-  template <typename It> CfiFunctionIndex(It B, It E) {
-    for (; B != E; ++B)
-      emplace(*B);
-  }
-
-  std::vector<StringRef> symbols() const {
-    std::vector<StringRef> Symbols;
-    for (auto &[GUID, Syms] : Index) {
-      (void)GUID;
-      llvm::append_range(Symbols, Syms);
-    }
+  CfiFunctionIndex(const CfiFunctionIndex &) = delete;
+  CfiFunctionIndex(CfiFunctionIndex &&) = default;
+
+  /// API used for serialization, e.g. YAML.
+  std::vector<std::pair<StringRef, GlobalValue::GUID>>
+  getSortedSymbols() const {
+    std::vector<std::pair<StringRef, GlobalValue::GUID>> Symbols;
+    for (auto &[GUID, Names] : ThinLTOToNamesIndex)
+      for (auto Name : Names)
+        Symbols.emplace_back(Name, GUID);
+    llvm::sort(Symbols);
     return Symbols;
   }
 
-  GUIDIterator guid_begin() const { return GUIDIterator(Index.begin()); }
-  GUIDIterator guid_end() const { return GUIDIterator(Index.end()); }
-  iterator_range<GUIDIterator> guids() const {
-    return make_range(guid_begin(), guid_end());
+  /// get the set of GUIDs that should also be exported because they are the
+  /// GUIDs of the cfi functions encapsulated here.
+  auto getExportedThinLTOGUIDs() const {
+    return map_range(ThinLTOToNamesIndex, [](auto I) { return I.first; });
   }
 
-  iterator_range<NestedIterator> forGuid(GlobalValue::GUID GUID) const {
-    auto I = Index.find(GUID);
-    if (I == Index.end())
+  /// get the name(s) associated with a given ThinLTO GUID. This enables
+  /// efficient identification of the subset of names that should be included in
+  /// a module summary.
+  auto getNamesForGUID(GlobalValue::GUID GUID) const {
+    auto I = ThinLTOToNamesIndex.find(GUID);
+    if (I == ThinLTOToNamesIndex.end())
       return make_range(NestedIterator{}, NestedIterator{});
     return make_range(I->second.begin(), I->second.end());
   }
 
-  template <typename... Args> void emplace(Args &&...A) {
-    StringRef S(std::forward<Args>(A)...);
-    GlobalValue::GUID GUID = GlobalValue::getGUIDAssumingExternalLinkage(
-        GlobalValue::dropLLVMManglingEscape(S));
-    Index[GUID].emplace(S);
+  /// Add the function name and the GUID that ThinLTO uses for it.
+  void addSymbolWithThinLTOGUID(StringRef Name, GlobalValue::GUID GUID) {
+    auto [Iter, _] = Names.insert(Name);
+    ThinLTOToNamesIndex[GUID].insert(Iter->first());
   }
 
-  size_t count(StringRef S) const {
-    GlobalValue::GUID GUID = GlobalValue::getGUIDAssumingExternalLinkage(
-        GlobalValue::dropLLVMManglingEscape(S));
-    auto I = Index.find(GUID);
-    if (I == Index.end())
-      return 0;
-    return I->second.count(S);
+  bool contains(StringRef Name) const {
+    return Names.find(Name) != Names.end();
   }
 
-  bool empty() const { return Index.empty(); }
+  bool empty() const {
+    assert(Names.empty() == ThinLTOToNamesIndex.empty());
+    return Names.empty();
+  }
 };
 
 /// 160 bits SHA1

llvm/include/llvm/IR/ModuleSummaryIndexYAML.h

@@ -327,6 +327,24 @@ template <> struct CustomMappingTraits<TypeIdSummaryMapTy> {
   }
 };
 
+template <> struct MappingTraits<std::pair<StringRef, GlobalValue::GUID>> {
+  static void mapping(IO &io,
+                      std::pair<StringRef, GlobalValue::GUID> &NameAndGUID) {
+    io.mapRequired("Name", NameAndGUID.first);
+    io.mapRequired("GUID", NameAndGUID.second);
+  }
+};
+
+using StringAndGUID = std::pair<llvm::StringRef, llvm::GlobalValue::GUID>;
+
+} // namespace yaml
+} // namespace llvm
+
+LLVM_YAML_IS_SEQUENCE_VECTOR(llvm::yaml::StringAndGUID)
+
+namespace llvm {
+namespace yaml {
+
 template <> struct MappingTraits<ModuleSummaryIndex> {
   static void mapping(IO &io, ModuleSummaryIndex& index) {
     io.mapOptional("GlobalValueMap", index.GlobalValueMap);
@@ -352,25 +370,24 @@ template <> struct MappingTraits<ModuleSummaryIndex> {
                    index.WithGlobalValueDeadStripping);
 
     if (io.outputting()) {
-      auto CfiFunctionDefs = index.CfiFunctionDefs.symbols();
-      llvm::sort(CfiFunctionDefs);
+      auto CfiFunctionDefs = index.CfiFunctionDefs.getSortedSymbols();
       io.mapOptional("CfiFunctionDefs", CfiFunctionDefs);
-      auto CfiFunctionDecls(index.CfiFunctionDecls.symbols());
-      llvm::sort(CfiFunctionDecls);
+      auto CfiFunctionDecls(index.CfiFunctionDecls.getSortedSymbols());
       io.mapOptional("CfiFunctionDecls", CfiFunctionDecls);
     } else {
-      std::vector<std::string> CfiFunctionDefs;
+      std::vector<std::pair<StringRef, GlobalValue::GUID>> CfiFunctionDefs;
       io.mapOptional("CfiFunctionDefs", CfiFunctionDefs);
-      index.CfiFunctionDefs = {CfiFunctionDefs.begin(), CfiFunctionDefs.end()};
-      std::vector<std::string> CfiFunctionDecls;
+      for (auto &[S, G] : CfiFunctionDefs)
+        index.CfiFunctionDefs.addSymbolWithThinLTOGUID(S, G);
+      std::vector<std::pair<StringRef, GlobalValue::GUID>> CfiFunctionDecls;
       io.mapOptional("CfiFunctionDecls", CfiFunctionDecls);
-      index.CfiFunctionDecls = {CfiFunctionDecls.begin(),
-                                CfiFunctionDecls.end()};
+      for (auto &[S, G] : CfiFunctionDecls)
+        index.CfiFunctionDecls.addSymbolWithThinLTOGUID(S, G);
     }
   }
 };
 
-} // End yaml namespace
-} // End llvm namespace
+} // namespace yaml
+} // namespace llvm
 
 #endif

llvm/lib/Bitcode/Reader/BitcodeReader.cpp

@@ -8146,17 +8146,25 @@ Error ModuleSummaryIndexBitcodeReader::parseEntireSummary(unsigned ID) {
 
     case bitc::FS_CFI_FUNCTION_DEFS: {
       auto &CfiFunctionDefs = TheIndex.cfiFunctionDefs();
-      for (unsigned I = 0; I != Record.size(); I += 2)
-        CfiFunctionDefs.emplace(Strtab.data() + Record[I],
-                                static_cast<size_t>(Record[I + 1]));
+      for (unsigned I = 0; I != Record.size(); I += 2) {
+        StringRef Name(Strtab.data() + Record[I],
+                       static_cast<size_t>(Record[I + 1]));
+        GlobalValue::GUID GUID = GlobalValue::getGUIDAssumingExternalLinkage(
+            GlobalValue::dropLLVMManglingEscape(Name));
+        CfiFunctionDefs.addSymbolWithThinLTOGUID(Name, GUID);
+      }
       break;
     }
 
     case bitc::FS_CFI_FUNCTION_DECLS: {
       auto &CfiFunctionDecls = TheIndex.cfiFunctionDecls();
-      for (unsigned I = 0; I != Record.size(); I += 2)
-        CfiFunctionDecls.emplace(Strtab.data() + Record[I],
-                                 static_cast<size_t>(Record[I + 1]));
+      for (unsigned I = 0; I != Record.size(); I += 2) {
+        StringRef Name(Strtab.data() + Record[I],
+                       static_cast<size_t>(Record[I + 1]));
+        GlobalValue::GUID GUID = GlobalValue::getGUIDAssumingExternalLinkage(
+            GlobalValue::dropLLVMManglingEscape(Name));
+        CfiFunctionDecls.addSymbolWithThinLTOGUID(Name, GUID);
+      }
       break;
     }
 

llvm/lib/Bitcode/Writer/BitcodeWriter.cpp

@@ -5322,8 +5322,8 @@ void IndexBitcodeWriter::writeCombinedGlobalValueSummary() {
     if (CfiIndex.empty())
       return;
     for (GlobalValue::GUID GUID : DefOrUseGUIDs) {
-      auto Defs = CfiIndex.forGuid(GUID);
-      llvm::append_range(Functions, Defs);
+      auto Names = CfiIndex.getNamesForGUID(GUID);
+      llvm::append_range(Functions, Names);
     }
     if (Functions.empty())
       return;

llvm/lib/LTO/LTO.cpp

@@ -1585,9 +1585,9 @@ class CGThinBackend : public ThinBackendProc {
                         OnWrite, ShouldEmitImportsFiles, ThinLTOParallelism),
         ShouldEmitIndexFiles(ShouldEmitIndexFiles) {
     auto &Defs = CombinedIndex.cfiFunctionDefs();
-    CfiFunctionDefs.insert_range(Defs.guids());
+    CfiFunctionDefs.insert_range(Defs.getExportedThinLTOGUIDs());
     auto &Decls = CombinedIndex.cfiFunctionDecls();
-    CfiFunctionDecls.insert_range(Decls.guids());
+    CfiFunctionDecls.insert_range(Decls.getExportedThinLTOGUIDs());
   }
 };
 
@@ -1948,10 +1948,10 @@ class WriteIndexesThinBackend : public ThinBackendProc {
         OldPrefix(OldPrefix), NewPrefix(NewPrefix),
         NativeObjectPrefix(NativeObjectPrefix),
         LinkedObjectsFile(LinkedObjectsFile) {
-    auto &Defs = CombinedIndex.cfiFunctionDefs();
-    CfiFunctionDefs.insert(Defs.guid_begin(), Defs.guid_end());
-    auto &Decls = CombinedIndex.cfiFunctionDecls();
-    CfiFunctionDecls.insert(Decls.guid_begin(), Decls.guid_end());
+    auto Defs = CombinedIndex.cfiFunctionDefs().getExportedThinLTOGUIDs();
+    CfiFunctionDefs.insert(Defs.begin(), Defs.end());
+    auto Decls = CombinedIndex.cfiFunctionDecls().getExportedThinLTOGUIDs();
+    CfiFunctionDecls.insert(Decls.begin(), Decls.end());
   }
 
   Error start(
@@ -2181,10 +2181,11 @@ Error LTO::runThinLTO(AddStreamFn AddStream, FileCache Cache,
 
   // Any functions referenced by the jump table in the regular LTO object must
   // be exported.
-  auto &Defs = ThinLTO.CombinedIndex.cfiFunctionDefs();
-  ExportedGUIDs.insert(Defs.guid_begin(), Defs.guid_end());
-  auto &Decls = ThinLTO.CombinedIndex.cfiFunctionDecls();
-  ExportedGUIDs.insert(Decls.guid_begin(), Decls.guid_end());
+  auto Defs = ThinLTO.CombinedIndex.cfiFunctionDefs().getExportedThinLTOGUIDs();
+  ExportedGUIDs.insert(Defs.begin(), Defs.end());
+  auto Decls =
+      ThinLTO.CombinedIndex.cfiFunctionDecls().getExportedThinLTOGUIDs();
+  ExportedGUIDs.insert(Decls.begin(), Decls.end());
 
   auto isExported = [&](StringRef ModuleIdentifier, ValueInfo VI) {
     const auto &ExportList = ExportLists.find(ModuleIdentifier);

llvm/lib/Transforms/IPO/LowerTypeTests.cpp

@@ -1788,10 +1788,15 @@ void LowerTypeTestsModule::buildBitSetsFromFunctionsNative(
     }
 
     if (IsExported) {
+      // TODO: use F->getGUID() once #184065 is relanded.
+      GlobalValue::GUID GUID = GlobalValue::getGUIDAssumingExternalLinkage(
+          GlobalValue::dropLLVMManglingEscape(F->getName()));
       if (IsJumpTableCanonical)
-        ExportSummary->cfiFunctionDefs().emplace(F->getName());
+        ExportSummary->cfiFunctionDefs().addSymbolWithThinLTOGUID(F->getName(),
+                                                                  GUID);
       else
-        ExportSummary->cfiFunctionDecls().emplace(F->getName());
+        ExportSummary->cfiFunctionDecls().addSymbolWithThinLTOGUID(F->getName(),
+                                                                   GUID);
     }
 
     if (!IsJumpTableCanonical) {
@@ -2123,9 +2128,9 @@ bool LowerTypeTestsModule::lower() {
       // have the same name, but it's not the one we are looking for.
       if (F.hasLocalLinkage())
         continue;
-      if (ImportSummary->cfiFunctionDefs().count(F.getName()))
+      if (ImportSummary->cfiFunctionDefs().contains(F.getName()))
         Defs.push_back(&F);
-      else if (ImportSummary->cfiFunctionDecls().count(F.getName()))
+      else if (ImportSummary->cfiFunctionDecls().contains(F.getName()))
         Decls.push_back(&F);
     }
 

llvm/test/Transforms/LowerTypeTests/Inputs/blockaddr-import.yaml

@@ -4,6 +4,6 @@ TypeIdMap:
     TTRes:
       Kind:            AllOnes
 CfiFunctionDefs:
-  - m
-CfiFunctionDecls:
-...
+- Name: m
+  GUID: 2799708793537531759
+CfiFunctionDecls: null

llvm/test/Transforms/LowerTypeTests/Inputs/cfi-direct-call.yaml

@@ -4,9 +4,12 @@ TypeIdMap:
     TTRes:
       Kind:            AllOnes
 CfiFunctionDefs:
-  - internal_default_def
-  - internal_hidden_def
-  - dsolocal_default_def
+- Name: internal_default_def
+  GUID: 2377677245844927262
+- Name: internal_hidden_def
+  GUID: 12222722662536945321
+- Name: dsolocal_default_def
+  GUID: 11010557671720506420
 CfiFunctionDecls:
-  - external_decl
-...
+- Name: external_decl
+  GUID: 2828901536631366483

llvm/test/Transforms/LowerTypeTests/Inputs/cfi-direct-call1.yaml

@@ -4,10 +4,14 @@ TypeIdMap:
     TTRes:
       Kind:            AllOnes
 CfiFunctionDefs:
-  - local_func1
-  - local_func2
-  - local_func3
+- Name: local_func1
+  GUID: 15974708163032388986
+- Name: local_func2
+  GUID: 10016086976061028907
+- Name: local_func3
+  GUID: 161148068062889559
 CfiFunctionDecls:
-  - extern_decl
-  - extern_weak
-...
+- Name: extern_decl
+  GUID: 15001569853432741044
+- Name: extern_weak
+  GUID: 9772975595203342681

llvm/test/Transforms/LowerTypeTests/Inputs/import-alias.yaml

@@ -6,6 +6,6 @@ TypeIdMap:
       SizeM1BitWidth:  7
 WithGlobalValueDeadStripping: false
 CfiFunctionDefs:
-  - f
-CfiFunctionDecls:
-...
+- Name: f
+  GUID: 14740650423002898831
+CfiFunctionDecls: null