[Metadata] Return the valid DebugLoc if one of them is null with -pick-merged-source-locations.

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

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) {