[GC] Support bidirectional iterator in GCStrategyMap

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.
@@ -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

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;
 }
 

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;