[llvm] [GC] Support bidirectional iterator in GCStrategyMap (PR #99316)

[via llvm-commits]

https://github.com/paperchalice created https://github.com/llvm/llvm-project/pull/99316

Convert this result into a flat-map like class, so `AsmPrinter` can iterate over this bidirectionally.


>From d99e03e4627877c5aed8c2194c23b1d99610f9f4 Mon Sep 17 00:00:00 2001
From: PaperChalice <liujunchang97 at outlook.com>
Date: Sun, 14 Jul 2024 21:12:07 +0800
Subject: [PATCH] [GC] Support bidirectional iterator in GCStrategyMap

---
 llvm/include/llvm/CodeGen/GCMetadata.h        | 33 ++++++++++++++++++-
 llvm/include/llvm/CodeGen/GCMetadataPrinter.h |  5 +++
 llvm/lib/CodeGen/GCMetadata.cpp               | 22 +++++++++----
 llvm/lib/CodeGen/ShadowStackGCLowering.cpp    |  2 +-
 4 files changed, 53 insertions(+), 9 deletions(-)

diff --git a/llvm/include/llvm/CodeGen/GCMetadata.h b/llvm/include/llvm/CodeGen/GCMetadata.h
index ca6a511185c7c..f22aa2c54f401 100644
--- a/llvm/include/llvm/CodeGen/GCMetadata.h
+++ b/llvm/include/llvm/CodeGen/GCMetadata.h
@@ -151,15 +151,46 @@ class GCFunctionInfo {
   size_t live_size(const iterator &p) const { return roots_size(); }
 };
 
-struct GCStrategyMap {
+class GCStrategyMap {
   StringMap<std::unique_ptr<GCStrategy>> StrategyMap;
+  SmallVector<GCStrategy *, 1> StrategyList; // For bidirectional iterator.
+  using StrategyListT = SmallVector<GCStrategy *, 1>;
 
+  FunctionAnalysisManager *FAM = nullptr;
+
+public:
   GCStrategyMap() = default;
   GCStrategyMap(GCStrategyMap &&) = default;
+  GCStrategyMap(FunctionAnalysisManager &FAM) : FAM(&FAM) {}
 
   /// Handle invalidation explicitly.
   bool invalidate(Module &M, const PreservedAnalyses &PA,
                   ModuleAnalysisManager::Invalidator &Inv);
+
+  GCFunctionInfo &getFunctionInfo(Function &F);
+
+  bool empty() const { return StrategyMap.empty(); }
+
+  bool contains(StringRef Name) const { return StrategyMap.contains(Name); }
+
+  /// Insert a new strategy if it is not existed, otherwise do nothing.
+  void insert(StringRef Name, std::unique_ptr<GCStrategy> Strategy) {
+    auto &S = StrategyMap[Name];
+    if (!S) {
+      S = std::move(Strategy);
+      StrategyList.push_back(S.get());
+    }
+  }
+
+  GCStrategy &at(StringRef Name) { return *StrategyMap.at(Name); }
+
+  // This class must support bidirectional iterator which is used by AsmPrinter.
+  using iterator = StrategyListT::iterator;
+  iterator begin() { return StrategyList.begin(); }
+  iterator end() { return StrategyList.end(); }
+  using reverse_iterator = StrategyListT::reverse_iterator;
+  reverse_iterator rbegin() { return StrategyList.rbegin(); }
+  reverse_iterator rend() { return StrategyList.rend(); }
 };
 
 /// An analysis pass which caches information about the entire Module.
diff --git a/llvm/include/llvm/CodeGen/GCMetadataPrinter.h b/llvm/include/llvm/CodeGen/GCMetadataPrinter.h
index f9527c9f8752e..a7cf6ec914b9c 100644
--- a/llvm/include/llvm/CodeGen/GCMetadataPrinter.h
+++ b/llvm/include/llvm/CodeGen/GCMetadataPrinter.h
@@ -27,8 +27,11 @@ class AsmPrinter;
 class GCMetadataPrinter;
 class GCModuleInfo;
 class GCStrategy;
+class GCStrategyMap;
 class Module;
 class StackMaps;
+template <typename IRUnitT, typename... ExtraArgTs> class AnalysisManager;
+using ModuleAnalysisManager = AnalysisManager<Module>;
 
 /// GCMetadataPrinterRegistry - The GC assembly printer registry uses all the
 /// defaults from Registry.
@@ -56,10 +59,12 @@ class GCMetadataPrinter {
   /// Called before the assembly for the module is generated by
   /// the AsmPrinter (but after target specific hooks.)
   virtual void beginAssembly(Module &M, GCModuleInfo &Info, AsmPrinter &AP) {}
+  virtual void beginAssembly(Module &M, GCStrategyMap &Map, AsmPrinter &AP) {}
 
   /// Called after the assembly for the module is generated by
   /// the AsmPrinter (but before target specific hooks)
   virtual void finishAssembly(Module &M, GCModuleInfo &Info, AsmPrinter &AP) {}
+  virtual void finishAssembly(Module &M, GCStrategyMap &Map, AsmPrinter &AP) {}
 
   /// Called when the stack maps are generated. Return true if
   /// stack maps with a custom format are generated. Otherwise
diff --git a/llvm/lib/CodeGen/GCMetadata.cpp b/llvm/lib/CodeGen/GCMetadata.cpp
index e1af457c9b9d7..abc7c7ddfb377 100644
--- a/llvm/lib/CodeGen/GCMetadata.cpp
+++ b/llvm/lib/CodeGen/GCMetadata.cpp
@@ -27,6 +27,10 @@ using namespace llvm;
 
 bool GCStrategyMap::invalidate(Module &M, const PreservedAnalyses &PA,
                                ModuleAnalysisManager::Invalidator &) {
+  auto PAC = PA.getChecker<CollectorMetadataAnalysis>();
+  if (PAC.preserved() || PAC.preservedSet<AllAnalysesOn<Function>>())
+    return false;
+
   for (const auto &F : M) {
     if (F.isDeclaration() || !F.hasGC())
       continue;
@@ -40,17 +44,22 @@ AnalysisKey CollectorMetadataAnalysis::Key;
 
 CollectorMetadataAnalysis::Result
 CollectorMetadataAnalysis::run(Module &M, ModuleAnalysisManager &MAM) {
-  Result R;
-  auto &Map = R.StrategyMap;
+  auto &FAM = MAM.getResult<FunctionAnalysisManagerModuleProxy>(M).getManager();
+  Result R(FAM);
   for (auto &F : M) {
     if (F.isDeclaration() || !F.hasGC())
       continue;
-    if (auto GCName = F.getGC(); !Map.contains(GCName))
-      Map[GCName] = getGCStrategy(GCName);
+    auto GCName = F.getGC();
+    R.insert(GCName, getGCStrategy(GCName));
   }
   return R;
 }
 
+GCFunctionInfo &llvm::GCStrategyMap::getFunctionInfo(Function &F) {
+  assert(FAM && "Need initialize!");
+  return FAM->getResult<GCFunctionAnalysis>(F);
+}
+
 AnalysisKey GCFunctionAnalysis::Key;
 
 GCFunctionAnalysis::Result
@@ -63,9 +72,8 @@ GCFunctionAnalysis::run(Function &F, FunctionAnalysisManager &FAM) {
       MAMProxy.cachedResultExists<CollectorMetadataAnalysis>(*F.getParent()) &&
       "This pass need module analysis `collector-metadata`!");
   auto &Map =
-      MAMProxy.getCachedResult<CollectorMetadataAnalysis>(*F.getParent())
-          ->StrategyMap;
-  GCFunctionInfo Info(F, *Map[F.getGC()]);
+      *MAMProxy.getCachedResult<CollectorMetadataAnalysis>(*F.getParent());
+  GCFunctionInfo Info(F, Map.at(F.getGC()));
   return Info;
 }
 
diff --git a/llvm/lib/CodeGen/ShadowStackGCLowering.cpp b/llvm/lib/CodeGen/ShadowStackGCLowering.cpp
index 232e5e2bb886d..8fd3359391106 100644
--- a/llvm/lib/CodeGen/ShadowStackGCLowering.cpp
+++ b/llvm/lib/CodeGen/ShadowStackGCLowering.cpp
@@ -110,7 +110,7 @@ class ShadowStackGCLowering : public FunctionPass {
 PreservedAnalyses ShadowStackGCLoweringPass::run(Module &M,
                                                  ModuleAnalysisManager &MAM) {
   auto &Map = MAM.getResult<CollectorMetadataAnalysis>(M);
-  if (Map.StrategyMap.contains("shadow-stack"))
+  if (Map.contains("shadow-stack"))
     return PreservedAnalyses::all();
 
   ShadowStackGCLoweringImpl Impl;