[GC] Support bidirectional iterator in GCStrategyMap

Author: paperchalice

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.

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.
@@ -58,10 +61,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

llvm/lib/CodeGen/GCMetadata.cpp

@@ -23,6 +23,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;
@@ -36,17 +40,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
@@ -59,9 +68,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;
 }
 

llvm/lib/CodeGen/ShadowStackGCLowering.cpp

@@ -109,7 +109,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;