shader: hand IR uses over at once on replacement and skip use lists at program teardown

src/graphics/shader/recompiler/ir/Program.cpp

@@ -118,15 +118,14 @@ ResourcePlan& ResourcePlan::operator=(ResourcePlan&& other) noexcept {
 }
 
 Program::~Program() {
-	// Planning expressions can refer to block values but outlive block storage in the base class.
+	// Planning expressions can refer to block values but outlive block storage in the base class,
+	// and values cross block boundaries: detach everything before any storage is destroyed.
 	for (auto& inst: value_storage) {
-		inst.Invalidate();
+		inst.DropArgs();
 	}
-	// Values may cross block boundaries. Detach all arguments before any block starts destroying
-	// its instruction storage so reverse-use links always point to live definitions.
 	for (auto* block: blocks) {
 		for (auto& inst: *block) {
-			inst.Invalidate();
+			inst.DropArgs();
 		}
 	}
 }

src/graphics/shader/recompiler/ir/Value.cpp

@@ -186,12 +186,17 @@ size_t Inst::NumPhiBlocks() const {
 	return num_args == PhiArity ? phi_args.size() : 0;
 }
 
-Value Inst::Arg(size_t index) const {
+template <typename Self>
-	EXIT_IF(index >= NumArgs());
+auto& Inst::ArgSlot(Self& self, size_t index) {
-	if (num_args <= InlineArity) {
+	EXIT_IF(index >= self.NumArgs());
-		return fixed_args[index];
+	if (self.num_args <= InlineArity) {
+		return self.fixed_args[index];
 	}
-	return num_args == PhiArity ? phi_args[index].second : large_args[index];
+	return self.num_args == PhiArity ? self.phi_args[index].second : self.large_args[index];
+}
+
+Value Inst::Arg(size_t index) const {
+	return ArgSlot(*this, index);
 }
 
 Block* Inst::PhiBlock(size_t index) const {
@@ -212,17 +217,11 @@ void Inst::SetParent(Block* block) {
 }
 
 void Inst::SetArg(size_t index, Value value) {
-	const auto old = Arg(index);
+	auto& arg = ArgSlot(*this, index);
-	if (auto* old_inst = old.TryInstruction(); old_inst != nullptr) {
+	if (auto* old_inst = arg.TryInstruction(); old_inst != nullptr) {
 		RemoveUse(old_inst, index);
 	}
-	if (num_args <= InlineArity) {
+	arg = value;
-		fixed_args[index] = value;
-	} else if (num_args == PhiArity) {
-		phi_args[index].second = value;
-	} else {
-		large_args[index] = value;
-	}
 	if (auto* new_inst = value.TryInstruction(); new_inst != nullptr) {
 		AddUse(new_inst, index);
 	}
@@ -238,11 +237,15 @@ void Inst::AddPhiOperand(Block* predecessor, Value value) {
 }
 
 void Inst::ReplaceUsesWith(Value replacement, bool preserve) {
-	const auto old_uses = uses;
+	// Detach the arguments first (a phi can use itself), then move every use over in order.
-	for (const auto& use: old_uses) {
-		use.user->SetArg(use.operand, replacement);
-	}
 	Invalidate();
+	const auto old_uses = std::exchange(uses, {});
+	for (const auto& use: old_uses) {
+		ArgSlot(*use.user, use.operand) = replacement;
+	}
+	if (auto* target = replacement.TryInstruction(); target != nullptr) {
+		target->uses.insert(target->uses.end(), old_uses.begin(), old_uses.end());
+	}
 	if (preserve) {
 		opcode = ValueOpcode::Identity;
 		num_args = 1;
@@ -272,6 +275,10 @@ void Inst::ClearArgs() {
 			RemoveUse(value_inst, index);
 		}
 	}
+	DropArgs();
+}
+
+void Inst::DropArgs() {
 	if (num_args == PhiArity) {
 		std::destroy_at(&phi_args);
 		std::construct_at(&fixed_args);

src/graphics/shader/recompiler/ir/Value.h

@@ -138,6 +138,8 @@ public:
 	void AddPhiOperand(Block* predecessor, Value value);
 	void ReplaceUsesWith(Value replacement, bool preserve = true);
 	void Invalidate();
+	// Drops the arguments without updating their use lists, for teardown that destroys both sides.
+	void DropArgs();
 
 	template <typename T>
 	requires(sizeof(T) <= sizeof(uint64_t) && std::is_trivially_copyable_v<T>)
@@ -157,6 +159,8 @@ public:
 private:
 	friend void EliminateDeadCode(const std::vector<Block*>& blocks);
 
+	template <typename Self>
+	static auto& ArgSlot(Self& self, size_t index);
 	void AddUse(Inst* used, size_t operand);
 	void RemoveUse(Inst* used, size_t operand);
 	void ClearArgs();

tests/ScalarProvenanceTests.cpp

@@ -353,6 +353,35 @@ void TestOperandMutationLifecycle() {
         "Phi replacement retained incoming uses or changed the consumer value");
 }
 
+void TestReplacementKeepsUseOrder() {
+  Fixture fixture(2);
+  const auto base = fixture.Emit(ValueOpcode::IAdd32, {Value(1u), Value(2u)});
+  auto *before = fixture.Emit(ValueOpcode::IAdd32, {base, Value(3u)}).TryInstruction();
+  auto *identity = fixture.Emit(ValueOpcode::Identity, {base}).TryInstruction();
+  auto *pair = fixture.Emit(ValueOpcode::IAdd32, {Value(identity), Value(identity)}).TryInstruction();
+  auto *after = fixture.Emit(ValueOpcode::IAdd32, {Value(4u), base}).TryInstruction();
+  auto &large = fixture.program.value_storage.emplace_back(ValueOpcode::GetImageResource);
+  const auto last = large.NumArgs() - 1;
+  large.SetArg(last, Value(identity));
+  identity->ReplaceUsesWith(base, false);
+  Check(!identity->HasUses() && identity->NumArgs() == 0 && pair->Arg(0) == base &&
+            pair->Arg(1) == base && large.Arg(last) == base &&
+            base.TryInstruction()->Uses() ==
+                std::vector<Use>({{before, 0}, {after, 1}, {pair, 0}, {pair, 1}, {&large, last}}),
+        "identity removal lost or reordered the uses of its argument");
+
+  auto &phi = fixture.BlockAt(1).AppendNewInst(ValueOpcode::Phi);
+  phi.SetFlags(Type::U32);
+  phi.AddPhiOperand(&fixture.BlockAt(0), Value(after));
+  phi.AddPhiOperand(&fixture.BlockAt(1), Value(&phi));
+  auto *user = fixture.Emit(ValueOpcode::IAdd32, {Value(&phi), Value(&phi)}, 0, 1).TryInstruction();
+  phi.ReplaceUsesWith(Value(after));
+  Check(phi.GetOpcode() == ValueOpcode::Identity && phi.Arg(0) == Value(after) && !phi.HasUses() &&
+            user->Arg(0) == Value(after) && user->Arg(1) == Value(after) &&
+            after->Uses() == std::vector<Use>({{user, 0}, {user, 1}, {&phi, 0}}),
+        "self-referencing Phi replacement kept a use of itself or reordered the moved uses");
+}
+
 void TestControlDependentStandaloneLoadStaysTyped() {
   Fixture fixture(3);
   const auto memory = fixture.AddMemory(ResourceKind::ScalarAddress);
@@ -916,6 +945,7 @@ int main() {
     TestCarryAndBitFields();
     TestInvariantAndDivergentPhi();
     TestOperandMutationLifecycle();
+    TestReplacementKeepsUseOrder();
     TestControlDependentStandaloneLoadStaysTyped();
     TestRuntime64BitDescriptorOps();
     TestUniformFirstLaneSamplerLod();