From ce98f8303309f9c65b40fe313bcb2b4760a83e44 Mon Sep 17 00:00:00 2001
From: Blazing Cook <>
Date: Sun, 11 Oct 2026 00:28:16 -0500
Subject: [PATCH] fix(ampr): retain pooled pages while remapping aliases

(cherry picked from commit a2408b5244c663849ae4f7022e7aff01f2ba5abc)
---
 core/libs/prx/libkernel/Apr/src/Apr.cpp | 57 ++++++++++++++++-------
 core/libs/tests/GuestAmpr.cpp           | 60 +++++++++++++++++++++++++
 2 files changed, 101 insertions(+), 16 deletions(-)

diff --git a/core/libs/prx/libkernel/Apr/src/Apr.cpp b/core/libs/prx/libkernel/Apr/src/Apr.cpp
index 631f5553e..4e9bf1288 100644
--- a/core/libs/prx/libkernel/Apr/src/Apr.cpp
+++ b/core/libs/prx/libkernel/Apr/src/Apr.cpp
@@ -347,15 +347,38 @@ void _addPage(AmmState& state, std::uint64_t address, const AmmPage& page) {
     if (page.pooled) ++state.pooledUses[page.physical];
 }
 
-void _forgetPages(AmmState& state, std::uint64_t address, std::uint64_t size, bool release) {
-    for (auto it = state.pages.lower_bound(address); it != state.pages.end() && it->first < address + size;) {
-        if (it->second.pooled) {
-            const auto uses = state.pooledUses.find(it->second.physical);
-            if (--uses->second == 0 && release) {
-                state.pooledUses.erase(uses);
-                _returnPoolPages(state, it->second.physical, PS5_PAGE_SIZE);
-            }
+void _releasePage(AmmState& state, const AmmPage& page) {
+    if (!page.pooled) return;
+    const auto uses = state.pooledUses.find(page.physical);
+    if (--uses->second == 0) {
+        state.pooledUses.erase(uses);
+        _returnPoolPages(state, page.physical, PS5_PAGE_SIZE);
+    }
+}
+
+class AmmPageReferences {
+public:
+    AmmPageReferences(AmmState& state, const std::vector<AmmPage>& pages) : state(state), pages(pages) {
+        for (const auto& page : pages) {
+            if (page.pooled) ++state.pooledUses.find(page.physical)->second;
         }
+    }
+
+    AmmPageReferences(const AmmPageReferences&) = delete;
+    AmmPageReferences& operator=(const AmmPageReferences&) = delete;
+
+    ~AmmPageReferences() {
+        for (const auto& page : pages) _releasePage(state, page);
+    }
+
+private:
+    AmmState& state;
+    const std::vector<AmmPage>& pages;
+};
+
+void _forgetPages(AmmState& state, std::uint64_t address, std::uint64_t size) {
+    for (auto it = state.pages.lower_bound(address); it != state.pages.end() && it->first < address + size;) {
+        _releasePage(state, it->second);
         it = state.pages.erase(it);
     }
 }
@@ -385,9 +408,9 @@ void _mapPages(AmmState& state, std::uint64_t address, const std::vector<AmmPage
         std::size_t last = first + 1;
         while (last < pages.size() && pages[last].physical == pages[last - 1].physical + PS5_PAGE_SIZE && pages[last].protection == pages[first].protection) ++last;
         _mapRun(address + first * PS5_PAGE_SIZE, pages[first].physical, (last - first) * PS5_PAGE_SIZE, pages[first].protection);
+        for (std::size_t index = first; index < last; ++index) _addPage(state, address + index * PS5_PAGE_SIZE, pages[index]);
         first = last;
     }
-    for (std::size_t index = 0; index < pages.size(); ++index) _addPage(state, address + index * PS5_PAGE_SIZE, pages[index]);
 }
 
 void _clearPrt(AmmState& state, std::uint64_t address, std::uint64_t size) {
@@ -410,7 +433,7 @@ void _ammUnmap(const Apr::AmmUnmapCommand& command) {
     auto& state = _amm();
     const std::lock_guard lock(state.lock);
     _ammCheckRange(state, command.address, command.size);
-    _forgetPages(state, command.address, command.size, true);
+    _forgetPages(state, command.address, command.size);
     _clearPrt(state, command.address, command.size);
     _reserve(command.address, command.size);
 }
@@ -419,7 +442,7 @@ void _ammMapAsPrt(const Apr::AmmUnmapCommand& command) {
     auto& state = _amm();
     const std::lock_guard lock(state.lock);
     _ammCheckRange(state, command.address, command.size);
-    _forgetPages(state, command.address, command.size, true);
+    _forgetPages(state, command.address, command.size);
     _zeroPrt(command.address, command.size);
     for (std::uint64_t offset = 0; offset < command.size; offset += PS5_PAGE_SIZE) state.prt.insert(command.address + offset);
 }
@@ -429,7 +452,7 @@ void _ammUnmapToPrt(const Apr::AmmUnmapCommand& command) {
     const std::lock_guard lock(state.lock);
     _ammCheckRange(state, command.address, command.size);
     _requirePrt(state, command.address, command.size);
-    _forgetPages(state, command.address, command.size, true);
+    _forgetPages(state, command.address, command.size);
     _zeroPrt(command.address, command.size);
 }
 
@@ -469,7 +492,7 @@ void _ammMap(const Apr::AmmMapCommand& command, bool direct) {
     auto& state = _amm();
     const std::lock_guard lock(state.lock);
     _ammCheckRange(state, command.address, command.size);
-    _forgetPages(state, command.address, command.size, true);
+    _forgetPages(state, command.address, command.size);
     _clearPrt(state, command.address, command.size);
     const auto physical = direct ? command.directOffset : _takePoolPages(state, command.size);
     try {
@@ -490,12 +513,14 @@ void _ammRemap(const Apr::AmmRemapCommand& command, bool alias, bool intoPrt) {
     if (intoPrt) _requirePrt(state, command.address, command.size);
     auto pages = _mappedPages(state, command.source, command.size);
     for (auto& page : pages) page.protection = command.protection;
+    const AmmPageReferences references(state, pages);
     if (!alias) {
-        _forgetPages(state, command.source, command.size, false);
-        _clearPrt(state, command.source, command.size);
         _reserve(command.source, command.size);
+        _forgetPages(state, command.source, command.size);
+        _clearPrt(state, command.source, command.size);
     }
-    _forgetPages(state, command.address, command.size, true);
+    _reserve(command.address, command.size);
+    _forgetPages(state, command.address, command.size);
     if (!intoPrt) _clearPrt(state, command.address, command.size);
     _mapPages(state, command.address, pages);
 }
diff --git a/core/libs/tests/GuestAmpr.cpp b/core/libs/tests/GuestAmpr.cpp
index 2bb8baefe..f11fd75e5 100644
--- a/core/libs/tests/GuestAmpr.cpp
+++ b/core/libs/tests/GuestAmpr.cpp
@@ -1013,6 +1013,65 @@ void SubmitAmm(Recorder& recorder) {
     Require(sceAmprAmmWaitCommandBufferCompletion(id) == 0);
 }
 
+void TestAmmRemapKeepsPoolReferences() {
+    std::uint64_t start = 0;
+    std::uint64_t end = 0;
+    std::uint64_t multimapStart = 0;
+    std::uint64_t multimapEnd = 0;
+    Require(sceAmprAmmGetVirtualAddressRanges(&start, &end, &multimapStart, &multimapEnd) == 0);
+    std::int64_t pool = -1;
+    Require(sceAmprAmmGiveDirectMemory(0, static_cast<std::int64_t>(sceKernelGetDirectMemorySize()), page, page, 1, &pool) == 0);
+
+    const auto source = start + 256 * page;
+    const auto alias = multimapStart + 256 * page;
+    const auto fresh = start + 257 * page;
+    const auto lastAlias = multimapStart + 257 * page;
+    const auto overlappingAlias = multimapStart + 258 * page;
+    Recorder maps;
+    Require(sceAmprAmmCommandBufferMap(&maps.buffer, source, page, 0, cpuReadWrite) == 0);
+    Require(sceAmprAmmCommandBufferMultiMap(&maps.buffer, alias, source, page, cpuReadWrite) == 0);
+    SubmitAmm(maps);
+    At(source) = 0xA1;
+    Require(At(alias) == 0xA1);
+
+    Recorder move;
+    Require(sceAmprAmmCommandBufferRemap(&move.buffer, alias, source, page, cpuReadWrite) == 0);
+    SubmitAmm(move);
+    Require(Protection(source) == 0 && At(alias) == 0xA1);
+
+    Recorder exhausted;
+    Require(sceAmprAmmCommandBufferMap(&exhausted.buffer, fresh, page, 0, cpuReadWrite) == 0);
+    const auto requireExhausted = [&] {
+        bool rejected = false;
+        try {
+            SubmitAmm(exhausted);
+        } catch (const std::runtime_error& error) {
+            rejected = std::string(error.what()).starts_with("AMM: no free run of ");
+        }
+        Require(rejected);
+    };
+    requireExhausted();
+    At(alias) = 0xB1;
+    Require(At(alias) == 0xB1);
+
+    Recorder retain;
+    Require(sceAmprAmmCommandBufferRemap(&retain.buffer, alias, alias, page, cpuReadWrite) == 0);
+    Require(sceAmprAmmCommandBufferMultiMap(&retain.buffer, lastAlias, alias, page, cpuReadWrite) == 0);
+    Require(sceAmprAmmCommandBufferRemap(&retain.buffer, lastAlias, alias, 2 * page, cpuReadWrite) == 0);
+    Require(sceAmprAmmCommandBufferUnmap(&retain.buffer, lastAlias, page) == 0);
+    SubmitAmm(retain);
+    Require(At(overlappingAlias) == 0xB1);
+    requireExhausted();
+
+    Recorder release;
+    Require(sceAmprAmmCommandBufferUnmap(&release.buffer, overlappingAlias, page) == 0);
+    Require(sceAmprAmmCommandBufferMap(&release.buffer, fresh, page, 0, cpuReadWrite) == 0);
+    SubmitAmm(release);
+    Require(At(fresh) == 0xB1);
+    At(fresh) = 0xC1;
+    Require(At(fresh) == 0xC1);
+}
+
 void TestAmmRemapAndProtect() {
     std::uint64_t start = 0;
     std::uint64_t end = 0;
@@ -1313,6 +1372,7 @@ int main() {
     TestGatherScatter();
     TestAmmRangesAboveHintedMaps();
     TestAmmUsageStats();
+    TestAmmRemapKeepsPoolReferences();
     TestAmm();
     TestAmmRemapAndProtect();
     TestAmmPrt();