[NFCI][Globals]Add GlobalObjects::updateSectionPrefix and change setSectionPrefix to handle empty strings

Author: mingmingl-llvm

llvm/include/llvm/IR/GlobalObject.h

@@ -121,9 +121,14 @@ class GlobalObject : public GlobalValue {
   /// appropriate default object file section.
   LLVM_ABI void setSection(StringRef S);
 
-  /// Set the section prefix for this global object.
+  /// Set the section prefix for this global object. If \p Prefix is empty,
+  /// the section prefix metadata will be cleared if it exists.
   LLVM_ABI void setSectionPrefix(StringRef Prefix);
 
+  /// If \p Prefix is different from existing prefix, update section prefix.
+  /// Returns true if an update happens and false otherwise.
+  LLVM_ABI bool updateSectionPrefix(StringRef Prefix);
+
   /// Get the section prefix for this global object.
   LLVM_ABI std::optional<StringRef> getSectionPrefix() const;
 

llvm/lib/IR/Globals.cpp

@@ -289,11 +289,28 @@ void GlobalObject::setSection(StringRef S) {
 }
 
 void GlobalObject::setSectionPrefix(StringRef Prefix) {
+  if (Prefix.empty()) {
+    setMetadata(LLVMContext::MD_section_prefix, nullptr);
+    return;
+  }
   MDBuilder MDB(getContext());
   setMetadata(LLVMContext::MD_section_prefix,
               MDB.createGlobalObjectSectionPrefix(Prefix));
 }
 
+bool GlobalObject::updateSectionPrefix(StringRef Prefix) {
+  auto MD = getMetadata(LLVMContext::MD_section_prefix);
+  StringRef ExistingPrefix; // Empty by default.
+  if (MD != nullptr)
+    ExistingPrefix = cast<MDString>(MD->getOperand(1))->getString();
+
+  if (ExistingPrefix != Prefix) {
+    setSectionPrefix(Prefix);
+    return true;
+  }
+  return false;
+}
+
 std::optional<StringRef> GlobalObject::getSectionPrefix() const {
   if (MDNode *MD = getMetadata(LLVMContext::MD_section_prefix)) {
     [[maybe_unused]] StringRef MDName =

llvm/unittests/IR/CMakeLists.txt

@@ -28,6 +28,7 @@ add_llvm_unittest(IRTests
   DominatorTreeBatchUpdatesTest.cpp
   DroppedVariableStatsIRTest.cpp
   FunctionTest.cpp
+  GlobalObjectTest.cpp
   PassBuilderCallbacksTest.cpp
   IRBuilderTest.cpp
   InstructionsTest.cpp

llvm/unittests/IR/GlobalObjectTest.cpp

@@ -0,0 +1,81 @@
+//===- GlobalObjectTest.cpp - Global object unit tests --------------------===//
+//
+// 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/IR/GlobalObject.h"
+#include "llvm-c/Core.h"
+#include "llvm/AsmParser/Parser.h"
+#include "llvm/IR/Module.h"
+#include "llvm/Support/SourceMgr.h"
+#include "gtest/gtest.h"
+#include "gmock/gmock.h"
+using namespace llvm;
+namespace {
+
+using testing::Eq;
+using testing::Optional;
+using testing::StrEq;
+
+static std::unique_ptr<Module> parseIR(LLVMContext &C, const char *IR) {
+  SMDiagnostic Err;
+  std::unique_ptr<Module> Mod = parseAssemblyString(IR, Err, C);
+  if (!Mod)
+    Err.print("GlobalObjectTests", errs());
+  return Mod;
+}
+
+static LLVMContext C;
+static std::unique_ptr<Module> M;
+
+class GlobalObjectTest : public testing::Test {
+public:
+static void SetUpTestSuite()  {
+   
+    M = parseIR(C, R"(
+@foo = global i32 3, !section_prefix !0
+@bar = global i32 0
+
+!0 = !{!"section_prefix", !"hot"}
+)");
+
+  }
+
+};
+
+TEST_F(GlobalObjectTest, SectionPrefix) {
+    GlobalVariable* Foo = M->getGlobalVariable("foo");
+    
+    // Initial section prefix is hot.
+    ASSERT_NE(Foo, nullptr);
+    ASSERT_THAT(Foo->getSectionPrefix(), Optional(StrEq("hot")));
+
+    // No actual update.
+    EXPECT_FALSE(Foo->updateSectionPrefix("hot"));
+
+    // Update prefix from hot to unlikely.
+    Foo->setSectionPrefix("unlikely");
+    EXPECT_THAT(Foo->getSectionPrefix(), Optional(StrEq("unlikely")));
+    
+    // Update prefix to empty is the same as clear.
+    Foo->setSectionPrefix("");
+    EXPECT_THAT(Foo->getSectionPrefix(), Eq(std::nullopt));
+
+    GlobalVariable* Bar = M->getGlobalVariable("bar");
+
+    // Initial section prefix is empty.
+    ASSERT_NE(Bar, nullptr);
+    ASSERT_THAT(Bar->getSectionPrefix(),  Eq(std::nullopt));
+ 
+    // No actual update.
+    EXPECT_FALSE(Bar->updateSectionPrefix(""));
+
+    // Update from empty to hot.
+    EXPECT_TRUE(Bar->updateSectionPrefix("hot"));
+    EXPECT_THAT(Bar->getSectionPrefix(), Optional(StrEq("hot")));
+}
+
+} // namespace