From a8b2946e939ee87d1a627b49ca86db41728cc22a Mon Sep 17 00:00:00 2001
From: Richard-Qin-X <richard.qin.001@gmail.com>
Date: Tue, 11 Aug 2026 18:53:06 +0800
Subject: [RISCV] Preserve folded constant placement

---
 backend/lib/Target/RISCV/RISCVCodeGen.cpp     | 18 ++++++++++++++++++
 backend/unittests/Target/RISCVCodeGenTest.cpp | 23 ++++++++++++++++++++---
 2 files changed, 38 insertions(+), 3 deletions(-)

(limited to 'backend')

diff --git a/backend/lib/Target/RISCV/RISCVCodeGen.cpp b/backend/lib/Target/RISCV/RISCVCodeGen.cpp
index 2f7eca9..e790edf 100644
--- a/backend/lib/Target/RISCV/RISCVCodeGen.cpp
+++ b/backend/lib/Target/RISCV/RISCVCodeGen.cpp
@@ -128,6 +128,7 @@ extendIntegerImmediate(const IntegerImmediate &Value, std::uint32_t DestinationW
                           "RISC-V constant extension requires virtual-register operands");
   }
   const Register SourceRegister = *Source;
+  const Register DestinationRegister = *Destination;
   const MachineOperand *SourceDefinition = Function.registerInfo().getUniqueVRegDef(SourceRegister);
   const MachineInstr *SourceDef = SourceDefinition == nullptr ? nullptr : SourceDefinition->parent();
   if (SourceDef == nullptr || SourceDef->opcode() != opcode(GenericOpcode::G_CONSTANT))
@@ -150,6 +151,23 @@ extendIntegerImmediate(const IntegerImmediate &Value, std::uint32_t DestinationW
   if (!Extended)
     return std::unexpected(Extended.error());
 
+  if (Function.registerInfo().hasOneUse(SourceRegister)) {
+    MachineInstr *Def = const_cast<MachineInstr *>(SourceDef);
+    if (auto Replaced = Def->operandAt(0)->setRegister(DestinationRegister); !Replaced)
+      return std::unexpected(Replaced.error());
+    auto Erased = Def->eraseOperand(1);
+    if (!Erased)
+      return std::unexpected(Erased.error());
+    auto Added = Def->addOperand(MachineOperand::integerImmediate(std::move(*Extended)));
+    if (!Added)
+      return std::unexpected(Added.error());
+    MachineBasicBlock *Parent = Instruction.parent();
+    if (Parent == nullptr)
+      return makeUnexpected(ErrorCode::VerificationError,
+                            "RISC-V constant extension is detached during folding");
+    return Parent->removeInstruction(Instruction);
+  }
+
   auto SetOpcode = Instruction.setOpcode(opcode(GenericOpcode::G_CONSTANT));
   if (!SetOpcode)
     return std::unexpected(SetOpcode.error());
diff --git a/backend/unittests/Target/RISCVCodeGenTest.cpp b/backend/unittests/Target/RISCVCodeGenTest.cpp
index 9665373..40f84f2 100644
--- a/backend/unittests/Target/RISCVCodeGenTest.cpp
+++ b/backend/unittests/Target/RISCVCodeGenTest.cpp
@@ -1090,11 +1090,16 @@ void testIntegerCompareReusesConstantExtension(TestContext &Context) {
     return;
   auto Constant = Builder.buildInstr(rcc::opcode(rcc::GenericOpcode::G_CONSTANT));
   if (!Constant || !Constant->addDef(*Right) ||
-      !Constant->addIntegerImmediate(std::move(*Immediate)) || !Constant->build())
+      !Constant->addIntegerImmediate(std::move(*Immediate)))
+    return;
+  auto BuiltConstant = Constant->build();
+  if (!BuiltConstant)
     return;
   auto Extension = Builder.buildInstr(rcc::opcode(rcc::GenericOpcode::G_ANYEXT));
-  if (!Extension || !Extension->addDef(*WideRight) || !Extension->addUse(*Right) ||
-      !Extension->build())
+  if (!Extension || !Extension->addDef(*WideRight) || !Extension->addUse(*Right))
+    return;
+  auto BuiltExtension = Extension->build();
+  if (!BuiltExtension)
     return;
 
   auto Action = rcc::riscv::legalizerInfo().classify(*Compare, *Function, *Target);
@@ -1107,6 +1112,18 @@ void testIntegerCompareReusesConstantExtension(TestContext &Context) {
                  "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");
+
+  auto ExtensionAction = rcc::riscv::legalizerInfo().classify(**BuiltExtension, *Function, *Target);
+  auto Folded = ExtensionAction ? rcc::riscv::legalizerInfo().lower(**BuiltExtension, *Function,
+                                                                    *Target, *ExtensionAction)
+                                : rcc::Expected<void>(std::unexpected(ExtensionAction.error()));
+  const rcc::Register *FoldedResult = (*BuiltConstant)->operandAt(0)->getIf<rcc::Register>();
+  Context.expect(ExtensionAction && *ExtensionAction == rcc::LegalizationAction::Custom && Folded &&
+                     (*BuiltExtension)->parent() == nullptr && FoldedResult != nullptr &&
+                     *FoldedResult == *WideRight &&
+                     Function->firstBlock()->firstInstruction() == *BuiltConstant &&
+                     Function->registerInfo().verifyDefUseChains(),
+                 "single-use constant extension folding must preserve the source layout position");
 }
 
 void testAtomicSelection(TestContext &Context) {
-- 
cgit v1.2.3