[GC] Use MapVector for GCStrategyMap

llvm/include/llvm/CodeGen/GCMetadata.h

@@ -33,6 +33,7 @@
 #define LLVM_CODEGEN_GCMETADATA_H
 
 #include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/MapVector.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringMap.h"
 #include "llvm/ADT/StringRef.h"
@@ -151,15 +152,49 @@ class GCFunctionInfo {
   size_t live_size(const iterator &p) const { return roots_size(); }
 };
 
-struct GCStrategyMap {
-  StringMap<std::unique_ptr<GCStrategy>> StrategyMap;
+class GCStrategyMap {
+  using MapT =
+      MapVector<StringRef, std::unique_ptr<GCStrategy>, StringMap<unsigned>>;
+  MapT Strategies;
 
+public:
   GCStrategyMap() = default;
   GCStrategyMap(GCStrategyMap &&) = default;
 
   /// Handle invalidation explicitly.
   bool invalidate(Module &M, const PreservedAnalyses &PA,
                   ModuleAnalysisManager::Invalidator &Inv);
+
+  using iterator = MapT::iterator;
+  using const_iterator = MapT::const_iterator;
+  using reverse_iterator = MapT::reverse_iterator;
+  using const_reverse_iterator = MapT::const_reverse_iterator;
+
+  iterator begin() { return Strategies.begin(); }
+  const_iterator begin() const { return Strategies.begin(); }
+  iterator end() { return Strategies.end(); }
+  const_iterator end() const { return Strategies.end(); }
+
+  reverse_iterator rbegin() { return Strategies.rbegin(); }
+  const_reverse_iterator rbegin() const { return Strategies.rbegin(); }
+  reverse_iterator rend() { return Strategies.rend(); }
+  const_reverse_iterator rend() const { return Strategies.rend(); }
+
+  bool empty() const { return Strategies.empty(); }
+
+  const GCStrategy &operator[](StringRef GCName) const {
+    auto I = Strategies.find(GCName);
+    assert(I != Strategies.end() && "Required strategy doesn't exist!");
+    return *I->second;
+  }
+
+  std::pair<iterator, bool> try_emplace(StringRef GCName) {
+    return Strategies.try_emplace(GCName);
+  }
+
+  bool contains(StringRef GCName) const {
+    return Strategies.find(GCName) != Strategies.end();
+  }
 };
 
 /// An analysis pass which caches information about the entire Module.

llvm/include/llvm/IR/GCStrategy.h

@@ -63,6 +63,7 @@ class Type;
 class GCStrategy {
 private:
   friend class GCModuleInfo;
+  friend class CollectorMetadataAnalysis;
 
   std::string Name;
 

llvm/lib/CodeGen/GCMetadata.cpp

@@ -26,7 +26,7 @@ bool GCStrategyMap::invalidate(Module &M, const PreservedAnalyses &PA,
   for (const auto &F : M) {
     if (F.isDeclaration() || !F.hasGC())
       continue;
-    if (!StrategyMap.contains(F.getGC()))
+    if (!contains(F.getGC()))
       return true;
   }
   return false;
@@ -36,17 +36,18 @@ AnalysisKey CollectorMetadataAnalysis::Key;
 
 CollectorMetadataAnalysis::Result
 CollectorMetadataAnalysis::run(Module &M, ModuleAnalysisManager &MAM) {
-  Result R;
-  auto &Map = R.StrategyMap;
+  Result StrategyMap;
   for (auto &F : M) {
     if (F.isDeclaration() || !F.hasGC())
       continue;
-    auto GCName = F.getGC();
-    auto [It, Inserted] = Map.try_emplace(GCName);
-    if (Inserted)
+    StringRef GCName = F.getGC();
+    auto [It, Inserted] = StrategyMap.try_emplace(GCName);
+    if (Inserted) {
       It->second = getGCStrategy(GCName);
+      It->second->Name = GCName;
+    }
   }
-  return R;
+  return StrategyMap;
 }
 
 AnalysisKey GCFunctionAnalysis::Key;
@@ -61,9 +62,9 @@ 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());
+  GCStrategy &S = *Map.try_emplace(F.getGC()).first->second;
+  GCFunctionInfo Info(F, S);
   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;

llvm/unittests/CodeGen/CMakeLists.txt

@@ -29,6 +29,7 @@ add_llvm_unittest(CodeGenTests
   DIETest.cpp
   DroppedVariableStatsMIRTest.cpp
   DwarfStringPoolEntryRefTest.cpp
+  GCMetadata.cpp
   InstrRefLDVTest.cpp
   LowLevelTypeTest.cpp
   LexicalScopesTest.cpp

llvm/unittests/CodeGen/GCMetadata.cpp

@@ -0,0 +1,76 @@
+//===- llvm/unittest/CodeGen/GCMetadata.cpp -------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/CodeGen/GCMetadata.h"
+#include "llvm/Analysis/CGSCCPassManager.h"
+#include "llvm/Analysis/LoopAnalysisManager.h"
+#include "llvm/AsmParser/Parser.h"
+#include "llvm/IR/Analysis.h"
+#include "llvm/IR/LLVMContext.h"
+#include "llvm/IR/Module.h"
+#include "llvm/IR/PassManager.h"
+#include "llvm/Passes/PassBuilder.h"
+#include "llvm/Support/SourceMgr.h"
+#include "gtest/gtest.h"
+
+using namespace llvm;
+
+namespace {
+
+std::unique_ptr<Module> parseIR(LLVMContext &Context, const char *IR) {
+  SMDiagnostic Err;
+  return parseAssemblyString(IR, Err, Context);
+}
+
+class GCMetadataTest : public ::testing::Test {
+protected:
+  LLVMContext Context;
+  std::unique_ptr<Module> M;
+
+public:
+  GCMetadataTest()
+      : M(parseIR(Context, R"(
+%Env = type ptr
+
+define void @.main(%Env) gc "shadow-stack" {
+	%Root = alloca %Env
+	call void @llvm.gcroot( ptr %Root, %Env null )
+	unreachable
+}
+
+define void @g() gc "erlang" {
+entry:
+	ret void
+}
+
+declare void @llvm.gcroot(ptr, %Env)
+)")) {}
+};
+
+TEST_F(GCMetadataTest, Basic) {
+  LoopAnalysisManager LAM;
+  FunctionAnalysisManager FAM;
+  CGSCCAnalysisManager CGAM;
+  ModuleAnalysisManager MAM;
+  PassBuilder PB;
+  PB.registerModuleAnalyses(MAM);
+  PB.registerCGSCCAnalyses(CGAM);
+  PB.registerFunctionAnalyses(FAM);
+  PB.registerLoopAnalyses(LAM);
+  PB.crossRegisterProxies(LAM, FAM, CGAM, MAM);
+
+  ModulePassManager MPM;
+  FunctionPassManager FPM;
+  GCStrategyMap &StrategyMap = MAM.getResult<CollectorMetadataAnalysis>(*M);
+  for (auto &[GCName, Strategy] : StrategyMap)
+    EXPECT_EQ(GCName, Strategy->getName());
+  for (auto &[GCName, Strategy] : llvm::reverse(StrategyMap))
+    EXPECT_EQ(GCName, Strategy->getName());
+}
+
+} // namespace