[MDBuilder] `mergeCallbackEncodings` fails due to inspecting the wrong node (#92466)

Given the following metadata as, with `!6` as `ExistingCallbacks` and
`!8` as `NewCB`:
```
!6 = !{!7}
!7 = !{i64 0, i1 false}
!8 = !{i64 2, i64 3, i1 false}
```
The merge function should add `!8` to the list of `!6`, i.e. `!6 =
!{!7,!8}`. However, at the moment the check if this is legal, tries to
interpret `!7` as integer instead of it's operand.

2 files changed, 46 insertions, 7 deletions

diff --git a/llvm/lib/IR/MDBuilder.cpp b/llvm/lib/IR/MDBuilder.cpp
index 0bf41d7..bd68db3 100644
--- a/llvm/lib/IR/MDBuilder.cpp
+++ b/llvm/lib/IR/MDBuilder.cpp
@@ -86,9 +86,8 @@ MDNode *MDBuilder::createFunctionEntryCount(
 }
 
 MDNode *MDBuilder::createFunctionSectionPrefix(StringRef Prefix) {
-  return MDNode::get(Context,
-                     {createString("function_section_prefix"),
-                      createString(Prefix)});
+  return MDNode::get(
+      Context, {createString("function_section_prefix"), createString(Prefix)});
 }
 
 MDNode *MDBuilder::createRange(const APInt &Lo, const APInt &Hi) {
@@ -148,9 +147,10 @@ MDNode *MDBuilder::mergeCallbackEncodings(MDNode *ExistingCallbacks,
   for (unsigned u = 0; u < NumExistingOps; u++) {
     Ops[u] = ExistingCallbacks->getOperand(u);
 
-    auto *OldCBCalleeIdxAsCM = cast<ConstantAsMetadata>(Ops[u]);
+    auto *OldCBCalleeIdxAsCM =
+        cast<ConstantAsMetadata>(cast<MDNode>(Ops[u])->getOperand(0));
     uint64_t OldCBCalleeIdx =
-      cast<ConstantInt>(OldCBCalleeIdxAsCM->getValue())->getZExtValue();
+        cast<ConstantInt>(OldCBCalleeIdxAsCM->getValue())->getZExtValue();
     (void)OldCBCalleeIdx;
     assert(NewCBCalleeIdx != OldCBCalleeIdx &&
            "Cannot map a callback callee index twice!");
@@ -339,8 +339,8 @@ MDNode *MDBuilder::createMutableTBAAAccessTag(MDNode *Tag) {
 
 MDNode *MDBuilder::createIrrLoopHeaderWeight(uint64_t Weight) {
   Metadata *Vals[] = {
-    createString("loop_header_weight"),
-    createConstant(ConstantInt::get(Type::getInt64Ty(Context), Weight)),
+      createString("loop_header_weight"),
+      createConstant(ConstantInt::get(Type::getInt64Ty(Context), Weight)),
   };
   return MDNode::get(Context, Vals);
 }
diff --git a/llvm/unittests/IR/MDBuilderTest.cpp b/llvm/unittests/IR/MDBuilderTest.cpp
index 2b5ab81..4656c70 100644
--- a/llvm/unittests/IR/MDBuilderTest.cpp
+++ b/llvm/unittests/IR/MDBuilderTest.cpp
@@ -127,4 +127,43 @@ TEST_F(MDBuilderTest, createPCSections) {
   EXPECT_EQ(mdconst::extract<ConstantInt>(Aux->getOperand(1))->getValue(),
             C2->getValue());
 }
+TEST_F(MDBuilderTest, createCallbackAndMerge) {
+  MDBuilder MDHelper(Context);
+  auto *CB1 = MDHelper.createCallbackEncoding(0, {1, -1}, false);
+  auto *CB2 = MDHelper.createCallbackEncoding(2, {-1}, false);
+  ASSERT_EQ(CB1->getNumOperands(), 4U);
+  ASSERT_TRUE(isa<ConstantAsMetadata>(CB1->getOperand(0)));
+  ASSERT_TRUE(isa<ConstantAsMetadata>(CB1->getOperand(1)));
+  ASSERT_TRUE(isa<ConstantAsMetadata>(CB1->getOperand(2)));
+  ASSERT_TRUE(isa<ConstantAsMetadata>(CB1->getOperand(3)));
+  EXPECT_EQ(mdconst::extract<ConstantInt>(CB1->getOperand(0))->getValue(), 0);
+  EXPECT_EQ(mdconst::extract<ConstantInt>(CB1->getOperand(1))->getValue(), 1);
+  EXPECT_EQ(mdconst::extract<ConstantInt>(CB1->getOperand(2))->getValue(), -1);
+  EXPECT_EQ(mdconst::extract<ConstantInt>(CB1->getOperand(3))->getValue(),
+            false);
+  ASSERT_EQ(CB2->getNumOperands(), 3U);
+  ASSERT_TRUE(isa<ConstantAsMetadata>(CB2->getOperand(0)));
+  ASSERT_TRUE(isa<ConstantAsMetadata>(CB2->getOperand(1)));
+  ASSERT_TRUE(isa<ConstantAsMetadata>(CB2->getOperand(2)));
+  EXPECT_EQ(mdconst::extract<ConstantInt>(CB2->getOperand(0))->getValue(), 2);
+  EXPECT_EQ(mdconst::extract<ConstantInt>(CB2->getOperand(1))->getValue(), -1);
+  EXPECT_EQ(mdconst::extract<ConstantInt>(CB2->getOperand(2))->getValue(),
+            false);
+  auto *CBList = MDNode::get(Context, {CB1, CB2});
+  auto *CB3 = MDHelper.createCallbackEncoding(4, {5}, false);
+  auto *NewCBList = MDHelper.mergeCallbackEncodings(CBList, CB3);
+  ASSERT_EQ(NewCBList->getNumOperands(), 3U);
+  EXPECT_TRUE(NewCBList->getOperand(0) == CB1);
+  EXPECT_TRUE(NewCBList->getOperand(1) == CB2);
+  EXPECT_TRUE(NewCBList->getOperand(2) == CB3);
+
+  ASSERT_EQ(CB3->getNumOperands(), 3U);
+  ASSERT_TRUE(isa<ConstantAsMetadata>(CB3->getOperand(0)));
+  ASSERT_TRUE(isa<ConstantAsMetadata>(CB3->getOperand(1)));
+  ASSERT_TRUE(isa<ConstantAsMetadata>(CB3->getOperand(2)));
+  EXPECT_EQ(mdconst::extract<ConstantInt>(CB3->getOperand(0))->getValue(), 4);
+  EXPECT_EQ(mdconst::extract<ConstantInt>(CB3->getOperand(1))->getValue(), 5);
+  EXPECT_EQ(mdconst::extract<ConstantInt>(CB3->getOperand(2))->getValue(),
+            false);
+}
 } // namespace