2 files changed, 71 insertions, 0 deletions

diff --git a/backend/lib/Target/RISCV/RISCVCodeGen.cpp b/backend/lib/Target/RISCV/RISCVCodeGen.cpp
index b856476..2f7eca9 100644
--- a/backend/lib/Target/RISCV/RISCVCodeGen.cpp
+++ b/backend/lib/Target/RISCV/RISCVCodeGen.cpp
@@ -698,6 +698,34 @@ lowerAtomicCompareExchangeWithSuccess(MachineInstr &Instruction, MachineFunction
   const MachineOperand *Definition = Function.registerInfo().getUniqueVRegDef(Source);
   MachineInstr *Def =
       Definition == nullptr ? nullptr : const_cast<MachineInstr *>(Definition->parent());
+  std::optional<Register> ExistingExtension;
+  bool AmbiguousExtension = false;
+  auto &RegisterInfo = Function.registerInfo();
+  for (MachineOperand &Use : RegisterInfo.useOperands(Source)) {
+    MachineInstr *User = Use.parent();
+    if (User == nullptr || User == &Instruction || User->parent() == nullptr ||
+        (User->opcode() != opcode(GenericOpcode::G_ANYEXT) &&
+         User->opcode() != opcode(GenericOpcode::G_SEXT))) {
+      continue;
+    }
+    const MachineOperand *ExtendedOperand = User->operandAt(0);
+    const Register *Extended =
+        ExtendedOperand == nullptr ? nullptr : ExtendedOperand->getIf<Register>();
+    const std::optional<LLT> ExtendedType =
+        ExtendedOperand == nullptr ? std::nullopt : registerType(Function, *ExtendedOperand);
+    if (Extended == nullptr || !ExtendedType || *ExtendedType != WideType ||
+        (Def == nullptr || (Def->opcode() != opcode(GenericOpcode::G_CONSTANT) &&
+                            User->opcode() != opcode(GenericOpcode::G_SEXT)))) {
+      continue;
+    }
+    if (ExistingExtension) {
+      AmbiguousExtension = true;
+      break;
+    }
+    ExistingExtension = *Extended;
+  }
+  if (ExistingExtension && !AmbiguousExtension)
+    return *ExistingExtension;
   if (Def != nullptr && Def->opcode() == opcode(GenericOpcode::G_CONSTANT) &&
       singleRegisterUse(Function, Source) == &Instruction) {
     auto Widened = widenConstantToRV64(*Def, Function);
diff --git a/backend/unittests/Target/RISCVCodeGenTest.cpp b/backend/unittests/Target/RISCVCodeGenTest.cpp
index 7c0ec33..9665373 100644
--- a/backend/unittests/Target/RISCVCodeGenTest.cpp
+++ b/backend/unittests/Target/RISCVCodeGenTest.cpp
@@ -1067,6 +1067,48 @@ void testConstantExtensionPreservesLiveSource(TestContext &Context) {
                  "constant-extension folding must preserve a source definition with live uses");
 }
 
+void testIntegerCompareReusesConstantExtension(TestContext &Context) {
+  auto Target = rcc::riscv::createTargetDescription();
+  auto I32 = rcc::LLT::makeAnyScalar(32);
+  auto I64 = rcc::LLT::makeAnyScalar(64);
+  if (!Target || !I32 || !I64)
+    return;
+  auto Function =
+      createCompareFunction(Context, *Target, rcc::GenericOpcode::G_ICMP, *I32, false,
+                            rcc::IntegerPredicate::Equal, rcc::FloatingPredicate::AlwaysFalse);
+  if (!Function)
+    return;
+  rcc::MachineInstr *Compare = Function->firstBlock()->firstInstruction();
+  const rcc::Register *Right = Compare->operandAt(3)->getIf<rcc::Register>();
+  auto WideRight = Function->registerInfo().createVirtualRegister(*I64);
+  auto Immediate = rcc::IntegerImmediate::fromUnsigned(32, 0);
+  if (Right == nullptr || !WideRight || !Immediate)
+    return;
+
+  rcc::MachineIRBuilder Builder(*Function, *Target);
+  if (!Builder.setInsertPointBefore(*Compare))
+    return;
+  auto Constant = Builder.buildInstr(rcc::opcode(rcc::GenericOpcode::G_CONSTANT));
+  if (!Constant || !Constant->addDef(*Right) ||
+      !Constant->addIntegerImmediate(std::move(*Immediate)) || !Constant->build())
+    return;
+  auto Extension = Builder.buildInstr(rcc::opcode(rcc::GenericOpcode::G_ANYEXT));
+  if (!Extension || !Extension->addDef(*WideRight) || !Extension->addUse(*Right) ||
+      !Extension->build())
+    return;
+
+  auto Action = rcc::riscv::legalizerInfo().classify(*Compare, *Function, *Target);
+  auto Lowered = Action ? rcc::riscv::legalizerInfo().lower(*Compare, *Function, *Target, *Action)
+                        : rcc::Expected<void>(std::unexpected(Action.error()));
+  const rcc::Register *LegalRight = Compare->operandAt(3)->getIf<rcc::Register>();
+  Context.expect(Action && *Action == rcc::LegalizationAction::Custom && Lowered,
+                 "RV64 integer comparison extension-reuse fixture must legalize");
+  Context.expect(LegalRight != nullptr && *LegalRight == *WideRight,
+                 "RV64 integer comparison legalization must reuse an existing constant extension");
+  Context.expect(Function->registerInfo().verifyDefUseChains().has_value(),
+                 "integer comparison extension reuse must preserve exact def-use chains");
+}
+
 void testAtomicSelection(TestContext &Context) {
   auto Target = rcc::riscv::createTargetDescription();
   if (!Target)
@@ -1346,6 +1388,7 @@ int main() {
   testOverflowSelection(Context);
   testAtomicCompareExchangeLegalization(Context);
   testConstantExtensionPreservesLiveSource(Context);
+  testIntegerCompareReusesConstantExtension(Context);
   testAtomicSelection(Context);
   testUnsupportedTypes(Context);
   testMemoryOffsetSelection(Context);