2 files changed, 33 insertions, 6 deletions

diff --git a/llvm/lib/IR/DebugInfoMetadata.cpp b/llvm/lib/IR/DebugInfoMetadata.cpp
index bfe956d..e8a58f1 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);
@@ -228,17 +228,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());
@@ -248,6 +249,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 e710f78..45860d6 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) {