[llvm] [Metadata] Return the valid DebugLoc if one of them is null with -pick-merged-source-locations. (PR #138148)

[Snehasish Kumar via llvm-commits]

https://github.com/snehasish updated https://github.com/llvm/llvm-project/pull/138148

>From a37c5ea512677523bb6e0ee36bee094da1095f5c Mon Sep 17 00:00:00 2001
From: Snehasish Kumar <snehasishk at google.com>
Date: Wed, 30 Apr 2025 23:50:45 +0000
Subject: [PATCH 1/2] [Metadata] Return the valid DebugLoc if one of them is
 null with -pick-merged-source-locations.

Previously when getMergedLocation was passed nullptr as one of the
parameters we returned nullptr. Change that behaviour to instead return
a valid DebugLoc. This is beneficial for binaries from which SamplePGO
profiles are obtained.
---
 llvm/lib/IR/DebugInfoMetadata.cpp  | 16 ++++++++++------
 llvm/unittests/IR/MetadataTest.cpp | 23 +++++++++++++++++++++++
 2 files changed, 33 insertions(+), 6 deletions(-)

diff --git a/llvm/lib/IR/DebugInfoMetadata.cpp b/llvm/lib/IR/DebugInfoMetadata.cpp
index 4f1b9e836120e..e0cc95ac0ca79 100644
--- a/llvm/lib/IR/DebugInfoMetadata.cpp
+++ b/llvm/lib/IR/DebugInfoMetadata.cpp
@@ -33,13 +33,13 @@ namespace llvm {
 cl::opt<bool> EnableFSDiscriminator(
     "enable-fs-discriminator", cl::Hidden,
     cl::desc("Enable adding flow sensitive discriminators"));
-} // namespace llvm
 
 // When true, preserves line and column number by picking one of the merged
 // location info in a deterministic manner to assist sample based PGO.
-static cl::opt<bool> PickMergedSourceLocations(
+cl::opt<bool> PickMergedSourceLocations(
     "pick-merged-source-locations", cl::init(false), cl::Hidden,
     cl::desc("Preserve line and column number when merging locations."));
+} // namespace llvm
 
 uint32_t DIType::getAlignInBits() const {
   return (getTag() == dwarf::DW_TAG_LLVM_ptrauth_type ? 0 : SubclassData32);
@@ -218,17 +218,18 @@ struct ScopeLocationsMatcher {
 };
 
 DILocation *DILocation::getMergedLocation(DILocation *LocA, DILocation *LocB) {
-  if (!LocA || !LocB)
-    return nullptr;
-
   if (LocA == LocB)
     return LocA;
 
   // For some use cases (SamplePGO), it is important to retain distinct source
   // locations. When this flag is set, we choose arbitrarily between A and B,
   // rather than computing a merged location using line 0, which is typically
-  // not useful for PGO.
+  // not useful for PGO. If one of them is null, then try to return one which is
+  // valid.
   if (PickMergedSourceLocations) {
+    if (!LocA || !LocB)
+      return LocA ? LocA : LocB;
+
     auto A = std::make_tuple(LocA->getLine(), LocA->getColumn(),
                              LocA->getDiscriminator(), LocA->getFilename(),
                              LocA->getDirectory());
@@ -238,6 +239,9 @@ DILocation *DILocation::getMergedLocation(DILocation *LocA, DILocation *LocB) {
     return A < B ? LocA : LocB;
   }
 
+  if (!LocA || !LocB)
+    return nullptr;
+
   LLVMContext &C = LocA->getContext();
 
   using LocVec = SmallVector<const DILocation *>;
diff --git a/llvm/unittests/IR/MetadataTest.cpp b/llvm/unittests/IR/MetadataTest.cpp
index e710f7845df5e..e0d96aab5bfa3 100644
--- a/llvm/unittests/IR/MetadataTest.cpp
+++ b/llvm/unittests/IR/MetadataTest.cpp
@@ -25,6 +25,10 @@
 #include <optional>
 using namespace llvm;
 
+namespace llvm {
+  extern cl::opt<bool> PickMergedSourceLocations;
+} // namespace llvm
+
 namespace {
 
 TEST(ContextAndReplaceableUsesTest, FromContext) {
@@ -1444,6 +1448,25 @@ TEST_F(DILocationTest, Merge) {
     auto *M2 = DILocation::getMergedLocation(A2, B);
     EXPECT_EQ(M1, M2);
   }
+
+  {
+    // If PickMergedSourceLocation is enabled, when one source location is null
+    // we should return the valid location.
+    PickMergedSourceLocations = true;
+    auto *A = DILocation::get(Context, 2, 7, N);
+    auto *M1 = DILocation::getMergedLocation(A, nullptr);
+    ASSERT_NE(nullptr, M1);
+    EXPECT_EQ(2u, M1->getLine());
+    EXPECT_EQ(7u, M1->getColumn());
+    EXPECT_EQ(N,  M1->getScope());
+
+    auto *M2 = DILocation::getMergedLocation(nullptr, A);
+    ASSERT_NE(nullptr, M2);
+    EXPECT_EQ(2u, M2->getLine());
+    EXPECT_EQ(7u, M2->getColumn());
+    EXPECT_EQ(N,  M2->getScope());
+    PickMergedSourceLocations = false;
+  }
 }
 
 TEST_F(DILocationTest, getDistinct) {

>From f62b0f07f8d7833bbc270aad9b982366cb39050f Mon Sep 17 00:00:00 2001
From: Snehasish Kumar <snehasishk at google.com>
Date: Thu, 1 May 2025 15:32:22 +0000
Subject: [PATCH 2/2] Fix formatting.

---
 llvm/unittests/IR/MetadataTest.cpp | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/llvm/unittests/IR/MetadataTest.cpp b/llvm/unittests/IR/MetadataTest.cpp
index e0d96aab5bfa3..45860d60d8f4c 100644
--- a/llvm/unittests/IR/MetadataTest.cpp
+++ b/llvm/unittests/IR/MetadataTest.cpp
@@ -26,7 +26,7 @@
 using namespace llvm;
 
 namespace llvm {
-  extern cl::opt<bool> PickMergedSourceLocations;
+extern cl::opt<bool> PickMergedSourceLocations;
 } // namespace llvm
 
 namespace {
@@ -1458,13 +1458,13 @@ TEST_F(DILocationTest, Merge) {
     ASSERT_NE(nullptr, M1);
     EXPECT_EQ(2u, M1->getLine());
     EXPECT_EQ(7u, M1->getColumn());
-    EXPECT_EQ(N,  M1->getScope());
+    EXPECT_EQ(N, M1->getScope());
 
     auto *M2 = DILocation::getMergedLocation(nullptr, A);
     ASSERT_NE(nullptr, M2);
     EXPECT_EQ(2u, M2->getLine());
     EXPECT_EQ(7u, M2->getColumn());
-    EXPECT_EQ(N,  M2->getScope());
+    EXPECT_EQ(N, M2->getScope());
     PickMergedSourceLocations = false;
   }
 }