fix(agc): retain registered headers sharing shader code

Keep each registered header and its prepared shader state instead of
replacing every registration for the same code address. Resolve explicit
ABIs by header and programs by type, preserving re-registration order.

Co-Authored-By: GPT-6.1-Sol <noreply@openai.com>

core/libs/prx/libSceAgcDriver/Execution/include/Driver/Shaders/ShaderRegistry.hpp

@@ -10,6 +10,7 @@
 #include <cstddef>
 #include <cstdint>
 #include <exception>
+#include <initializer_list>
 #include <map>
 #include <memory>
 #include <mutex>
@@ -27,7 +28,7 @@ class VulkanDevice;
 namespace AgcDriver::DriverDetail {
 
 struct ShaderSnapshot;
-using ShaderRegistry = std::map<std::uint64_t, std::shared_ptr<const ShaderSnapshot>>;
+using ShaderRegistry = std::map<std::uint64_t, std::vector<std::shared_ptr<const ShaderSnapshot>>>;
 
 struct PreparedShaderState {
     struct Entry {
@@ -89,6 +90,7 @@ std::shared_ptr<const ShaderRecompiler::SourceHandle> PrepareShaderWithDiagnosti
 std::uint64_t NullPixelProgramAddress();
 ShaderSnapshot PrepareNullPixelProgram(const VulkanDevice& device);
 std::optional<ShaderRecompiler::ShaderFloatMode> RegisteredFloatMode(const ShaderSnapshot& snapshot);
+std::shared_ptr<const ShaderSnapshot> RegisteredProgram(const ShaderRegistry& registry, std::uint64_t address, std::initializer_list<std::uint8_t> types);
 void PublishRegisteredShader(std::shared_ptr<ShaderRegistry>& registry, const std::shared_ptr<const ShaderSnapshot>& snapshot);
 
 void ResolvePreparedGraphics(const ShaderSnapshot& front, const std::shared_ptr<const ShaderSnapshot>& fragment, std::uint32_t primitiveType, const ShaderRecompiler::SpirvTarget& target);

core/libs/prx/libSceAgcDriver/Execution/src/Driver/Dispatch/Dispatch.cpp

@@ -16,12 +16,7 @@ namespace AgcDriver::DriverDetail {
 void Driver::dispatch(QueueState& queue, std::span<const std::uint32_t> packet, const Submission& submission, std::uint64_t indirectArguments) {
     PerformanceTimer timing("Driver.Dispatch");
     const auto address = (static_cast<std::uint64_t>(readRegister(queue.shader, 0x20c)) << 8u) | (static_cast<std::uint64_t>(readRegister(queue.shader, 0x20d) & 0xffu) << 40u);
-    auto it = submission.shaders->upper_bound(address);
-    std::shared_ptr<const ShaderSnapshot> registeredShader;
-    if (it != submission.shaders->begin()) {
-        --it;
-        if (address - it->second->codeAddress < it->second->code.size() * sizeof(std::uint32_t)) registeredShader = it->second;
-    }
+    auto registeredShader = RegisteredProgram(*submission.shaders, address, {0});
     if (!registeredShader) registeredShader = ReadRawComputeShader(address);
     const auto& snapshot = *registeredShader;
     require(snapshot.type == 0, "compute program refers to a non-compute shader");

core/libs/prx/libSceAgcDriver/Execution/src/Driver/Draw/DrawDecode.cpp

@@ -28,13 +28,10 @@ void DecodeGraphicsPrograms(DrawDecode& decoded, const QueueState& queue, const
     const auto prepare = [&](std::uint64_t address, std::uint8_t type, Stage stage, std::uint32_t rsrc2, std::uint32_t userDataBase) {
         const bool nullPixel = stage == Stage::Fragment && (address == 0 || pixelSkipped);
         if (nullPixel) address = NullPixelProgramAddress();
-        auto it = registry.upper_bound(address);
-        require(it != registry.begin(), "graphics program does not belong to a registered shader");
-        --it;
-        const auto& snapshot = *it->second;
-        require(address - snapshot.codeAddress < snapshot.code.size() * sizeof(std::uint32_t), "graphics program is outside registered shader code");
+        const auto registered = RegisteredProgram(registry, address, {type});
+        require(registered != nullptr, "graphics program does not belong to a compatible registered shader");
+        const auto& snapshot = *registered;
         require((address - snapshot.codeAddress) % sizeof(std::uint32_t) == 0, "graphics entry point is not dword aligned");
-        require(snapshot.type == type, "graphics program refers to an incompatible shader binary type");
         if (!nullPixel) Graphics::NoteRegisterRead(Graphics::RegisterBank::Shader, rsrc2);
         const auto resources = nullPixel ? 0u : ReadGraphicsRegister(queue.shader, rsrc2);
         const auto userCount = ((resources >> 1u) & 0x1fu) | (((resources >> 27u) & 1u) << 5u);
@@ -46,7 +43,7 @@ void DecodeGraphicsPrograms(DrawDecode& decoded, const QueueState& queue, const
             8,
             {},
             {{{snapshot.codeAddress, std::as_bytes(std::span(snapshot.code))}, {snapshot.headerAddress, snapshot.header}}},
-            it->second,
+            registered,
             codeOffset
         };
         for (std::uint32_t i = 0; i < userCount; ++i) {
@@ -89,10 +86,9 @@ void DecodeGraphicsPrograms(DrawDecode& decoded, const QueueState& queue, const
             append(0x0c8, 2, Stage::TessellationEvaluation, 0x08b, 0x08c, Role::Domain);
         } else if (graphics.stages.path == Graphics::ShaderPath::Geometry) {
             const auto frontAddress = programAddress(0xc8);
-            auto snapshot = registry.upper_bound(frontAddress);
-            require(snapshot != registry.begin(), "geometry front program is not registered");
-            --snapshot;
-            const auto type = snapshot->second->type;
+            const auto snapshot = RegisteredProgram(registry, frontAddress, {2, 4});
+            require(snapshot != nullptr, "geometry front program is not registered");
+            const auto type = snapshot->type;
             require(type == 2 || type == 4, "invalid geometry front binary type");
             append(0xc8, type, Stage::Mesh, 0x8b, 0x8c, Role::Main);
             initializeMerged(programs.back(), 0x82, type == 4);

core/libs/prx/libSceAgcDriver/Execution/src/Driver/Draw/DrawRegisterKey.cpp

@@ -50,8 +50,8 @@ std::uint64_t Driver::drawRegisterKey(const QueueState& queue, const ShaderRegis
             continue;
         }
         --it;
-        mix(reinterpret_cast<std::uintptr_t>(it->second.get()));
-        mix(address - it->second->codeAddress);
+        for (const auto& snapshot : it->second) mix(reinterpret_cast<std::uintptr_t>(snapshot.get()));
+        mix(address - it->first);
     }
     return key;
 }

core/libs/prx/libSceAgcDriver/Execution/src/Driver/Shaders/ShaderRegistry.cpp

@@ -207,20 +207,57 @@ bool SameHeader(const ShaderSnapshot& snapshot, const Shader* shader) {
     return std::memcmp(&used, &registered, offsetof(Shader, num_sh_registers) + sizeof(Shader::num_sh_registers)) == 0;
 }
 
+std::shared_ptr<const ShaderSnapshot> RegisteredHeader(const ShaderRegistry& registry, const Shader* shader) {
+    const auto found = registry.find(reinterpret_cast<std::uintptr_t>(const_cast<const void*>(shader->code)));
+    if (found == registry.end()) return {};
+    for (const auto& snapshot : found->second) {
+        if (snapshot->headerAddress == reinterpret_cast<std::uintptr_t>(shader)) return snapshot;
+    }
+    for (const auto& snapshot : found->second) {
+        if (SameHeader(*snapshot, shader)) return snapshot;
+    }
+    return {};
+}
+
+std::shared_ptr<const ShaderSnapshot> RegisteredProgram(const ShaderRegistry& registry, std::uint64_t address, std::initializer_list<std::uint8_t> types) {
+    auto found = registry.upper_bound(address);
+    if (found == registry.begin()) return {};
+    --found;
+    std::shared_ptr<const ShaderSnapshot> result;
+    bool registered = false;
+    for (const auto& snapshot : found->second) {
+        if (address - snapshot->codeAddress >= snapshot->code.size() * sizeof(std::uint32_t)) continue;
+        registered = true;
+        if (std::ranges::find(types, snapshot->type) == types.end()) continue;
+        require(result == nullptr || result->type == snapshot->type, "registered program has ambiguous shader binary types");
+        result = snapshot;
+    }
+    require(!registered || result != nullptr, "registered program refers to an incompatible shader binary type");
+    return result;
+}
+
 void PublishRegisteredShader(std::shared_ptr<ShaderRegistry>& registry, const std::shared_ptr<const ShaderSnapshot>& snapshot) {
     ShaderPreparationTransaction transaction;
+    auto published = snapshot;
     if (registry != nullptr) {
         const auto found = registry->find(snapshot->codeAddress);
         if (found != registry->end()) {
-            const auto& current = *found->second;
-            if (current.headerAddress == snapshot->headerAddress && current.type == snapshot->type && current.code == snapshot->code && current.header == snapshot->header) {
-                transaction.Commit();
-                return;
+            for (const auto& current : found->second) {
+                if (current->headerAddress == snapshot->headerAddress && current->type == snapshot->type && current->code == snapshot->code && current->header == snapshot->header) {
+                    if (current == found->second.back()) {
+                        transaction.Commit();
+                        return;
+                    }
+                    published = current;
+                    break;
+                }
             }
         }
     }
     auto next = registry != nullptr ? std::make_shared<ShaderRegistry>(*registry) : std::make_shared<ShaderRegistry>();
-    next->insert_or_assign(snapshot->codeAddress, snapshot);
+    auto& headers = (*next)[snapshot->codeAddress];
+    std::erase_if(headers, [&](const auto& current) { return current->headerAddress == snapshot->headerAddress || current->code != snapshot->code; });
+    headers.push_back(published);
     registry = std::move(next);
     transaction.Commit();
 }
@@ -681,9 +718,10 @@ void Driver::ResolveGraphicsStagesAbi(std::span<const Shader* const> stages, std
             GuestMemory::CheckRange(shader, sizeof(Shader), alignof(Shader));
             const auto address = reinterpret_cast<std::uintptr_t>(const_cast<const void*>(shader->code));
             require(registry != nullptr && registry->contains(address), "graphics ABI refers to an unregistered shader");
-            if (owner == nullptr) owner = registry->at(address);
-            const auto& snapshot = *registry->at(address);
-            require(SameHeader(snapshot, shader), "graphics ABI refers to a replaced shader header");
+            const auto selected = RegisteredHeader(*registry, shader);
+            require(selected != nullptr, "graphics ABI refers to a replaced shader header");
+            if (owner == nullptr) owner = selected;
+            const auto& snapshot = *selected;
             require(snapshot.registeredState != nullptr, "registered shader state is missing");
             const auto& registered = *snapshot.registeredState;
             for (const auto& [offset, value] : registered.shader) state.shader.insert_or_assign(offset, value);
@@ -747,8 +785,8 @@ void Driver::ResolveShaderAbi(const Shader* shader, std::span<const ShaderRegist
         std::lock_guard lock(mutex);
         const auto address = reinterpret_cast<std::uintptr_t>(const_cast<const void*>(shader->code));
         require(shaders != nullptr && shaders->contains(address), "static ABI refers to an unregistered shader");
-        snapshot = shaders->at(address);
-        require(SameHeader(*snapshot, shader), "static ABI refers to a replaced shader header");
+        snapshot = RegisteredHeader(*shaders, shader);
+        require(snapshot != nullptr, "static ABI refers to a replaced shader header");
     }
     require(snapshot->registeredState != nullptr, "registered shader state is missing");
     QueueState state{};
@@ -809,8 +847,8 @@ void Driver::ResolveGraphicsAbi(const Shader* vertex, const Shader* pixel, std::
         const auto lookup = [&](const Shader* shader) {
             const auto address = reinterpret_cast<std::uintptr_t>(const_cast<const void*>(shader->code));
             require(shaders != nullptr && shaders->contains(address), "rectangle ABI refers to an unregistered shader");
-            const auto snapshot = shaders->at(address);
-            require(SameHeader(*snapshot, shader), "rectangle ABI refers to a replaced shader header");
+            const auto snapshot = RegisteredHeader(*shaders, shader);
+            require(snapshot != nullptr, "rectangle ABI refers to a replaced shader header");
             return snapshot;
         };
         front = lookup(vertex);
@@ -926,10 +964,13 @@ void Driver::RegisterShader(const Shader* shader) {
         if (shaders != nullptr) {
             const auto found = shaders->find(snapshot.codeAddress);
             if (found != shaders->end()) {
-                const auto& current = *found->second;
-                if (current.headerAddress == snapshot.headerAddress && current.type == snapshot.type && current.code == snapshot.code && current.header == snapshot.header) {
-                    transaction.Commit();
-                    return;
+                for (const auto& current : found->second) {
+                    if (current->headerAddress == snapshot.headerAddress && current->type == snapshot.type && current->code == snapshot.code && current->header == snapshot.header) {
+                        const auto identical = current;
+                        PublishRegisteredShader(shaders, identical);
+                        transaction.Commit();
+                        return;
+                    }
                 }
             }
         }

core/libs/prx/libSceAgcDriver/tests/execution/PreparedGraphics.cpp

@@ -80,7 +80,7 @@ void Check(AgcDriver::VulkanDevice& device, AgcDriver::Graphics::ShaderPath path
     fragment.snapshot->code.resize(64);
     fragment.snapshot->code.insert(fragment.snapshot->code.end(), pixelCode.begin(), pixelCode.end());
     ShaderRegistry registry;
-    for (const auto* fixture : {&front, &back, &domain, &fragment}) registry.emplace(fixture->snapshot->codeAddress, fixture->snapshot);
+    for (const auto* fixture : {&front, &back, &domain, &fragment}) registry.emplace(fixture->snapshot->codeAddress, std::vector<std::shared_ptr<const ShaderSnapshot>>{fixture->snapshot});
     AgcDriver::QueueState queue{};
     queue.context[0x8e] = 0xfu;
     queue.context[0x8f] = 0xfu;
@@ -256,8 +256,8 @@ void NullPixelMatchesRegistration(AgcDriver::VulkanDevice& device) {
     const auto registered = null->prepared->entries;
     Require(!registered.empty(), "the null pixel program was not prepared at registration");
     ShaderRegistry registry;
-    registry.emplace(front.snapshot->codeAddress, front.snapshot);
-    registry.emplace(null->codeAddress, null);
+    registry.emplace(front.snapshot->codeAddress, std::vector<std::shared_ptr<const ShaderSnapshot>>{front.snapshot});
+    registry.emplace(null->codeAddress, std::vector<std::shared_ptr<const ShaderSnapshot>>{null});
     AgcDriver::QueueState queue{};
     queue.context[0x8e] = 0xfu;
     queue.context[0x8f] = 0xfu;

core/libs/prx/libSceAgcDriver/tests/execution/PreparedPublication.cpp

@@ -1,4 +1,5 @@
 #include "VulkanTestDevice.hpp"
+#include "SceShaders.hpp"
 #include "prx/libSceAgcDriver/Execution/include/ShaderPreparation.hpp"
 #include "prx/libSceAgcDriver/Execution/include/Driver/Shaders/ShaderRegistry.hpp"
 #include "Optimization/ResourceProgram.hpp"
@@ -6,6 +7,7 @@
 #include <barrier>
 #include <chrono>
 #include <future>
+#include <cstring>
 #include <iostream>
 #include <thread>
 
@@ -106,6 +108,45 @@ void Transactions() {
     Require(!partial && front.prepared->entries.size() == 200 && pixel->prepared->entries.size() == 200, "concurrent group publication was partial or lost an update");
 }
 
+void HeaderSelection() {
+    const auto makeSnapshot = [](std::uint64_t headerAddress, std::uint8_t type) {
+        auto snapshot = std::make_shared<ShaderSnapshot>();
+        snapshot->codeAddress = 0x10000;
+        snapshot->headerAddress = headerAddress;
+        snapshot->type = type;
+        snapshot->code = {0xbf810000u};
+        Shader header{};
+        header.code = reinterpret_cast<const void*>(snapshot->codeAddress);
+        header.type = type;
+        snapshot->header.resize(sizeof(Shader));
+        std::memcpy(snapshot->header.data(), &header, sizeof(Shader));
+        return snapshot;
+    };
+    const auto vertex = makeSnapshot(0x20000, 2);
+    const auto compute = makeSnapshot(0x30000, 0);
+    std::shared_ptr<ShaderRegistry> registry;
+    PublishRegisteredShader(registry, vertex);
+    PublishRegisteredShader(registry, compute);
+    Require(RegisteredProgram(*registry, 0x10000, {2}) == vertex && RegisteredProgram(*registry, 0x10000, {0}) == compute, "shared code selected another registered type");
+    Reject([&] { static_cast<void>(RegisteredProgram(*registry, 0x10000, {1})); });
+    Require(RegisteredProgram(*registry, 0x10100, {2}) == nullptr, "an out-of-range entry selected registered code");
+    const auto copy = makeSnapshot(0x40000, 2);
+    PublishRegisteredShader(registry, copy);
+    Require(RegisteredProgram(*registry, 0x10000, {2}) == copy, "equivalent registered headers became ambiguous");
+    PublishRegisteredShader(registry, makeSnapshot(vertex->headerAddress, 2));
+    Require(RegisteredProgram(*registry, 0x10000, {2}) == vertex && registry->at(0x10000).size() == 3, "identical earlier registration did not restore its prepared snapshot as current");
+    const auto replacement = makeSnapshot(copy->headerAddress, 2);
+    Shader changed;
+    std::memcpy(&changed, replacement->header.data(), sizeof(Shader));
+    changed.target = 1;
+    std::memcpy(replacement->header.data(), &changed, sizeof(Shader));
+    PublishRegisteredShader(registry, replacement);
+    Require(RegisteredProgram(*registry, 0x10000, {2}) == replacement, "same-type registration did not keep the latest program metadata");
+    const auto geometry = makeSnapshot(0x60000, 4);
+    PublishRegisteredShader(registry, geometry);
+    Reject([&] { static_cast<void>(RegisteredProgram(*registry, 0x10000, {2, 4})); });
+}
+
 void Registration(AgcDriver::VulkanDevice& device) {
     const std::array<std::uint32_t, 1> code{0xbf810000u};
     RecompileRequest request{{ShaderStage::Compute, 0x10000, code, 0, {}}, {32, 0, {}, ShaderComputeStageInfo{{1, 1, 1}, 0, {false, false, false}, false, 0, {}}, {}, {}, {}}, device.ComputeTarget(32), {0, 0, 0, 128}};
@@ -130,7 +171,7 @@ void Registration(AgcDriver::VulkanDevice& device) {
             for (unsigned iteration = 0; iteration < 40; ++iteration) {
                 ShaderPreparationTransaction transaction;
                 PublishRegisteredShader(registry, makeSnapshot());
-                const auto current = registry->at(0x10000);
+                const auto current = registry->at(0x10000).back();
                 Require(current == original && current->prepared->rectangleRequested, "identical registration replaced a prepared snapshot");
                 auto& entries = transaction.Edit(*current).entries;
                 entries.push_back({iteration + 1, handle});
@@ -160,11 +201,11 @@ void Registration(AgcDriver::VulkanDevice& device) {
         transaction.Commit();
     }
     replacementWriter.get();
-    Require(registry->at(0x10000) == replacement && submission->at(0x10000) == original, "replacement modified an in-flight registry");
+    Require(registry->at(0x10000).back() == replacement && submission->at(0x10000).back() == original, "replacement modified an in-flight registry");
     Require(original->prepared->entries.size() == 162 && replacement->prepared->entries.empty(), "stale preparation was published into a replacement snapshot");
     original.reset();
-    request.shader.code = submission->at(0x10000)->code;
-    const auto invocation = InvocationFor(*submission->at(0x10000), 0, request);
+    request.shader.code = submission->at(0x10000).back()->code;
+    const auto invocation = InvocationFor(*submission->at(0x10000).back(), 0, request);
     const auto capture = invocation.Capture({});
     const auto result = invocation.Materialize(*capture);
     Require(result->variantId == GetPreparedArtifact(*handle).variantId, "old submission lost its prepared artifact");
@@ -172,19 +213,19 @@ void Registration(AgcDriver::VulkanDevice& device) {
     device.WaitIdle();
     Require(!old.expired(), "in-flight submission released its snapshot early");
     submission.reset();
-    Require(old.expired(), "completed submission retained an obsolete snapshot");
+    Require(!old.expired() && registry->at(0x10000).size() == 2, "registration discarded an earlier header for shared code");
     Require(invocation.Materialize(*capture)->spirv.data() == result->spirv.data(), "prepared handle did not retain its artifact");
     auto changed = makeSnapshot();
     changed->headerAddress = replacement->headerAddress;
     changed->code[0] ^= 1;
     PublishRegisteredShader(registry, changed);
-    Require(registry->at(0x10000) == changed, "changed shader code was treated as identical");
+    Require(registry->at(0x10000).back() == changed && registry->at(0x10000).size() == 1 && old.expired(), "changed shader code retained obsolete headers");
     auto metadata = makeSnapshot();
     metadata->headerAddress = changed->headerAddress;
     metadata->code = changed->code;
     metadata->header.push_back(std::byte{1});
     PublishRegisteredShader(registry, metadata);
-    Require(registry->at(0x10000) == metadata, "changed shader metadata was treated as identical");
+    Require(registry->at(0x10000).back() == metadata, "changed shader metadata was treated as identical");
 }
 
 }
@@ -192,6 +233,7 @@ void Registration(AgcDriver::VulkanDevice& device) {
 int main() {
     try {
         Transactions();
+        HeaderSelection();
         auto device = OpenVulkanTestDevice();
         if (!device) return VulkanTestSkipped;
         Registration(*device);

core/libs/prx/libSceAgcDriver/tests/execution/PreparedShaders.cpp

@@ -6,6 +6,7 @@
 #include "prx/libSceAgcDriver/Execution/include/Driver.hpp"
 #include "prx/libSceAgcDriver/Execution/include/ShaderPreparation.hpp"
 #include "prx/libSceAgcDriver/Submit/include/Acb.hpp"
+#include "prx/libSceAgcDriver/Submit/include/Dcb.hpp"
 #include <spirv/unified1/spirv.hpp>
 #include <algorithm>
 #include <array>
@@ -334,6 +335,67 @@ void UnsupportedTypeRegistration() {
     AgcDriverRegisterShader_nid_postfix(&header.shader);
 }
 
+void SharedCodeHeaders() {
+    alignas(256) static const std::array<std::uint32_t, 11> code{0x7e0002ffu, 0, 0x7e0202ffu, 0, 0x7e0402ffu, 0, 0x7e0602ffu, 0x3f800000u, 0xf80008cfu, 0x03020100u, 0xbf810000u};
+    struct Header {
+        Shader shader{};
+        std::array<ShaderRegister, 7> registers{};
+        ShaderSpecialRegs specials{};
+        ShaderUserData users{};
+        std::array<ShaderRegister, 1> context{{{0x2d5, 0x2000}}};
+    } vertex, mesh;
+    const auto address = reinterpret_cast<std::uintptr_t>(code.data());
+    for (auto* header : {&vertex, &mesh}) {
+        header->shader.file_header = 0x34333231u;
+        header->shader.version = 0x18;
+        header->shader.header_size = sizeof(Header);
+        header->shader.shader_size = sizeof(code);
+        header->shader.code = code.data();
+        header->shader.sh_registers = header->registers.data();
+        header->shader.num_sh_registers = header->registers.size();
+        header->shader.specials = &header->specials;
+        header->shader.user_data = &header->users;
+        header->specials.dispatch_modifier = 0x8000;
+    }
+    vertex.shader.cx_registers = vertex.context.data();
+    vertex.shader.num_cx_registers = vertex.context.size();
+    vertex.shader.type = 2;
+    vertex.registers = {{{0xc8, static_cast<std::uint32_t>(address >> 8u)}, {0xc9, static_cast<std::uint32_t>(address >> 40u)}, {0x8b, 0}, {0x8a, 0}, {0xca, 0}, {0xcb, 0}, {0xcc, 0}}};
+    mesh.shader.type = 4;
+    mesh.registers = vertex.registers;
+    const std::array<ShaderRegister, 1> primitive{{{0x242, 4}}};
+    AgcDriverRegisterShader_nid_postfix(&vertex.shader);
+    AgcDriverRegisterShader_nid_postfix(&mesh.shader);
+    AgcDriverResolveShaderAbi_nid_postfix(&vertex.shader, {}, primitive);
+    AgcDriverRegisterShader_nid_postfix(&mesh.shader);
+    Shader copy = vertex.shader;
+    copy.user_data = nullptr;
+    AgcDriverResolveShaderAbi_nid_postfix(&copy, {}, primitive);
+    copy.target ^= 1u;
+    ExpectFailure([&] { AgcDriverResolveShaderAbi_nid_postfix(&copy, {}, primitive); }, "replaced shader header");
+    std::vector<std::uint32_t> commands;
+    for (const auto reg : vertex.registers) commands.insert(commands.end(), {0xc0017600u, reg.offset, reg.value});
+    for (const auto reg : vertex.context) commands.insert(commands.end(), {0xc0016900u, reg.offset, reg.value});
+    commands.insert(commands.end(), {0xc0017600u, 0x8, 0, 0xc0016900u, 0xd, 0x00100010u, 0xc0016900u, 0x1c3, 4, 0xc0017900u, 0x242, 4, 0xc0012d00u, 3, 2});
+    Packet packet{commands.data(), static_cast<std::uint32_t>(commands.size()), 0, {}};
+    for (unsigned iteration = 0; iteration < 2; ++iteration) {
+        Require(sceAgcDriverSubmitDcb(&packet) == 0, "shared-code draw submission failed");
+        AgcDriverWaitIdle_nid_postfix();
+    }
+    mesh.shader.type = 2;
+    mesh.shader.cx_registers = mesh.context.data();
+    mesh.shader.num_cx_registers = mesh.context.size();
+    AgcDriverRegisterShader_nid_postfix(&mesh.shader);
+    AgcDriverResolveShaderAbi_nid_postfix(&mesh.shader, {}, primitive);
+    AgcDriverResolveShaderAbi_nid_postfix(&vertex.shader, {}, primitive);
+    Require(sceAgcDriverSubmitDcb(&packet) == 0, "same-type registered header draw failed");
+    AgcDriverWaitIdle_nid_postfix();
+    AgcDriverRegisterShader_nid_postfix(&vertex.shader);
+    AgcDriverResolveShaderAbi_nid_postfix(&vertex.shader, {}, primitive);
+    Require(sceAgcDriverSubmitDcb(&packet) == 0, "re-registered earlier header draw failed");
+    AgcDriverWaitIdle_nid_postfix();
+}
+
 void RegistrationWithoutSpecials() {
     alignas(256) std::array<std::uint32_t, 1> code{0xbf810000u};
     struct Header {
@@ -521,6 +583,7 @@ int main(int argc, char** argv) {
         Require(argc != 2 || std::string_view(argv[1]) != "--fail-before-registration", "injected failure before registration");
         PrepareMultisampledStorage(*device);
         device.reset();
+        SharedCodeHeaders();
         RegistrationWithoutSpecials();
         Registration(argc == 2);
         std::cout << "prepared shader and transactional registration tests passed\n";