fix(agc): save the guest bytes of results that become pending over a surface that is no longer fully write-watched

core/libs/prx/libSceAgcDriver/CMakeLists.txt

@@ -1733,6 +1733,12 @@ if(BUILD_TESTING)
         set_tests_properties(agc_storage_idle_writeback_race agc_storage_kept_unit_settle PROPERTIES ENVIRONMENT_MODIFICATION
             "PATH=path_list_prepend:$<TARGET_FILE_DIR:libc>")
     endif()
+    add_test(NAME agc_storage_pending_after_unwatch COMMAND agc_driver_recorder_tests --pending-after-unwatch-only)
+    set_tests_properties(agc_storage_pending_after_unwatch PROPERTIES SKIP_RETURN_CODE 77 TIMEOUT 30)
+    if(WIN32)
+        set_tests_properties(agc_storage_pending_after_unwatch PROPERTIES ENVIRONMENT_MODIFICATION
+            "PATH=path_list_prepend:$<TARGET_FILE_DIR:libc>")
+    endif()
     add_test(NAME agc_storage_baseline_settled COMMAND agc_driver_recorder_tests --settled-baseline-only)
     add_test(NAME agc_storage_clear_hand_over COMMAND agc_driver_recorder_tests --clear-hand-over-only)
     set_tests_properties(agc_storage_clear_hand_over PROPERTIES SKIP_RETURN_CODE 77)

core/libs/prx/libSceAgcDriver/Execution/include/GuestMemory.hpp

@@ -65,6 +65,8 @@ void Unwatch(std::uint64_t address, std::size_t bytes);
 void Rewatch(std::span<const std::pair<std::uint64_t, std::uint64_t>> ranges);
 bool ImportWatched(std::uint64_t address, std::size_t bytes, const std::function<bool()>& import, bool staysWatched = true, bool importMayWriteGuest = true);
 void SetUnwatchHook(void (*hook)(std::uint64_t address, std::size_t bytes));
+bool UnwatchWritesProven();
+void SetImportWalkFailureForTest(int walk, std::uint64_t span = 0);
 #ifdef _WIN32
 using FreshCollectHook = bool (*)(std::uintptr_t address, std::size_t bytes, void (*written)(void*, std::uintptr_t, std::uintptr_t), void* context);
 void SetFreshCollectHook(FreshCollectHook hook);

core/libs/prx/libSceAgcDriver/Execution/src/GuestMemory.cpp

@@ -1357,6 +1357,23 @@ class TrackerLock {
         lock.lock();
     }
 
+    class Suspension {
+    public:
+        explicit Suspension(TrackerLock& owner) : owner(owner) {
+            owner.Suspend();
+        }
+
+        ~Suspension() {
+            owner.Resume();
+        }
+
+        Suspension(const Suspension&) = delete;
+        Suspension& operator=(const Suspension&) = delete;
+
+    private:
+        TrackerLock& owner;
+    };
+
 private:
     void publish() {
         const auto current = tracker.generation.load(std::memory_order_relaxed);
@@ -1758,6 +1775,10 @@ namespace {
 
 std::atomic<void (*)(std::uint64_t, std::size_t)> unwatchHook{nullptr};
 
+thread_local bool unwatchWritesProven = true;
+std::atomic<int> importWalkFailureForTest{0};
+std::atomic<std::uint64_t> importWalkFailureSpanForTest{0};
+
 void beforeUnwatch(std::uint64_t address, std::size_t bytes) {
     const auto hook = unwatchHook.load(std::memory_order_acquire);
     if (hook == nullptr) return;
@@ -1776,6 +1797,16 @@ void beforeUnwatch(std::uint64_t address, std::size_t bytes) {
 #endif
 }
 
+void beforeUnwatchUnproven(std::uint64_t address, std::size_t bytes) {
+    struct Unproven {
+        Unproven() { unwatchWritesProven = false; }
+        ~Unproven() { unwatchWritesProven = true; }
+        Unproven(const Unproven&) = delete;
+        Unproven& operator=(const Unproven&) = delete;
+    } unproven;
+    beforeUnwatch(address, bytes);
+}
+
 void unwatchLocked(WriteTracker& tracker, std::uint64_t address, std::size_t bytes) {
     if (!tracker.watched) return;
 #ifdef _WIN32
@@ -1792,6 +1823,15 @@ void SetUnwatchHook(void (*hook)(std::uint64_t, std::size_t)) {
     unwatchHook.store(hook, std::memory_order_release);
 }
 
+bool UnwatchWritesProven() {
+    return unwatchWritesProven;
+}
+
+void SetImportWalkFailureForTest(int walk, std::uint64_t span) {
+    importWalkFailureSpanForTest.store(span, std::memory_order_release);
+    importWalkFailureForTest.store(walk, std::memory_order_release);
+}
+
 #ifdef _WIN32
 void SetFreshCollectHook(FreshCollectHook hook) {
     freshCollectHook.store(hook, std::memory_order_release);
@@ -1884,29 +1924,39 @@ bool ImportWatched(std::uint64_t address, std::size_t bytes, const std::function
     std::vector<std::uint64_t> hashes(static_cast<std::size_t>((stop - first) / page), 0);
     for (const auto& [begin, end] : spans) CheckRange(reinterpret_cast<const void*>(begin), static_cast<std::size_t>(end - begin), 1, false);
     HashPages(first, spans, hashes);
-    const auto saveBytes = [&](std::uint64_t begin, std::uint64_t end) {
-        lock.Suspend();
-        beforeUnwatch(begin, static_cast<std::size_t>(end - begin));
-        lock.Resume();
+    const auto failWalk = importWalkFailureForTest.load(std::memory_order_acquire);
+    const auto failSpan = importWalkFailureSpanForTest.load(std::memory_order_acquire);
+    const auto failsWalk = [&](int walk, std::uint64_t begin, std::uint64_t end) {
+        return failWalk == walk && (failSpan == 0 || (begin <= failSpan && failSpan < end));
+    };
+    const auto leave = [&](const std::vector<std::pair<std::uint64_t, std::uint64_t>>& ranges) {
+        if (ranges.empty()) return;
+        const TrackerLock::Suspension suspended{lock};
+        for (const auto& [begin, end] : ranges) beforeUnwatchUnproven(begin, static_cast<std::size_t>(end - begin));
     };
     std::vector<std::pair<std::uint64_t, std::uint64_t>> leaving;
     std::erase_if(spans, [&](const auto& span) {
-        if (walkWrites(tracker, span.first, span.second, StampKind::Cpu, true)) return false;
+        if (!failsWalk(1, span.first, span.second) && walkWrites(tracker, span.first, span.second, StampKind::Cpu, true)) return false;
         leaving.push_back(span);
         return true;
     });
-    for (const auto& [begin, end] : leaving) saveBytes(begin, end);
+    leave(leaving);
     if (!runImport()) return false;
     const WindowCopy window{first, hashes};
-    auto unwatchedFrom = first;
-    for (const auto& [begin, end] : spans) {
-        if (!walkWrites(tracker, begin, end, StampKind::Cpu, true, &window)) {
-            saveBytes(begin, end);
-            continue;
+    leaving.clear();
+    std::erase_if(spans, [&](const auto& span) {
+        if (failsWalk(2, span.first, span.second) || !walkWrites(tracker, span.first, span.second, StampKind::Cpu, true, &window)) {
+            leaving.push_back(span);
+            return true;
         }
 #ifdef _WIN32
-        GuestArena::GuestArenaCompareResidentWhenShared_nid_no_patch(static_cast<std::uintptr_t>(begin), static_cast<std::size_t>(end - begin));
+        GuestArena::GuestArenaCompareResidentWhenShared_nid_no_patch(static_cast<std::uintptr_t>(span.first), static_cast<std::size_t>(span.second - span.first));
 #endif
+        return false;
+    });
+    leave(leaving);
+    auto unwatchedFrom = first;
+    for (const auto& [begin, end] : spans) {
         if (unwatchedFrom < begin) unwatchLocked(tracker, unwatchedFrom, static_cast<std::size_t>(begin - unwatchedFrom));
         unwatchedFrom = end;
     }
@@ -2058,9 +2108,10 @@ static std::uint64_t markWritten(std::uint64_t address, std::size_t bytes, bool
         const bool collected = GuestWriteWatch::GuestWriteWatchCollectFresh_nid_postfix(static_cast<std::uintptr_t>(blockFirst), static_cast<std::size_t>(blockStop - blockFirst), &stampWrittenRun, &runs);
 #endif
         if (!collected) {
-            lock.Suspend();
-            beforeUnwatch(blockFirst, static_cast<std::size_t>(blockStop - blockFirst));
-            lock.Resume();
+            {
+                const TrackerLock::Suspension suspended{lock};
+                beforeUnwatchUnproven(blockFirst, static_cast<std::size_t>(blockStop - blockFirst));
+            }
             unwatchLocked(tracker, blockFirst, static_cast<std::size_t>(blockStop - blockFirst));
             tracker.noteUnstampedStore(address, address + bytes);
             return 0;

core/libs/prx/libSceAgcDriver/Graphics/include/Texture.hpp

@@ -401,7 +401,26 @@ class StorageTexture : public std::enable_shared_from_this<StorageTexture> {
     std::vector<VkBufferImageCopy> CopyRegions(const std::vector<bool>* layers = nullptr) const;
     void upload(const std::vector<bool>* layers = nullptr);
     std::vector<bool> cpuUploadLayers(const std::vector<bool>& units) const;
-    void captureGuestBytes(const std::vector<bool>* layers);
+    void captureGuestBytes(const std::vector<bool>* layers, bool keepWritten = false);
+    bool keepWrittenSince(std::uint32_t unit);
+    enum class BaselineMiss : std::uint8_t {
+        None,
+        LeftWatchWhilePending,
+        SaveInFlight,
+        SavedBytesDropped,
+        NoUploadCopy,
+        UploadCopyDiffers,
+        NoUnwatchHash,
+        UnwatchUnproven,
+        UnwatchSawCpuWrites,
+        UnwatchWhileSaved,
+        HashDiffers,
+    };
+    static const char* BaselineMissText(BaselineMiss miss);
+    BaselineMiss baselineMiss(std::uint32_t unit) const;
+    std::array<std::uint64_t, 2> unitContentHash(std::uint32_t unit) const;
+    void hashBeforeUnwatch(std::uint64_t address, std::size_t bytes, bool proven);
+    void forgetUnwatchHash(std::uint32_t first, std::uint32_t count);
     bool originalHolds(std::uint64_t address, std::size_t bytes) const;
     bool compareUntracked(std::uint64_t address, std::size_t bytes, std::span<std::uint8_t> changed, bool memoize = false) const;
     // Stores the pending tracked layers overlapping [address, address + bytes) to guest memory; in
@@ -606,6 +625,14 @@ class StorageTexture : public std::enable_shared_from_this<StorageTexture> {
     mutable std::array<ComparedRange, 4> comparedRanges{};
     mutable std::uint32_t nextComparedRange = 0;
     std::vector<std::byte> generationBaseline;
+    struct UnwatchHash {
+        std::array<std::uint64_t, 2> hash{};
+        std::uint64_t generation = 0;
+        BaselineMiss skipped = BaselineMiss::None;
+        std::uint64_t skippedGeneration = 0;
+    };
+    std::vector<UnwatchHash> unwatchHash;
+    std::vector<BaselineMiss> pendingMiss;
     // DCC keys the image content was uploaded under: a fast-cleared surface starts as its clear value.
     DccKeys uploadedKeys = DccKeys::Uncompressed;
     mutable DccKeys filledKeys = DccKeys::Uncompressed;