1 files changed, 91 insertions, 0 deletions
diff --git a/backend/unittests/CodeGen/VirtRegMapTest.cpp b/backend/unittests/CodeGen/VirtRegMapTest.cpp
new file mode 100644
index 0000000..a3bdfc1
--- /dev/null
+++ b/backend/unittests/CodeGen/VirtRegMapTest.cpp
@@ -0,0 +1,91 @@
+// SPDX-License-Identifier: GPL-3.0-or-later
+
+#include "CodeGen/RegisterAllocationTestSupport.h"
+#include "rcc/CodeGen/LiveRangeEdit.h"
+
+#include <iostream>
+
+namespace {
+
+using namespace rcc;
+using namespace rcc::test;
+using namespace rcc::test::ra;
+
+void testSplitMetadata(TestContext &Context, const TargetDescription &Target) {
+  auto Fixture = createFunction("virt-map", 10);
+  if (!Fixture)
+    return;
+  auto Parent = createVReg(*Fixture->Function);
+  if (!Parent)
+    return;
+  VirtRegMap Mapping(*Fixture->Function);
+  LiveRangeEdit Edit(*Fixture->Function, Target, Mapping);
+  Context.expect(Mapping.add(*Parent).has_value(), "VirtRegMap must accept a known vreg");
+  auto Slot = Edit.getOrCreateSpillSlot(*Parent);
+  auto Child = Edit.createFrom(*Parent, false);
+  Context.expect(
+      Slot && Child && Mapping.verify().has_value() && Mapping.stackSlot(*Child) == Slot &&
+          Mapping.entry(*Child)->SplitParent == *Parent && Mapping.entry(*Child)->Origin == *Parent,
+      "LiveRangeEdit must preserve split origin and spill metadata");
+  if (!Child)
+    return;
+  const PhysicalRegister G0 = physical(TestRegister::G0);
+  Context.expect(Mapping.assign(*Child, G0).has_value() && Mapping.physicalRegister(*Child) == G0 &&
+                     Mapping.entry(*Child)->State == VirtRegAllocationState::Assigned,
+                 "VirtRegMap must transition a split child to an assigned state");
+}
+
+void testRewriteAndAtomicValidation(TestContext &Context, const TargetDescription &Target) {
+  auto Fixture = createFunction("virt-rewrite", 11);
+  if (!Fixture)
+    return;
+  auto Value = createVReg(*Fixture->Function);
+  if (!Value ||
+      !buildDef(*Fixture->Function, Target, *Fixture->Block, *Value, 1, TestSubReg::Low32) ||
+      !buildUse(*Fixture->Function, Target, *Fixture->Block, *Value, TestSubReg::Low32) ||
+      !prepareForAllocation(*Fixture->Function, Target))
+    return;
+  VirtRegMap Mapping(*Fixture->Function);
+  const PhysicalRegister G0 = physical(TestRegister::G0);
+  Context.expect(Mapping.add(*Value).has_value() && Mapping.assign(*Value, G0).has_value(),
+                 "VirtRegMap must accept a physical assignment before rewrite");
+  auto Rewritten = rewriteVirtualRegisters(*Fixture->Function, Target, Mapping);
+  const MachineOperand *Def = Fixture->Block->firstInstruction()->operandAt(0);
+  Context.expect(
+      Rewritten && *Rewritten && Def != nullptr && Def->registerValue() &&
+          Def->registerValue() == Register::fromPhysical(physical(TestRegister::G0Low32)) &&
+          Def->subReg().value() == 0 &&
+          std::ranges::none_of(Fixture->Function->registerInfo().virtualRegisters(),
+                               [](const VirtualRegisterInfo &Info) { return !Info.IsTombstone; }),
+      "VirtualRegisterRewrite must project and clear subregister operands");
+
+  auto InvalidFixture = createFunction("virt-rewrite-invalid", 12);
+  if (!InvalidFixture)
+    return;
+  auto Invalid = createVReg(*InvalidFixture->Function);
+  if (!Invalid ||
+      !buildDef(*InvalidFixture->Function, Target, *InvalidFixture->Block, *Invalid, 2) ||
+      !buildUse(*InvalidFixture->Function, Target, *InvalidFixture->Block, *Invalid) ||
+      !prepareForAllocation(*InvalidFixture->Function, Target))
+    return;
+  VirtRegMap InvalidMapping(*InvalidFixture->Function);
+  const PhysicalRegister F0 = physical(TestRegister::F0);
+  if (!InvalidMapping.add(*Invalid) || !InvalidMapping.assign(*Invalid, F0))
+    return;
+  auto Failed = rewriteVirtualRegisters(*InvalidFixture->Function, Target, InvalidMapping);
+  Context.expect(!Failed && Failed.error().code() == ErrorCode::RegisterAllocationFailed,
+                 "VirtualRegisterRewrite must reject class-incompatible mappings atomically");
+}
+
+} // namespace
+
+int main() {
+  TestContext Context;
+  auto Target = createTarget();
+  Context.expect(Target.has_value(), "VirtRegMap test target must be constructible");
+  if (!Target)
+    return Context.result();
+  testSplitMetadata(Context, *Target);
+  testRewriteAndAtomicValidation(Context, *Target);
+  return Context.result();
+}