Index: llvm/trunk/lib/CodeGen/RegisterCoalescer.cpp
===================================================================
--- llvm/trunk/lib/CodeGen/RegisterCoalescer.cpp
+++ llvm/trunk/lib/CodeGen/RegisterCoalescer.cpp
@@ -1910,6 +1910,13 @@
     }
     LI.removeEmptySubRanges();
   }
+
+  // CP.getSrcReg()'s live interval has been merged into CP.getDstReg's live
+  // interval. Since CP.getSrcReg() is in ToBeUpdated set and its live interval
+  // is not up-to-date, need to update the merged live interval here.
+  if (ToBeUpdated.count(CP.getSrcReg()))
+    ShrinkMainRange = true;
+
   if (ShrinkMainRange) {
     LiveInterval &LI = LIS->getInterval(CP.getDstReg());
     shrinkToUses(&LI);
Index: llvm/trunk/test/CodeGen/X86/late-remat-update-2.mir
===================================================================
--- llvm/trunk/test/CodeGen/X86/late-remat-update-2.mir
+++ llvm/trunk/test/CodeGen/X86/late-remat-update-2.mir
@@ -0,0 +1,63 @@
+# RUN: llc -mtriple=x86_64-- -run-pass=simple-register-coalescing -run-pass=regallocbasic -run-pass=virtregrewriter -late-remat-update-threshold=0 %s -o - | FileCheck %s
+#
+# PR40061: %t2 = %t1 is rematerialized and %t1 is added into toBeUpdated set
+# to postpone its live interval update. After the rematerialization, the live
+# interval of %t1 is larger than necessary. Then %t1 is merged into %t3 and %t1
+# gets removed. After the merge, %t3 contains live interval larger than
+# necessary. Because %t3 is not in toBeUpdated set so its live interval is not
+# updated after register coalescing, and it will break some assumption in
+# regalloc. This test wants to check the live interval is up-to-date after
+# register coalescing.
+#
+# To prevent the test from taking effect only in assert enabled mode, we want
+# to achieve the test goal without dumping regalloc trace. We add strong hint
+# to allocate both %t1 and %t2 to $rax register. If the %t1's live interval is
+# not shrinked properly after register coalescing, %t1 and %t2 will not be
+# both allocated to $rax because of inference, and we utilize the fact to
+# achieve the test goal. But note that the assumption only holds when we use
+# regallocbasic instead of greedy because greedy can update the live interval
+# in the process of splitting.
+#
+# CHECK-LABEL: name: foo
+# CHECK: bb.0.entry:
+# CHECK: $rax = MOV64ri32 -11
+# CHECK: bb.1:
+# CHECK: $rax = MOV64ri32 -11
+# CHECK: $rax = ADD64ri8 killed renamable $rax, 5
+# CHECK: CMP64ri8 renamable $rax
+# CHECK: RET 0, $rax
+# CHECK: bb.2:
+# CHECK: $rax = ADD64ri8 killed renamable $rax, 10
+# CHECK: bb.3:
+# CHECK: RET 0, $rax
+---
+name:            foo
+body:             |
+  bb.0.entry:
+    successors: %bb.1(0x15555555), %bb.2(0x6aaaaaab)
+
+    %t1:gr64 = MOV64ri32 -11
+    CMP64ri8 %t1, 1, implicit-def $eflags
+    JE_1 %bb.2, implicit killed $eflags
+    JMP_1 %bb.1
+
+  bb.1:
+    successors: %bb.1(0x80000000)
+
+    %t2:gr64 = COPY %t1
+    %t2:gr64 = ADD64ri8 %t2, 5, implicit-def $eflags
+    $rax = COPY %t2
+    CMP64ri8 %t2, 1, implicit-def $eflags
+    JE_1 %bb.1, implicit killed $eflags
+    RET 0, $rax
+
+  bb.2:
+    successors: %bb.3(0x80000000)
+    %t3:gr64 = COPY %t1
+    %t3:gr64 = ADD64ri8 %t3, 10, implicit-def $eflags
+
+  bb.3:
+    $rax = COPY %t3
+    RET 0, $rax
+
+...